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

# # Tolman-Oppenheimer-Volkoff equations
# 
# This Jupyter/SageMath worksheet is relative to the lectures
# [General relativity computations with SageManifolds](https://indico.cern.ch/event/505595/) given at the NewCompStar School 2016 (Coimbra, Portugal).
#   
# These computations are based on [SageManifolds](http://sagemanifolds.obspm.fr) (curently version 1.1, as included in SageMath 8.1).
# 
# Click [here](https://raw.githubusercontent.com/sagemanifolds/SageManifolds/master/Worksheets/v1.1/SM_TOV.ipynb) to download the worksheet file (ipynb format). To run it, you must start SageMath with the Jupyter notebook, with the command `sage -n jupyter`

# This worksheet is divided in two parts:
# 
# 1. Deriving the TOV system from the Einstein equation
# 2. Solving the TOV system to get stellar models
# 

# *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')


# ## 1. Deriving the TOV system from the Einstein equation
# 
# ### Spacetime
# 
# We declare the spacetime manifold $M$:

# In[3]:


M = Manifold(4, 'M')
print(M)


# To get some information about the object `M`, we use the question mark:

# In[4]:


get_ipython().run_line_magic('pinfo', 'M')


# Using a double question mark, we get the Python source code (SageMath is **open source**, isn't it?):

# In[5]:


get_ipython().run_line_magic('pinfo2', 'M')


# We declare the chart of spherical coordinates $(t,r,\theta,\phi)$, via the method `chart` acting on `M`; to see how to use it, we again use the question mark:

# In[6]:


get_ipython().run_line_magic('pinfo', 'M.chart')


# In[7]:


X.<t,r,th,ph> = M.chart(r't r:(0,+oo) th:(0,pi):\theta ph:(0,2*pi):\phi')
X


# ### Metric tensor
# 
# The static and spherically symmetric metric ansatz, with the unknown functions $\nu(r)$ and $m(r)$:

# In[8]:


g = M.lorentzian_metric('g')
nu = function('nu')
m = function('m')
g[0,0] = -exp(2*nu(r))
g[1,1] = 1/(1-2*m(r)/r)
g[2,2] = r^2
g[3,3] = (r*sin(th))^2
g.display()


# One can display the metric components as a list:

# In[9]:


g.display_comp()


# By default, only the nonzero components are shown; to get all the components, set the option `only_nonzero` to `False`:

# In[10]:


g.display_comp(only_nonzero=False)


# We can also display the metric components as a matrix, via the `[]` operator:

# In[11]:


g[:]


# The `[]` operator can also be used to access to individual elements:

# In[12]:


g[0,0]


# ## Einstein equation
# 
# Let us start by evaluating the Ricci tensor of $g$:

# In[13]:


Ric = g.ricci()
print(Ric)


# In[14]:


Ric.display()


# The Ricci scalar is naturally obtained by the method `ricci_scalar()`:

# In[15]:


get_ipython().run_line_magic('pinfo', 'g.ricci_scalar')


# In[16]:


g.ricci_scalar().display()


# As a check we can also compute it by taking the trace of the Ricci tensor with respect to $g$:

# In[17]:


g.ricci_scalar() == g.inverse()['^{ab}']*Ric['_{ab}']


# The Einstein tensor
# $$ G_{ab} := R_{ab} - \frac{1}{2} R \, g_{ab},$$
# or in index-free notation:
# $$ G := \mathrm{Ric}(g) - \frac{1}{2} r(g) \, g $$

# In[18]:


G = Ric - 1/2*g.ricci_scalar() * g
G.set_name('G')
print(G)


# In[19]:


G.display()


# ### The energy-momentum tensor
# 
# We consider a perfect fluid matter model. 
# Let us first defined the fluid 4-velocity $u$: 

# In[20]:


u = M.vector_field('u')
u[0] = exp(-nu(r))
u.display()


# In[21]:


u[:]


# In[22]:


print(u.parent())


# Let us check that $u$ is a normalized timelike vector, i.e. that $g_{ab} u^a u^b = -1$, or, in index-free notation, $g(u,u)=-1$:

# In[23]:


g(u,u)


# In[24]:


print(g(u,u))


# In[25]:


g(u,u).display()


# In[26]:


g(u,u).parent()


# In[27]:


print(g(u,u).parent())


# In[28]:


g(u,u) == -1


# To form the energy-momentum tensor, we need the 1-form $\underline{u}$ that is metric-dual to the vector $u$, i.e. 
# $u_a = g_{ab} u^b$:

# In[29]:


u_form = u.down(g)
print(u_form)


# In[30]:


u_form.display()


# The energy-momentum tensor is then
# $$ T_{ab} = (\rho + p) u_a u_b + p \, g_{ab},$$
# or in index-free notation:
# $$ T = (\rho + p) \underline{u}\otimes\underline{u} + p \, g$$

# Since the tensor product $\otimes$ is taken with the `*` operator, we write:

# In[31]:


