NRPyPN: Validated Post-Newtonian Expressions for Numerical Relativity
"""
As documented in the NRPyPN notebook
NRPyPN_shortcuts.ipynb, this Python script
provides useful shortcuts for inputting
post-Newtonian expressions into SymPy/NRPy+

Basic functions:
dot(a,b): 3-vector dot product
cross(a,b): 3-vector cross product
div(a,b): a shortcut for SymPy's sp.Rational(a,b), to declare rational numbers
num_eval(expr): Numerically evaluates NRPyPN expressions

Author:  Zachariah B. Etienne
         zachetie **at** gmail **dot* com
"""
from typing import List, cast, Optional
import sympy as sp  # SymPy: The Python computer algebra package upon which NRPy+ depends
import nrpy.indexedexp as ixp  # NRPy+: Symbolic indexed expression (e.g., tensors, vectors, etc.) support

# Step 1: Declare several global variables used
#         throughout NRPyPN
m1, m2 = sp.symbols("m1 m2", real=True)
S1U = cast(List[sp.Expr], ixp.declarerank1("S1U"))
S2U = cast(List[sp.Expr], ixp.declarerank1("S2U"))
pU = cast(List[sp.Expr], ixp.declarerank1("pU"))
nU = cast(List[sp.Expr], ixp.declarerank1("nU"))

drdt = sp.symbols("drdt", real=True)
Pt, Pr = sp.symbols("Pt Pr", real=True)
# Some references use r, others use q to represent the
#   distance between the two point masses. This is rather
#   confusing since q is also used to represent the
#   mass ratio m2/m1. However, q is the canonical position
#   variable name in Hamiltonian mechanics, so both are
#   well justified. It should be obvious which is which
#   throughout NRPyPN.
r, q = sp.symbols("r q", real=True)
chi1U = cast(List[sp.Expr], ixp.declarerank1("chi1U"))
chi2U = cast(List[sp.Expr], ixp.declarerank1("chi2U"))

# Euler-Mascheroni gamma constant:
gamma_EulerMascheroni = sp.EulerGamma

# Derived quantities used in Damour et al papers:
n12U = ixp.zerorank1()
n21U = ixp.zerorank1()
p1U = ixp.zerorank1()
p2U = ixp.zerorank1()
for _ in range(3):
    n12U[_] = +nU[_]
    n21U[_] = -nU[_]
    p1U[_] = +pU[_]
    p2U[_] = -pU[_]


# Step 2.a: Define dot and cross product of vectors
def dot(vec1: List[sp.Expr], vec2: List[sp.Expr]) -> sp.Expr:
    """
    Compute the dot product of two vectors.

    :param vec1: First vector
    :param vec2: Second vector
    :return: Dot product of vec1 and vec2
    """
    vec1_dot_vec2: sp.Expr = sp.sympify(0)
    for i in range(3):
        vec1_dot_vec2 += vec1[i] * vec2[i]
    return vec1_dot_vec2


def cross(vec1: List[sp.Expr], vec2: List[sp.Expr]) -> List[sp.Expr]:
    """
    Compute the cross product of two vectors.

    :param vec1: First vector
    :param vec2: Second vector
    :return: Cross product of vec1 and vec2
    """
    vec1_cross_vec2: List[sp.Expr] = ixp.zerorank1()
    LeviCivitaSymbol = ixp.LeviCivitaSymbol_dim3_rank3()
    for i in range(3):
        for j in range(3):
            for k in range(3):
                vec1_cross_vec2[i] += LeviCivitaSymbol[i][j][k] * vec1[j] * vec2[k]
    return vec1_cross_vec2


def div(a: int, b: int) -> sp.Rational:
    """
    Create a SymPy Rational number a/b=Rational(a,b) from two integers a and b.

    :param a: Numerator
    :param b: Denominator
    :return: Rational number a/b
    """
    return cast(sp.Rational, sp.Rational(a, b))


# Step 3: num_eval(expr), a means to numerically evaluate SymPy/NRPyPN
#         expressions
def num_eval(
    expr: sp.Expr,
    qmassratio: float = 1.0,  # must be >= 1
    nr: float = 12.0,  # Orbital separation
    nchi1x: float = +0.0,
    nchi1y: float = +0.0,
    nchi1z: float = +0.0,
    nchi2x: float = +0.0,
    nchi2y: float = +0.0,
    nchi2z: float = +0.0,
    nPt: Optional[float] = None,
    ndrdt: Optional[float] = None,
) -> float:
    """
    Numerically evaluate SymPy/NRPyPN expressions.

    :param expr: Expression to evaluate
    :param qmassratio: Mass ratio, must be >= 1
    :param nr: Orbital separation
    :param nchi1x: x-component of the first spin
    :param nchi1y: y-component of the first spin
    :param nchi1z: z-component of the first spin
    :param nchi2x: x-component of the second spin
    :param nchi2y: y-component of the second spin
    :param nchi2z: z-component of the second spin
    :param nPt: Optional Pt value
    :param ndrdt: Optional drdt value
    :return: Numerical value of the expression
    """
    # DERIVED QUANTITIES BELOW
    # We want m1+m2 = 1, so that
    #         m2/m1 = qmassratio
    # We find below:
    nm1 = 1 / (1 + qmassratio)
    nm2 = qmassratio / (1 + qmassratio)
    # This way nm1+nm2 = (qmassratio+1)/(1+qmassratio) = 1 CHECK
    #      and nm2/nm1 = qmassratio                        CHECK

    nS1U0 = nchi1x * nm1**2
    nS1U1 = nchi1y * nm1**2
    nS1U2 = nchi1z * nm1**2
    nS2U0 = nchi2x * nm2**2
    nS2U1 = nchi2y * nm2**2
    nS2U2 = nchi2z * nm2**2

    if nPt:
        expr2 = expr.subs(Pt, nPt)
        expr = expr2
    if ndrdt:
        expr2 = expr.subs(drdt, ndrdt)
        expr = expr2
    return cast(
        float,
        expr.subs(m1, nm1)
        .subs(m2, nm2)
        .subs(S1U[0], nS1U0)
        .subs(S1U[1], nS1U1)
        .subs(S1U[2], nS1U2)
        .subs(S2U[0], nS2U0)
        .subs(S2U[1], nS2U1)
        .subs(S2U[2], nS2U2)
        .subs(chi1U[0], nchi1x)
        .subs(chi1U[1], nchi1y)
        .subs(chi1U[2], nchi1z)
        .subs(chi2U[0], nchi2x)
        .subs(chi2U[1], nchi2y)
        .subs(chi2U[2], nchi2z)
        .subs(r, nr)
        .subs(q, nr)
        .subs(sp.pi, sp.N(sp.pi))
        .subs(gamma_EulerMascheroni, sp.N(sp.EulerGamma)),
    )

nrpypn/NRPyPN_shortcuts.py
nrpy/nrpypn
"""
As documented in the NRPyPN notebook
PN-Hamiltonian-Nonspinning.ipynb, this Python script
generates nonspinning pieces of the post-Newtonian (PN)
Hamiltonian, up to and including third PN order.

Author:  Zachariah B. Etienne
         zachetie **at** gmail **dot* com
"""

from typing import List, cast
import sympy as sp  # SymPy: The Python computer algebra package upon which NRPy+ depends
import nrpy.indexedexp as ixp  # NRPy+: Symbolic indexed expression (e.g., tensors, vectors, etc.) support
import nrpy.validate_expressions.validate_expressions as ve

# NRPyPN: shortcuts for e.g., vector operations
from nrpypn.NRPyPN_shortcuts import div, dot


class PN_Hamiltonian_NS:
    """
    Class to handle Hamiltonian calculations in Newtonian and Post-Newtonian (PN) formalism for non-spinning objects.

    :param m1: Mass of the first object
    :param m2: Mass of the second object
    :param PU: Momentum 3-vector
    :param nU: Unit vector pointing from m1 to m2
    :param q: Scalar related to the separation between m1 and m2

    Attributes:
    - H_Newt (sp.Expr): Newtonian Hamiltonian
    - H_NS_1PN (sp.Expr): 1st order Post-Newtonian Hamiltonian
    - H_NS_2PN (sp.Expr): 2nd order Post-Newtonian Hamiltonian
    - H_NS_3PN (sp.Expr): 3rd order Post-Newtonian Hamiltonian

    Basic functions:
    f_H_Newt__H_NS_1PN__H_NS_2PN(m1,m2, pU, nU, q): Compute H_Newt,
                                                    H_NS_1PN, and H_NS_2PN
                                                    and store to class
                                                    variables of the same
                                                    names.
    f_H_NS_3PN: Compute H_NS_3PN, and store to global variable of same name
    """

    def __init__(
        self, m1: sp.Expr, m2: sp.Expr, PU: List[sp.Expr], nU: List[sp.Expr], q: sp.Expr
    ) -> None:
        self.H_Newt: sp.Expr = sp.sympify(0)
        self.H_NS_1PN: sp.Expr = sp.sympify(0)
        self.H_NS_2PN: sp.Expr = sp.sympify(0)
        self.H_NS_3PN: sp.Expr = sp.sympify(0)

        self.f_H_Newt__H_NS_1PN__H_NS_2PN(m1, m2, PU, nU, q)
        self.f_H_NS_3PN(m1, m2, PU, nU, q)

    def f_H_Newt__H_NS_1PN__H_NS_2PN(
        self, m1: sp.Expr, m2: sp.Expr, PU: List[sp.Expr], nU: List[sp.Expr], q: sp.Expr
    ) -> None:
        """
        Compute H_Newt, H_NS_1PN, and H_NS_2PN and store them in the class variables of the same names.
        """
        mu = m1 * m2 / (m1 + m2)
        eta = m1 * m2 / (m1 + m2) ** 2
        pU = ixp.zerorank1()
        for i in range(3):
            pU[i] = PU[i] / mu

        self.H_Newt = mu * (+div(1, 2) * dot(pU, pU) - 1 / q)

        # fmt: off
        self.H_NS_1PN = mu*(+div(1,8)*(3*eta-1)*dot(pU,pU)**2
                       -div(1,2)*((3+eta)*dot(pU,pU) + eta*dot(nU,pU)**2)/q
                       +div(1,2)/q**2)

        self.H_NS_2PN = mu*(+div(1,16)*(1 -  5*eta + 5*eta**2)*dot(pU,pU)**3
                       +div(1,8)*(+(5 - 20*eta - 3*eta**2)*dot(pU,pU)**2
                                  -2*eta**2*dot(nU,pU)**2*dot(pU,pU)
                                  -3*eta**2*dot(nU,pU)**4)/q
                       +div(1,2)*((5+8*eta)*dot(pU,pU) + 3*eta*dot(nU,pU)**2)/q**2
                       -div(1,4)*(1+3*eta)/q**3)
        # fmt: on

    def f_H_NS_3PN(
        self, m1: sp.Expr, m2: sp.Expr, PU: List[sp.Expr], nU: List[sp.Expr], q: sp.Expr
    ) -> None:
        """
        Compute H_NS_3PN and store it in the class variable of the same name.
        """
        mu = m1 * m2 / (m1 + m2)
        eta = m1 * m2 / (m1 + m2) ** 2
        pU = ixp.zerorank1()
        for i in range(3):
            pU[i] = PU[i] / mu

        # fmt: off
        self.H_NS_3PN = mu*(+div(1,128)*(-5 + 35*eta - 70*eta**2 + 35*eta**3)*dot(pU,pU)**4
                       +div(1, 16)*(+(-7 + 42*eta - 53*eta**2 -  5*eta**3)*dot(pU,pU)**3
                                    +(2-3*eta)*eta**2*dot(nU,pU)**2*dot(pU,pU)**2
                                    +3*(1-eta)*eta**2*dot(nU,pU)**4*dot(pU,pU) - 5*eta**3*dot(nU,pU)**6)/q
                       +(+div(1,16)*(-27 + 136*eta + 109*eta**2)*dot(pU,pU)**2
                         +div(1,16)*(+17 +  30*eta)*eta*dot(nU,pU)**2*dot(pU,pU)
                         +div(1,12)*(+ 5 +  43*eta)*eta*dot(nU,pU)**4)/q**2
                       +(+(-div(25, 8) + (div(1,64)*sp.pi**2 - div(335,48))*eta - div(23,8)*eta**2)*dot(pU,pU)
                         +(-div(85,16) - div(3,64)*sp.pi**2 - div(7,4)*eta)*eta*dot(nU,pU)**2)/q**3
                       +(+div(1,8)+(div(109,12) - div(21,32)*sp.pi**2)*eta)/q**4)
        # fmt: on


if __name__ == "__main__":
    import doctest
    import os
    import sys

    results = doctest.testmod()
    if results.failed > 0:
        print(f"Doctest failed: {results.failed} of {results.attempted} test(s)")
        sys.exit(1)
    else:
        print(f"Doctest passed: All {results.attempted} test(s) passed")

    in_m1, in_m2, in_q = sp.symbols("m1 m2 q")
    in_PU: List[sp.Expr] = cast(List[sp.Expr], ixp.declarerank1("PU"))
    in_nU: List[sp.Expr] = cast(List[sp.Expr], ixp.declarerank1("nU"))
    non_spinning_Hamiltonian = PN_Hamiltonian_NS(in_m1, in_m2, in_PU, in_nU, in_q)
    results_dict = ve.process_dictionary_of_expressions(
        non_spinning_Hamiltonian.__dict__, fixed_mpfs_for_free_symbols=True
    )
    ve.compare_or_generate_trusted_results(
        os.path.abspath(__file__),
        os.getcwd(),
        # File basename. If this is set to "trusted_module_test1", then
        #   trusted results_dict will be stored in tests/trusted_module_test1.py
        f"{os.path.splitext(os.path.basename(__file__))[0]}",
        results_dict,
    )

nrpypn/PN_Hamiltonian_NS.py
nrpy/nrpypn
"""
As documented in the NRPyPN notebook
PN-Hamiltonian-Spin-Orbit.ipynb, this Python script
generates spin-orbit coupling pieces of the
post-Newtonian (PN) Hamiltonian, up to and
including 3.5PN order.

Author:  Zachariah B. Etienne
         zachetie **at** gmail **dot* com
"""
from typing import List, cast
import sympy as sp
import nrpy.indexedexp as ixp  # NRPy+: Symbolic indexed expression (e.g., tensors, vectors, etc.) support
import nrpy.validate_expressions.validate_expressions as ve

# NRPyPN: shortcuts for e.g., vector operations
from nrpypn.NRPyPN_shortcuts import div, dot, cross


class PN_Hamiltonian_SO:
    """
    Class to compute post-Newtonian Hamiltonians for Spin-Orbit coupling terms.

    The Hamiltonians are calculated up to 3.5PN order based on Damour, Jaranowski, and Schäfer (2008).
    Reference: https://arxiv.org/abs/0711.1048

    Attributes:
    H_SO_1p5PN : sympy.Expr
        1.5PN Hamiltonian term
    H_SO_2p5PN : sympy.Expr
        2.5PN Hamiltonian term
    H_SO_3p5PN : sympy.Expr
        3.5PN Hamiltonian term

    :param m1: Mass of the first object
    :param m2: Mass of the second object
    :param n12U: Unit vector from object 1 to object 2
    :param n21U: Unit vector from object 2 to object 1
    :param S1U: Spin vector of object 1
    :param S2U: Spin vector of object 2
    :param p1U: Momentum vector of object 1
    :param p2U: Momentum vector of object 2
    :param r12: Scalar distance between object 1 and object 2
    """

    def __init__(
        self,
        m1: sp.Expr,
        m2: sp.Expr,
        n12U: List[sp.Expr],
        n21U: List[sp.Expr],
        S1U: List[sp.Expr],
        S2U: List[sp.Expr],
        p1U: List[sp.Expr],
        p2U: List[sp.Expr],
        r12: sp.Expr,
    ):
        self.H_SO_1p5PN: sp.Expr = sp.sympify(0)
        self.H_SO_2p5PN: sp.Expr = sp.sympify(0)
        self.H_SO_3p5PN: sp.Expr = sp.sympify(0)

        self.f_H_SO_1p5PN(m1, m2, n12U, n21U, S1U, S2U, p1U, p2U, r12)
        self.f_H_SO_2p5PN(m1, m2, n12U, n21U, S1U, S2U, p1U, p2U, r12)
        self.f_H_SO_3p5PN(m1, m2, n12U, n21U, S1U, S2U, p1U, p2U, r12)

    def f_H_SO_1p5PN(
        self,
        m1: sp.Expr,
        m2: sp.Expr,
        n12U: List[sp.Expr],
        n21U: List[sp.Expr],
        S1U: List[sp.Expr],
        S2U: List[sp.Expr],
        p1U: List[sp.Expr],
        p2U: List[sp.Expr],
        r12: sp.Expr,
    ) -> None:
        """
        1.5PN spin-orbit coupling term, from Eq. 4.11a of
           Damour, Jaranowski, and Schäfer (2008)
             https://arxiv.org/abs/0711.1048
        """

        def f_Omega1(
            m1: sp.Expr,
            m2: sp.Expr,
            n12U: List[sp.Expr],
            p1U: List[sp.Expr],
            p2U: List[sp.Expr],
            r12: sp.Expr,
        ) -> List[sp.Expr]:
            """Compute Omega vector"""
            Omega1 = ixp.zerorank1()
            for i in range(3):
                Omega1[i] = (
                    div(3, 2) * m2 / m1 * cross(n12U, p1U)[i] - 2 * cross(n12U, p2U)[i]
                ) / r12**2
            return Omega1

        Omega1 = f_Omega1(m1, m2, n12U, p1U, p2U, r12)
        Omega2 = f_Omega1(m2, m1, n21U, p2U, p1U, r12)
        self.H_SO_1p5PN = dot(Omega1, S1U) + dot(Omega2, S2U)

    #################################
    #################################

    def f_H_SO_2p5PN(
        self,
        m1: sp.Expr,
        m2: sp.Expr,
        n12U: List[sp.Expr],
        n21U: List[sp.Expr],
        S1U: List[sp.Expr],
        S2U: List[sp.Expr],
        p1U: List[sp.Expr],
        p2U: List[sp.Expr],
        r12: sp.Expr,
    ) -> None:
        """
        2.5PN spin-orbit coupling term, from Eq. 4.11b of
           Damour, Jaranowski, and Schäfer (2008)
             https://arxiv.org/abs/0711.1048
        """

        def f_Omega_SO_2p5PN(
            m1: sp.Expr,
            m2: sp.Expr,
            n12U: List[sp.Expr],
            p1U: List[sp.Expr],
            p2U: List[sp.Expr],
            r12: sp.Expr,
        ) -> List[sp.Expr]:
            """Compute Omega vector"""
            Omega1 = ixp.zerorank1()
            # fmt: off
            for i in range(3):
                Omega1[i] = (+(+(-div(11, 2) * m2 - 5 * m2 ** 2 / m1) * cross(n12U, p1U)[i]
                               + (6 * m1 + div(15, 2) * m2) * cross(n12U, p2U)[i]) / r12 ** 3
                             + (+(-div(5, 8) * m2 * dot(p1U, p1U) / m1 ** 3
                                  - div(3, 4) * dot(p1U, p2U) / m1 ** 2
                                  + div(3, 4) * dot(p2U, p2U) / (m1 * m2)
                                  - div(3, 4) * dot(n12U, p1U) * dot(n12U, p2U) / m1 ** 2
                                  - div(3, 2) * dot(n12U, p2U) ** 2 / (m1 * m2)) * cross(n12U, p1U)[i]
                                + (dot(p1U, p2U) / (m1 * m2) + 3 * dot(n12U, p1U) * dot(n12U, p2U) / (m1 * m2)) *
                                cross(n12U, p2U)[i]
                                + (div(3, 4) * dot(n12U, p1U) / m1 ** 2 - 2 * dot(n12U, p2U) / (m1 * m2)) * cross(p1U, p2U)[
                                    i]) / r12 ** 2)
            # fmt: on
            return Omega1

        Omega1_2p5PNU = f_Omega_SO_2p5PN(m1, m2, n12U, p1U, p2U, r12)
        Omega2_2p5PNU = f_Omega_SO_2p5PN(m2, m1, n21U, p2U, p1U, r12)

        self.H_SO_2p5PN = dot(Omega1_2p5PNU, S1U) + dot(Omega2_2p5PNU, S2U)

    #################################
    #################################
    def HS2011_Omega_SO_3p5PN_pt1(
        self,
        m1: sp.Expr,
        m2: sp.Expr,
        n12U: List[sp.Expr],
        p1U: List[sp.Expr],
        p2U: List[sp.Expr],
        r12: sp.Expr,
    ) -> List[sp.Expr]:
        """3.5PN H_SO:  Omega_1, part 1:"""

        Omega1 = ixp.zerorank1()
        # fmt: off
        for i in range(3):
            Omega1[i] = ((+7*m2*dot(p1U,p1U)**2/(16*m1**5)
                          +9*dot(n12U,p1U)*dot(n12U,p2U)*dot(p1U,p1U)/(16*m1**4)
                          +3*dot(p1U,p1U)*dot(n12U,p2U)**2/(4*m1**3*m2)
                          +45*dot(n12U,p1U)*dot(n12U,p2U)**3/(16*m1**2*m2**2)
                          +9*dot(p1U,p1U)*dot(p1U,p2U)/(16*m1**4)
                          -3*dot(n12U,p2U)**2*dot(p1U,p2U)/(16*m1**2*m2**2)
                          -3*dot(p1U,p1U)*dot(p2U,p2U)/(16*m1**3*m2)
                          -15*dot(n12U,p1U)*dot(n12U,p2U)*dot(p2U,p2U)/(16*m1**2*m2**2)
                          +3*dot(n12U,p2U)**2*dot(p2U,p2U)/(4*m1*m2**3)
                          -3*dot(p1U,p2U)*dot(p2U,p2U)/(16*m1**2*m2**2)
                          -3*dot(p2U,p2U)**2/(16*m1*m2**3))*cross(n12U,p1U)[i])/r12**2
        # fmt: on
        return Omega1

    def HS2011_Omega_SO_3p5PN_pt2(
        self,
        m1: sp.Expr,
        m2: sp.Expr,
        n12U: List[sp.Expr],
        p1U: List[sp.Expr],
        p2U: List[sp.Expr],
        r12: sp.Expr,
    ) -> List[sp.Expr]:
        """3.5PN H_SO:  Omega_1, part 2:"""
        Omega1 = ixp.zerorank1()
        # fmt: off
        for i in range(3):
            Omega1[i] = (+(-3*dot(n12U,p1U)*dot(n12U,p2U)*dot(p1U,p1U)/(2*m1**3*m2)
                           -15*dot(n12U,p1U)**2*dot(n12U,p2U)**2/(4*m1**2*m2**2)
                           +3*dot(p1U,p1U)*dot(n12U,p2U)**2/(4*m1**2*m2**2)
                           -dot(p1U,p1U)*dot(p1U,p2U)/(2*m1**3*m2)
                           +dot(p1U,p2U)**2/(2*m1**2*m2**2)
                           +3*dot(n12U,p1U)**2*dot(p2U,p2U)/(4*m1**2*m2**2)
                           -dot(p1U,p1U)*dot(p2U,p2U)/(4*m1**2*m2**2)
                           -3*dot(n12U,p1U)*dot(n12U,p2U)*dot(p2U,p2U)/(2*m1*m2**3)
                           -dot(p1U,p2U)*dot(p2U,p2U)/(2*m1*m2**3))*cross(n12U,p2U)[i])/r12**2
        # fmt: on
        return Omega1

    def HS2011_Omega_SO_3p5PN_pt3(
        self,
        m1: sp.Expr,
        m2: sp.Expr,
        n12U: List[sp.Expr],
        p1U: List[sp.Expr],
        p2U: List[sp.Expr],
        r12: sp.Expr,
    ) -> List[sp.Expr]:
        """3.5PN H_SO:  Omega_1, part 3:"""
        Omega1 = ixp.zerorank1()
        # fmt: off
        for i in range(3):
            Omega1[i] = (+(-9*dot(n12U,p1U)*dot(p1U,p1U)/(16*m1**4)
                           +dot(p1U,p1U)*dot(n12U,p2U)/(m1**3*m2)
                           +27*dot(n12U,p1U)*dot(n12U,p2U)**2/(16*m1**2*m2**2)
                           -dot(n12U,p2U)*dot(p1U,p2U)/(8*m1**2*m2**2)
                           -5*dot(n12U,p1U)*dot(p2U,p2U)/(16*m1**2*m2**2))*cross(p1U,p2U)[i])/r12**2
        # fmt: on
        return Omega1

    def HS2011_Omega_SO_3p5PN_pt4(
        self,
        m1: sp.Expr,
        m2: sp.Expr,
        n12U: List[sp.Expr],
        p1U: List[sp.Expr],
        p2U: List[sp.Expr],
        r12: sp.Expr,
    ) -> List[sp.Expr]:
        """3.5PN H_SO:  Omega_1, part 4:"""
        Omega1 = ixp.zerorank1()
        # fmt: off
        for i in range(3):
            Omega1[i] = (+(-3*m2*dot(n12U,p1U)**2/(2*m1**2)
                           +((-3*m2)/(2*m1**2) + 27*m2**2/(8*m1**3))*dot(p1U,p1U)
                           +(177/(16*m1) + 11/m2)*dot(n12U,p2U)**2
                           +(11/(2*m1) + 9*m2/(2*m1**2))*dot(n12U,p1U)*dot(n12U,p2U)
                           +(23/(4*m1) + 9*m2/(2*m1**2))*dot(p1U,p2U)
                           -(159/(16*m1) + 37/(8*m2))*dot(p2U,p2U))*cross(n12U,p1U)[i])/r12**3
        # fmt: on
        return Omega1

    def HS2011_Omega_SO_3p5PN_pt5(
        self,
        m1: sp.Expr,
        m2: sp.Expr,
        n12U: List[sp.Expr],
        p1U: List[sp.Expr],
        p2U: List[sp.Expr],
        r12: sp.Expr,
    ) -> List[sp.Expr]:
        """3.5PN H_SO:  Omega_1, part 5:"""
        Omega1 = ixp.zerorank1()
        # fmt: off
        for i in range(3):
            Omega1[i] = (+(+4*dot(n12U,p1U)**2/m1
                           +13*dot(p1U,p1U)/(2*m1)
                           +5*dot(n12U,p2U)**2/m2
                           +53*dot(p2U,p2U)/(8*m2)
                           -(211/(8*m1) + 22/m2)*dot(n12U,p1U)*dot(n12U,p2U)
                           -(47/(8*m1) + 5/m2)*dot(p1U,p2U))*cross(n12U,p2U)[i])/r12**3
        # fmt: on
        return Omega1

    def HS2011_Omega_SO_3p5PN_pt6(
        self,
        m1: sp.Expr,
        m2: sp.Expr,
        n12U: List[sp.Expr],
        p1U: List[sp.Expr],
        p2U: List[sp.Expr],
        r12: sp.Expr,
    ) -> List[sp.Expr]:
        """3.5PN H_SO:  Omega_1, part 6:"""
        Omega1 = ixp.zerorank1()
        # fmt: off
        for i in range(3):
            Omega1[i] = (+(-(8/m1 + 9*m2/(2*m1**2))*dot(n12U,p1U)
                           +(59/(4*m1) + 27/(2*m2))*dot(n12U,p2U))*cross(p1U,p2U)[i])/r12**3
        return Omega1

    def HS2011_Omega_SO_3p5PN_pt7(
        self,
        m1: sp.Expr,
        m2: sp.Expr,
        n12U: List[sp.Expr],
        p1U: List[sp.Expr],
        p2U: List[sp.Expr],
        r12: sp.Expr,
    ) -> List[sp.Expr]:
        """3.5PN H_SO:  Omega_1, part 7:"""
        Omega1 = ixp.zerorank1()
        # fmt: off
        for i in range(3):
            Omega1[i] = (+(181*m1*m2/16 + 95*m2**2/4 + 75*m2**3/(8*m1))*cross(n12U,p1U)[i]
                         -(21*m1**2/2 + 473*m1*m2/16 + 63*m2**2/4)*cross(n12U,p2U)[i])/r12**4
        # fmt: on
        return Omega1

    def f_H_SO_3p5PN(
        self,
        m1: sp.Expr,
        m2: sp.Expr,
        n12U: List[sp.Expr],
        n21U: List[sp.Expr],
        S1U: List[sp.Expr],
        S2U: List[sp.Expr],
        p1U: List[sp.Expr],
        p2U: List[sp.Expr],
        r12: sp.Expr,
    ) -> None:
        """
        3.5PN spin-orbit coupling term, from Eq. 5 of
        Hartung and Steinhoff (2011)
        https://arxiv.org/abs/1104.3079

        3.5PN H_SO: Combining all the above Omega terms
        into the full 3.5PN S-O Hamiltonian
        expression
        """
        Omega1_3p5PNU = ixp.zerorank1()
        Omega2_3p5PNU = ixp.zerorank1()
        for i in range(3):
            Omega1_3p5PNU[i] = self.HS2011_Omega_SO_3p5PN_pt1(
                m1, m2, n12U, p1U, p2U, r12
            )[i]
            Omega1_3p5PNU[i] += self.HS2011_Omega_SO_3p5PN_pt2(
                m1, m2, n12U, p1U, p2U, r12
            )[i]
            Omega1_3p5PNU[i] += self.HS2011_Omega_SO_3p5PN_pt3(
                m1, m2, n12U, p1U, p2U, r12
            )[i]
            Omega1_3p5PNU[i] += self.HS2011_Omega_SO_3p5PN_pt4(
                m1, m2, n12U, p1U, p2U, r12
            )[i]
            Omega1_3p5PNU[i] += self.HS2011_Omega_SO_3p5PN_pt5(
                m1, m2, n12U, p1U, p2U, r12
            )[i]
            Omega1_3p5PNU[i] += self.HS2011_Omega_SO_3p5PN_pt6(
                m1, m2, n12U, p1U, p2U, r12
            )[i]
            Omega1_3p5PNU[i] += self.HS2011_Omega_SO_3p5PN_pt7(
                m1, m2, n12U, p1U, p2U, r12
            )[i]

            Omega2_3p5PNU[i] = self.HS2011_Omega_SO_3p5PN_pt1(
                m2, m1, n21U, p2U, p1U, r12
            )[i]
            Omega2_3p5PNU[i] += self.HS2011_Omega_SO_3p5PN_pt2(
                m2, m1, n21U, p2U, p1U, r12
            )[i]
            Omega2_3p5PNU[i] += self.HS2011_Omega_SO_3p5PN_pt3(
                m2, m1, n21U, p2U, p1U, r12
            )[i]
            Omega2_3p5PNU[i] += self.HS2011_Omega_SO_3p5PN_pt4(
                m2, m1, n21U, p2U, p1U, r12
            )[i]
            Omega2_3p5PNU[i] += self.HS2011_Omega_SO_3p5PN_pt5(
                m2, m1, n21U, p2U, p1U, r12
            )[i]
            Omega2_3p5PNU[i] += self.HS2011_Omega_SO_3p5PN_pt6(
                m2, m1, n21U, p2U, p1U, r12
            )[i]
            Omega2_3p5PNU[i] += self.HS2011_Omega_SO_3p5PN_pt7(
                m2, m1, n21U, p2U, p1U, r12
            )[i]

        self.H_SO_3p5PN = dot(Omega1_3p5PNU, S1U) + dot(Omega2_3p5PNU, S2U)


if __name__ == "__main__":
    import doctest
    import os
    import sys

    results = doctest.testmod()
    if results.failed > 0:
        print(f"Doctest failed: {results.failed} of {results.attempted} test(s)")
        sys.exit(1)
    else:
        print(f"Doctest passed: All {results.attempted} test(s) passed")

    in_m1, in_m2, in_r12 = sp.symbols("m1 m2 r12")
    in_n12U: List[sp.Expr] = cast(List[sp.Expr], ixp.declarerank1("n12U"))
    in_n21U: List[sp.Expr] = cast(List[sp.Expr], ixp.declarerank1("n21U"))
    in_S1U: List[sp.Expr] = cast(List[sp.Expr], ixp.declarerank1("S1U"))
    in_S2U: List[sp.Expr] = cast(List[sp.Expr], ixp.declarerank1("S2U"))
    in_p1U: List[sp.Expr] = cast(List[sp.Expr], ixp.declarerank1("p1U"))
    in_p2U: List[sp.Expr] = cast(List[sp.Expr], ixp.declarerank1("p2U"))
    in_nU: List[sp.Expr] = cast(List[sp.Expr], ixp.declarerank1("nU"))
    non_spinning_Hamiltonian = PN_Hamiltonian_SO(
        in_m1, in_m2, in_n12U, in_n21U, in_S1U, in_S2U, in_p1U, in_p2U, in_r12
    )
    results_dict = ve.process_dictionary_of_expressions(
        non_spinning_Hamiltonian.__dict__, fixed_mpfs_for_free_symbols=True
    )
    ve.compare_or_generate_trusted_results(
        os.path.abspath(__file__),
        os.getcwd(),
        # File basename. If this is set to "trusted_module_test1", then
        #   trusted results_dict will be stored in tests/trusted_module_test1.py
        f"{os.path.splitext(os.path.basename(__file__))[0]}",
        results_dict,
    )

nrpypn/PN_Hamiltonian_SO.py
nrpy/nrpypn
"""
As documented in the NRPyPN notebook
PN-Hamiltonian-Spin-Spin.ipynb, this Python script
generates spin-spin coupling pieces of the
post-Newtonian (PN) Hamiltonian, up to and
including 3PN order.

Author:  Zachariah B. Etienne
         zachetie **at** gmail **dot* com
"""
from typing import List, cast
import sympy as sp
import nrpy.indexedexp as ixp  # NRPy+: Symbolic indexed expression (e.g., tensors, vectors, etc.) support
import nrpy.validate_expressions.validate_expressions as ve

# NRPyPN: shortcuts for e.g., vector operations
from nrpypn.NRPyPN_shortcuts import div, dot, cross


class PN_Hamiltonian_SS:
    """
    This Python class generates spin-spin coupling pieces of the
    post-Newtonian (PN) Hamiltonian, up to and including 3PN order.

    :param m1: Mass of object 1.
    :param m2: Mass of object 2.
    :param n12U: Unit vector from object 1 to object 2.
    :param n21U: Unit vector from object 2 to object 1.
    :param S1U: Spin vector of object 1.
    :param S2U: Spin vector of object 2.
    :param p1U: Momentum vector of object 1.
    :param p2U: Momentum vector of object 2.
    :param r12: Distance between object 1 and object 2.
    """

    def __init__(
        self,
        m1: sp.Expr,
        m2: sp.Expr,
        n12U: List[sp.Expr],
        n21U: List[sp.Expr],
        S1U: List[sp.Expr],
        S2U: List[sp.Expr],
        p1U: List[sp.Expr],
        p2U: List[sp.Expr],
        r12: sp.Expr,
    ) -> None:
        # f_H_SS_2PN(m1,m2, S1U,S2U, nU, r12):
        #       Compute the complete H_SS_2PN term and store to
        #                     global variable of the same name.
        # nU: as defined in NRPyPN_shortcuts:
        #  n12U[i] = +nU[i]
        self.f_H_SS_2PN(m1, m2, S1U, S2U, n12U, r12)
        # f_H_SS_S1S2_3PN(m1,m2, n12U, S1U,S2U, p1U,p2U, r12):
        #       Compute HS1S2_3PN and store to global variable
        #                     of the same name.
        self.f_H_SS_S1S2_3PN(m1, m2, n12U, S1U, S2U, p1U, p2U, r12)
        # f_H_SS_S1sq_S2sq_3PN(m1,m2, n12U,n21U, S1U,S2U, p1U,p2U, r12):
        #       Compute H_SS_S1sq_S2sq_3PN and store to global
        #                     variable of the same name.
        self.f_H_SS_S1sq_S2sq_3PN(m1, m2, n12U, n21U, S1U, S2U, p1U, p2U, r12)

    def f_H_SS_2PN(
        self,
        m1: sp.Expr,
        m2: sp.Expr,
        S1U: List[sp.Expr],
        S2U: List[sp.Expr],
        nU: List[sp.Expr],
        q: sp.Expr,
    ) -> None:
        """
        2PN spin-spin term, from Eqs. 2.18 and 2.19 of
             Buonanno, Chen, and Damour (2006):
           https://arxiv.org/abs/gr-qc/0508067
        """
        S0U = ixp.zerorank1()
        for i in range(3):
            S0U[i] = (1 + m2 / m1) * S1U[i] + (1 + m1 / m2) * S2U[i]
        mu = m1 * m2 / (m1 + m2)
        self.H_SS_2PN = (
            mu / (m1 + m2) * (3 * dot(S0U, nU) ** 2 - dot(S0U, S0U)) / (2 * q**3)
        )

    #################################
    #################################
    def f_H_SS_S1S2_3PN(
        self,
        m1: sp.Expr,
        m2: sp.Expr,
        n12U: List[sp.Expr],
        S1U: List[sp.Expr],
        S2U: List[sp.Expr],
        p1U: List[sp.Expr],
        p2U: List[sp.Expr],
        r12: sp.Expr,
    ) -> None:
        """
        3PN spin-spin S_1,S_2 coupling term, from Eq. 2.11 of
              Steinhoff, Hergt, and Sch\"afer (2008a)
                https://arxiv.org/abs/0712.1716
        """
        # fmt: off
        H_SS_S1S2_3PN = (+div(3,2)*(dot(cross(p1U,S1U),n12U)*dot(cross(p2U,S2U),n12U))
                         +       6*(dot(cross(p2U,S1U),n12U)*dot(cross(p1U,S2U),n12U))
                         -15*dot(S1U,n12U)*dot(S2U,n12U)*dot(p1U,n12U)*dot(p2U,n12U)
                         -3*dot(S1U,n12U)*dot(S2U,n12U)*dot(p1U,p2U)
                         +3*dot(S1U,p2U)*dot(S2U,n12U)*dot(p1U,n12U)
                         +3*dot(S2U,p1U)*dot(S1U,n12U)*dot(p2U,n12U)
                         +3*dot(S1U,p1U)*dot(S2U,n12U)*dot(p2U,n12U)
                         +3*dot(S2U,p2U)*dot(S1U,n12U)*dot(p1U,n12U)
                         -div(1,2)*dot(S1U,p2U)*dot(S2U,p1U)
                         +dot(S1U,p1U)*dot(S2U,p2U)
                         -3*dot(S1U,S2U)*dot(p1U,n12U)*dot(p2U,n12U)
                         +div(1,2)*dot(S1U,S2U)*dot(p1U,p2U))/(2*m1*m2*r12**3)
        H_SS_S1S2_3PN +=(-dot(cross(p1U,S1U),n12U)*dot(cross(p1U,S2U),n12U)
                         +dot(S1U,S2U)*dot(p1U,n12U)**2
                         -dot(S1U,n12U)*dot(S2U,p1U)*dot(p1U,n12U))*3/(2*m1**2*r12**3)
        H_SS_S1S2_3PN +=(-dot(cross(p2U,S2U),n12U)*dot(cross(p2U,S1U),n12U)
                         +dot(S1U,S2U)*dot(p2U,n12U)**2
                         -dot(S2U,n12U)*dot(S1U,p1U)*dot(p2U,n12U))*3/(2*m2**2*r12**3)
        H_SS_S1S2_3PN +=( dot(S1U,S2U)-2*dot(S1U,n12U)*dot(S2U,n12U))*6*(m1+m2)/r12**4
        self.H_SS_S1S2_3PN = H_SS_S1S2_3PN
        # fmt: on

    #################################
    #################################
    def f_H_SS_S1sq_S2sq_3PN(
        self,
        m1: sp.Expr,
        m2: sp.Expr,
        n12U: List[sp.Expr],
        n21U: List[sp.Expr],
        S1U: List[sp.Expr],
        S2U: List[sp.Expr],
        p1U: List[sp.Expr],
        p2U: List[sp.Expr],
        r12: sp.Expr,
    ) -> None:
        """
        # 3PN spin-orbit coupling term, from Eq. 9 of
        #    Steinhoff, Hergt, and Sch\"afer (2008b)
        #       https://arxiv.org/abs/0809.2200
        """

        def f_H_SS_particle(
            m1: sp.Expr,
            m2: sp.Expr,
            n12U: List[sp.Expr],
            S1U: List[sp.Expr],
            _S2U: List[sp.Expr],  # _S2U unused.
            p1U: List[sp.Expr],
            p2U: List[sp.Expr],
            r12: sp.Expr,
        ) -> sp.Expr:
            """Per-object H_SS contribution, call once to compute 1 on 2 and again for 2 on 1"""
            # fmt: off
            H_SS_S1sq_S2sq_3PN_particle = (
                +  m2/(4*m1**3)*dot(p1U,S1U)**2
                +3*m2/(8*m1**3)*dot(p1U,n12U)**2*dot(S1U,S1U)
                -3*m2/(8*m1**3)*dot(p1U,p1U)*dot(S1U,n12U)**2
                -3*m2/(4*m1**3)*dot(p1U,n12U)*dot(S1U,n12U)*dot(p1U,S1U)
                -3/(4*m1*m2)*dot(p2U,p2U)*dot(S1U,S1U)
                +9/(4*m1*m2)*dot(p2U,p2U)*dot(S1U,n12U)**2
                +3/(4*m1**2)*dot(p1U,p2U)*dot(S1U,S1U)
                -9/(4*m1**2)*dot(p1U,p2U)*dot(S1U,n12U)**2
                -3/(2*m1**2)*dot(p1U,n12U)*dot(p2U,S1U)*dot(S1U,n12U)
                +3/(m1**2)  *dot(p2U,n12U)*dot(p1U,S1U)*dot(S1U,n12U)
                +3/(4*m1**2)*dot(p1U,n12U)*dot(p2U,n12U)*dot(S1U,S1U)
                -15/(4*m1**2)*dot(p1U,n12U)*dot(p2U,n12U)*dot(S1U,n12U)**2)/r12**3
            H_SS_S1sq_S2sq_3PN_particle+= -(+div(9,2)*dot(S1U,n12U)**2
                                            -div(5,2)*dot(S1U,S1U)
                                            +7*m2/m1*dot(S1U,n12U)**2
                                            -3*m2/m1*dot(S1U,S1U))*m2/r12**4
            # fmt: on
            return cast(sp.Expr, H_SS_S1sq_S2sq_3PN_particle)

        # fmt: off
        self.H_SS_S1sq_S2sq_3PN = (+f_H_SS_particle(m1,m2, n12U, S1U,S2U, p1U,p2U, r12)
                                   +f_H_SS_particle(m2,m1, n21U, S2U,S1U, p2U,p1U, r12))
        # fmt: on


if __name__ == "__main__":
    import doctest
    import os
    import sys

    results = doctest.testmod()
    if results.failed > 0:
        print(f"Doctest failed: {results.failed} of {results.attempted} test(s)")
        sys.exit(1)
    else:
        print(f"Doctest passed: All {results.attempted} test(s) passed")

    in_m1, in_m2, in_r12 = sp.symbols("m1 m2 r12")
    in_n12U: List[sp.Expr] = cast(List[sp.Expr], ixp.declarerank1("n12U"))
    in_n21U: List[sp.Expr] = cast(List[sp.Expr], ixp.declarerank1("n21U"))
    in_S1U: List[sp.Expr] = cast(List[sp.Expr], ixp.declarerank1("S1U"))
    in_S2U: List[sp.Expr] = cast(List[sp.Expr], ixp.declarerank1("S2U"))
    in_p1U: List[sp.Expr] = cast(List[sp.Expr], ixp.declarerank1("p1U"))
    in_p2U: List[sp.Expr] = cast(List[sp.Expr], ixp.declarerank1("p2U"))
    spin_spin_Hamiltonian = PN_Hamiltonian_SS(
        in_m1, in_m2, in_n12U, in_n21U, in_S1U, in_S2U, in_p1U, in_p2U, in_r12
    )
    results_dict = ve.process_dictionary_of_expressions(
        spin_spin_Hamiltonian.__dict__, fixed_mpfs_for_free_symbols=True
    )
    ve.compare_or_generate_trusted_results(
        os.path.abspath(__file__),
        os.getcwd(),
        # File basename. If this is set to "trusted_module_test1", then
        #   trusted results_dict will be stored in tests/trusted_module_test1.py
        f"{os.path.splitext(os.path.basename(__file__))[0]}",
        results_dict,
    )

nrpypn/PN_Hamiltonian_SS.py
nrpy/nrpypn
"""
As documented in the NRPyPN notebook
PN-Hamiltonian-SSS.ipynb, this Python script
generates spin-spin-spin coupling pieces of the
post-Newtonian (PN) Hamiltonian, up to and
including 3PN order.

Author:  Zachariah B. Etienne
         zachetie **at** gmail **dot* com
"""
from typing import List, cast
import sympy as sp
import nrpy.indexedexp as ixp  # NRPy+: Symbolic indexed expression (e.g., tensors, vectors, etc.) support
import nrpy.validate_expressions.validate_expressions as ve

# NRPyPN: shortcuts for e.g., vector operations
from nrpypn.NRPyPN_shortcuts import div, dot, cross


#################################
class PN_Hamiltonian_SSS:
    """
    Implement 3PN spin-spin-spin term, from Eq. 3.12 of
            Levi and Steinhoff (2015):
         https://arxiv.org/abs/1410.2601

    :param m1: Mass of the first object
    :param m2: Mass of the second object
    :param n12U: Unit vector from the first to the second object
    :param n21U: Unit vector from the second to the first object
    :param S1U: Spin vector of the first object
    :param S2U: Spin vector of the second object
    :param p1U: Momentum of the first object
    :param p2U: Momentum of the second object
    :param r12: Separation between the two objects
    """

    def __init__(
        self,
        m1: sp.Expr,
        m2: sp.Expr,
        n12U: List[sp.Expr],
        n21U: List[sp.Expr],
        S1U: List[sp.Expr],
        S2U: List[sp.Expr],
        p1U: List[sp.Expr],
        p2U: List[sp.Expr],
        r12: sp.Expr,
    ) -> None:
        """Initialize class by calling f_H_SSS_3PN."""
        self.f_H_SSS_3PN(m1, m2, n12U, n21U, S1U, S2U, p1U, p2U, r12)

    def f_H_SSS_3PN(
        self,
        m1: sp.Expr,
        m2: sp.Expr,
        n12U: List[sp.Expr],
        n21U: List[sp.Expr],
        S1U: List[sp.Expr],
        S2U: List[sp.Expr],
        p1U: List[sp.Expr],
        p2U: List[sp.Expr],
        r12: sp.Expr,
    ) -> None:
        """3PN spin-spin-spin term, from Eq. 3.12 of Levi and Steinhoff (2015):
        https://arxiv.org/abs/1410.2601"""

        def f_H_SSS_3PN_pt(
            m1: sp.Expr,
            m2: sp.Expr,
            nU: List[sp.Expr],
            S1U: List[sp.Expr],
            S2U: List[sp.Expr],
            p1U: List[sp.Expr],
            p2U: List[sp.Expr],
            r: sp.Expr,
        ) -> sp.Expr:
            """Implement one part of the spin-spin-spin, call once for particle 1 acting on 2,
            and again for vice-versa"""
            p2_minus_m2_over_4m1_p1 = ixp.zerorank1()
            for i in range(3):
                p2_minus_m2_over_4m1_p1[i] = p2U[i] - m2 / (4 * m1) * p1U[i]
            # fmt: off
            H_SSS_3PN_pt = (div(3,2)*(dot(S1U,S1U)*dot(S2U,cross(nU,p1U))
                                       +dot(S1U,nU)*dot(S2U,cross(S1U,p1U))
                                       -5*dot(S1U,nU)**2*dot(S2U,cross(nU,p1U))
                                       +dot(nU,cross(S1U,S2U))*( dot(S1U,p1U)
                                                                -5*dot(S1U,nU)*dot(p1U,nU)))
                            -3*m1/(2*m2)*(   dot(S1U,S1U)  *dot(S2U,cross(nU,p2U))
                                          +2*dot(S1U,nU)   *dot(S2U,cross(S1U,p2U))
                                          -5*dot(S1U,nU)**2*dot(S2U,cross(nU,p2U)))
                            -dot(cross(S1U,nU),p2_minus_m2_over_4m1_p1)*(dot(S1U,S1U) - 5*dot(S1U,nU)**2))/(m1**2*r**4)
            # fmt: on
            return cast(sp.Expr, H_SSS_3PN_pt)

        self.H_SSS_3PN = +f_H_SSS_3PN_pt(
            m1, m2, n12U, S1U, S2U, p1U, p2U, r12
        ) + f_H_SSS_3PN_pt(m2, m1, n21U, S2U, S1U, p2U, p1U, r12)


if __name__ == "__main__":
    import doctest
    import os
    import sys

    results = doctest.testmod()
    if results.failed > 0:
        print(f"Doctest failed: {results.failed} of {results.attempted} test(s)")
        sys.exit(1)
    else:
        print(f"Doctest passed: All {results.attempted} test(s) passed")

    in_m1, in_m2, in_r12 = sp.symbols("m1 m2 r12")
    in_n12U: List[sp.Expr] = cast(List[sp.Expr], ixp.declarerank1("n12U"))
    in_n21U: List[sp.Expr] = cast(List[sp.Expr], ixp.declarerank1("n21U"))
    in_S1U: List[sp.Expr] = cast(List[sp.Expr], ixp.declarerank1("S1U"))
    in_S2U: List[sp.Expr] = cast(List[sp.Expr], ixp.declarerank1("S2U"))
    in_p1U: List[sp.Expr] = cast(List[sp.Expr], ixp.declarerank1("p1U"))
    in_p2U: List[sp.Expr] = cast(List[sp.Expr], ixp.declarerank1("p2U"))
    spin_spin_spin_Hamiltonian = PN_Hamiltonian_SSS(
        in_m1, in_m2, in_n12U, in_n21U, in_S1U, in_S2U, in_p1U, in_p2U, in_r12
    )
    results_dict = ve.process_dictionary_of_expressions(
        spin_spin_spin_Hamiltonian.__dict__, fixed_mpfs_for_free_symbols=True
    )
    ve.compare_or_generate_trusted_results(
        os.path.abspath(__file__),
        os.getcwd(),
        # File basename. If this is set to "trusted_module_test1", then
        #   trusted results_dict will be stored in tests/trusted_module_test1.py
        f"{os.path.splitext(os.path.basename(__file__))[0]}",
        results_dict,
    )

nrpypn/PN_Hamiltonian_SSS.py
nrpy/nrpypn
"""
As documented in the NRPyPN notebook
PN-MOmega.ipynb, this Python script
generates the expression for the orbital
angular frequency M*Omega up to and
including terms at 3.5PN order.
It largely follows the expression of
 Ramos-Buades, Husa, and Pratten (2018)
   https://arxiv.org/abs/1810.00036
but fixes obvious typos in the given
equation.

Author:  Zachariah B. Etienne
         zachetie **at** gmail **dot* com
"""
from typing import List, cast
import sympy as sp  # SymPy: The Python computer algebra package upon which NRPy+ depends
import nrpy.indexedexp as ixp  # NRPy+: Symbolic indexed expression (e.g., tensors, vectors, etc.) support
import nrpy.validate_expressions.validate_expressions as ve

from nrpypn.NRPyPN_shortcuts import div  # NRPyPN: shortcuts for e.g., vector operations


class PN_MOmega:
    """
    Basic equation for MOmega can be written in the
      form:
    MOmega = 1/r^(3/2) * (1 + sum_{k=2}^7 (a_k/r^{k/2}))
    where we construct the a_k terms in the sum below:
    """

    def __init__(
        self,
        m1: sp.Expr,
        m2: sp.Expr,
        chi1U: List[sp.Expr],
        chi2U: List[sp.Expr],
        r: sp.Expr,
    ):
        self.a_2: sp.Expr = sp.sympify(0)
        self.a_3: sp.Expr = sp.sympify(0)
        self.a_4: sp.Expr = sp.sympify(0)
        self.a_5: sp.Expr = sp.sympify(0)
        self.a_6: sp.Expr = sp.sympify(0)
        self.a_7: sp.Expr = sp.sympify(0)
        self.MOmega: sp.Expr = sp.sympify(0)
        self.f_MOmega(m1, m2, chi1U, chi2U, r)

    def MOmega__a_2_thru_a_4(
        self,
        m1: sp.Expr,
        m2: sp.Expr,
        chi1x: sp.Expr,
        chi1y: sp.Expr,
        chi1z: sp.Expr,
        chi2x: sp.Expr,
        chi2y: sp.Expr,
        chi2z: sp.Expr,
    ) -> None:
        """
        Construct terms a_2, a_3, and a_4, from
          Eq A2 of Ramos-Buades, Husa, and Pratten (2018)
        https://arxiv.org/abs/1810.00036
        These terms have been independently validated
        against the same terms in Eq 6 of
          Healy, Lousto, Nakano, and Zlochower (2017)
          https://arxiv.org/abs/1702.00872
        """
        q = m2 / m1  # It is assumed that q >= 1, so m2 >= m1.
        # fmt: off
        self.a_2 = -((3*q**2+5*q+3)/(2*(q+1)**2))
        self.a_3 = (-(3*q+4)*chi1z/(4*(q+1)**2) - q*(4*q+3)*chi2z/(4*(q+1)**2))
        self.a_4 = (-3*q**2*chi2x**2/(2*(q+1)**2)
                    +3*q**2*chi2y**2/(4*(q+1)**2)
                    +3*q**2*chi2z**2/(4*(q+1)**2)
                    +(+24*q**4 + 103*q**3 + 164*q**2 + 103*q + 24)/(16*(q+1)**4)
                    -3*chi1x**2/(2*(q+1)**2)
                    -3*q*chi1x*chi2x/(q+1)**2
                    +3*chi1y**2/(4*(q+1)**2)
                    +3*q*chi1y*chi2y/(2*(q+1)**2)
                    +3*chi1z**2/(4*(q+1)**2)
                    +3*q*chi1z*chi2z/(2*(q+1)**2))
        # fmt: on

    def MOmega__a_5_thru_a_6(
        self,
        m1: sp.Expr,
        m2: sp.Expr,
        chi1x: sp.Expr,
        chi1y: sp.Expr,
        chi1z: sp.Expr,
        chi2x: sp.Expr,
        chi2y: sp.Expr,
        chi2z: sp.Expr,
    ) -> None:
        """
        Construct terms a_5 and a_6, from
         Eq A1 of Ramos-Buades, Husa, and Pratten (2018)
           https://arxiv.org/abs/1810.00036
         These terms have been independently validated
           against the same terms in Eq 6 of
         Healy, Lousto, Nakano, and Zlochower (2017)
           https://arxiv.org/abs/1702.00872
        """
        q = m2 / m1  # It is assumed that q >= 1, so m2 >= m1.
        # fmt: off
        self.a_5 = (+3*  (13*q**3 + 34*q**2 + 30*q + 16)*chi1z/(16*(q+1)**4)
                    +3*q*(16*q**3 + 30*q**2 + 34*q + 13)*chi2z/(16*(q+1)**4))
        self.a_6 = (+(+155*q**2 + 180*q + 76)*chi1x**2/(16*(q+1)**4)
                    +q*(+120*q**2 + 187*q + 120)*chi1x*chi2x/(8*(q+1)**4)
                    -(+55*q**2 + 85*q + 43)*chi1y**2/(8*(q+1)**4)
                    -q*(+54*q**2 + 95*q + 54)*chi1y*chi2y/( 4*(q+1)**4)
                    -q*(+96*q**2 +127*q + 96)*chi1z*chi2z/(16*(q+1)**4)
                    +q**2*(+76*q**2 + 180*q + 155)*chi2x**2/(16*(q+1)**4)
                    -q**2*(+43*q**2 +  85*q +  55)*chi2y**2/( 8*(q+1)**4)
                    -q**2*(+2*q+5)*(+14*q+27)*chi2z**2/(32*(q+1)**4)
                    -     (+5*q+2)*(+27*q+14)*chi1z**2/(32*(q+1)**4)
                    +(+501*sp.pi**2*q*(q+1)**4
                      -4*(120*q**6 + 2744*q**5 + 10049*q**4 + 14820*q**3 + 10049*q**2 + 2744*q + 120))/(384*(q+1)**6))
        # fmt: on

    def MOmega__a_7(
        self,
        m1: sp.Expr,
        m2: sp.Expr,
        chi1x: sp.Expr,
        chi1y: sp.Expr,
        chi1z: sp.Expr,
        chi2x: sp.Expr,
        chi2y: sp.Expr,
        chi2z: sp.Expr,
    ) -> None:
        """Construct term a_7, from Eq A1 of  Ramos-Buades, Husa, and Pratten (2018) https://arxiv.org/abs/1810.00036"""
        q = m2 / m1  # It is assumed that q >= 1, so m2 >= m1.
        # fmt: off
        self.a_7 = (+3*(4*q+1)*q**3*chi2x**2*chi2z/(2*(q+1)**4)
               -3*(4*q+1)*q**3*chi2y**2*chi2z/(8*(q+1)**4)
               -3*(4*q+1)*q**3*chi2z**3      /(8*(q+1)**4)
               +chi1x*(+9*(2*q+1)*q**2*chi2x*chi2z/(4*(q+1)**4)
                       +9*(1*q+2)*q   *chi2x*chi1z/(4*(q+1)**4))
               +chi1y*(+9*q**2*chi2y*chi1z/(4*(q+1)**4)
                       +9*q**2*chi2y*chi2z/(4*(q+1)**4))
               +chi1z*(+9*q**2*(2*q+3)*chi2x**2/(4*(q+1)**4)
                       -9*q**2*(  q+2)*chi2y**2/(4*(q+1)**4)
                       -9*q**2        *chi2z**2/(4*(q+1)**3)
                       -(135*q**5 + 385*q**4 + 363*q**3 + 377*q**2 + 387*q + 168)/(32*(q+1)**6))
               -(+168*q**5 + 387*q**4 + 377*q**3 + 363*q**2 + 385*q + 135)*q*chi2z/(32*(q+1)**6)
               +chi1x**2*(+3*(q+4)*chi1z/(2*(q+1)**4)
                          +9*q*(3*q+2)*chi2z/(4*(q+1)**4))
               +chi1y**2*(-3*(q+4)*chi1z/(8*(q+1)**4)
                          -9*q*(2*q+1)*chi2z/(4*(q+1)**4))
               -9*q*chi1z**2*chi2z/(4*(q+1)**3)
               -3*(q+4)*chi1z**3/(8*(q+1)**4))
        # fmt: on

    def f_MOmega(
        self,
        m1: sp.Expr,
        m2: sp.Expr,
        chi1U: List[sp.Expr],
        chi2U: List[sp.Expr],
        r: sp.Expr,
    ) -> None:
        """Sum the expressions for a_k to construct p_t as prescribed:
        MOmega = 1/r^(3/2) * (1 + sum_{k=2}^7 (a_k/r^{k/2}))"""
        a: List[sp.Expr] = ixp.zerorank1(dimension=10)
        self.MOmega__a_2_thru_a_4(
            m1, m2, chi1U[0], chi1U[1], chi1U[2], chi2U[0], chi2U[1], chi2U[2]
        )
        a[2] = self.a_2
        a[3] = self.a_3
        a[4] = self.a_4
        self.MOmega__a_5_thru_a_6(
            m1, m2, chi1U[0], chi1U[1], chi1U[2], chi2U[0], chi2U[1], chi2U[2]
        )
        a[5] = self.a_5
        a[6] = self.a_6
        self.MOmega__a_7(
            m1, m2, chi1U[0], chi1U[1], chi1U[2], chi2U[0], chi2U[1], chi2U[2]
        )
        a[7] = self.a_7
        self.MOmega = sp.sympify(1)  # Term prior to the sum in parentheses
        for k in range(8):
            self.MOmega += a[k] / r ** div(k, 2)
        self.MOmega *= 1 / r ** div(3, 2)


if __name__ == "__main__":
    import doctest
    import os
    import sys

    results = doctest.testmod()
    if results.failed > 0:
        print(f"Doctest failed: {results.failed} of {results.attempted} test(s)")
        sys.exit(1)
    else:
        print(f"Doctest passed: All {results.attempted} test(s) passed")

    in_m1, in_m2, in_r = sp.symbols("m1 m2 r")
    in_chi1U: List[sp.Expr] = cast(List[sp.Expr], ixp.declarerank1("chi1U"))
    in_chi2U: List[sp.Expr] = cast(List[sp.Expr], ixp.declarerank1("chi2U"))
    m_omega = PN_MOmega(in_m1, in_m2, in_chi1U, in_chi2U, in_r)
    results_dict = ve.process_dictionary_of_expressions(
        m_omega.__dict__, fixed_mpfs_for_free_symbols=True
    )
    ve.compare_or_generate_trusted_results(
        os.path.abspath(__file__),
        os.getcwd(),
        # File basename. If this is set to "trusted_module_test1", then
        #   trusted results_dict will be stored in tests/trusted_module_test1.py
        f"{os.path.splitext(os.path.basename(__file__))[0]}",
        results_dict,
    )

nrpypn/PN_MOmega.py
nrpy/nrpypn
"""
As documented in the NRPyPN notebook
PN-dE_GW_dt.ipynb, this Python script
generates dE_GW/dt at highest known
post-Newtonian order (as of 2015, at
least).

Author:  Zachariah B. Etienne
         zachetie **at** gmail **dot* com
"""
from typing import List, Tuple, cast
import sympy as sp  # SymPy: The Python computer algebra package upon which NRPy+ depends
import nrpy.indexedexp as ixp  # NRPy+: Symbolic indexed expression (e.g., tensors, vectors, etc.) support
import nrpy.validate_expressions.validate_expressions as ve

# NRPyPN: shortcuts for e.g., vector operations
from nrpypn.NRPyPN_shortcuts import div, dot, gamma_EulerMascheroni


class PN_dE_GW_dt_and_dM_dt:
    """
    This class calculates gravitational-wave energy flux and mass flux.
    Implements Eqs: A1-13 of https://arxiv.org/abs/1502.01747
    Eqs A22-28 of https://arxiv.org/abs/1502.01747, with
    Eq A.14 of https://arxiv.org/abs/0709.0093 for Mdot
           and correction on b[7] term by comparison with
     https://link.springer.com/content/pdf/10.12942/lrr-2014-2.pdf

    :param mOmega: Angular frequency
    :param m1: Mass of first object
    :param m2: Mass of second object
    :param n12U: List of symbolic expressions
    :param S1U: List of symbolic expressions for S1U
    :param S2U: List of symbolic expressions for S2U
    """

    def __init__(
        self,
        mOmega: sp.Expr,
        m1: sp.Expr,
        m2: sp.Expr,
        n12U: List[sp.Expr],
        S1U: List[sp.Expr],
        S2U: List[sp.Expr],
    ):
        self.dE_GW_dt_plus_dM_dt: sp.Expr = sp.sympify(0)
        self.dE_GW_dt: sp.Expr = sp.sympify(0)
        self.dM_dt: sp.Expr = sp.sympify(0)

        self.f_dE_GW_dt_and_dM_dt(mOmega, m1, m2, n12U, S1U, S2U)

    # Constants given in Eqs A1-13 of https://arxiv.org/abs/1502.01747
    def dE_GW_dt_OBKPSS2015_consts(
        self,
        m1: sp.Expr,
        m2: sp.Expr,
        _n12U: List[sp.Expr],
        S1U: List[sp.Expr],
        S2U: List[sp.Expr],
    ) -> Tuple[
        sp.Expr, sp.Expr, List[sp.Expr], List[sp.Expr], List[sp.Expr], sp.Expr, sp.Expr
    ]:  # _n12U unused.
        """
        This method calculates and returns several constants based on given parameters.

        :param m1: Mass of first object
        :param m2: Mass of second object
        :param _n12U: Unused parameter for compatibility
        :param S1U: List of symbolic expressions for S1U
        :param S2U: List of symbolic expressions for S2U
        :return: Tuple of calculated constants and lists of symbolic expressions: nu, delta, l, chi_a, chi_s, s_l, sigma_l
        """
        # define scalars:
        m = m1 + m2
        nu = m1 * m2 / m**2
        delta = (m1 - m2) / m
        # define vectors:
        Stot = ixp.zerorank1()
        Sigma = ixp.zerorank1()
        l = ixp.zerorank1()
        l[2] = sp.sympify(1)
        chi1U = ixp.zerorank1()
        chi2U = ixp.zerorank1()
        chi_s = ixp.zerorank1()
        chi_a = ixp.zerorank1()
        for i in range(3):
            Stot[i] = S1U[i] + S2U[i]
            Sigma[i] = (m1 + m2) / m2 * S2U[i] - (m1 + m2) / m1 * S1U[i]
            chi1U[i] = S1U[i] / m1**2
            chi2U[i] = S2U[i] / m2**2
            chi_s[i] = div(1, 2) * (chi1U[i] + chi2U[i])
            chi_a[i] = div(1, 2) * (chi1U[i] - chi2U[i])
        # define scalars that depend on vectors
        s_l = dot(Stot, l) / m**2
        # s_n = dot(Stot,n12U)/m**2
        sigma_l = dot(Sigma, l) / m**2
        # sigma_n = dot(Sigma,n12U)/m**2
        return nu, delta, l, chi_a, chi_s, s_l, sigma_l

    def f_dE_GW_dt_and_dM_dt(
        self,
        mOmega: sp.Expr,
        m1: sp.Expr,
        m2: sp.Expr,
        n12U: List[sp.Expr],
        S1U: List[sp.Expr],
        S2U: List[sp.Expr],
    ) -> None:
        """
        This method calculates and updates the class attributes `dE_GW_dt_plus_dM_dt`, `dE_GW_dt`, `dM_dt`.
        Based on Eqs A22-28 of https://arxiv.org/abs/1502.01747, with
           Eq A.14 of https://arxiv.org/abs/0709.0093 for Mdot
           and correction on b[7] term by comparison with
        https://link.springer.com/content/pdf/10.12942/lrr-2014-2.pdf

        :param mOmega: Angular frequency
        :param m1: Mass of first object
        :param m2: Mass of second object
        :param n12U: List of symbolic expressions for n12U
        :param S1U: List of symbolic expressions for S1U
        :param S2U: List of symbolic expressions for S2U
        """

        def f_compute_quantities(
            mOmega: sp.Expr,
            m1: sp.Expr,
            m2: sp.Expr,
            n12U: List[sp.Expr],
            S1U: List[sp.Expr],
            S2U: List[sp.Expr],
            which_quantity: str,
        ) -> sp.Expr:
            if which_quantity not in ("dM_dt", "dE_GW_dt", "dE_GW_dt_plus_dM_dt"):
                raise ValueError(f"which_quantity == {which_quantity} not supported!")

            nu, delta, l, chi_a, chi_s, s_l, sigma_l = self.dE_GW_dt_OBKPSS2015_consts(
                m1, m2, n12U, S1U, S2U
            )
            x = (mOmega) ** div(2, 3)

            # fmt: off
            # Compute b_5_Mdot:
            b_5_Mdot = (-div(1,4)*(+(1-3*nu)*dot(chi_s,l)*(1+3*dot(chi_s,l)**2+9*dot(chi_a,l)**2)
                                +(1-  nu)*delta*dot(chi_a,l)*(1+3*dot(chi_a,l)**2+9*dot(chi_s,l)**2)))
            if which_quantity == "dM_dt":
                return cast(sp.Expr, div(32,5)*nu**2*x**5*b_5_Mdot*x**div(5,2))

            b = ixp.zerorank1(dimension=10)
            b[2] = -div(1247,336) - div(35,12)*nu
            b[3] = +4*sp.pi - 4*s_l - div(5,4)*delta*sigma_l
            b[4] =(-div(44711,9072) + div(9271,504)*nu + div(65,18)*nu**2
                   +(+div(287,96) + div( 1,24)*nu)*dot(chi_s,l)**2
                   -(+div( 89,96) + div( 7,24)*nu)*dot(chi_s,chi_s)
                   +(+div(287,96) -         12*nu)*dot(chi_a,l)**2
                   +(-div( 89,96) +          4*nu)*dot(chi_a,chi_a)
                   +div(287,48)*delta*dot(chi_s,l)*dot(chi_a,l) - div(89,48)*delta*dot(chi_s,chi_a))
            b[5] =(-div(8191,672)*sp.pi - div(9,2)*s_l - div(13,16)*delta*sigma_l
                   +nu*(-div(583,24)*sp.pi + div(272,9)*s_l + div(43,4)*delta*sigma_l))
            if which_quantity == "dE_GW_dt_plus_dM_dt":
                b[5]+= b_5_Mdot
            b[6] =(+div(6643739519,69854400) + div(16,3)*sp.pi**2 - div(1712,105)*gamma_EulerMascheroni
                   -div(856,105)*sp.log(16*x) + (-div(134543,7776) + div(41,48)*sp.pi**2)*nu
                   -div(94403,3024)*nu**2 - div(775,324)*nu**3 - 16*sp.pi*s_l - div(31,6)*sp.pi*delta*sigma_l)
            b[7] =(+(+div(476645,6804) + div(6172,189)*nu - div(2810,27)*nu**2)*s_l
                   +(+div(9535,336) + div(1849,126)*nu - div(1501,36)*nu**2)*delta*sigma_l
                   +(-div(16285,504) + div(214745,1728)*nu + div(193385,3024)*nu**2)*sp.pi)
            b[8] =(+(-div(3485,96)*sp.pi + div(13879,72)*sp.pi*nu)*s_l
                   +(-div(7163,672)*sp.pi + div(130583,2016)*sp.pi*nu)*delta*sigma_l)
            b_sum = sp.sympify(1)
            for k in range(9):
                b_sum += b[k]*x**div(k,2)
            return cast(sp.Expr, div(32,5)*nu**2*x**5*b_sum)
            # fmt: on

        self.dE_GW_dt_plus_dM_dt = f_compute_quantities(
            mOmega, m1, m2, n12U, S1U, S2U, "dE_GW_dt_plus_dM_dt"
        )
        self.dE_GW_dt = f_compute_quantities(mOmega, m1, m2, n12U, S1U, S2U, "dE_GW_dt")
        self.dM_dt = f_compute_quantities(mOmega, m1, m2, n12U, S1U, S2U, "dM_dt")


if __name__ == "__main__":
    import doctest
    import os
    import sys

    results = doctest.testmod()
    if results.failed > 0:
        print(f"Doctest failed: {results.failed} of {results.attempted} test(s)")
        sys.exit(1)
    else:
        print(f"Doctest passed: All {results.attempted} test(s) passed")

    in_mOmega, in_m1, in_m2 = sp.symbols("mOmega m1 m2")
    in_n12U: List[sp.Expr] = cast(List[sp.Expr], ixp.declarerank1("n12U"))
    in_S1U: List[sp.Expr] = cast(List[sp.Expr], ixp.declarerank1("S1U"))
    in_S2U: List[sp.Expr] = cast(List[sp.Expr], ixp.declarerank1("S2U"))
    mass_energy_fluxes = PN_dE_GW_dt_and_dM_dt(
        in_mOmega, in_m1, in_m2, in_n12U, in_S1U, in_S2U
    )
    results_dict = ve.process_dictionary_of_expressions(
        mass_energy_fluxes.__dict__, fixed_mpfs_for_free_symbols=True
    )
    ve.compare_or_generate_trusted_results(
        os.path.abspath(__file__),
        os.getcwd(),
        # File basename. If this is set to "trusted_module_test1", then
        #   trusted results_dict will be stored in tests/trusted_module_test1.py
        f"{os.path.splitext(os.path.basename(__file__))[0]}",
        results_dict,
    )

nrpypn/PN_dE_GW_dt_and_dM_dt.py
nrpy/nrpypn
"""
As documented in the NRPyPN notebook
PN-p_r.ipynb, this Python script
generates the expression for the radial
component of momentum p_r up to and
including terms at 3.5PN order.
It largely follows the technique of
 Ramos-Buades, Husa, and Pratten (2018)
   https://arxiv.org/abs/1810.00036

Author:  Zach Etienne
         zachetie **at** gmail **dot* com
"""

from typing import List, cast
import sympy as sp  # SymPy: The Python computer algebra package upon which NRPy+ depends
import nrpy.indexedexp as ixp
import nrpy.validate_expressions.validate_expressions as ve

# NRPyPN: shortcuts for e.g., vector operations
from nrpypn.NRPyPN_shortcuts import Pt, Pr, nU, div
from nrpypn.PN_p_t import PN_p_t
from nrpypn.PN_MOmega import PN_MOmega
from nrpypn.PN_dE_GW_dt_and_dM_dt import PN_dE_GW_dt_and_dM_dt
from nrpypn.PN_Hamiltonian_NS import PN_Hamiltonian_NS
from nrpypn.PN_Hamiltonian_SS import PN_Hamiltonian_SS
from nrpypn.PN_Hamiltonian_SSS import PN_Hamiltonian_SSS
from nrpypn.PN_Hamiltonian_SO import PN_Hamiltonian_SO


class PN_p_r:
    """

    Core class functions:
    f_dr_dt(Htot_xyplane_binary, m1,m2, n12U,n21U, chi1U,chi2U, S1U,S2U, p1U,p2U, r)
        Given Htot_xyplane_binary and other standard input parameters, compute
        dr_dt = dr/dt and store to class variable of the same name.

    f_p_r_fullHam(m1,m2, n12U,n21U, chi1U,chi2U, S1U,S2U, p1U,p2U, r)
          Compute p_r using the full Hamiltonian, without truncating
          higher-order terms self-consistently.

    f_p_r(m1,m2, chi1U,chi2U, r)
          Compute p_r and store to class variable of the same name.
    """

    def __init__(
        self,
        m1: sp.Expr,
        m2: sp.Expr,
        n12U: List[sp.Expr],
        n21U: List[sp.Expr],
        chi1U: List[sp.Expr],
        chi2U: List[sp.Expr],
        S1U: List[sp.Expr],
        S2U: List[sp.Expr],
        p1U: List[sp.Expr],
        p2U: List[sp.Expr],
        r: sp.Expr,
        use_alternative_approach_for_p_r: bool = False,
    ):
        self.Htot_xyplane_binary: sp.Expr = sp.sympify(0)
        self.dr_dt: sp.Expr = sp.sympify(0)
        self.p_r: sp.Expr = sp.sympify(0)

        if use_alternative_approach_for_p_r:
            self.f_p_r_alternative__use_fullHam(
                m1, m2, n12U, n21U, chi1U, chi2U, S1U, S2U, p1U, p2U, r
            )
        else:
            self.f_p_r(m1, m2, n12U, n21U, chi1U, chi2U, S1U, S2U, p1U, p2U, r)

    #################################
    #################################

    # Step 1: Construct full Hamiltonian
    #         expression for a binary instantaneously
    #         orbiting on the xy plane, store
    #         result to Htot_xyplane_binary
    def f_Htot_xyplane_binary(
        self,
        m1: sp.Expr,
        m2: sp.Expr,
        n12U: List[sp.Expr],
        n21U: List[sp.Expr],
        S1U: List[sp.Expr],
        S2U: List[sp.Expr],
        p1U: List[sp.Expr],
        p2U: List[sp.Expr],
        r: sp.Expr,
    ) -> None:
        """
        Given standard input parameters, compute
        the Hamiltonian for a binary system
        orbiting instantaneously on the xy plane,
        and store to the class variable
        Htot_xyplane_binary
        """

        def make_replacements(expr: sp.Expr) -> sp.Expr:
            zero = sp.sympify(0)
            one = sp.sympify(1)
            return cast(
                sp.Expr,
                expr.subs(p1U[1], Pt)
                .subs(p2U[1], -Pt)
                .subs(p1U[2], zero)
                .subs(p2U[2], zero)
                .subs(p1U[0], -Pr)
                .subs(p2U[0], Pr)
                .subs(nU[0], one)
                .subs(nU[1], zero)
                .subs(nU[2], zero),
            )

        H_NS = PN_Hamiltonian_NS(m1, m2, p1U, n12U, r)

        self.Htot_xyplane_binary = make_replacements(
            H_NS.H_Newt + H_NS.H_NS_1PN + H_NS.H_NS_2PN + H_NS.H_NS_3PN
        )

        H_SO = PN_Hamiltonian_SO(m1, m2, n12U, n21U, S1U, S2U, p1U, p2U, r)
        self.Htot_xyplane_binary += make_replacements(
            H_SO.H_SO_1p5PN + H_SO.H_SO_2p5PN + H_SO.H_SO_3p5PN
        )

        H_SS = PN_Hamiltonian_SS(m1, m2, n12U, n21U, S1U, S2U, p1U, p2U, r)
        self.Htot_xyplane_binary += make_replacements(
            H_SS.H_SS_2PN + H_SS.H_SS_S1S2_3PN + H_SS.H_SS_S1sq_S2sq_3PN
        )

        H_SSS = PN_Hamiltonian_SSS(m1, m2, n12U, n21U, S1U, S2U, p1U, p2U, r)
        self.Htot_xyplane_binary += make_replacements(H_SSS.H_SSS_3PN)

    # Function for computing dr/dt
    def f_dr_dt(
        self,
        Htot_xyplane_binary: sp.Expr,
        m1: sp.Expr,
        m2: sp.Expr,
        n12U: List[sp.Expr],
        chi1U: List[sp.Expr],
        chi2U: List[sp.Expr],
        S1U: List[sp.Expr],
        S2U: List[sp.Expr],
        r: sp.Expr,
    ) -> None:
        """
        Given Htot_xyplane_binary (computed
        above) and other standard input
        parameters, compute
        dr_dt = dr/dt and store to class
        variable of the same name.
        """
        # First compute p_t
        pt = PN_p_t(m1, m2, chi1U, chi2U, r)

        # Then compute dH_{circ}/dr = partial_H(p_r=0)/partial_r
        #                                  + partial_H(p_r=0)/partial_{p_t} partial_{p_t}/partial_r
        dHcirc_dr = +sp.diff(Htot_xyplane_binary.subs(Pr, sp.sympify(0)), r) + sp.diff(
            Htot_xyplane_binary.subs(Pr, sp.sympify(0)), Pt
        ) * sp.diff(pt.p_t, r)

        # Then compute M Omega
        MOm = PN_MOmega(m1, m2, chi1U, chi2U, r)

        # Next compute dE_GW_dt_plus_dM_dt
        dE_GW_dt_and_dM_dt = PN_dE_GW_dt_and_dM_dt(MOm.MOmega, m1, m2, n12U, S1U, S2U)
        dE_GW_dt_plus_dM_dt = dE_GW_dt_and_dM_dt.dE_GW_dt_plus_dM_dt

        # Finally, compute dr/dt
        self.dr_dt = dE_GW_dt_plus_dM_dt / dHcirc_dr

    # Next we compute p_r as a function of dr_dt (unknown) and known quantities using
    # p_r  approx [dr/dt - (partial_H/partial_{p_r})|_{p_r=0}] * [(partial^2_{H}/partial_{p_r^2})|_{p_r=0}]^{-1}
    def f_p_r_alternative__use_fullHam(
        self,
        m1: sp.Expr,
        m2: sp.Expr,
        n12U: List[sp.Expr],
        n21U: List[sp.Expr],
        chi1U: List[sp.Expr],
        chi2U: List[sp.Expr],
        S1U: List[sp.Expr],
        S2U: List[sp.Expr],
        p1U: List[sp.Expr],
        p2U: List[sp.Expr],
        r: sp.Expr,
    ) -> None:
        """
        Basic equation for p_r can be written in the
          form:
         p_r  approx [dr/dt - (partial_H/partial_{p_r})|_{p_r=0}] * [(partial^2_{H}/partial_{p_r^2})|_{p_r=0}]^{-1},
          where
        dr/dt = [dE_{rm GW}/dt + dM/dt] * [dH_{circ} / dr]^{-1},
         and
        H_{circ} = Htot_xyplane_binary|_{p_r=0}
         -> [dH_{circ}(r,p_t(r)) / dr] = partial_{H(p_r=0)}/partial_r
                   + partial_{H(p_r=0)}/partial_{p_t} partial_{p_t}/partial_r.
         Here,
         * the expression for p_t is given by PN_p_t.py
         * the expression for [dE_{rm GW}/dt + dM/dt] is given by
                   PN_dE_GW_dt_and_dM_dt.py
           + Since [dE_{rm GW}/dt + dM/dt] is a function of MOmega,
           we also need input from the PN_MOmega.py Python module.
        """
        self.f_Htot_xyplane_binary(m1, m2, n12U, n21U, S1U, S2U, p1U, p2U, r)
        self.f_dr_dt(self.Htot_xyplane_binary, m1, m2, n12U, chi1U, chi2U, S1U, S2U, r)

        dHdpr_przero = sp.diff(self.Htot_xyplane_binary, Pr).subs(Pr, sp.sympify(0))
        d2Hdpr2_przero = sp.diff(sp.diff(self.Htot_xyplane_binary, Pr), Pr).subs(
            Pr, sp.sympify(0)
        )
        self.p_r = (self.dr_dt - dHdpr_przero) / (d2Hdpr2_przero)

    def f_p_r(
        self,
        m1: sp.Expr,
        m2: sp.Expr,
        n12U: List[sp.Expr],
        n21U: List[sp.Expr],
        chi1U: List[sp.Expr],
        chi2U: List[sp.Expr],
        S1U: List[sp.Expr],
        S2U: List[sp.Expr],
        p1U: List[sp.Expr],
        p2U: List[sp.Expr],
        r: sp.Expr,
    ) -> None:
        """
        Ramos-Buades, Husa, and Pratten (2018)
          approach for computing p_r.
        Transcribed from Eq 2.18 of
        Ramos-Buades, Husa, and Pratten (2018),
          https://arxiv.org/abs/1810.00036
        """
        q = m2 / m1  # It is assumed that q >= 1, so m2 >= m1.
        self.f_Htot_xyplane_binary(m1, m2, n12U, n21U, S1U, S2U, p1U, p2U, r)
        self.f_dr_dt(self.Htot_xyplane_binary, m1, m2, n12U, chi1U, chi2U, S1U, S2U, r)
        chi1x = chi1U[0]
        chi1y = chi1U[1]
        chi1z = chi1U[2]
        chi2x = chi2U[0]
        chi2y = chi2U[1]
        chi2z = chi2U[2]
        # fmt: off
        p_r_num = (-self.dr_dt
                   +(-(6*q+13)*q**2*chi1x*chi2y/(4*(q+1)**4)
                     -(6*q+ 1)*q**2*chi2x*chi2y/(4*(q+1)**4)
                     +chi1y*(-q*(   q+6)*chi1x/(4*(q+1)**4)
                             -q*(13*q+6)*chi2x/(4*(q+1)**4)))/r**div(7,2)
                   +(+chi1z*(+3*q   *(5*q+2)*chi1x*chi2y/(2*(q+1)**4)
                             -3*q**2*(2*q+5)*chi2x*chi2y/(2*(q+1)**4))
                     +chi1y*chi2z*(+3*q**2*(2*q+5)*chi2x/(2*(q+1)**4)
                                   -3*q   *(5*q+2)*chi1x/(2*(q+1)**4)))/r**4)
        p_r_den = (-(q+1)**2/q - (-7*q**2-15*q-7)/(2*q*r)
                   -(47*q**4 + 229*q**3 + 363*q**2 + 229*q + 47)/(8*q*(q+1)**2*r**2)
                   -(+( 4*q**2 + 11*q + 12)*chi1z/(4*q*(q+1))
                     +(12*q**2 + 11*q +  4)*chi2z/(4*  (q+1)))/r**div(5,2)
                   -(+(- 53*q**5 - 357*q**4 - 1097*q**3 - 1486*q**2 - 842*q - 144)*chi1z/(16*q*(q+1)**4)
                     +(-144*q**5 - 842*q**4 - 1486*q**3 - 1097*q**2 - 357*q -  53)*chi2z/(16  *(q+1)**4))/r**div(7,2)
                   -(+(  q**2 + 9*q + 9)*chi1x**2/(2*q*(q+1)**2)
                     +(3*q**2 + 5*q + 3)*chi2x*chi1x/((q+1)**2)
                     +(3*q**2 + 8*q + 3)*chi1y*chi2y/(2*(q+1)**2)
                     -9*q**2*chi2y**2/(4*(q+1))
                     +(3*q**2 + 8*q + 3)*chi1z*chi2z/(2*(q+1)**2)
                     -9*q**2*chi2z**2/(4*(q+1))
                     +(9*q**3 + 9*q**2 + q)*chi2x**2/(2*(q+1)**2)
                     +(-363*q**6 - 2608*q**5 - 7324*q**4 - 10161*q**3 - 7324*q**2 - 2608*q - 363)/(48*q*(q+1)**4)
                     -9*chi1y**2/(4*q*(q+1))
                     -9*chi1z**2/(4*q*(q+1)) - sp.pi**2/16)/r**3)
        # fmt: on
        self.p_r = p_r_num / p_r_den


if __name__ == "__main__":
    import doctest
    import os
    import sys

    results = doctest.testmod()
    if results.failed > 0:
        print(f"Doctest failed: {results.failed} of {results.attempted} test(s)")
        sys.exit(1)
    else:
        print(f"Doctest passed: All {results.attempted} test(s) passed")

    in_m1, in_m2, in_r12 = sp.symbols("m1 m2 r12")
    in_n12U: List[sp.Expr] = cast(List[sp.Expr], ixp.declarerank1("n12U"))
    in_n21U: List[sp.Expr] = cast(List[sp.Expr], ixp.declarerank1("n21U"))
    in_chi1U: List[sp.Expr] = cast(List[sp.Expr], ixp.declarerank1("chi1U"))
    in_chi2U: List[sp.Expr] = cast(List[sp.Expr], ixp.declarerank1("chi2U"))
    in_S1U: List[sp.Expr] = cast(List[sp.Expr], ixp.declarerank1("S1U"))
    in_S2U: List[sp.Expr] = cast(List[sp.Expr], ixp.declarerank1("S2U"))
    in_p1U: List[sp.Expr] = cast(List[sp.Expr], ixp.declarerank1("p1U"))
    in_p2U: List[sp.Expr] = cast(List[sp.Expr], ixp.declarerank1("p2U"))
    in_nU: List[sp.Expr] = cast(List[sp.Expr], ixp.declarerank1("nU"))
    for Alternative in [False, True]:
        p_r = PN_p_r(
            in_m1,
            in_m2,
            in_n12U,
            in_n21U,
            in_chi1U,
            in_chi2U,
            in_S1U,
            in_S2U,
            in_p1U,
            in_p2U,
            in_r12,
            use_alternative_approach_for_p_r=Alternative,
        )
        results_dict = ve.process_dictionary_of_expressions(
            p_r.__dict__, fixed_mpfs_for_free_symbols=True
        )
        ve.compare_or_generate_trusted_results(
            os.path.abspath(__file__),
            os.getcwd(),
            # File basename. If this is set to "trusted_module_test1", then
            #   trusted results_dict will be stored in tests/trusted_module_test1.py
            f"{os.path.splitext(os.path.basename(__file__))[0]}_Alternative{Alternative}",
            results_dict,
        )

nrpypn/PN_p_r.py
nrpy/nrpypn
"""
As documented in the NRPyPN notebook
PN-p_t.ipynb, this Python script
generates the expression for the transverse
component of momentum p_t up to and
including terms at 3.5PN order.
This is an implementation of the equations of
 Ramos-Buades, Husa, and Pratten (2018)
   https://arxiv.org/abs/1810.00036
but validates against the relevant equation
  in Healy, Lousto, Nakano, and Zlochower (2017)
   https://arxiv.org/abs/1702.00872

Author:  Zachariah B. Etienne
         zachetie **at** gmail **dot* com
"""
from typing import List, cast
import sympy as sp  # SymPy: The Python computer algebra package upon which NRPy+ depends
import nrpy.indexedexp as ixp  # NRPy+: Symbolic indexed expression (e.g., tensors, vectors, etc.) support
import nrpy.validate_expressions.validate_expressions as ve

from nrpypn.NRPyPN_shortcuts import div  # NRPyPN: shortcuts for e.g., vector operations


class PN_p_t:
    """
    Basic equation for p_t can be written in the
    form:
    p_t = q / (sqrt(r) * (1 + q)^2) * (1 + sum_{k=2}^7 (a_k / r^{k/2}))
    where we construct the a_k terms in the sum below.

    :param m1: mass of object 1
    :param m2: mass of object 2
    :param chi1U: spin of object 1
    :param chi2U: spin of object 2
    :param r: radius of separation
    """

    def __init__(
        self,
        m1: sp.Expr,
        m2: sp.Expr,
        chi1U: List[sp.Expr],
        chi2U: List[sp.Expr],
        r: sp.Expr,
    ):
        self.a_2: sp.Expr = sp.sympify(0)
        self.a_3: sp.Expr = sp.sympify(0)
        self.a_4: sp.Expr = sp.sympify(0)
        self.a_5: sp.Expr = sp.sympify(0)
        self.a_6: sp.Expr = sp.sympify(0)
        self.a_7: sp.Expr = sp.sympify(0)

        self.p_t: sp.Expr = sp.sympify(0)

        self.f_p_t(m1, m2, chi1U, chi2U, r)

    #################################
    #################################
    def p_t__a_2_thru_a_4(
        self,
        m1: sp.Expr,
        m2: sp.Expr,
        chi1x: sp.Expr,
        chi1y: sp.Expr,
        chi1z: sp.Expr,
        chi2x: sp.Expr,
        chi2y: sp.Expr,
        chi2z: sp.Expr,
    ) -> None:
        """
        Step 1: Construct terms a_2, a_3, and a_4, from
        Eq A2 of Ramos-Buades, Husa, and Pratten (2018)
        https://arxiv.org/abs/1810.00036
        These terms have been independently validated
        against the same terms in Eq 7 of
        Healy, Lousto, Nakano, and Zlochower (2017)
        https://arxiv.org/abs/1702.00872

        :param m1: mass of object 1
        :param m2: mass of object 2
        :param chi1x: x-component of spin of object 1
        :param chi1y: y-component of spin of object 1
        :param chi1z: z-component of spin of object 1
        :param chi2x: x-component of spin of object 2
        :param chi2y: y-component of spin of object 2
        :param chi2z: z-component of spin of object 2
        """
        # fmt: off
        q = m2/m1 # It is assumed that q >= 1, so m2 >= m1.
        self.a_2 = sp.sympify(2)
        self.a_3 = (-3*(4*q**2+3*q)*chi2z/(4*(q+1)**2) - 3*(3*q+4)*chi1z/(4*(q+1)**2))
        self.a_4 = (-3*q**2*chi2x**2/(2*(q+1)**2)
                    +3*q**2*chi2y**2/(4*(q+1)**2)
                    +3*q**2*chi2z**2/(4*(q+1)**2)
                    +(+42*q**2 + 41*q + 42)/(16*(q+1)**2)
                    -3*chi1x**2/(2*(q+1)**2)
                    -3*q*chi1x*chi2x/(q+1)**2
                    +3*chi1y**2/(4*(q+1)**2)
                    +3*q*chi1y*chi2y/(2*(q+1)**2)
                    +3*chi1z**2/(4*(q+1)**2)
                    +3*q*chi1z*chi2z/(2*(q+1)**2))
        # fmt: on

    def p_t__a_5_thru_a_6(
        self,
        m1: sp.Expr,
        m2: sp.Expr,
        chi1x: sp.Expr,
        chi1y: sp.Expr,
        chi1z: sp.Expr,
        chi2x: sp.Expr,
        chi2y: sp.Expr,
        chi2z: sp.Expr,
        FixSignError: bool = True,
    ) -> None:
        """
        Construct terms a_5 and a_6, from
        Eq A2 of Ramos-Buades, Husa, and Pratten (2018)
        https://arxiv.org/abs/1810.00036
        These terms have been independently validated
        against the same terms in Eq 7 of
        Healy, Lousto, Nakano, and Zlochower (2017)
        https://arxiv.org/abs/1702.00872
        and a sign error was corrected in the a_5
        expression.

        :param m1: mass of object 1
        :param m2: mass of object 2
        :param chi1x: x-component of spin of object 1
        :param chi1y: y-component of spin of object 1
        :param chi1z: z-component of spin of object 1
        :param chi2x: x-component of spin of object 2
        :param chi2y: y-component of spin of object 2
        :param chi2z: z-component of spin of object 2
        :param FixSignError: boolean to determine if sign error should be fixed
        """
        # Initialize the SignFix based on FixSignError
        SignFix = sp.sympify(-1) if FixSignError else sp.sympify(1)

        # It is assumed that q >= 1, so m2 >= m1.
        q = m2 / m1
        # fmt: off
        self.a_5 = (SignFix*(13*q**3 + 60*q**2 + 116*q + 72)*chi1z/(16*(q+1)**4)
                    +(-72*q**4 - 116*q**3 - 60*q**2 - 13*q)*chi2z/(16*(q+1)**4))
        self.a_6 = (+(+472*q**2 - 640)*chi1x**2/(128*(q+1)**4)
                    +(-512*q**2 - 640*q - 64)*chi1y**2/(128*(q+1)**4)
                    +(-108*q**2 + 224*q +512)*chi1z**2/(128*(q+1)**4)
                    +(+472*q**2 - 640*q**4)*chi2x**2/(128*(q+1)**4)
                    +(+192*q**3 + 560*q**2 + 192*q)*chi1x*chi2x/(128*(q+1)**4)
                    +(-864*q**3 -1856*q**2 - 864*q)*chi1y*chi2y/(128*(q+1)**4)
                    +(+480*q**3 +1064*q**2 + 480*q)*chi1z*chi2z/(128*(q+1)**4)
                    +( -64*q**4 - 640*q**3 - 512*q**2)*chi2y**2/(128*(q+1)**4)
                    +(+512*q**4 + 224*q**3 - 108*q**2)*chi2z**2/(128*(q+1)**4)
                    +(+480*q**4 + 163*sp.pi**2*q**3 - 2636*q**3 + 326*sp.pi**2*q**2 - 6128*q**2 + 163*sp.pi**2*q-2636*q+480)
                     /(128*(q+1)**4))
        # fmt: on

    def p_t__a_7(
        self,
        m1: sp.Expr,
        m2: sp.Expr,
        chi1x: sp.Expr,
        chi1y: sp.Expr,
        chi1z: sp.Expr,
        chi2x: sp.Expr,
        chi2y: sp.Expr,
        chi2z: sp.Expr,
    ) -> None:
        """
        Construct term a_7, from Eq A2 of
         Ramos-Buades, Husa, and Pratten (2018)
           https://arxiv.org/abs/1810.00036

        :param m1: mass of object 1
        :param m2: mass of object 2
        :param chi1x: x-component of spin of object 1
        :param chi1y: y-component of spin of object 1
        :param chi1z: z-component of spin of object 1
        :param chi2x: x-component of spin of object 2
        :param chi2y: y-component of spin of object 2
        :param chi2z: z-component of spin of object 2
        """
        q = m2 / m1  # It is assumed that q >= 1, so m2 >= m1.
        # fmt: off
        self.a_7 = (+5*(4*q+1)*q**3*chi2x**2*chi2z/(2*(q+1)**4)
                    -5*(4*q+1)*q**3*chi2y**2*chi2z/(8*(q+1)**4)
                    -5*(4*q+1)*q**3*chi2z**3      /(8*(q+1)**4)
                    +chi1x*(+15*(2*q+1)*q**2*chi2x*chi2z/(4*(q+1)**4)
                            +15*(1*q+2)*q   *chi2x*chi1z/(4*(q+1)**4))
                    +chi1y*(+15*q**2*chi2y*chi1z/(4*(q+1)**4)
                            +15*q**2*chi2y*chi2z/(4*(q+1)**4))
                    +chi1z*(+15*q**2*(2*q+3)*chi2x**2/(4*(q+1)**4)
                            -15*q**2*(  q+2)*chi2y**2/(4*(q+1)**4)
                            -15*q**2        *chi2z**2/(4*(q+1)**3)
                            -(103*q**5 + 145*q**4 - 27*q**3 + 252*q**2 + 670*q + 348)/(32*(q+1)**6))
                    -(+348*q**5 + 670*q**4 + 252*q**3 - 27*q**2 + 145*q + 103)*q*chi2z/(32*(q+1)**6)
                    +chi1x**2*(+5*(q+4)*chi1z/(2*(q+1)**4)
                               +15*q*(3*q+2)*chi2z/(4*(q+1)**4))
                    +chi1y**2*(-5*(q+4)*chi1z/(8*(q+1)**4)
                               -15*q*(2*q+1)*chi2z/(4*(q+1)**4))
                    -15*q*chi1z**2*chi2z/(4*(q+1)**3)
                    -5*(q+4)*chi1z**3/(8*(q+1)**4))
        # fmt: on

    def f_p_t(
        self,
        m1: sp.Expr,
        m2: sp.Expr,
        chi1U: List[sp.Expr],
        chi2U: List[sp.Expr],
        r: sp.Expr,
    ) -> None:
        """
        Finally, sum the expressions for a_k to construct p_t as prescribed:
        p_t = q/(sqrt(r)*(1+q)^2) (1 + sum_{k=2}^7 (a_k/r^{k/2}))

        :param m1: mass of object 1
        :param m2: mass of object 2
        :param chi1U: spin of object 1
        :param chi2U: spin of object 2
        :param r: radius of separation
        """
        q = m2 / m1  # It is assumed that q >= 1, so m2 >= m1.
        a = ixp.zerorank1(dimension=10)
        self.p_t__a_2_thru_a_4(
            m1, m2, chi1U[0], chi1U[1], chi1U[2], chi2U[0], chi2U[1], chi2U[2]
        )
        self.p_t__a_5_thru_a_6(
            m1, m2, chi1U[0], chi1U[1], chi1U[2], chi2U[0], chi2U[1], chi2U[2]
        )
        self.p_t__a_7(
            m1, m2, chi1U[0], chi1U[1], chi1U[2], chi2U[0], chi2U[1], chi2U[2]
        )
        a[2] = self.a_2
        a[3] = self.a_3
        a[4] = self.a_4
        a[5] = self.a_5
        a[6] = self.a_6
        a[7] = self.a_7
        self.p_t = sp.sympify(1)  # Term prior to the sum in parentheses
        for k in range(8):
            self.p_t += a[k] / r ** div(k, 2)
        self.p_t *= q / (1 + q) ** 2 * 1 / r ** div(1, 2)


if __name__ == "__main__":
    import doctest
    import os
    import sys

    results = doctest.testmod()
    if results.failed > 0:
        print(f"Doctest failed: {results.failed} of {results.attempted} test(s)")
        sys.exit(1)
    else:
        print(f"Doctest passed: All {results.attempted} test(s) passed")

    in_m1, in_m2, in_r = sp.symbols("m1 m2, r")
    in_chi1U: List[sp.Expr] = cast(List[sp.Expr], ixp.declarerank1("chi1U"))
    in_chi2U: List[sp.Expr] = cast(List[sp.Expr], ixp.declarerank1("chi2U"))
    tangential_momentum = PN_p_t(in_m1, in_m2, in_chi1U, in_chi2U, in_r)
    results_dict = ve.process_dictionary_of_expressions(
        tangential_momentum.__dict__, fixed_mpfs_for_free_symbols=True
    )
    ve.compare_or_generate_trusted_results(
        os.path.abspath(__file__),
        os.getcwd(),
        # File basename. If this is set to "trusted_module_test1", then
        #   trusted results_dict will be stored in tests/trusted_module_test1.py
        f"{os.path.splitext(os.path.basename(__file__))[0]}",
        results_dict,
    )

nrpypn/PN_p_t.py
nrpy/nrpypn
"""
As documented in the NRPyPN notebook
NRPyPN_shortcuts.ipynb, this Python script
provides useful shortcuts for inputting
post-Newtonian expressions into SymPy/NRPy+

Basic functions:

Author:  Zachariah B. Etienne
         zachetie **at** gmail **dot* com
"""
from typing import Tuple, Dict
from collections import namedtuple

from nrpypn.NRPyPN_shortcuts import (
    m1,
    m2,
    chi1U,
    chi2U,
    r,
    n12U,
    n21U,
    S1U,
    S2U,
    p1U,
    p2U,
    num_eval,
)
from nrpypn.PN_p_t import PN_p_t
from nrpypn.PN_p_r import PN_p_r


def eval__P_t__and__P_r(
    qmassratio: float,
    nr: float,
    nchi1x: float,
    nchi1y: float,
    nchi1z: float,
    nchi2x: float,
    nchi2y: float,
    nchi2z: float,
) -> Tuple[float, float]:
    """
    Core function to numerically evaluate P_t and P_r

    :param qmassratio: mass ratio q >= 1
    :param nr: Orbital separation (total coordinate distance between black holes)
    :param nchi1x: x-component of the dimensionless spin vector for BH 1
    :param nchi1y: y-component of the dimensionless spin vector for BH 1
    :param nchi1z: z-component of the dimensionless spin vector for BH 1
    :param nchi2x: x-component of the dimensionless spin vector for BH 2
    :param nchi2y: y-component of the dimensionless spin vector for BH 2
    :param nchi2z: z-component of the dimensionless spin vector for BH 2
    :return: Tuple containing the numerical values for P_t (the tangential momentum) and P_r (the radial momentum)
    """
    # Compute p_t, the tangential component of momentum
    pt = PN_p_t(m1, m2, chi1U, chi2U, r)

    # Compute p_r, the radial component of momentum
    pr = PN_p_r(m1, m2, n12U, n21U, chi1U, chi2U, S1U, S2U, p1U, p2U, r)

    nPt = num_eval(
        pt.p_t,
        qmassratio=qmassratio,
        nr=nr,
        nchi1x=nchi1x,
        nchi1y=nchi1y,
        nchi1z=nchi1z,
        nchi2x=nchi2x,
        nchi2y=nchi2y,
        nchi2z=nchi2z,
    )

    nPr = num_eval(
        pr.p_r,
        qmassratio=qmassratio,
        nr=nr,
        nchi1x=nchi1x,
        nchi1y=nchi1y,
        nchi1z=nchi1z,
        nchi2x=nchi2x,
        nchi2y=nchi2y,
        nchi2z=nchi2z,
        nPt=nPt,
    )
    return float(nPt), float(nPr)


