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

# # Friedmann equations
# 
# This worksheet demonstrates a few capabilities of SageMath in computations regarding cosmological spacetimes with Friedmann-Lemaître-Robertson-Walker (FLRW) metrics.
# The corresponding tools have been developed within the  [SageManifolds](http://sagemanifolds.obspm.fr) project (version 1.2, as included in SageMath 8.2).
# 
# Click [here](https://raw.githubusercontent.com/sagemanifolds/SageManifolds/master/Worksheets/v1.2/SM_Friedmann_equations.ipynb) to download the worksheet file (ipynb format). To run it, you must start SageMath within the Jupyter notebook, via the command `sage -n jupyter`

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

# In[1]:


version()


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

# In[2]:


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


# We declare the spacetime M as a 4-dimensional Lorentzian manifold:

# In[3]:


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


# We introduce the standard FLRW coordinates, via the method `chart()`, the argument of which is a string expressing the coordinates names, their ranges (the default is $(-\infty,+\infty)$) and their LaTeX symbols:

# In[4]:


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


# <p>Assuming that the speed of light c=1, let us define a few variables: Newton's constant $G$, the cosmological constant $\Lambda$, the spatial curvature constant $k$, the scale factor $a(t)$, the fluid proper density $\rho(t)$ and the fluid pressure $p(t)$:</p>

# In[5]:


var('G, Lambda, k', domain='real')
a = M.scalar_field(function('a')(t), name='a')
rho = M.scalar_field(function('rho')(t), name='rho')
p = M.scalar_field(function('p')(t), name='p')


# The FLRW metric is defined by its components in the manifold's default frame, i.e. the frame associated with the FLRW coordinates:

# In[6]:


g = M.metric()
g[0,0] = -1
g[1,1] = a*a/(1 - k*r^2)
g[2,2] = a*a*r^2
g[3,3] = a*a*(r*sin(th))^2
g.display()


# <p>A matrix view of the metric components:</p>

# In[7]:


g[:]


# <p>The Levi-Civita connection associated with the metric is computed:</p>

# In[8]:


nabla = g.connection()
g.christoffel_symbols_display()


# Ricci tensor:

# In[9]:


Ricci = nabla.ricci()
Ricci.display()


# In[10]:


Ricci.display_comp()


# <p>Ricci scalar ($R^\mu_{\ \, \mu}$):</p>

# In[11]:


Ricci_scalar = g.ricci_scalar()
Ricci_scalar.display()


# <p>The fluid 4-velocity:</p>

# In[12]:


u = M.vector_field('u')
u[0] = 1
u.display()


# In[13]:


g(u,u).expr()


# `u.dot(u)` is equivalent to `g(u,u)`:

# In[14]:


u.dot(u).expr()


# <p>Perfect fluid energy-momentum tensor $T$:</p>

# In[15]:


u_form = u.down(g) # the 1-form associated to u by metric duality
T = (rho+p)*(u_form*u_form) + p*g
T.set_name('T')
print(T)
T.display()


# <p>The trace of $T$ (we use index notation to denote the double contraction $g^{ab} T_{ab}$):</p>

# In[16]:


Ttrace = g.inverse()['^ab']*T['_ab']
Ttrace.display()


# <p>Einstein equation: $R_{\mu \nu} - {1 \over 2} R g_{\mu \nu} + \Lambda g_{\mu \nu} = {8 \pi G} T_{\mu \nu}$</p>

# In[17]:


E1 = Ricci - Ricci_scalar/2*g + Lambda*g - (8*pi*G)*T
print("First Friedmann equation:\n")
E1[0,0].expr().expand() == 0


# <p>Trace-reversed version of the Einstein equation: $R_{\mu \nu} - \Lambda g_{\mu \nu} = {8 \pi G} \left(T_{\mu \nu} - {1 \over 2}T\,g_{\mu \nu}\right)$</p>

# In[18]:


E2 = Ricci - Lambda*g - (8*pi*G)*(T - Ttrace/2*g)
print("Second Friedmann equation:\n")
E2[0,0].expr().expand() == 0


# In[ ]: