#!/usr/bin/env python
# coding: utf-8

# <script async src="https://www.googletagmanager.com/gtag/js?id=UA-59152712-8"></script>
# <script>
#   window.dataLayer = window.dataLayer || [];
#   function gtag(){dataLayer.push(arguments);}
#   gtag('js', new Date());
# 
#   gtag('config', 'UA-59152712-8');
# </script>
# 
# # $p_t$, the tangential component of the momentum vector, up to and including 3.5 post-Newtonian order
# 
# ## Author: Zach Etienne
# 
# ## This notebook constructs the tangential component of the momentum vector, $p_t$, up to and including 3.5 PN order, and performs various validation tests. 
# 
# **Notebook Status:** <font color='green'><b> Validated </b></font>
# 
# **Validation Notes:** All expressions in this notebook were transcribed twice by hand on separate occasions, and expressions were corrected as needed to ensure consistency with published work. Published work was cross-validated and typo(s) in published work were corrected. In addition, this tutorial notebook has been confirmed to be self-consistent with its corresponding NRPy+ module, as documented [below](#code_validation). **Additional validation tests may have been performed, but are as yet, undocumented.**
# 
# ### This notebook exists as the following Python module:
# 1. [PN_p_t.py](../../edit/NRPyPN/PN_p_t.py)
# 
# ### This notebook & corresponding Python module depend on the following NRPy+/NRPyPN Python modules:
# 1. [indexedexp.py](../../edit/indexedexp.py): [**documentation+tutorial**](../Tutorial-Indexed_Expressions.ipynb)
# 1. [NRPyPN_shortcuts.py](../../edit/NRPyPN/NRPyPN_shortcuts.py): [**documentation**](NRPyPN_shortcuts.ipynb)

# <a id='toc'></a>
# 
# # Table of Contents
# $$\label{toc}$$
# 
# 1. Part 1: [$p_t$](#p_t), up to and including 3.5PN order, as derived in [Ramos-Buades, Husa, and Pratten (2018)](https://arxiv.org/abs/1810.00036)
# 1. Part 2: [Validation against second transcription and corresponding Python module](#code_validation)
# 1. Part 3: [Validation against trusted numerical values](#code_validationv2) (i.e., in Table V of [Ramos-Buades, Husa, and Pratten (2018)](https://arxiv.org/abs/1810.00036))
# 1. Part 4: [LaTeX PDF output](#latex_pdf_output): $\LaTeX$ PDF Output

# <a id='p_t'></a>
# 
# # Part 1: $p_t$, up to and including 3.5PN order, as derived in [Ramos-Buades, Husa, and Pratten (2018)](https://arxiv.org/abs/1810.00036) \[Back to [top](#toc)\]
# $$\label{p_t}$$ 
# 
# As described in the [nonspinning Hamiltonian notebook](PN-Hamiltonian-Nonspinning.ipynb), the basic physical system assumes two point particles of mass $m_1$ and $m_2$ with corresponding momentum vectors $\mathbf{P}_1$ and $\mathbf{P}_2$, and displacement vectors $\mathbf{X}_1$ and $\mathbf{X}_2$ with respect to the center of mass. Here we also consider the spin vectors of each point mass $\mathbf{S}_1$ and $\mathbf{S}_2$, respectively.
# 
# To reduce possibility of copying error, the equation for $p_t$ is taken directly from the arXiv LaTeX source code of Eq A2 in [Ramos-Buades, Husa, and Pratten (2018)](https://arxiv.org/abs/1810.00036), and only mildly formatted to (1) improve presentation in Jupyter notebooks, (2) to ensure some degree of consistency in notation across different terms in other NRPyPN notebooks, and (3) to correct any errors. In particular, the boxed negative sign at 2.5PN order ($a_5$ below) was missing in the original equation. We will later show that this negative sign is necessary for consistency with other expressions in the same paper, as well as with the expression up to 3PN order in [Healy, Lousto, Nakano, and Zlochower (2017)](https://arxiv.org/abs/1702.00872):
# 
# $$
# p_t = \frac{q}{(1+q)^2}\frac{1}{r^{1/2}}\left(1 + \sum_{k=2}^7 \frac{a_k}{r^{k/2}}\right),
# $$
# where
# 
# \begin{align}
# a_2 &= 2\\
# a_3 &= \left[-\frac{3 \left(4 q^2+3 q\right) \chi _{2z}}{4 (q+1)^2}-\frac{3 (3 q+4) \chi _{1z}}{4 (q+1)^2}\right]\\
# a_4 &= \left[ -\frac{3 q^2 \chi _{2x}^2}{2 (q+1)^2} +\frac{3 q^2 \chi _{2y}^2}{4 (q+1)^2}+\frac{3 q^2 \chi _{2z}^2}{4 (q+1)^2} 
#  +\frac{42 q^2+41 q+42}{16 (q+1)^2}-\frac{3 \chi _{1x}^2}{2 (q+1)^2} \right.\\
# &\quad\quad \left. -\frac{3 q  \chi _{1x} \chi _{2x}}{(q+1)^2}+\frac{3 \chi _{1y}^2}{4 (q+1)^2}+\frac{3 q \chi _{1y}\chi _{2y}}{2 (q+1)^2}+\frac{3 \chi _{1z}^2}{4 (q+1)^2}+\frac{3 q \chi _{1z} \chi _{2z}}{2 (q+1)^2}\right]\\
# a_5 &= \left[ \boxed{-1} \frac{\left(13 q^3+60 q^2+116 q+72\right) \chi _{1z}}{16 (q+1)^4}+\frac{\left(-72 q^4-116 q^3-60 q^2-13 q\right) \chi _{2z}}{16 (q+1)^4}  \right]\\
# a_6 &= \left[\frac{\left(472 q^2-640\right) \chi _{1x}^2}{128 (q+1)^4} + \frac{\left(-512 q^2-640 q-64\right) \chi _{1y}^2}{128 (q+1)^4}+\frac{\left(-108 q^2+224 q+512\right) \chi _{1z}^2}{128 (q+1)^4}\right.\\
# &\quad\quad \left.+\frac{\left(472 q^2-640 q^4\right) \chi _{2x}^2}{128 (q+1)^4}+\frac{\left(192 q^3+560 q^2+192 q\right) \chi _{1x} \chi _{2x}}{128 (q+1)^4} +\frac{\left(-864 q^3-1856 q^2-864 q\right) \chi _{1y} \chi _{2y}}{128 (q+1)^4}\right.\\
# &\quad\quad \left.+\frac{\left(480 q^3+1064 q^2+480 q\right) \chi _{1z} \chi _{2z}}{128 (q+1)^4}+\frac{\left(-64 q^4-640 q^3-512 q^2\right) \chi _{2y}^2}{128 (q+1)^4}+\frac{\left(512 q^4+224 q^3-108 q^2\right) \chi _{2z}^2}{128 (q+1)^4} \right. \nonumber
# \\
# &\quad\quad\left.+\frac{480 q^4+163 \pi ^2 q^3-2636 q^3+326 \pi ^2 q^2-6128 q^2+163 \pi ^2 q-2636 q+480}{128 (q+1)^4} \right]\\
# a_7 &= \left[ \frac{5 (4 q+1) q^3 \chi _{2 x}^2 \chi _{2 z}}{2 (q+1)^4}-\frac{5 (4 q+1) q^3 \chi _{2 y}^2 \chi _{2 z}}{8 (q+1)^4}-\frac{5 (4 q+1) q^3 \chi _{2 z}^3}{8 (q+1)^4}+\chi _{1x} \left(\frac{15 (2 q+1) q^2 \chi _{2 x} \chi _{2 z}}{4 (q+1)^4}+\frac{15 (q+2) q \chi _{2 x} \chi _{1z}}{4 (q+1)^4}\right)\right. \nonumber
# \\
# &\quad\quad \left.+\chi _{1y} \left(\frac{15 q^2 \chi _{2 y} \chi _{1z}}{4 (q+1)^4}+\frac{15 q^2 \chi _{2 y} \chi _{2 z}}{4 (q+1)^4}\right)+\chi _{1z} \left(\frac{15 q^2 (2 q+3) \chi _{2 x}^2}{4 (q+1)^4}-\frac{15 q^2 (q+2) \chi _{2 y}^2}{4 (q+1)^4}-\frac{15 q^2 \chi _{2 z}^2}{4 (q+1)^3} \right.\right. \nonumber
# \\
# &\quad\quad \left.\left. -\frac{103 q^5+145 q^4-27 q^3+252 q^2+670 q+348}{32 (q+1)^6}\right)-\frac{\left(348 q^5+670 q^4+252 q^3-27 q^2+145 q+103\right) q \chi _{2 z}}{32 (q+1)^6}\right.\nonumber
# \\
# &\quad\quad \left.+\chi _{1x}^2 \left(\frac{5 (q+4) \chi _{1z}}{2 (q+1)^4}+\frac{15 q (3 q+2) \chi _{2 z}}{4 (q+1)^4}\right)+\chi _{1y}^2 \left(-\frac{5 (q+4) \chi _{1z}}{8 (q+1)^4}-\frac{15 q (2 q+1) \chi _{2 z}}{4 (q+1)^4}\right)-\frac{15 q \chi _{1z}^2 \chi _{2 z}}{4 (q+1)^3}-\frac{5 (q+4) \chi _{1z}^3}{8 (q+1)^4} \right]
# \end{align}

# Let's divide and conquer, by tackling the coefficients one at a time:
# 
# \begin{align}
# a_2 &= 2\\
# a_3 &= \left[-\frac{3 \left(4 q^2+3 q\right) \chi _{2z}}{4 (q+1)^2}-\frac{3 (3 q+4) \chi _{1z}}{4 (q+1)^2}\right]\\
# a_4 &= \left[ -\frac{3 q^2 \chi _{2x}^2}{2 (q+1)^2} +\frac{3 q^2 \chi _{2y}^2}{4 (q+1)^2}+\frac{3 q^2 \chi _{2z}^2}{4 (q+1)^2} 
#  +\frac{42 q^2+41 q+42}{16 (q+1)^2}-\frac{3 \chi _{1x}^2}{2 (q+1)^2} \right.\\
# &\quad\quad \left. -\frac{3 q  \chi _{1x} \chi _{2x}}{(q+1)^2}+\frac{3 \chi _{1y}^2}{4 (q+1)^2}+\frac{3 q \chi _{1y}\chi _{2y}}{2 (q+1)^2}+\frac{3 \chi _{1z}^2}{4 (q+1)^2}+\frac{3 q \chi _{1z} \chi _{2z}}{2 (q+1)^2}\right]
# \end{align}

# In[1]:


# Step 0: Add NRPy's directory to the path
# https://stackoverflow.com/questions/16780014/import-file-from-parent-directory
import sympy as sp               # SymPy: The Python computer algebra package upon which NRPy+ depends
import indexedexpNRPyPN as ixp         # NRPy+: Symbolic indexed expression (e.g., tensors, vectors, etc.) support
from NRPyPN_shortcuts import div # NRPyPN: shortcuts for e.g., vector operations

# 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
def p_t__a_2_thru_a_4(m1,m2, chi1x,chi1y,chi1z, chi2x,chi2y,chi2z):
    q = m2/m1 # It is assumed that q >= 1, so m2 >= m1.
    global a_2,a_3,a_4
    a_2 = 2
    a_3 = (-3*(4*q**2+3*q)*chi2z/(4*(q+1)**2) - 3*(3*q+4)*chi1z/(4*(q+1)**2))
    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))


# In[2]:


# Second version, for validation purposes only.
def p_t__a_2_thru_a_4v2(m1,m2, chi1x,chi1y,chi1z, chi2x,chi2y,chi2z):
    q = m2/m1 # It is assumed that q >= 1, so m2 >= m1.
    global a_2v2,a_3v2,a_4v2
    # Validated against HLNZ2017 version
    a_2v2 = 2
    # Validated against HLNZ2017 version
    a_3v2 = (-(3*(4*q**2+3*q)*chi2z)/(4*(q+1)**2)-(3*(3*q+4)*chi1z)/(4*(q+1)**2))
    # Validated against HLNZ2017 version
    a_4v2 = -(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)


# Next, $a_5$ and $a_6$:
# 
# \begin{align}
# a_5 &= \left[ \boxed{-1} \frac{\left(13 q^3+60 q^2+116 q+72\right) \chi _{1z}}{16 (q+1)^4}+\frac{\left(-72 q^4-116 q^3-60 q^2-13 q\right) \chi _{2z}}{16 (q+1)^4}  \right]\\
# a_6 &= \left[\frac{\left(472 q^2-640\right) \chi _{1x}^2}{128 (q+1)^4} + \frac{\left(-512 q^2-640 q-64\right) \chi _{1y}^2}{128 (q+1)^4}+\frac{\left(-108 q^2+224 q+512\right) \chi _{1z}^2}{128 (q+1)^4}\right.\\
# &\quad\quad \left.+\frac{\left(472 q^2-640 q^4\right) \chi _{2x}^2}{128 (q+1)^4}+\frac{\left(192 q^3+560 q^2+192 q\right) \chi _{1x} \chi _{2x}}{128 (q+1)^4} +\frac{\left(-864 q^3-1856 q^2-864 q\right) \chi _{1y} \chi _{2y}}{128 (q+1)^4}\right.\\
# &\quad\quad \left.+\frac{\left(480 q^3+1064 q^2+480 q\right) \chi _{1z} \chi _{2z}}{128 (q+1)^4}+\frac{\left(-64 q^4-640 q^3-512 q^2\right) \chi _{2y}^2}{128 (q+1)^4}+\frac{\left(512 q^4+224 q^3-108 q^2\right) \chi _{2z}^2}{128 (q+1)^4} \right. \nonumber
# \\
# &\quad\quad\left.+\frac{480 q^4+163 \pi ^2 q^3-2636 q^3+326 \pi ^2 q^2-6128 q^2+163 \pi ^2 q-2636 q+480}{128 (q+1)^4} \right]\\
# \end{align}

# In[3]:


# 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.
def p_t__a_5_thru_a_6(m1,m2, chi1x,chi1y,chi1z, chi2x,chi2y,chi2z, FixSignError=True):
    SignFix = sp.sympify(-1)
    if FixSignError == False:
        SignFix = sp.sympify(+1)
    q = m2/m1 # It is assumed that q >= 1, so m2 >= m1.
    global a_5,a_6
    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))
    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))


# In[4]:


# Second version, for validation purposes only.
def p_t__a_5_thru_a_6v2(m1,m2, chi1x,chi1y,chi1z, chi2x,chi2y,chi2z, FixSignError=True):
    SignFix = sp.sympify(-1)
    if FixSignError == False:
        SignFix = sp.sympify(+1)
    q = m2/m1 # It is assumed that q >= 1, so m2 >= m1.
    pi = sp.pi
    global a_5v2,a_6v2
    # Validated (separately) against HLNZ2017, as well as row 3 of Table V in RHP2018
    a_5v2 =  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)
    # Validated (separately) against HLNZ2017 version
    a_6v2 = (+(+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*pi**2*q**3 - 2636*q**3 + 326*pi**2*q**2 - 6128*q**2 + 163*pi**2*q - 2636*q + 480)
               /(128*(q+1)**4))


# Next, we compare the expression for $a_5$ with Eq. 7 of [Healy, Lousto, Nakano, and Zlochower (2017)](https://arxiv.org/abs/1702.00872), as additional validation that there is at least is a sign inconsistency:
# 
# To reduce the possibility of copying errors, the following equation for $a_5$ is taken directly from the arXiv LaTeX source code of Eq. 7 of [Healy, Lousto, Nakano, and Zlochower (2017)](https://arxiv.org/abs/1702.00872), and only mildly formatted to (1) improve presentation in Jupyter notebooks and (2) to ensure some degree of consistency in notation across different terms in other NRPyPN notebooks.
# 
# **Important: Note that [Healy, Lousto, Nakano, and Zlochower (2017)](https://arxiv.org/abs/1702.00872) adopts notation such that particle labels are interchanged: $1\leftrightarrow 2$, with respect to [Ramos-Buades, Husa, and Pratten (2018)](https://arxiv.org/abs/1810.00036)**
# 
# \begin{align}
# a_5 &= + \left( -\frac{1}{16}\,{\frac {q \left( 72\,{q}^{3}+116\,{q}^{2}+60\,q+13 \right) 
# {\chi_{1z}}}{ \left( 1+q \right) ^{4}}}
# -\frac{1}{16}\,{\frac { \left( 13\,{q}^{3}+60\,{q}^{2}+116\,q+72 \right) {\chi_{2z}}}
# { \left( 1+q \right) ^{4}}} \right)\\
# \end{align}

# In[5]:


# Third version, for addtional validation.
def p_t__a_5_thru_a_6_HLNZ2017(m1,m2, chi1x,chi1y,chi1z, chi2x,chi2y,chi2z):
    q = m2/m1 # It is assumed that q >= 1, so m2 >= m1.
    global a_5_HLNZ2017
    a_5_HLNZ2017 = (-div(1,16)*(q*(72*q**3 + 116*q**2 + 60*q + 13)*chi1z/(1+q)**4)
                    -div(1,16)*(  (13*q**3 +  60*q**2 +116*q + 72)*chi2z/(1+q)**4))


# Finally, we validate that all 3 expressions for $a_5$ agree. (At the bottom, we confirm that all v2 expressions for $a_i$ match.)

# In[6]:


from NRPyPN_shortcuts import m1,m2, chi1U,chi2U # Import needed input variables

p_t__a_5_thru_a_6(         m1,m2, chi1U[0],chi1U[1],chi1U[2], chi2U[0],chi2U[1],chi2U[2])
p_t__a_5_thru_a_6v2(       m1,m2, chi1U[0],chi1U[1],chi1U[2], chi2U[0],chi2U[1],chi2U[2])
# Again, the particle labels are interchanged in Healy, Lousto, Nakano, and Zlochower (2017):
p_t__a_5_thru_a_6_HLNZ2017(m1,m2, chi2U[0],chi2U[1],chi2U[2], chi1U[0],chi1U[1],chi1U[2])

def error(varname):
    print("ERROR: When comparing Python module & notebook, "+varname+" was found not to match.")
    sys.exit(1)

if sp.simplify(a_5 - a_5v2)        != 0: error("a_5v2")
if sp.simplify(a_5 - a_5_HLNZ2017) != 0: error("a_5_HLNZ2017")


# Finally $a_7$:
# 
# \begin{align}
# a_7 &= \left[ \frac{5 (4 q+1) q^3 \chi _{2 x}^2 \chi _{2 z}}{2 (q+1)^4}-\frac{5 (4 q+1) q^3 \chi _{2 y}^2 \chi _{2 z}}{8 (q+1)^4}-\frac{5 (4 q+1) q^3 \chi _{2 z}^3}{8 (q+1)^4}+\chi _{1x} \left(\frac{15 (2 q+1) q^2 \chi _{2 x} \chi _{2 z}}{4 (q+1)^4}+\frac{15 (q+2) q \chi _{2 x} \chi _{1z}}{4 (q+1)^4}\right)\right. \nonumber
# \\
# &\quad\quad \left.+\chi _{1y} \left(\frac{15 q^2 \chi _{2 y} \chi _{1z}}{4 (q+1)^4}+\frac{15 q^2 \chi _{2 y} \chi _{2 z}}{4 (q+1)^4}\right)+\chi _{1z} \left(\frac{15 q^2 (2 q+3) \chi _{2 x}^2}{4 (q+1)^4}-\frac{15 q^2 (q+2) \chi _{2 y}^2}{4 (q+1)^4}-\frac{15 q^2 \chi _{2 z}^2}{4 (q+1)^3} \right.\right. \nonumber
# \\
# &\quad\quad \left.\left. -\frac{103 q^5+145 q^4-27 q^3+252 q^2+670 q+348}{32 (q+1)^6}\right)-\frac{\left(348 q^5+670 q^4+252 q^3-27 q^2+145 q+103\right) q \chi _{2 z}}{32 (q+1)^6}\right.\nonumber
# \\
# &\quad\quad \left.+\chi _{1x}^2 \left(\frac{5 (q+4) \chi _{1z}}{2 (q+1)^4}+\frac{15 q (3 q+2) \chi _{2 z}}{4 (q+1)^4}\right)+\chi _{1y}^2 \left(-\frac{5 (q+4) \chi _{1z}}{8 (q+1)^4}-\frac{15 q (2 q+1) \chi _{2 z}}{4 (q+1)^4}\right)-\frac{15 q \chi _{1z}^2 \chi _{2 z}}{4 (q+1)^3}-\frac{5 (q+4) \chi _{1z}^3}{8 (q+1)^4} \right]
# \end{align}

# In[7]:


# Construct term a_7, from Eq A2 of
#  Ramos-Buades, Husa, and Pratten (2018)
#    https://arxiv.org/abs/1810.00036
def p_t__a_7(m1,m2, chi1x,chi1y,chi1z, chi2x,chi2y,chi2z):
    q = m2/m1 # It is assumed that q >= 1, so m2 >= m1.
    global a_7
    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))


# In[8]:


# Second version, for validation purposes only.
def p_t__a_7v2(m1,m2, chi1x,chi1y,chi1z, chi2x,chi2y,chi2z):
    q = m2/m1 # It is assumed that q >= 1, so m2 >= m1.
    global a_7v2
    a_7v2  = (+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*(  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))


# Putting it all together, recall that
# 
# $$
# p_t = \frac{q}{(1+q)^2}\frac{1}{r^{1/2}}\left(1 + \sum_{k=2}^7 \frac{a_k}{r^{k/2}}\right),
# $$
# where $k/2$ is the post-Newtonian order.

# In[9]:


# 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}))
def f_p_t(m1,m2, chi1U,chi2U, r):
    q = m2/m1 # It is assumed that q >= 1, so m2 >= m1.
    a = ixp.zerorank1(DIM=10)
    p_t__a_2_thru_a_4(m1,m2, chi1U[0],chi1U[1],chi1U[2], chi2U[0],chi2U[1],chi2U[2])
    a[2] = a_2
    a[3] = a_3
    a[4] = a_4
    p_t__a_5_thru_a_6(m1,m2, chi1U[0],chi1U[1],chi1U[2], chi2U[0],chi2U[1],chi2U[2])
    a[5] = a_5
    a[6] = a_6
    p_t__a_7(         m1,m2, chi1U[0],chi1U[1],chi1U[2], chi2U[0],chi2U[1],chi2U[2])
    a[7] = a_7
    global p_t
    p_t = 1 # Term prior to the sum in parentheses
    for k in range(8):
        p_t += a[k]/r**div(k,2)
    p_t *= q / (1+q)**2 * 1/r**div(1,2)


# In[10]:


# Second version, for validation purposes only.
def f_p_tv2(m1,m2, chi1U,chi2U, r):
    q = m2/m1 # It is assumed that q >= 1, so m2 >= m1.
    a = ixp.zerorank1(DIM=10)
    p_t__a_2_thru_a_4v2(m1,m2, chi1U[0],chi1U[1],chi1U[2], chi2U[0],chi2U[1],chi2U[2])
    a[2] = a_2v2
    a[3] = a_3v2
    a[4] = a_4v2
    p_t__a_5_thru_a_6v2(m1,m2, chi1U[0],chi1U[1],chi1U[2], chi2U[0],chi2U[1],chi2U[2])
    a[5] = a_5v2
    a[6] = a_6v2
    p_t__a_7v2(         m1,m2, chi1U[0],chi1U[1],chi1U[2], chi2U[0],chi2U[1],chi2U[2])
    a[7] = a_7v2
    global p_tv2
    p_tv2 = 1 # Term prior to the sum in parentheses
    for k in range(8):
        p_tv2 += a[k]/r**div(k,2)
    p_tv2 *= q / (1+q)**2 * 1/r**div(1,2)


# <a id='code_validation'></a>
# 
# # Part 2: Validation against second transcription and corresponding Python module \[Back to [top](#toc)\]
# $$\label{code_validation}$$ 
# 
# As a code validation check, we verify agreement between 
# * the SymPy expressions transcribed from the cited published work on two separate occasions, and
# * the SymPy expressions generated in this notebook, and the corresponding Python module.

