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

# # 3+1 Simon-Mars tensor in the $\delta=2$ Tomimatsu-Sato spacetime
# 
# This worksheet demonstrates a few capabilities of [SageManifolds](http://sagemanifolds.obspm.fr/) (version 1.0, as included in SageMath 7.5) in computations regarding the 3+1 decomposition of the Simon-Mars tensor in the $\delta=2$ Tomimatsu-Sato spacetime. The results obtained here are used in the article [arXiv:1412.6542](http://arxiv.org/abs/1412.6542).
# 
# Click [here](https://raw.githubusercontent.com/sagemanifolds/SageManifolds/master/Worksheets/v1.0/SM_Simon-Mars_3p1_TS2.ipynb) to download the worksheet file (ipynb format). To run it, you must start SageMath with the Jupyter notebook, via the command `sage -n jupyter`

# *NB:* a version of SageMath at least equal to 7.5 is required to run this worksheet:

# In[1]:


version()


# First we set up the notebook to display mathematical objects using LaTeX rendering:

# In[2]:


get_ipython().run_line_magic('display', 'latex')


# Since some computations are quite long, we ask for running them in parallel on 8 cores:

# In[3]:


Parallelism().set(nproc=8)


# <h2>Tomimatsu-Sato spacetime</h2>
# <p>The Tomimatsu-Sato metric is an exact stationary and axisymmetric  solution of the vacuum Einstein equation, which is asymptotically flat and has a non-zero angular momentum. It has been found in 1972 by A. Tomimatsu and H. Sato [<a href="http://dx.doi.org/10.1103/PhysRevLett.29.1344">Phys. Rev. Lett. <strong>29</strong>, 1344 (1972)</a>], as a solution of the Ernst equation. It is actually the member $\delta=2$ of a larger family of solutions parametrized by a positive integer $\delta$ and exhibited by Tomimatsu and Sato in 1973 <a href="http://dx.doi.org/10.1143/PTP.50.95">[Prog. Theor. Phys. <strong>50</strong>, 95 (1973)</a>], the member $\delta=1$ being nothing but the Kerr metric. We refer to [<a href="http://dx.doi.org/10.1143/PTP.127.1057">Manko, Prog. Theor. Phys.<strong> 127</strong>, 1057 (2012)</a>] for a discussion of the properties of this solution.</p>
# <h2>Spacelike hypersurface</h2>
# <p>We consider some hypersurface $\Sigma$ of a spacelike foliation $(\Sigma_t)_{t\in\mathbb{R}}$ of $\delta=2$ Tomimatsu-Sato spacetime; we declare $\Sigma_t$ as a 3-dimensional manifold:</p>

# In[4]:


Sig = Manifold(3, 'Sigma', r'\Sigma', start_index=1)


# <p>On $\Sigma$, we consider the prolate spheroidal <span id="cell_outer_1">coordinates</span> $(x,y,\phi)$, with $x\in(1,+\infty)$, $y\in(-1,1)$ and $\phi\in(0,2\pi)$ :</p>

# In[5]:


X.<r,y,ph> = Sig.chart(r'x:(1,+oo) y:(-1,1) ph:(0,2*pi):\phi')
print X ; X


# <h3>Riemannian metric on $\Sigma$</h3>
# <p>The Tomimatsu-Sato metric depens on three parameters: the integer $\delta$, the real number $p\in[0,1]$, and the total mass $m$:</p>

# In[6]:


var('d, p, m')
assume(m>0)
assumptions()


# <p>We set $\delta=2$ and choose a specific value for $p$, namely $p=1/5$:</p>

# In[7]:


d = 2
p = 1/5


# <p>Furthermore, without any loss of generality, we may set $m=1$ (this simply fixes some length scale):</p>

# In[8]:


m=1


# <p>The parameter $q$ is related to $p$ by $p^2+q^2=1$:</p>

# In[9]:


q = sqrt(1-p^2)


# <p>Some shortcut notations:</p>

# In[10]:


AA2 = (p^2*(x^2-1)^2+q^2*(1-y^2)^2)^2 \
      - 4*p^2*q^2*(x^2-1)*(1-y^2)*(x^2-y^2)^2
BB2 = (p^2*x^4+2*p*x^3-2*p*x+q^2*y^4-1)^2 \
      + 4*q^2*y^2*(p*x^3-p*x*y^2-y^2+1)^2
CC2 = p^3*x*(1-x^2)*(2*(x^4-1)+(x^2+3)*(1-y^2)) \
      + p^2*(x^2-1)*((x^2-1)*(1-y^2)-4*x^2*(x^2-y^2)) \
      + q^2*(1-y^2)^3*(p*x+1)


# <p>The Riemannian metric $\gamma$ induced by the spacetime metric $g$ on $\Sigma$:</p>

# In[11]:


gam = Sig.riemannian_metric('gam', latex_name=r'\gamma') 
gam[1,1] = m^2*BB2/(p^2*d^2*(x^2-1)*(x^2-y^2)^3)
gam[2,2] = m^2*BB2/(p^2*d^2*(y^2-1)*(-x^2+y^2)^3)
gam[3,3] = - m^2*(y^2-1)*(p^2*BB2^2*(x^2-1)
                          + 4*q^2*d^2*CC2^2*(y^2-1))/(AA2*BB2*d^2)
gam.display()


# A view of the non-vanishing components of $\gamma$ w.r.t. coordinates $(x,y,\phi)$:

# In[12]:


gam.display_comp()


# The expression of the metric determinant with respect to the default chart (coordinates $(x,y,\phi)$):

# In[13]:


gam.determinant().expr()


# <h3>Lapse function and shift vector</h3>

# In[14]:


N2 = AA2/BB2 - 2*m*q*CC2*(y^2-1)/BB2*(2*m*q*CC2*(y^2-1)
                /(BB2*(m^2*(y^2-1)*(p^2*BB2^2*(x^2-1)
                +4*q^2*d^2*CC2^2*(y^2-1))/(AA2*BB2*d^2)))) 
N2.simplify_full()


# In[15]:


N = Sig.scalar_field(sqrt(N2.simplify_full()), name='N')
print N
N.display()


# The coordinate expression of the scalar field $N$:

# In[16]:


N.expr()


# In[17]:


b3 = 2*m*q*CC2*(y^2-1)/(BB2*(m^2*(y^2-1)*(p^2*BB2^2*(x^2-1)
        +4*q^2*d^2*CC2^2*(y^2-1))/(AA2*BB2*d^2)))
b = Sig.vector_field('beta', latex_name=r'\beta') 
b[3] = b3.simplify_full()
# unset components are zero 
b.display_comp(only_nonzero=False)


# <h3>Extrinsic curvature of $\Sigma$</h3>
# <p>We use the formula $$K_{ij} = \frac{1}{2N} \mathcal{L}_{\beta} \gamma_{ij}, $$ which is valid for any stationary spacetime:</p>

# In[18]:


K = gam.lie_derivative(b) / (2*N)
K.set_name('K')
print K


# <p>The component $K_{13} = K_{x\phi}$:</p>

# In[19]:


K[1,3]


# <p>The type-(1,1) tensor $K^\sharp$ of components $K^i_{\ \, j} = \gamma^{ik} K_{kj}$:</p>

# In[20]:


Ku = K.up(gam, 0)
print Ku


# <p>We may check that the hypersurface $\Sigma$ is maximal, i.e. that $K^k_{\ \, k} = 0$:</p>

# In[21]:


trK = Ku.trace()
trK


# <h3>Connection and curvature</h3>
# <p>Let us call $D$ the Levi-Civita connection associated with $\gamma$: </p>

# In[22]:


D = gam.connection(name='D')
print D


# <p>The Ricci tensor associated with $\gamma$:</p>

# In[23]:


Ric = gam.ricci()
print Ric


# <p>The scalar curvature $R = \gamma^{ij} R_{ij}$:</p>

# In[24]:


R = gam.ricci_scalar(name='R')
print R


# <h3>Terms related to the extrinsic curvature</h3>
# <p>Let us first evaluate the term $K_{ij} K^{ij}$:</p>

# In[25]:


Kuu = Ku.up(gam, 1)
trKK = K['_ij']*Kuu['^ij']
print trKK


# <p>Then we compute the symmetric bilinear form $k_{ij} := K_{ik} K^k_{\ \, j}$:</p>

# In[26]:


KK = K['_ik']*Ku['^k_j']
print KK


# <p>We check that this tensor field is symmetric:</p>

# In[27]:


KK1 = KK.symmetrize()
KK == KK1


# <p>Accordingly, we work with the explicitly symmetric version:</p>

# In[28]:


KK = KK1
print KK


# <p> </p>
# <h2>Electric and magnetic parts of the Weyl tensor</h2>
# <p>The electric part is the bilinear form $E$ given by $$ E_{ij} = R_{ij} + K K_{ij} - K_{ik} K^k_{\ \, j} $$</p>

# In[29]:


E = Ric + trK*K - KK
print E


# <p>The magnetic part is the bilinear form $B$ defined by $$ B_{ij} = \epsilon^k_{\ \, l i} D_k K^l_{\ \, j}, $$</p>
# <p>where $\epsilon^k_{\ \, l i}$ are the components of the type-(1,2) tensor $\epsilon^\sharp$, related to the Levi-Civita alternating tensor $\epsilon$ associated with $\gamma$ by $\epsilon^k_{\ \, l i} = \gamma^{km} \epsilon_{m l i}$. In SageManifolds, $\epsilon$ is obtained by the command <span style="font-family: courier new,courier;">volume_form()</span> and $\epsilon^\sharp$ by the command <span style="font-family: courier new,courier;">volume_form(1)</span> (1 = 1 index raised):</p>

# In[30]:


eps = gam.volume_form() 
print eps


# In[31]:


epsu = gam.volume_form(1)
print epsu


# In[32]:


DKu = D(Ku)
B = epsu['^k_li']*DKu['^l_jk'] 
print B


# <p>Let us check that $B$ is symmetric:</p>

# In[33]:


B1 = B.symmetrize()
B == B1


# <p>Accordingly, we set</p>

# In[34]:


B = B1
B.set_name('B')
print B


# <h2>3+1 decomposition of the Simon-Mars tensor</h2>
# <p>We proceed according to the computation presented in <a href="http://arxiv.org/abs/1412.6542">arXiv:1412.6542</a>.</p>
# <p>Tensor $E^\sharp$ of components $E^i_ {\ \, j}$:</p>

# In[35]:


Eu = E.up(gam, 0) 
print Eu


# <p>Tensor $B^\sharp$ of components $B^i_{\ \, j}$:</p>

# In[36]:


Bu = B.up(gam, 0)
print Bu


# <p>1-form $\beta^\flat$ of components $\beta_i$ and its exterior derivative:</p>

# In[37]:


bd = b.down(gam)
xdb = bd.exterior_derivative()
print xdb


# <p>Scalar square of shift $\beta_i \beta^i$:</p>

# In[38]:


b2 = bd(b)
print b2


# <p>Scalar $Y = E(\beta,\beta) = E_{ij} \beta^i \beta^j$:</p>

# In[39]:


Ebb = E(b,b)
Y = Ebb
print Y


# <p>Scalar $\bar Y = B(\beta,\beta) = B_{ij}\beta^i \beta^j$:</p>

# In[40]:


Bbb = B(b,b)
Y_bar = Bbb
print Y_bar


# <p>1-form of components $Eb_i = E_{ij} \beta^j$:</p>

# In[41]:


Eb = E.contract(b)
print Eb


# <p>Vector field of components $Eub^i = E^i_{\ \, j} \beta^j$:</p>

# In[42]:


Eub = Eu.contract(b)
print Eub


# <p>1-form of components $Bb_i = B_{ij} \beta^j$:</p>

# In[43]:


Bb = B.contract(b)
print Bb


# <p>Vector field of components $Bub^i = B^i_{\ \, j} \beta^j$:</p>

# In[44]:


Bub = Bu.contract(b)
print Bub


# <p>Vector field of components $Kub^i = K^i_{\ \, j} \beta^j$:</p>

# In[45]:


Kub = Ku.contract(b)
print Kub


# In[46]:


T = 2*b(N) - 2*K(b,b)
print T ; T.display()


# In[47]:


Db = D(b)  #  Db^i_j = D_j b^i
Dbu = Db.up(gam, 1)  # Dbu^{ij} = D^j b^i
bDb = b*Dbu  # bDb^{ijk} = b^i D^k b^j
T_bar = eps['_ijk']*bDb['^ikj']
print T_bar ; T_bar.display()


# In[48]:


epsb = eps.contract(b) 
print epsb


# In[49]:


epsB = eps['_ijl']*Bu['^l_k']
print epsB


# In[50]:


Z = 2*N*( D(N)  -K.contract(b)) + b.contract(xdb)
print Z


# In[51]:


DNu = D(N).up(gam)
A = 2*(DNu - Ku.contract(b))*b + N*Dbu
Z_bar = eps['_ijk']*A['^kj']
print Z_bar


# In[52]:


W = N*Eb + epsb.contract(Bub)
print W


# In[53]:


W_bar = N*Bb - epsb.contract(Eub)
print W_bar


# In[54]:


M = - 4*Eb(Kub - DNu) - 2*(epsB['_ij.']*Dbu['^ji'])(b)
print M ; M.display()


# In[55]:


M_bar = 2*(eps.contract(Eub))['_ij']*Dbu['^ji'] - 4*Bb(Kub - DNu)
print M_bar ; M_bar.display()


# In[56]:


F = (N^2 - b2)*gam + bd*bd
print F


# In[57]:


A = epsB['_ilk']*b['^l'] + epsB['_ikl']*b['^l'] \
    + Bu['^m_i']*epsb['_mk'] - 2*N*E
xdbE = xdb['_kl']*Eu['^k_i']
L = 2*N*epsB['_kli']*Dbu['^kl'] + 2*xdb['_ij']*Eub['^j'] \
    + 2*xdbE['_li']*b['^l'] + 2*A['_ik']*(Kub - DNu)['^k']
print L


# In[58]:


N2pbb = N^2 + b2
V = N2pbb*E - 2*(b.contract(E)*bd).symmetrize() + Ebb*gam \
    + 2*N*(b.contract(epsB).symmetrize())
print V


# In[59]:


beps = b.contract(eps)
V_bar = N2pbb*B - 2*(b.contract(B)*bd).symmetrize() + Bbb*gam \
        -2*N*(beps['_il']*Eu['^l_j']).symmetrize()
print V_bar


# In[60]:


F = (N^2 - b2)*gam + bd*bd
print F


# In[61]:


R1 = (4*(V*Z - V_bar*Z_bar) + F*L).antisymmetrize(1,2)
print R1


# In[62]:


R2 = 4*(T*V - T_bar*V_bar - W*Z + W_bar*Z_bar) + M*F - N*bd*L
print R2


# In[63]:


R3 = (4*(W*Z - W_bar*Z_bar) + N*bd*L).antisymmetrize()
print R3


# In[64]:


R2[3,1] == -2*R3[3,1]


# In[65]:


R2[3,2] == -2*R3[3,2]


# In[66]:


R4 = 4*(T*W - T_bar*W_bar) -4*(Y*Z - Y_bar*Z_bar) + N*M*bd - b2*L
print R4


# In[67]:


epsE = eps['_ijl']*Eu['^l_k']
print epsE


# In[68]:


A = - epsE['_ilk']*b['^l'] - epsE['_ikl']*b['^l'] \
    - Eu['^m_i']*epsb['_mk'] - 2*N*B
xdbB = xdb['_kl']*Bu['^k_i']
L_bar = - 2*N*epsE['_kli']*Dbu['^kl'] + 2*xdb['_ij']*Bub['^j'] \
        + 2*xdbB['_li']*b['^l'] + 2*A['_ik']*(Kub - DNu)['^k']
