Riemann_Calculations.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b944d7e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np #We import all our libraries\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "from itertools import product\n",
    "from scipy import optimize # for fits\n",
    "\n",
    "from sympy import latex, symbols, sin, cos, pi, simplify, lambdify, Matrix\n",
    "\n",
    "%matplotlib inline\n",
    "rng = np.random.RandomState(0)\n",
    "from IPython.display import display, Math\n",
    "import torch\n",
    "import torch.autograd.functional as functional"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "12871fae",
   "metadata": {},
   "outputs": [],
   "source": [
    "def fgrad(inputs): #This function allow us to create the hessian by calling it later.\n",
    "    return torch.autograd.functional.jacobian(metric_tensor,inputs, create_graph=True)\n",
    "def christoffelderiv(coord1, coord2, deriv, alpha, beta): #It's the part of the Riemann tensor with the derivative of the Christoffel.\n",
    "    christoffelderiv = 0.5*((jacobinv_l[alpha][beta][deriv] * (jacob_l[alpha][coord1][coord2] + jacob_l[alpha][coord2][coord1] - jacob_l[coord1][coord2][alpha]))\\\n",
    "                                    + g_inv_l[alpha][beta] * (hessian_l[alpha][coord1][coord2][deriv] + hessian_l[alpha][coord2][coord1][deriv] - hessian_l[coord1][coord2][alpha][deriv]))\n",
    "    return christoffelderiv \n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "756453c5",
   "metadata": {},
   "source": [
    "## Metric tensor: 2D unit sphere\n",
    "\n",
    "Coordinates $x^\\alpha = (\\theta, \\phi)$.\n",
    "\n",
    "\\begin{equation}\n",
    "g_{\\mu\\nu} = \n",
    "\\begin{pmatrix}\n",
    "1 & 0 \\\\\n",
    "0 & \\sin^2 \\theta \\\\\n",
    "\\end{pmatrix}\n",
    "\\end{equation}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "2c96fb01",
   "metadata": {},
   "outputs": [],
   "source": [
    "#We define a metric tensor. \n",
    "dim=2 #Tell the dimension.\n",
    "\n",
    "def metric_tensor(coordinate,*arg): #This function create the metric tensor, we need to put in a fonction to be able to\n",
    "                                    #calculate its derivatives with our jacobian later.\n",
    "    \"\"\"\n",
    "    Here, we have as arg for our function : \n",
    "    coordinates : (theta, phi) of a point \n",
    "    \"\"\"\n",
    "    theta=coordinate[0] \n",
    "    phi=coordinate[1]\n",
    "    tensor = torch.zeros(dim,dim)\n",
    "    tensor[0,0] = 1\n",
    "    tensor[1,1] = torch.pow(torch.sin(theta),2)\n",
    "    return tensor\n",
    "\n",
    "def inverse(coordinate, *arg): #Gives us the inverse of the metric_tensor.\n",
    "    theta=coordinate[0]\n",
    "    phi=coordinate[1]\n",
    "    tensor = torch.zeros(dim,dim)\n",
    "    tensor[0,0] = 1\n",
    "    tensor[1,1] = torch.pow(torch.sin(theta),2)\n",
    "    tensor = torch.inverse(tensor)\n",
    "    return tensor\n",
    "def christoffelproduct_liste(coord1, coord2, coord3): #This function gives us the product of two christoffels.\n",
    "        christoffelproduct = christoffel_l[beta][coord1][eta]*christoffel_l[eta][coord2][coord3]\n",
    "        return christoffelproduct"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "4efb8f99",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "theta,phi coordinates:  tensor([ 3., 15.])\n"
     ]
    }
   ],
   "source": [
    "#Let's set the initial coordinate to start with.  \n",
    "coord = torch.tensor([3., 15.],requires_grad=True)\n",
    "print(\"theta,phi coordinates: \",coord.data)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e2ddba31",
   "metadata": {},
   "source": [
    "## Metric tensor: 4D flat space, spherical coordinates\n",
    "\n",
    "Coordinates $x^\\alpha = (ct, R, \\theta, \\phi)$.\n",
    "\n",
    "\\begin{equation}\n",
    "g_{\\mu\\nu} = \n",
    "\\begin{pmatrix}\n",
    "1 & 0 & 0 & 0  \\\\\n",
    "0 & -1 & 0 & 0 \\\\\n",
    "0 & 0 & -R^2 & 0 \\\\\n",
    "0 & 0 & 0 & -R^2 \\sin^2 \\theta \\\\\n",
    "\\end{pmatrix}\n",
    "\\end{equation}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8058fda6",
   "metadata": {},
   "outputs": [],
   "source": [
    "#We define a metric tensor. \n",
    "dim=4 #Tell the dimension.\n",
    "\n",
    "def metric_tensor(coordinate,*arg): \n",
    "                                    #This function create the metric tensor, we need to put in a fonction to be able to\n",
    "                                    #calculate its derivatives with our jacobian later.\n",
    "    \"\"\"\n",
    "    Args: \n",
    "        coordinates (ct,r,theta,phi) of a point\n",
    "        arg: not used\n",
    "    \"\"\"\n",
    "    t=coordinate[0] # c=1\n",
    "    r=coordinate[1]\n",
    "    theta=coordinate[2]\n",
    "    phi=coordinate[3]\n",
    "    tensor = torch.zeros(dim,dim)\n",
    "    tensor[0,0] = 1\n",
    "    tensor[1,1] = -1\n",
    "    tensor[2,2] = -1 * torch.pow(r,2)\n",
    "    tensor[3,3] = -1 * torch.pow(r*torch.sin(theta),2)\n",
    "    return tensor\n",
    "\n",
    "def inverse(coordinate, *arg): #And the inverse of the metric tensor. \n",
    "    tensor = metric_tensor(coordinate, *arg)\n",
    "    tensor = torch.inverse(tensor)\n",
    "    return tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "66df97d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Let's set the initial coordinate to start with.  \n",
    "coord = torch.tensor([0, 10, 30, 0.5],requires_grad=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3ad0085",
   "metadata": {},
   "source": [
    "## Metric tensor: Schwarzschid metric\n",
    "\n",
    "Coordinates $x^\\alpha = (ct, R, \\theta, \\phi)$.\n",
    "\n",
    "$$g_{\\mu\\nu} = \n",
    "\\begin{pmatrix}\n",
    "1 - \\frac{2GM}{Rc^2} & 0 & 0 & 0 \\\\\n",
    "0 & -\\left(1 - \\frac{2GM}{Rc^2}\\right)^{-1} & 0 & 0 \\\\\n",
    "0 & 0 & -R^2 & 0 \\\\\n",
    "0 & 0 & 0 & -R^2 \\sin^2 \\theta \\\\\n",
    "\\end{pmatrix}\n",
    "= \n",
    "\\begin{pmatrix}\n",
    "1 - \\frac{R_s}{R} & 0 & 0 & 0 \\\\\n",
    "0 & -\\left(1 - \\frac{R_s}{R}\\right)^{-1} & 0 & 0 \\\\\n",
    "0 & 0 & -R^2 & 0 \\\\\n",
    "0 & 0 & 0 & -R^2 \\sin^2 \\theta \\\\\n",
    "\\end{pmatrix}\n",
    "$$\n",
    "\n",
    "With $R_s = \\frac{2 G M}{c^2}$, Schwarzschild radius."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "0cc9413f",
   "metadata": {},
   "outputs": [],
   "source": [
    "#We define a metric tensor. \n",
    "dim=4 #Tell the dimension.\n",
    "Rs = 1 #Here, we have arguments. Rs is the one argument of this metric tensor.\n",
    "def metric_tensor(coordinate,*arg):#This function create the metric tensor, we need to put in a fonction to be able to\n",
    "                                    #calculate its derivatives with our jacobian later.\n",
    "    \"\"\"\n",
    "    Args: \n",
    "        coordinates (ct,r,theta,phi) of a point\n",
    "        arg[0]: Rs, Schwarzschid radius\n",
    "    \"\"\"\n",
    "    t=coordinate[0] # c=1\n",
    "    r=coordinate[1]\n",
    "    theta=coordinate[2]\n",
    "    phi=coordinate[3]\n",
    "    #if (len(arg)>0):\n",
    "    #    Rs = arg[0] # Schwarzschid radius\n",
    "    tensor = torch.zeros(dim,dim)\n",
    "    tensor[0,0] = 1 - Rs/r\n",
    "    tensor[1,1] = -1 * torch.pow(1 - Rs/r,-1)\n",
    "    tensor[2,2] = -1 * torch.pow(r,2)\n",
    "    tensor[3,3] = -1 * torch.pow(r*torch.sin(theta),2)\n",
    "    return tensor\n",
    "\n",
    "def inverse(coordinate, *arg):\n",
    "    tensor = metric_tensor(coordinate, *arg)\n",
    "    tensor = torch.inverse(tensor)\n",
    "    return tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "b89c7b5e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "t,r,theta,phi coordinates:  tensor([ 1.0000, 10.0000,  0.5000,  0.2000])\n"
     ]
    }
   ],
   "source": [
    "coord = torch.tensor([1,10,0.5,0.2],requires_grad=True)\n",
    "print(\"t,r,theta,phi coordinates: \",coord.data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "d40b8c32",
   "metadata": {},
   "outputs": [],
   "source": [
    "#here, we calculate all that is needed to do the further calculations such as the Christofell and the jacobian.\n",
    "\n",
    "g = metric_tensor(coord)\n",
    "g_l = g.tolist() #We also convert it into a list.\n",
    "g_inv = torch.inverse(g)\n",
    "g_inv_l = g_inv.tolist()\n",
    "#We do all the preliminary calculations : \n",
    "jacob = functional.jacobian(metric_tensor,(coord)) #We take the jacobian of our metric tensor at the desirated coordinate\n",
    "jacob_l = jacob.tolist()\n",
    "jacobinv = functional.jacobian(inverse, (coord)) #And we take alsothe jacobian of the inverse metric tensor as we will need ti.\n",
    "jacobinv_l = jacobinv.tolist()\n",
    "hessian = torch.autograd.functional.jacobian(fgrad,(coord), create_graph=True) #fgrad here is the jacobian put in a function. \n",
    "#By doing the jacobian of the jacobian of the metric tensor at a given coordinate, one gets the hessian at those coordinates.\n",
    "hessian_l = hessian.tolist()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "022f1dbd",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "What do you want to display ? Choose a number : \n",
      " 1 : Christoffel_symbols \n",
      " 2 : Riemann_tensor \n",
      " 3 : Ricci_tensor \n",
      " 4 : Scalar_curvature \n",
      "1\n"
     ]
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\Gamma^{0}_{0 0} = 0.0$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\Gamma^{0}_{0 1} = 0.0$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\Gamma^{0}_{1 0} = 0.0$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\Gamma^{0}_{1 1} = 0.139707744121552$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\Gamma^{1}_{0 0} = 0.0$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\Gamma^{1}_{0 1} = -7.01525354385376$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\Gamma^{1}_{1 0} = -7.01525354385376$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\Gamma^{1}_{1 1} = 0.0$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Do you want to display an other value ? y/n\n",
      "n\n"
     ]
    }
   ],
   "source": [
    "def Relat_calc():\n",
    "    def christoffelproduct_liste(coord1, coord2, coord3):\n",
    "        christoffelproduct = christoffel_l[beta][coord1][eta]*christoffel_l[eta][coord2][coord3]\n",
    "        return christoffelproduct\n",
    "    # Initialize Christoffel\n",
    "    christoffel = torch.zeros(dim,dim,dim)\n",
    "    for beta in range(0,dim):\n",
    "        for mu in range(0,dim):\n",
    "            for nu in range(0,dim): #Here the loop are done in respect the einstein summation convention.\n",
    "                cricri = 0\n",
    "                for alpha in range(0,dim):\n",
    "\n",
    "                    cricri += 0.5 * g_inv_l[alpha][beta] * (jacob_l[alpha][mu][nu] + jacob_l[alpha][nu][mu] - jacob_l[mu][nu][alpha])\n",
    "                christoffel[beta,mu,nu] = cricri\n",
    "            \n",
    "    christoffel_l = christoffel.tolist()\n",
    "    # Initialize Riemannian tensor\n",
    "    riemann = torch.zeros(dim,dim,dim,dim)\n",
    "    riemann_l = riemann.tolist()\n",
    "    for beta in range(0,dim):\n",
    "        for epsilon in range(0,dim):\n",
    "            for mu in range(0,dim):\n",
    "                for nu in range(0,dim):\n",
    "                    #Here as you can see, the riemann tensor is divided in the addition of those two functions\n",
    "                    riemannscalar = 0\n",
    "                    #First the christoffel derivatives\n",
    "                    for alpha in range(0,dim):\n",
    "                        riemannscalar += christoffelderiv(mu, epsilon, nu, alpha, beta)-christoffelderiv(mu, nu, epsilon, alpha, beta)\n",
    "                    for eta in range(0,dim):\n",
    "                        #Then the christoffelproduct.\n",
    "                        riemannscalar += christoffelproduct_liste(nu, mu, epsilon) - christoffelproduct_liste(epsilon, mu, nu) \n",
    "                    \n",
    "                    riemann_l[beta][mu][nu][epsilon] = riemannscalar\n",
    "\n",
    "    #Initialize Ricci tensor\n",
    "    ricci = torch.zeros(dim,dim)\n",
    "    ricci_l = ricci.tolist()\n",
    "    for mu in range(0, dim):\n",
    "        for nu in range(0, dim):\n",
    "            riccisca = 0\n",
    "            for alpha in range(0, dim):\n",
    "                riccisca += riemann_l[alpha][mu][alpha][nu]\n",
    "                ricci_l[mu][nu]= riccisca\n",
    "    #scalarcurvature \n",
    "    scalarcurvature = 0\n",
    "    for mu in range(dim):\n",
    "        for nu in range(dim):\n",
    "            scalarcurvature += g_inv_l[mu][nu]*ricci_l[mu][nu]\n",
    "    i = 0\n",
    "    while i < 1:\n",
    "        print(\"What do you want to display ? Choose a number : \\n 1 : Christoffel_symbols \\n 2 : Riemann_tensor \\n 3 : Ricci_tensor \\n 4 : Scalar_curvature \")\n",
    "        detail = input()\n",
    "        if detail== '1':\n",
    "            for beta in range(0,dim):\n",
    "                for mu in range(0,dim):\n",
    "                    for nu in range(0,dim):\n",
    "                        display(Math(f'\\Gamma^{{{beta}}}_{{{mu} {nu}}} = ' + latex(christoffel_l[beta][mu][nu])))\n",
    "        elif detail=='2':\n",
    "            for beta in range(0,dim):\n",
    "                for epsilon in range(0,dim):\n",
    "                    for mu in range(0,dim):\n",
    "                        for nu in range(0,dim):\n",
    "                            display(Math(f'R^{{{beta}}}_{{{mu} {nu} {epsilon}}} = ' + latex(riemann_l[beta][mu][nu][epsilon])))\n",
    "        elif detail=='3':\n",
    "            for mu in range(0, dim):\n",
    "                for nu in range(0, dim):\n",
    "                    display(Math(f'R_{{{mu} {nu}}} = ' + latex(ricci_l[mu][nu])))\n",
    "        elif detail=='4':\n",
    "            display(Math(f'R = ' + latex(scalarcurvature)))\n",
    "        else :\n",
    "            print(\"Error : please choose a number between 1, 2, 3 and 4.\")\n",
    "        print(\"Do you want to display an other value ? y/n\")\n",
    "        x = input()\n",
    "        if x == 'n':\n",
    "            i += 1\n",
    "Relat_calc()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88e95641",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f18a1ec2",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}