# In[11]:


from NRPyPN_shortcuts import q, num_eval # Import needed input variable & numerical evaluation routine

f_p_t(m1,m2, chi1U,chi2U, q)

def error(varname):
    print("ERROR: When comparing Python module & notebook, "+varname+" was found not to match.")
    sys.exit(1)

# Validation against second transcription of the expressions:
f_p_tv2(m1,m2, chi1U,chi2U, q)
if sp.simplify(p_t - p_tv2) != 0: error("p_tv2")

# Validation against corresponding Python module:
import PN_p_t as pt
pt.f_p_t(m1,m2, chi1U,chi2U, q)

if sp.simplify(p_t - pt.p_t) != 0: error("pt.p_t")
print("ALL TESTS PASS")


# <a id='code_validationv2'></a>
# 
# # Part 3: Validation against trusted numerical values (i.e., in Table V of [Ramos-Buades, Husa, and Pratten (2018)](https://arxiv.org/abs/1810.00036)) \[Back to [top](#toc)\]
# $$\label{code_validationv2}$$ 

# In[12]:


# Useful function for comparing published & NRPyPN results
def compare_pub_NPN(desc, pub,NPN,NPN_with_a5_chi1z_sign_error):
    print("##################################################")
    print("  "+desc)
    print("##################################################")
    print(str(pub) + "          <- Expected result, from Table V of Ramos-Buades, Husa, and Pratten (2018)")
    print(str(NPN) + " <- Result from NRPyPN")
    relerror = abs(pub-NPN)/pub
    resultstring = "Relative error between NRPyPN & published: "+str(relerror*100)+"%"
    if relerror > 1e-3:
        resultstring += " <--- NOT GOOD! (see explanation below)"
    else:
        resultstring += " <--- EXCELLENT AGREEMENT!"
    print(resultstring+"\n")
    print(str(NPN_with_a5_chi1z_sign_error) + " <- Result from NRPyPN, with chi1z sign error in a_5 expression.")


# In[13]:


# 1. Let's consider the case:
# * Mass ratio q=1, chi1=chi2=(0,0,0), radial separation r=12
pub_result = 0.850941e-1 # Expected result, from Table V of Ramos-Buades, Husa, and Pratten (2018) https://arxiv.org/abs/1810.00036
NPN_with_a5_chi1z_sign_error = 0.0850940927209620 # should be unaffected by sign error, as chi1z=0.
NPN_result = num_eval(p_t,
                      qmassratio =  1.0,  # must be >= 1
                      nr         = 12.0,  # Orbital separation
                      nchi1x     = +0.,
                      nchi1y     = +0.,
                      nchi1z     = +0.,
                      nchi2x     = +0.,
                      nchi2y     = +0.,
                      nchi2z     = +0.)
compare_pub_NPN("Case: q=1, nonspinning, initial separation 12",
                pub_result,NPN_result,NPN_with_a5_chi1z_sign_error)


# In[14]:


# 2. Let's consider the case:
# * Mass ratio q=1.5, chi1= (0,0,-0.6); chi2=(0,0,0.6), radial separation r=10.8
pub_result = 0.868557e-1 # Expected result, from Table V of Ramos-Buades, Husa, and Pratten (2018) https://arxiv.org/abs/1810.00036
NPN_with_a5_chi1z_sign_error = 0.0867002374951143
NPN_result = num_eval(p_t,
                      qmassratio =  1.5,  # must be >= 1
                      nr         = 10.8,  # Orbital separation
                      nchi1x     = +0.,
                      nchi1y     = +0.,
                      nchi1z     = -0.6,
                      nchi2x     = +0.,
                      nchi2y     = +0.,
                      nchi2z     = +0.6)
compare_pub_NPN("Case: q=1.5, chi1z=-0.6, chi2z=0.6, initial separation 10.8",
                pub_result,NPN_result,NPN_with_a5_chi1z_sign_error)


# In[15]:


# 3. Let's consider the case:
# * Mass ratio q=4, chi1= (0,0,-0.8); chi2=(0,0,0.8), radial separation r=11
pub_result = 0.559207e-1 # Expected result, from Table V of Ramos-Buades, Husa, and Pratten (2018) https://arxiv.org/abs/1810.00036
NPN_with_a5_chi1z_sign_error = 0.0557629777874552
NPN_result = num_eval(p_t,
                      qmassratio =  4.0,  # must be >= 1
                      nr         = 11.0,  # Orbital separation
                      nchi1x     = +0.,
                      nchi1y     = +0.,
                      nchi1z     = -0.8,
                      nchi2x     = +0.,
                      nchi2y     = +0.,
                      nchi2z     = +0.8)
compare_pub_NPN("Case: q=4.0, chi1z=-0.8, chi2z=0.8, initial separation 11.0",
                pub_result,NPN_result,NPN_with_a5_chi1z_sign_error)
print("0.0558369          <- Second iteration value in pub result. Note that NRPyPN value is *closer* to this value.")


# In[16]:


# 4. Let's consider the case:
# * Mass ratio q=2, chi1= (0,0,0); chi2=(−0.3535, 0.3535, 0.5), radial separation r=10.8
pub_result = 0.7935e-1 # Expected result, from Table V of Ramos-Buades, Husa, and Pratten (2018) https://arxiv.org/abs/1810.00036
NPN_with_a5_chi1z_sign_error = 0.0793500403866190 # should be unaffected by sign error, as chi1z=0.
NPN_result = num_eval(p_t,
                      qmassratio =  2.0,  # must be >= 1
                      nr         = 10.8,  # Orbital separation
                      nchi1x     = +0.,
                      nchi1y     = +0.,
                      nchi1z     = +0.,
                      nchi2x     = -0.3535,
                      nchi2y     = +0.3535,
                      nchi2z     = +0.5)
compare_pub_NPN("Case: q=2.0, chi2x=-0.3535, chi2y=+0.3535, chi2z=+0.5, initial separation 10.8",
                pub_result,NPN_result,NPN_with_a5_chi1z_sign_error)


# In[17]:


# 5. Let's consider the case:
# * Mass ratio q=8, chi1= (0, 0, 0.5); chi2=(0, 0, 0.5), radial separation r=11
pub_result = 0.345755e-1 # Expected result, from Table V of Ramos-Buades, Husa, and Pratten (2018) https://arxiv.org/abs/1810.00036
NPN_with_a5_chi1z_sign_error = 0.0345584951081129 # should be unaffected by sign error, as chi1z=0.
NPN_result = num_eval(p_t,
                      qmassratio =  8.0,  # must be >= 1
                      nr         = 11.0,  # Orbital separation
                      nchi1x     = +0.,
                      nchi1y     = +0.,
                      nchi1z     = +0.5,
                      nchi2x     = +0.,
                      nchi2y     = +0.,
                      nchi2z     = +0.5)
compare_pub_NPN("""
Case: q=8.0, chi1z=chi2z=+0.5, initial separation 11

Note: This one is weird. Clearly the value in the table
      has a typo, such that the p_r and p_t values
      should be interchanged; p_t is about 20% the
      next smallest value in the table, and the
      parameters aren't that different. We therefore
      assume that this is the case, and find agreement
      with the published result to about 0.07%, which
      isn't the best, but given that the table values
      seem to be clearly wrong, it's an encouraging
      sign.
""",pub_result,NPN_result,NPN_with_a5_chi1z_sign_error)


# <a id='latex_pdf_output'></a>
# 
# # Part 4: Output this notebook to $\LaTeX$-formatted PDF file \[Back to [top](#toc)\]
# $$\label{latex_pdf_output}$$
# 
# The following code cell converts this Jupyter notebook into a proper, clickable $\LaTeX$-formatted PDF file. After the cell is successfully run, the generated PDF may be found in the root NRPy+ tutorial directory, with filename
# [PN-p_t.pdf](PN-p_t.pdf) (Note that clicking on this link may not work; you may need to open the PDF file through another means.)

# In[18]:


import os,sys                    # Standard Python modules for multiplatform OS-level functions
import cmdline_helperNRPyPN as cmd    # NRPy+: Multi-platform Python command-line interface
cmd.output_Jupyter_notebook_to_LaTeXed_PDF("PN-p_t")