print L_bar


# In[69]:


R1_bar = (4*(V*Z_bar + V_bar*Z) + F*L_bar).antisymmetrize(1,2)
print R1_bar


# In[70]:


R2_bar = 4*(T_bar*V + T*V_bar) - 4*(W*Z_bar + W_bar*Z) \
         + M_bar*F - N*bd*L_bar
print R2_bar


# In[71]:


R3_bar = (4*(W*Z_bar + W_bar*Z) + N*bd*L_bar).antisymmetrize()
print R3_bar


# In[72]:


R4_bar = 4*(T_bar*W + T*W_bar - Y*Z_bar - Y_bar*Z) \
         + M_bar*N*bd - b2*L_bar
print R4_bar


# In[73]:


R1u = R1.up(gam)
print R1u


# In[74]:


R2u = R2.up(gam)
print R2u


# In[75]:


R3u = R3.up(gam)
print R3u


# In[76]:


R4u = R4.up(gam)
print R4u


# In[77]:


R1_baru = R1_bar.up(gam)
print R1_baru


# In[78]:


R2_baru = R2_bar.up(gam)
print R2_baru


# In[79]:


R3_baru = R3_bar.up(gam)
print R3_baru


# In[80]:


R4_baru = R4_bar.up(gam)
print R4_baru


# <h3>Simon-Mars scalars</h3>

# In[81]:


S1 = 4*(R1['_ijk']*R1u['^ijk'] - R1_bar['_ijk']*R1_baru['^ijk'] \
     - 2*(R2['_ij']*R2u['^ij'] - R2_bar['_ij']*R2_baru['^ij']) \
     - R3['_ij']*R3u['^ij'] + R3_bar['_ij']*R3_baru['^ij'] \
     + 2*(R4['_i']*R4u['^i'] - R4_bar['_i']*R4_baru['^i']))
print S1


# In[82]:


S1E = S1.expr()


# In[83]:


S2 = 4*(R1['_ijk']*R1_baru['^ijk'] + R1_bar['_ijk']*R1u['^ijk'] \
     - 2*(R2['_ij']*R2_baru['^ij'] + R2_bar['_ij']*R2u['^ij']) \
     - R3['_ij']*R3_baru['^ij'] - R3_bar['_ij']*R3u['^ij'] \
     + 2*(R4['_i']*R4_baru['^i'] + R4_bar['_i']*R4u['^i']))
print S2


# In[84]:


S2E = S2.expr()


# In[85]:


lS1E = log(S1E,10).simplify_full()


# In[86]:


lS2E = log(S2E,10).simplify_full()


# <p>Simon-Mars scalars expressed in terms of the coordinates $X=-1/x,y$:</p>

# In[87]:


var('X')
S1EX = S1E.subs(x=-1/X).simplify_full()
S2EX = S2E.subs(x=-1/X).simplify_full()


# <p>Definition of the ergoregion:</p>

# In[88]:


g00 = - AA2/BB2
g00X = g00.subs(x=-1/X).simplify_full()


# In[89]:


ergXy = implicit_plot(g00X, (X,-1,0), (y,-1,1), plot_points=200, 
                      fill=False, linewidth=1, color='black', 
                      axes_labels=(r"$X\,\left[M^{-1}\right]$", 
                                   r"$y\,\left[M\right]$"), 
                      fontsize=14)


# <p>Due to the very high degree of the polynomials involved in the expression of the Simon-Mars scalars, the floating-point precision of Sage's <span style="font-family: courier new,courier;">contour_plot</span> function (53 bits) is not sufficient. Taking avantage that Sage is <strong>open-source</strong>, we modify the function to allow for an arbitrary precision. First, we define a sampling function with a floating-point precision specified by the user (argument <span style="font-family: courier new,courier;">precis</span>): </p>

# In[90]:


def array_precisXy(fXy, Xmin, Xmax, ymin, ymax, np, precis, tronc):
    RP = RealField(precis)
    Xmin = RP(Xmin)
    Xmax = RP(Xmax)
    ymin = RP(ymin)
    ymax = RP(ymax)
    dX = (Xmax - Xmin) / RP(np-1)
    dy = (ymax - ymin) / RP(np-1)
    resu = []
    for i in range(np):
        list_y = []
        yy = ymin + dy * RP(i)
        fyy = fXy.subs(y=yy)
        for j in range(np):
            XX = Xmin + dX * RP(j)
            fyyXX = fyy.subs(X = XX)
            val = RP(log(abs(fyyXX) + 1e-20, 10))
            if val < -tronc:
                val = -tronc
            elif val > tronc:
                val = tronc
            list_y.append(val)
        resu.append(list_y)
    return resu


# <p>Then we redefine <span style="font-family: courier new,courier;">contour_plot</span> so that it uses the sampling function with a floating-point precision of 200 bits:</p>

# In[91]:


from sage.misc.decorators import options, suboptions

@suboptions('colorbar', orientation='vertical', format=None, 
            spacing=None)
@suboptions('label', fontsize=9, colors='blue', inline=None, 
            inline_spacing=3, fmt="%1.2f")
@options(plot_points=100, fill=True, contours=None, linewidths=None, 
         linestyles=None, labels=False, frame=True, axes=False, 
         colorbar=False, legend_label=None, aspect_ratio=1)
def contour_plot_precisXy(f, xrange, yrange, **options):
    from sage.plot.all import Graphics
    from sage.plot.misc import setup_for_eval_on_grid
    from sage.plot.contour_plot import ContourPlot

    np = options['plot_points']
    precis = 200  # floating-point precision = 200 bits 
    tronc = 10
    xy_data_array = array_precisXy(f, xrange[0], xrange[1], 
                                   yrange[0], yrange[1], np, precis, 
                                   tronc)

    g = Graphics()

    # Reset aspect_ratio to 'automatic' in case scale is 'semilog[xy]'.
    # Otherwise matplotlib complains.
    scale = options.get('scale', None)
    if isinstance(scale, (list, tuple)):
        scale = scale[0]
    if scale == 'semilogy' or scale == 'semilogx':
        options['aspect_ratio'] = 'automatic'

    g._set_extra_kwds(Graphics._extract_kwds_for_show(options, 
                                    ignore=['xmin', 'xmax']))
    g.add_primitive(ContourPlot(xy_data_array, xrange, 
                                yrange, options))
    return g


# <p>Then we are able to draw the contour plot of the two Simon-Mars scalars, in terms of the coordinates $(X,y)$ (Figure 11 of <a href="http://arxiv.org/abs/1412.6542">arXiv:1412.6542</a>):</p>

# In[92]:


c1Xy = contour_plot_precisXy(S1EX, (-1,0), (-1,1), 
                             plot_points=200, 
                             fill=False, cmap='hsv', 
                             linewidths=1, 
                             contours=(-10,-9,-8,-7,-6,-5,-4,-3,
                                       -2,-1,0,1,2,3,4,5,6,7,8), 
                             colorbar=True, 
                             colorbar_spacing='uniform', 
                             colorbar_format='%1.f', 
                             axes_labels=(r"$X\,\left[M^{-1}\right]$", 
                                          r"$y\,\left[M\right]$"), 
                             fontsize=14)


# In[93]:


S1TSXy = c1Xy+ergXy
show(S1TSXy)


# In[94]:


c2Xy = contour_plot_precisXy(S2EX, (-1,0), (-1,1), 
                             plot_points=200, 
                             fill=False, cmap='hsv', 
                             linewidths=1, 
                             contours=(-10,-9,-8,-7,-6,-5,-4,-3,-2,
                                       -1,0,1,2,3,4,5,6,7,8,9,10), 
                             colorbar=True, 
                             colorbar_spacing='uniform', 
                             colorbar_format='%1.f', 
                             axes_labels=(r"$X\,\left[M^{-1}\right]$", 
                                          r"$y\,\left[M\right]$"), 
                             fontsize=14)


# In[95]:


S2TSXy = c2Xy + ergXy
show(S2TSXy)


# <p>Let us do the same in terms of the Weyl-Lewis-Papapetrou cylindrical coordinates $(\rho,z)$, which are related to the prolate spheroidal coordinates $(x,y)$ by $$ \rho = \sqrt{(x^2-1)(1-y^2)}  \quad\mbox{and}\quad z=xy . $$ </p>
# <p>For simplicity, we denote $\rho$ by $r$:</p>

# In[96]:


var('r z')


# In[97]:


S1Erz = S1E.subs(x=1/2*(sqrt(r^2+(z+1)^2)+sqrt(r^2+(z-1)^2)), 
                 y=1/2*(sqrt(r^2+(z+1)^2)-sqrt(r^2+(z-1)^2)))
S1Erz = S1Erz.simplify_full()


# In[98]:


S2Erz = S2E.subs(x=1/2*(sqrt(r^2+(z+1)^2)+sqrt(r^2+(z-1)^2)), 
                 y=1/2*(sqrt(r^2+(z+1)^2)-sqrt(r^2+(z-1)^2)))
S2Erz = S2Erz.simplify_full()


# In[99]:


def tab_precis(fz, zz, rmin, rmax, np, precis, tronc):
    RP = RealField(precis)
    rmin = RP(rmin)
    rmax = RP(rmax)
    zz = RP(zz)
    dr = (rmax - rmin) / RP(np-1)
    resu = []
    fzz = fz.subs(z=zz)
    for i in range(np):
        rr = rmin + dr * RP(i)
        val = RP(log(abs(fzz.subs(r = rr)), 10))
        if val < -tronc:
            val = -tronc
        elif val > tronc:
            val = tronc
        resu.append((rr, zz, val))
    return resu


# ### 3D plots

# We also a viewer for 3D plots (use `'threejs'` or `'jmol'` for interactive 3D graphics):

# In[100]:


viewer3D = 'threejs' # must be 'threejs', jmol', 'tachyon' or None (default)


# In[101]:


gg = Graphics()
rmin = 0.1
rmax = 3
zmin = -2
zmax = 2
npr = 200
npz = npr
precis = 200 # 200-bits floating-point precision
tronc = 5
dz = (zmax-zmin) / (npz-1)
for i in range(npz):
    zz = zmin + i*dz
    gg += line3d(tab_precis(S1Erz, zz, rmin, rmax, 
                            npr, precis, tronc))
show(gg, viewer=viewer3D, axes_labels=['rho', 'z', 'S_1'], 
     aspect_ratio=[1,1,0.3])


# In[102]:


gg2 = Graphics()
for i in range(npz):
    zz = zmin + i*dz
    gg2 += line3d(tab_precis(S2Erz, zz, rmin, rmax, 
                             npr, precis, tronc))
show(gg2, viewer=viewer3D, axes_labels=['rho', 'z', 'S_2'], 
     aspect_ratio=[1,1,0.3])


# ### 2D contour plots
# 
# Contour plots of the two Simon-Mars scalar fields in terms of coordinates $(\rho,z)$ (Figure 12 of <a href="http://arxiv.org/abs/1412.6542">arXiv:1412.6542</a>)

# In[103]:


def array_precis(frz, rmin, rmax, zmin, zmax, np, precis,
                 tronc):
    RP = RealField(precis)
    rmin = RP(rmin)
    rmax = RP(rmax)
    zmin = RP(zmin)
    zmax = RP(zmax)
    dr = (rmax - rmin) / RP(np-1)
    dz = (zmax - zmin) / RP(np-1)
    resu = []
    for i in range(np):
        list_z = []
        zz = zmin + dz * RP(i)
        fzz = frz.subs(z=zz)
        for j in range(np):
            rr = rmin + dr * RP(j)
            fzzrr = fzz.subs(r = rr)
            val = RP(log(abs(fzzrr) + 1e-20, 10))
            if val < -tronc:
                val = -tronc
            elif val > tronc:
                val = tronc
            list_z.append(val)
        resu.append(list_z)
    return resu