def test_results_against_trusted() -> None:
    """
    Test the function eval__P_t__and__P_r against trusted values.

    :raises ValueError: If the relative error exceeds the tolerance for either p_t or p_r.
    """
    bbh_params = namedtuple(
        "bbh_params",
        [
            "q",
            "d",
            "chi1x",
            "chi1y",
            "chi1z",
            "chi2x",
            "chi2y",
            "chi2z",
            "tr_p_t",
            "tr_p_r",
        ],
    )
    case_dict: Dict[str, bbh_params] = {
        "Mass ratio q=2, chi1= (0,0,0); chi2=(-0.3535, 0.3535, 0.5), radial separation r=10.8": bbh_params(
            q=2.0,
            d=10.8,
            chi1x=0,
            chi1y=0,
            chi1z=0,
            chi2x=-0.3535,
            chi2y=0.3535,
            chi2z=0.5,
            tr_p_t=0.0793500403866190,
            tr_p_r=0.0005426257166677216,
        ),
        "Mass ratio q=8, chi1= (0, 0, 0.5); chi2=(0, 0, 0.5), radial separation r=11": bbh_params(
            q=8.0,
            d=11.0,
            chi1x=0,
            chi1y=0,
            chi1z=0.5,
            chi2x=0,
            chi2y=0,
            chi2z=0.5,
            tr_p_t=0.0345503689803291,
            tr_p_r=0.0000975735300199125,
        ),
        "Mass ratio q=1.5, chi1= (0,0,-0.6); chi2=(0,0,0.6), radial separation r=10.8": bbh_params(
            q=1.5,
            d=10.8,
            chi1x=0,
            chi1y=0,
            chi1z=-0.6,
            chi2x=0,
            chi2y=0,
            chi2z=0.6,
            tr_p_t=0.0868556558764586,
            tr_p_r=0.0006770196671656516,
        ),
        "Mass ratio q=4, chi1= (0,0,-0.8); chi2=(0,0,0.8), radial separation r=11": bbh_params(
            q=4,
            d=11,
            chi1x=0,
            chi1y=0,
            chi1z=-0.8,
            chi2x=0,
            chi2y=0,
            chi2z=0.8,
            tr_p_t=0.0558077537453816,
            tr_p_r=0.0002511426409814753,
        ),
        "Mass ratio q=1, chi1=chi2=(0,0,0), radial separation r=12": bbh_params(
            q=1,
            d=12,
            chi1x=0,
            chi1y=0,
            chi1z=0,
            chi2x=0,
            chi2y=0,
            chi2z=0,
            tr_p_t=0.0850940927209620,
            tr_p_r=0.0005398602170955123,
        ),
        "Mass ratio q=1.391, chi1=(0.4381,-0.221,0.314), chi2=(-0.1,0.4,-0.26), radial separation r=8.124981": bbh_params(
            q=1.391,
            d=8.124981,
            chi1x=0.4381,
            chi1y=-0.221,
            chi1z=0.314,
            chi2x=-0.1,
            chi2y=0.4,
            chi2z=-0.26,
            tr_p_t=0.10894035623063517,
            tr_p_r=0.002093970161808933,
        ),
    }

    for key, value in sorted(case_dict.items()):
        p_t, p_r = eval__P_t__and__P_r(
            qmassratio=value.q,
            nr=value.d,
            nchi1x=value.chi1x,
            nchi1y=value.chi1y,
            nchi1z=value.chi1z,
            nchi2x=value.chi2x,
            nchi2y=value.chi2y,
            nchi2z=value.chi2z,
        )
        # rel_error = abs(value.tr_p_t - p_t) / value.tr_p_t < 1e-13 -> abs(value.tr_p_t - p_t) < 1e-13 * value.tr_p_t
        rel_error_p_t = abs(value.tr_p_t - p_t) / value.tr_p_t
        rel_error_p_r = abs(value.tr_p_r - p_r) / value.tr_p_r
        tol = 1.0e-13
        if rel_error_p_t > tol:
            raise ValueError(
                f"Error: p_t relative error ({rel_error_p_t}) too large, when compared with trusted value"
            )
        if rel_error_p_r > tol:
            raise ValueError(
                f"Error: p_r relative error ({rel_error_p_r}) too large, when compared with trusted value"
            )
        print(value.tr_p_t, value.tr_p_r, key)
        print(p_t, p_r, key)
        print(rel_error_p_t, rel_error_p_r)

nrpypn/eval_p_t_and_p_r.py
nrpy/nrpypn
"""
Driver Python script for generating quasicircular
parameters, accurate to 3.5PN, for tangential
and radial momenta for e.g., binary black holes.

Author:  Zachariah B. Etienne
         zachetie **at** gmail **dot* com
"""

import argparse
from nrpypn.eval_p_t_and_p_r import eval__P_t__and__P_r


def main() -> None:
    """
    Main function for evaluating P_t and P_r for in quasi-circular orbits.
    """
    parser = argparse.ArgumentParser(
        description="Evaluate P_t and P_r for e.g., black holes in quasi-circular orbits."
    )
    parser.add_argument("qmassratio", type=float, nargs="?", help="Mass ratio q >= 1")
    parser.add_argument(
        "nr",
        type=float,
        nargs="?",
        help="Orbital separation (total coordinate distance between black holes)",
    )
    parser.add_argument(
        "nchi1x",
        type=float,
        nargs="?",
        help="x-component of the dimensionless spin vector for BH 1",
    )
    parser.add_argument(
        "nchi1y",
        type=float,
        nargs="?",
        help="y-component of the dimensionless spin vector for BH 1",
    )
    parser.add_argument(
        "nchi1z",
        type=float,
        nargs="?",
        help="z-component of the dimensionless spin vector for BH 1",
    )
    parser.add_argument(
        "nchi2x",
        type=float,
        nargs="?",
        help="x-component of the dimensionless spin vector for BH 2",
    )
    parser.add_argument(
        "nchi2y",
        type=float,
        nargs="?",
        help="y-component of the dimensionless spin vector for BH 2",
    )
    parser.add_argument(
        "nchi2z",
        type=float,
        nargs="?",
        help="z-component of the dimensionless spin vector for BH 2",
    )

    args = parser.parse_args()

    # Check for missing arguments and raise appropriate errors
    if any(arg is None for arg in vars(args).values()):
        print(
            "****************************************************************************"
            "\nThis program evaluates P_t and P_r for black holes in quasi-circular orbits."
        )
        print("It assumes the black holes are:")
        print(
            "* Situated initially on the x-axis, with black hole 1 at x>0, and black hole 2 at x<0."
        )
        print(
            "* Orbiting instantaneously in the x-y plane, with the center-of-mass at x=y=z=0."
        )
        print(
            "If you plan to use these parameters with TwoPunctures, you will need to set:"
        )
        print("TwoPunctures::give_bare_mass = no")
        print("TwoPunctures::target_m_minus = q/(1+q)")
        print("TwoPunctures::target_m_plus = 1/(1+q)")
        print("TwoPunctures::par_P_plus[0] = -|P_r| (from output)")
        print("TwoPunctures::par_P_plus[1] = +|P_t| (from output)")
        print("TwoPunctures::par_P_minus[0] = +|P_r| (from output)")
        print("TwoPunctures::par_P_minus[1] = -|P_t| (from output)")
        print(
            "TwoPunctures::par_S_plus[0] = nchi1x*TwoPunctures::target_m_plus*TwoPunctures::target_m_plus (from input into code cell below)"
        )
        print(
            "TwoPunctures::par_S_plus[1] = nchi1y*TwoPunctures::target_m_plus*TwoPunctures::target_m_plus (from input into code cell below)"
        )
        print(
            "TwoPunctures::par_S_plus[2] = nchi1z*TwoPunctures::target_m_plus*TwoPunctures::target_m_plus (from input into code cell below)"
        )
        print("... similarly for TwoPunctures::par_S_minus[]\n")
        parser.print_help()
    else:
        pt, pr = eval__P_t__and__P_r(
            args.qmassratio,
            args.nr,
            args.nchi1x,
            args.nchi1y,
            args.nchi1z,
            args.nchi2x,
            args.nchi2y,
            args.nchi2z,
        )
        print(f"P_t: {pt}, P_r: {pr}")


if __name__ == "__main__":
    main()

nrpypn/quasicirc.py
nrpy/nrpypn
from mpmath import mpf, mpc  # type: ignore

trusted_dict = {
    "H_NS_1PN": mpf("-19.8257035427912225797300718768688"),
    "H_NS_2PN": mpf("570.950538354676501528950589026389"),
    "H_NS_3PN": mpc(real="-7086.01166408597866097681333750691", imag="0.0"),
    "H_Newt": mpf("-2.3040324979173311641201327255852"),
}

nrpypn/tests/PN_Hamiltonian_NS.py
nrpy/nrpypn
from mpmath import mpf  # type: ignore

trusted_dict = {
    "H_SO_1p5PN": mpf("1.47046378087759998473172957458297"),
    "H_SO_2p5PN": mpf("-16.3180355853968259575125774805119"),
    "H_SO_3p5PN": mpf("102.89022500849432889566059836349"),
}

nrpypn/tests/PN_Hamiltonian_SO.py
nrpy/nrpypn
from mpmath import mpf  # type: ignore

trusted_dict = {
    "H_SS_2PN": mpf("-2.37131331324167797657501460979584"),
    "H_SS_S1S2_3PN": mpf("8.60496364312357694579811897556445"),
    "H_SS_S1sq_S2sq_3PN": mpf("28.1951662859833824059995742363262"),
}

nrpypn/tests/PN_Hamiltonian_SS.py
nrpy/nrpypn
from mpmath import mpf  # type: ignore

trusted_dict = {"H_SSS_3PN": mpf("3.568445155230340913887981317855")}

nrpypn/tests/PN_Hamiltonian_SSS.py
nrpy/nrpypn
from mpmath import mpf, mpc  # type: ignore

trusted_dict = {
    "MOmega": mpc(real="-1186.56943466450256043085172488669", imag="0.0"),
    "a_2": mpf("-1.37788070533649696515252454675881"),
    "a_3": mpf("-0.505455357265858588781512877689063"),
    "a_4": mpf("2.02368122116637566316339871454775"),
    "a_5": mpf("1.26142199627990527483776741040888"),
    "a_6": mpc(real="-6.80607857297921310532087197953488", imag="0.0"),
    "a_7": mpf("-1.00062206284378827224708966596041"),
}

nrpypn/tests/PN_MOmega.py
nrpy/nrpypn
from mpmath import mpf, mpc  # type: ignore

trusted_dict = {
    "dE_GW_dt": mpc(real="0.384077195398243128300620588750028", imag="0.0"),
    "dE_GW_dt_plus_dM_dt": mpc(real="0.383755343832215091981825945626755", imag="0.0"),
    "dM_dt": mpf("-0.000321851566028036318794643123292238"),
}

nrpypn/tests/PN_dE_GW_dt_and_dM_dt.py
nrpy/nrpypn
from mpmath import mpc  # type: ignore

trusted_dict = {
    "Htot_xyplane_binary": mpc(real="34.3246183582895289098123690031488", imag="0.0"),
    "dr_dt": mpc(
        real="-186971493.425943392178665219836809",
        imag="169624446.398018278703267691599858",
    ),
    "p_r": mpc(
        real="679448.448389355606899440876802426",
        imag="-616409.839689796322172292512186887",
    ),
}

nrpypn/tests/PN_p_r_AlternativeFalse.py
nrpy/nrpypn
from mpmath import mpc  # type: ignore

trusted_dict = {
    "Htot_xyplane_binary": mpc(real="34.3246183582895289098123690031488", imag="0.0"),
    "dr_dt": mpc(
        real="-186971493.425943392178665219836809",
        imag="169624446.398018278703267691599858",
    ),
    "p_r": mpc(
        real="3173456.22633299297826350862165735",
        imag="-2879025.64931096048450287201403591",
    ),
}

nrpypn/tests/PN_p_r_AlternativeTrue.py
nrpy/nrpypn
from mpmath import mpf, mpc  # type: ignore

trusted_dict = {
    "a_2": mpf("2.0"),
    "a_3": mpf("-1.51636607179757576634453863306679"),
    "a_4": mpf("2.36306594632581447415902941581588"),
    "a_5": mpf("-1.22489199276253877238252208698009"),
    "a_6": mpc(real="-3.43460222742737298949919905455675", imag="0.0"),
    "a_7": mpf("-0.883596441737575244634834758138545"),
    "p_t": mpc(real="-62.8291361830855940468190931599334", imag="0.0"),
}

nrpypn/tests/PN_p_t.py
nrpy/nrpypn
"""
To install the nrpypn package, navigate to this directory and execute:

    pip install .

This will install nrpypn and its required dependencies.

Instructions for uploading latest release to PyPI:
    rm -rf build dist && python setup.py sdist bdist_wheel && twine check dist/*
    twine upload dist/*
"""
import os
import sys
from pathlib import Path
from typing import List
from setuptools import setup, find_packages  # type: ignore

# pylint: disable=consider-using-f-string


def check_python_version() -> None:
    """
    Check for the minimum Python version (3.6 or newer).

    :raises: SystemExit if the Python version is less than 3.6.
    """
    if sys.version_info < (3, 6):
        raise SystemExit(
            "This project requires Python 3.6 and newer. Python {0}.{1} detected.".format(
                sys.version_info[0], sys.version_info[1]
            )
        )


def read_requirements_file() -> List[str]:
    """
    Read the contents of the requirements.txt file.

    :return: List of strings containing the required packages.
    """
    with open("requirements.txt", "r", encoding="utf-8") as file:
        return file.read().splitlines()


def get_nrpypn_version(pkg_root_directory: str) -> str:
    """
    Fetches the version from the nrpypn package.

    :param pkg_root_directory: Root directory where 'release.txt' is located.
    :return: Version as a string.
    :raises ValueError: When version information could not be found.
    """
    with open(
        os.path.join(pkg_root_directory, "nrpypn", "release.txt"), encoding="utf-8"
    ) as file:
        for line in file:
            if line.startswith("version ="):
                return line.split("=")[1].strip().strip("\"'")
    raise ValueError("Version information could not be found in 'release.txt'")


if __name__ == "__main__":
    # Don't install NRPyPN if this is run from a doctest.
    if "DOCTEST_MODE" in os.environ:
        sys.exit(0)

    check_python_version()

    dir_setup = os.path.dirname(os.path.realpath(__file__))

    requirements = read_requirements_file()

    setup(
        name="nrpypn",
        version=get_nrpypn_version(dir_setup),
        license="BSD-2-Clause",
        data_files=[("license", ["LICENSE"])],
        description="Validated Post-Newtonian Expressions for Numerical Relativity",
        long_description=(Path(__file__).parent / "README.md").read_text("UTF-8"),
        long_description_content_type="text/markdown",
        python_requires=">=3.6",
        packages=find_packages(),
        package_data={
            "nrpypn": ["py.typed"],
        },
        classifiers=[
            "License :: OSI Approved :: BSD License",
            "Operating System :: OS Independent",
            "Programming Language :: Python",
            "Topic :: Scientific/Engineering",
            "Topic :: Scientific/Engineering :: Astronomy",
            "Topic :: Scientific/Engineering :: Mathematics",
            "Topic :: Scientific/Engineering :: Physics",
            "Programming Language :: Python :: 3",
            "Programming Language :: Python :: 3.6",
            "Programming Language :: Python :: 3.7",
            "Programming Language :: Python :: 3.8",
            "Programming Language :: Python :: 3.9",
            "Programming Language :: Python :: 3.10",
            "Programming Language :: Python :: 3.11",
            "Programming Language :: Python :: 3.12",
            "Programming Language :: Python :: 3 :: Only",
            "Programming Language :: Python :: Implementation :: CPython",
            "Programming Language :: Python :: Implementation :: PyPy",
        ],
        url="https://github.com/nrpy/nrpypn",
        author="Zachariah Etienne",
        install_requires=requirements,
    )

setup.py
nrpy/nrpypn