rho = function('rho')
p = function('p')
T = (rho(r)+p(r))* (u_form * u_form) + p(r) * g
T.set_name('T')
print(T)


# In[32]:


T.display()


# In[33]:


T(u,u)


# In[34]:


print(T(u,u))


# In[35]:


T(u,u).display()


# ### The Einstein equation
# 
# The Einstein equation is 
# $$ G = 8\pi T$$
# We rewrite it as $E = 0$ with $E := G - 8\pi T$:

# In[36]:


E = G - 8*pi*T
E.set_name('E')
print(E)


# In[37]:


E.display()


# In[38]:


E.display_comp()


# In[39]:


E[0,0]


# In[40]:


EE0_sol = solve(E[0,0].expr()==0, diff(m(r),r))
EE0_sol


# In[41]:


EE0 = EE0_sol[0]
EE0


# In[42]:


E[1,1]


# In[43]:


EE1_sol = solve(E[1,1].expr()==0, diff(nu(r),r))
EE1 = EE1_sol[0]
EE1


# In[44]:


E[3,3] == E[2,2]*sin(th)^2


# In[45]:


E[2,2]


# ### The energy-momentum conservation equation
# The energy-momentum tensor must obey
# $$ \nabla_b T^b_{\ \, a} = 0$$
# We first form the tensor $T^b_{\ \, a}$ by raising the first index of $T_{ab}$:

# In[46]:


Tu = T.up(g, 0)
print(Tu)


# We get the Levi-Civita connection $\nabla$ associated with the metric $g$:

# In[47]:


nabla = g.connection()
print(nabla)


# In[48]:


nabla.display()


# We apply `nabla` to `Tu` to get the tensor $(\nabla T)^b_{\ \, ac} = \nabla_c T^b_{\ \, a}$ (MTW index convention):

# In[49]:


dTu = nabla(Tu)
print(dTu)


# The divergence $\nabla_b T^b_{\ \, a}$ is then computed as the trace of the tensor $(\nabla T)^b_{\ \, ac}$ on the first index ($b$, position `0`) and last index ($c$, position `2`):

# In[50]:


divT = dTu.trace(0,2)
print(divT)


# We can also take the trace by using the index notation:

# In[51]:


divT == dTu['^b_{ab}']


# In[52]:


print(divT)


# In[53]:


divT.display()


# In[54]:


divT[:]


# The only non trivially vanishing components is thus

# In[55]:


divT[1]


# Hence the energy-momentum conservation equation $\nabla_b T^b_{\ \, a}=0$ reduces to

# In[56]:


EE2_sol = solve(divT[1].expr()==0, diff(p(r),r))
EE2 = EE2_sol[0]
EE2


# ### The TOV system
# 
# Let us collect all the independent equations obtained so far:

# In[57]:


for eq in [EE0, EE1, EE2]:
    show(eq)


# ## 2. Solving the TOV system
# 
# In order to solve the TOV system, we need to specify a fluid equation of state. For simplicity, we select a polytropic one:
# $$ p = k \rho^2$$

# In[58]:


var('k', domain='real')
p_eos(r) = k*rho(r)^2
p_eos(r)


# We substitute this expression for $p$ in the TOV equations:

# In[59]:


EE1_rho = EE1.substitute_function(p, p_eos)
EE1_rho


# In[60]:


EE2_rho = EE2.substitute_function(p, p_eos)
EE2_rho


# In[61]:


EE2_rho = (EE2_rho / (2*k*rho(r))).simplify_full()
EE2_rho


# We subsitute the expression of equation `EE1_rho` for $\partial \nu/\partial r$, in order to get rid of any derivative in the right-hand sides:

# In[62]:


EE2_rho = EE2_rho.subs({diff(nu(r),r): EE1_rho.rhs()}).simplify_full()
EE2_rho


# The system to solve for $(m(r), \nu(r), \rho(r))$ is thus

# In[63]:


for eq in [EE0, EE1_rho, EE2_rho]:
    show(eq)


# ### Numerical resolution
# 
# Let us use a standard 4th-order Runge-Kutta method:

# In[64]:


get_ipython().run_line_magic('pinfo', 'desolve_system_rk4')


# We gather all equations in a list for the ease of manipulation:

# In[65]:


eqs = [EE0, EE1_rho, EE2_rho]


# To get a numerical solution, we have of course to specify some numerical value for the EOS constant $k$; let us choose $k=1/4$:

# In[66]:


k0 = 1/4
rhs = [eq.rhs().subs(k=k0) for eq in eqs]
rhs


# The integration for $m(r)$ and $\rho(r)$ has to stop as soon as $\rho(r)$ become negative. An easy way to ensure this is to use the Heaviside function (actually SageMath's `unit_step`; for some reason, `heaviside` does not work for the RK4 numerical resolution):

# In[67]:


plot(unit_step(x), (x,-4,4), thickness=2, aspect_ratio=1)


# and to multiply the r.h.s. of the $dm/dr$ and $d\rho/dr$ equations by $h(\rho)$:

# In[68]:


rhs[0] = rhs[0] * unit_step(rho(r))
rhs[2] = rhs[2] * unit_step(rho(r))
rhs


# Let us add an extra equation, for the purpose of getting the star's radius as an output of the integration, via the equation $dR/dr = 1$, again multiplied by the Heaviside function of $\rho$:

# In[69]:


rhs.append(1 * unit_step(rho(r)))
rhs


# For the purpose of the numerical integration via `desolve_system_rk4`, we have to replace the symbolic functions $m(r)$, $\nu(r)$ and $\rho(r)$ by some symbolic variables, $m_1$, $\nu_1$ and
# $\rho_1$, say:

# In[70]:


var('m_1 nu_1 rho_1 r_1')
rhs = [y.subs({m(r): m_1, nu(r): nu_1, rho(r): rho_1}) for y in rhs]
rhs


# The integration parameters:

# In[71]:


rho_c = 1
r_min = 1e-8
r_max = 1
np = 200
delta_r = (r_max - r_min) / (np-1)


# The numerical resolution, with the initial conditions 
# $$(r_{\rm min}, m(r_{\rm min}), \nu(r_{\rm min}), \rho(r_{\rm min}), R(r_{\rm min})) = (r_{\rm min},0,0,\rho_{\rm c},r_{\rm min}) $$
# set in the parameter `ics`:

# In[72]:


sol = desolve_system_rk4(rhs, vars=(m_1, nu_1, rho_1, r_1), ivar=r, 
                         ics=[r_min, 0, 0, rho_c, r_min], 
                         end_points=(r_min, r_max), step=delta_r)


# The solution is returned as a list, the first 10 elements of which being:

# In[73]:


sol[:10]


# Each element is of the form $[r, m(r), \nu(r), \rho(r), R(r)]$. So to get the list of $(r, \rho(r))$ values, we write

# In[74]:


rho_sol = [(s[0], s[3]) for s in sol]
rho_sol[:10]


# We may then use this list to have some plot of $\rho(r)$, thanks to the function `line`, which transforms a list into a graphical object:

# In[75]:


graph = line(rho_sol, axes_labels=[r'$r$', r'$\rho$'], gridlines=True)
graph


# The solution for $m(r)$:

# In[76]:


m_sol = [(s[0], s[1]) for s in sol]
m_sol[:10]


# In[77]:


graph = line(m_sol, axes_labels=[r'$r$', r'$m$'], gridlines=True)
graph


# The solution for $\nu(r)$ (has to be rescaled by adding a constant to ensure $\nu(+\infty) = 1$):

# In[78]:


nu_sol = [(s[0], s[2]) for s in sol]
nu_sol[:10]


# In[79]:


graph = line(nu_sol, axes_labels=[r'$r$', r'$\tilde{\nu}$'], gridlines=True)
graph


# The solution for $R(r)$:

# In[80]:


r_sol = [(s[0], s[4]) for s in sol]
line(r_sol)


# The total gravitational mass of the star is obtained via the last element (index: `-1`) of the list of $(r,m(r))$ values:

# In[81]:


M_grav = m_sol[-1][1]
M_grav


# Similarly, the stellar radius is obtained through the last element of the list of $(r,R(r))$ values:

# In[82]:


R = r_sol[-1][1]
R


# The star's compactness:

# In[83]:


M_grav/R


# ### Sequence of stellar models
# 
# Let us perform a loop on the central density.
# First we set up a list of values for $\rho_c$:

# In[84]:


rho_c_min = 0.01
rho_c_max = 3
n_conf = 40
rho_c_list = [rho_c_min + i * (rho_c_max-rho_c_min)/(n_conf-1) for i in range(n_conf)]
rho_c_list


# The loop:

# In[85]:


M_list = list()
R_list = list()
for rho_c in rho_c_list:
    sol = desolve_system_rk4(rhs, vars=(m_1, nu_1, rho_1, r_1), ivar=r, 
                             ics=[r_min, 0, 0, rho_c, r_min], 
                             end_points=(r_min, r_max), step=delta_r)
    M_list.append( sol[-1][1] )
    R_list.append( sol[-1][4] )


# The mass along the sequence:

# In[86]:


M_list


# The radius along the sequence:

# In[87]:


R_list


# To draw $M$ as a function of $\rho_{\rm c}$, we use the Python function `zip` to construct a list of $(\rho_{\rm c}, M)$ values:

# In[88]:


zip(rho_c_list, M_list)


# In[89]:


graph = line(zip(rho_c_list, M_list), axes_labels=[r'$\rho_c$', r'$M$'], gridlines=True)
graph


# Similarly, we draw $M$ as a function of $R$:

# In[90]:


graph = line(zip(R_list, M_list), axes_labels=[r'$R$', r'$M$'], gridlines=True)
graph


# and we save the plot in a pdf file to use it in our next publication ;-)

# In[91]:


graph.save('plot_M_R.pdf')