# In[104]:


rmin = 0.1
rmax = 3
zmin = -2
zmax = 2
npr = 10
npz = npr
precis = 100
tronc = 5
val = array_precis(S1Erz, rmin, rmax, zmin, zmax, npr, 
                   precis, tronc)


# In[105]:


from sage.misc.decorators import options, suboptions

@suboptions('colorbar', orientation='vertical', format=None, 
            spacing=None)
@suboptions('label', fontsize=9, colors='blue', inline=None, 
            inline_spacing=3, fmt="%1.2f")
@options(plot_points=100, fill=True, contours=None, 
         linewidths=None, linestyles=None, labels=False, 
         frame=True, axes=False, colorbar=False, 
         legend_label=None, aspect_ratio=1)
def contour_plot_precis(f, xrange, yrange, **options):
    from sage.plot.all import Graphics
    from sage.plot.misc import setup_for_eval_on_grid
    from sage.plot.contour_plot import ContourPlot

    np = options['plot_points']
    precis = 200
    tronc = 10
    xy_data_array = array_precis(f, xrange[0], xrange[1], 
                                 yrange[0], yrange[1], np, 
                                 precis, tronc)

    g = Graphics()

    # Reset aspect_ratio to 'automatic' in case scale is 'semilog[xy]'.
    # Otherwise matplotlib complains.
    scale = options.get('scale', None)
    if isinstance(scale, (list, tuple)):
        scale = scale[0]
    if scale == 'semilogy' or scale == 'semilogx':
        options['aspect_ratio'] = 'automatic'

    g._set_extra_kwds(Graphics._extract_kwds_for_show(options, 
                                     ignore=['xmin', 'xmax']))
    g.add_primitive(ContourPlot(xy_data_array, xrange, yrange, 
                                options))
    return g


# In[106]:


c1rz = contour_plot_precis(S1Erz, (0.0001,10), (-5,5.001), 
                           plot_points=300, fill=False, 
                           cmap='hsv', linewidths=1, 
                           contours=(-10,-9,-8,-7,-6,-5.5,-5,-4.5,
                                     -4,-3.5,-3,-2.5,-2,-1.5,-1,
                                     -0.5,0,0.5,1,1.5,2,2.5,3,3.5,
                                     4,4.5,5), 
                           colorbar=True, 
                           colorbar_spacing='uniform', 
                           colorbar_format='%1.f', 
                           axes_labels=(r"$\rho\,\left[M\right]$", 
                                        r"$z\,\left[M\right]$"), 
                           fontsize=14)


# In[107]:


g00rz = g00.subs(x=1/2*(sqrt(r^2+(z+1)^2)+sqrt(r^2+(z-1)^2)), 
                 y=1/2*(sqrt(r^2+(z+1)^2)
                        -sqrt(r^2+(z-1)^2))).simplify_full()
c2 = implicit_plot(g00rz, (r,0.0001,10), (z,-5,5.001), 
                   plot_points=200, fill=False, 
                   linewidth=1, color='black', 
                   axes_labels=(r"$\rho\,\left[M\right]$", 
                                r"$z\,\left[M\right]$"), 
                   fontsize=14)
S1TSrz = c1rz+c2
show(S1TSrz)


# In[108]:


c2rz = contour_plot_precis(S2Erz, (0.0001,10), (-5,5.001), 
                           plot_points=300, fill=False, 
                           cmap='hsv', linewidths=1, 
                           contours=(-10,-9,-8,-7,-6,-5.5,-5,-4.5,
                                     -4,-3.5,-3,-2.5,-2,-1.5,-1,
                                     -0.5,0,0.5,1,1.5,2,2.5,3,3.5,
                                     4,4.5,5), 
                           colorbar=True, 
                           colorbar_spacing='uniform', 
                           colorbar_format='%1.f', 
                           axes_labels=(r"$\rho\,\left[M\right]$", 
                                        r"$z\,\left[M\right]$"), 
                           fontsize=14)
S2TSrz = c2rz+c2
show(S2TSrz)