Python_CRFD.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Solving Relativistic Fluid 2D\n",
    "\n",
    "### Introduction\n",
    "The goal of this notebook is to study turbulence from a perfect relativistic fluid. We consider the stress energy tensor for a perfect fluid in Minkowski space,\n",
    "\n",
    "\\begin{equation}\n",
    "T_{ab} = p\\eta_{ab} + (e+p)u_a u_b\n",
    "\\end{equation}\n",
    "\n",
    "The equations of motion for such fluid are,\n",
    "\n",
    "\\begin{equation}\n",
    "\\label{EOM}\n",
    "\\partial_a T^{ab} = 0. \\quad (1)\n",
    "\\end{equation}\n",
    "\n",
    "Here we focus on a conformal fluid, that is the stress energy tensor is traceless,\n",
    "\n",
    "$$T_a^a = 0 \\to e = -2p.$$\n",
    "\n",
    "To solve eq.1 numerically, we rewrite the EOM in conservative variables $U=(D,S^i)$ where,\n",
    "\n",
    "$$\\frac{D}{2} = T^{00} = \\frac{e}{2}\\left(-1+\\frac{3}{W^2}\\right)$$\n",
    "\n",
    "$$\\frac{S^i}{2} = T^{0i} = \\frac{3e}{2}\\frac{v^i}{W^2}$$\n",
    "\n",
    "Where we have the Lorentz factor $W = \\sqrt{1-v^iv_i}$. Now the EOM in conservative form,\n",
    "\n",
    "$$\\partial_t D + \\partial_i S^i = 0$$\n",
    "$$\\partial_t S^i + \\partial_j(e \\delta^{ij} + S^i v^j) = 0$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import packages\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits.mplot3d import Axes3D  \n",
    "import matplotlib.cm as cm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Number of grid points\n",
    "Nx = 201\n",
    "Ny = Nx\n",
    "Nt = 2000\n",
    "\n",
    "# Discretisation parameters\n",
    "cfl = 1.0\n",
    "dx  = 10./(Nx)\n",
    "dy  = 10./(Ny)\n",
    "dt  = cfl*min(dx,dy)\n",
    "\n",
    "# Spacial linspace\n",
    "x = np.linspace(0,dx*Nx,Nx)\n",
    "y = np.linspace(0,dy*Ny,Ny)\n",
    "\n",
    "# physical parameters\n",
    "gamma = 1.5\n",
    "\n",
    "# Numerical parameters (floor)\n",
    "epsilon = 1e-5\n",
    "dissip = -0.01"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Variables to store solutions\n",
    "v_x_old = np.zeros((Nx,Ny))\n",
    "\n",
    "v_y_old = np.zeros((Nx,Ny))\n",
    "\n",
    "rho_old = np.ones((Nx,Ny))\n",
    "\n",
    "D_old = np.zeros((Nx,Ny))\n",
    "\n",
    "s_x_old = np.zeros((Nx,Ny))\n",
    "\n",
    "s_y_old = np.zeros((Nx,Ny))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Variable conversion (C->P)\n",
    "\n",
    "To switch from conservative to primitive variable we first note that,\n",
    "\n",
    "$$S_iS^i/D^2 = 9(1-W^2)(3-W^2)^2$$\n",
    "\n",
    "Now let $\\Theta = S_iS^i/D^2$ and solving for $W^2$ we get,\n",
    "\n",
    "$$W^2 = \\frac{3}{2\\Theta} \\left(2 \\Theta - 3 + \\sqrt{9-8\\Theta}\\right)$$\n",
    "\n",
    "With this,\n",
    "\n",
    "$$e = \\frac{D}{-1+3W^{-2}}$$\n",
    "\n",
    "$$v^i = \\frac{W^2 S^i}{3e}$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def CtoP(D,Sx,Sy):\n",
    "    \n",
    "    Nx = S_x_old.shape[0]\n",
    "    Nx = S_x_old.shape[1]\n",
    "    prho = np.zeros((Nx,Ny))\n",
    "    v_x  = np.zeros((Nx,Ny))\n",
    "    v_y  = np.zeros((Nx,Ny))\n",
    "    W2   = np.zeros((Nx,Ny))\n",
    "    \n",
    "    # Compute theta\n",
    "    theta = (np.multiply(Sx,Sx)+np.multiply(Sy,Sy))/np.multiply(D,D)\n",
    "    \n",
    "    # Compute W^2\n",
    "    W2 = (1.5/theta)*(2*theta-3.0+np.sqrt(9.0-8.0*theta))\n",
    "    if np.any(np.isnan(W2)):\n",
    "        print(np.amin(9.0-8.0*theta))\n",
    "        index1 = np.unravel_index(W2.argmin(), W2.shape)[0]\n",
    "        index2 = np.unravel_index(W2.argmin(), W2.shape)[1]\n",
    "        print(\"Position: \",index1,index2)\n",
    "        print(\"Theta: \", theta[index1,index2])\n",
    "        print(\"Sx: \",Sx[index1,index2])\n",
    "        print(\"Sy: \",Sy[index1,index2])\n",
    "        print(\"D: \",D[index1,index2])\n",
    "    for i in range(Nx):\n",
    "        for j in range(Ny):\n",
    "            if (theta[i,j] < epsilon**2):\n",
    "                W2[i,j] = 1.0\n",
    "        \n",
    "    # Compute primitives\n",
    "    prho = D/(-1+(3/W2))\n",
    "    \n",
    "    v_x = (np.multiply(W2,Sx))/(3*prho)\n",
    "    v_y = (np.multiply(W2,Sy))/(3*prho)\n",
    "    \n",
    "    return prho, v_x, v_y\n",
    "\n",
    "def PtoC(rho,v_x,v_y,Nx,Ny):\n",
    "    \n",
    "    D = np.zeros((Nx,Ny))\n",
    "    S_x = np.zeros((Nx,Ny))\n",
    "    S_y = np.zeros((Nx,Ny))\n",
    "    \n",
    "    # Compute W^2\n",
    "    W2 = np.ones((Nx,Ny))-(np.multiply(v_x,v_x)+np.multiply(v_y,v_y))\n",
    "    \n",
    "    D   = np.multiply(rho,-np.ones((Nx,Ny)) + 3.0/W2)\n",
    "    S_x = 3.0*v_x/W2\n",
    "    S_y = 3.0*v_y/W2\n",
    "    \n",
    "    return D, S_x, S_y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initial Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def gen_init(dx,dy,Nx,Ny,x,y,plot=False):\n",
    "    # initial data assumes square grid\n",
    "    v_x = np.zeros((Nx,Ny))\n",
    "    v_y = np.zeros((Nx,Ny))\n",
    "    \n",
    "    S_x = np.zeros((Nx,Ny))\n",
    "    S_y = np.zeros((Nx,Ny))\n",
    "    \n",
    "    D = np.zeros((Nx,Ny))\n",
    "    \n",
    "    # Density of 1 everywhere\n",
    "    rho = np.ones((Nx,Ny))\n",
    "    \n",
    "    W = np.zeros((Nx,Ny))\n",
    "    W2 = np.zeros((Nx,Ny))\n",
    "    \n",
    "    # Generate velocities\n",
    "    for i in range(Nx):\n",
    "        if(y[i]>3 and y[i]<7):\n",
    "            v_x[:,i] = 0.5*np.sin((2*np.pi)*y[i])*(y[i]-3.0)**2*((y[i]-7.0)**2)/32\n",
    "        else:\n",
    "            # Set to floor\n",
    "            v_x[:,i] = epsilon\n",
    "        if(x[i]>3 and x[i] < 7):\n",
    "            v_y[i,:] = 0.5*np.sin((0.5*np.pi)*x[i])*(x[i]-3.0)**2*((x[i]-7.0)**2)/32\n",
    "        else:\n",
    "            # Set to floor\n",
    "            v_y[i,:] = epsilon\n",
    "\n",
    "    # Compute conservative variables\n",
    "    D, S_x, S_y = PtoC(rho,v_x,v_y,Nx,Ny)\n",
    "\n",
    "    if plot:\n",
    "        fig = plt.figure(figsize=(15,15))\n",
    "        ax = fig.add_subplot(111, projection='3d')\n",
    "        X, Y = np.meshgrid(x, y)\n",
    "        \n",
    "        ax.plot_surface(X, Y, S_x)\n",
    "        \n",
    "        ax.set_xlabel('X')\n",
    "        ax.set_ylabel('Y')\n",
    "        ax.set_zlabel('S_x')\n",
    "        \n",
    "        plt.show()\n",
    "    \n",
    "    return v_x, v_y, rho, D, S_x, S_y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate initial data\n",
    "v_x_old, v_y_old, rho_old, D_old, S_x_old, S_y_old = gen_init(dx,dy,Nx,Ny,x,y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define finite difference fn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 2nd order Central difference for 1st degree derivative \n",
    "def dudx(u,dx):\n",
    "    \n",
    "    du = np.zeros((len(u),len(u)))\n",
    "    \n",
    "    for i in range(1,len(u)-1):\n",
    "        du[i,:] = (u[i+1,:]-u[i-1,:])/(2*dx)\n",
    "        \n",
    "    # Assume periodic BC\n",
    "    du[0,:]  = (u[1,:]-u[-2,:])/(2*dx)\n",
    "    du[-1,:] = (u[1,:]-u[-2,:])/(2*dx)\n",
    "        \n",
    "    return du\n",
    "\n",
    "# 2nd order Central difference for 1st degree derivative \n",
    "def dudy(u,dy):\n",
    "    \n",
    "    du = np.zeros((len(u),len(u)))\n",
    "    \n",
    "    for i in range(1,len(u)-1):\n",
    "        du[:,i] = (u[:,i+1]-u[:,i-1])/(2*dy)\n",
    "    \n",
    "    # Assume periodic BC\n",
    "    du[:,0]  = (u[:,1]-u[:,-2])/(2*dy)\n",
    "    du[:,-1] = (u[:,1]-u[:,-2])/(2*dy)\n",
    "        \n",
    "    return du\n",
    "\n",
    "# Need some dissapation\n",
    "def dissip_u(u,dx):\n",
    "    du = np.zeros((len(u),len(u)))\n",
    "    du_temp = np.zeros((len(u),len(u)))\n",
    "    for i in range(len(u)):\n",
    "        du[:,i] = dissip4(u[:,i])\n",
    "    for j in range(len(u)):\n",
    "        du_temp[j,:] = dissip4(u[j,:])\n",
    "    \n",
    "    du = (du + du_temp)/dx\n",
    "    \n",
    "    return du\n",
    "\n",
    "def dissip4(u):\n",
    "    du = np.zeros(len(u))\n",
    "    N = len(u)-1\n",
    "    for i in range(2,N-1):\n",
    "        du[i] = u[i+2]+u[i-2]-4*(u[i+1]+u[i-1])+6*u[i]\n",
    "        \n",
    "    du[1] = u[3]+u[-1]-4*(u[2]+u[0])+6*u[1]\n",
    "    du[0] = u[2]+u[-2]-4*(u[1]+u[-1])+6*u[0]\n",
    "    du[-1] = du[0]\n",
    "    du[-2] = u[0]+u[-4]-4*(u[-1]+u[-3])+6*u[-2]\n",
    "    return du"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define fluxes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def floor(u,Nx):\n",
    "    for i in range(Nx):\n",
    "        for j in range(Nx):\n",
    "            if u[i,j] < epsilon:\n",
    "                u[i,j] = epsilon\n",
    "    return u\n",
    "\n",
    "def fluxes(rho,v_x,v_y,D,S_x,S_y):\n",
    "    \n",
    "    Fx_rho = S_x\n",
    "    Fy_rho = S_y\n",
    "    \n",
    "    Fx_Sx = rho+np.multiply(S_x,v_x)\n",
    "    Fy_Sx = np.multiply(S_x,v_y)\n",
    "    \n",
    "    Fx_Sy = np.multiply(S_y,v_x)\n",
    "    Fy_Sy = rho+np.multiply(S_y,v_y)\n",
    "    \n",
    "    return Fx_rho, Fy_rho, Fx_Sx, Fy_Sx, Fx_Sy, Fy_Sy\n",
    "\n",
    "def RHS(Z):\n",
    "    \n",
    "    D   = Z[0]\n",
    "    S_x = Z[1]\n",
    "    S_y = Z[2]\n",
    "    \n",
    "    D = floor(D,Nx)\n",
    "    \n",
    "    rho, v_x, v_y =  CtoP(D,S_x,S_y)\n",
    "    \n",
    "    # Compute fluxes\n",
    "    Fx_rho, Fy_rho, Fx_Sx, Fy_Sx, Fx_Sy, Fy_Sy = fluxes(rho,v_x,v_y,D,S_x,S_y)\n",
    "        \n",
    "    # Compute derivatives\n",
    "    dxFx_rho = dudx(Fx_rho,dx)\n",
    "    dyFy_rho = dudy(Fy_rho,dx)\n",
    "    \n",
    "    dxFx_Sx  = dudx(Fx_Sx,dx)\n",
    "    dyFy_Sx  = dudy(Fy_Sx,dx)\n",
    "    \n",
    "    dxFx_Sy  = dudx(Fx_Sy,dx)\n",
    "    dyFy_Sy  = dudy(Fy_Sy,dx)\n",
    "    \n",
    "    dD   = (-dxFx_rho-dyFy_rho) + dissip*dissip_u(D,dx)\n",
    "    dS_x = (-dxFx_Sx-dyFy_Sx) + dissip*dissip_u(S_x,dx)\n",
    "    dS_y = (-dxFx_Sy-dyFy_Sy) + dissip*dissip_u(S_y,dx)\n",
    "        \n",
    "    del dxFx_rho,dyFy_rho,dxFx_Sx,dyFy_Sx,dxFx_Sy,dyFy_Sy\n",
    "    return np.array([dD, dS_x, dS_y])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define RK4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rk4_step(f, y, h):\n",
    "    \n",
    "    k1 = h*f(y)\n",
    "    k2 = h*f(y + 0.5 * k1)\n",
    "    k3 = h*f(y + 0.5 * k2)\n",
    "    k4 = h*f(y + k3)\n",
    "    \n",
    "    new_data = y + (k1 + 2*k2 + 2*k3 + k4) / 6\n",
    "    del k1,k2,k3,k4\n",
    "    return new_data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Evolve"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1\n",
      "2\n",
      "3\n",
      "4\n",
      "5\n"
     ]
    }
   ],
   "source": [
    "# Create Z vector\n",
    "Z_old = np.array([D_old,S_x_old,S_y_old])\n",
    "t = [0]\n",
    "X, Y = np.meshgrid(x, y)\n",
    "\n",
    "# Save D figure\n",
    "figD, axD = plt.subplots(figsize=(15,15))\n",
    "imD = axD.imshow(Z_old[0], interpolation='bilinear', \n",
    "origin='lower', extent=[min(x), max(x), min(y), max(y)],cmap=cm.RdYlGn,\n",
    "              vmax=abs(Z_old[0]).max(), vmin=abs(Z_old[0]).min())\n",
    "\n",
    "figSx, axSx = plt.subplots(figsize=(15,15))\n",
    "imSx = axSx.imshow(Z_old[1], interpolation='bilinear', \n",
    "origin='lower', extent=[min(x), max(x), min(y), max(y)],cmap=cm.RdYlGn,\n",
    "              vmax=abs(Z_old[1]).max(), vmin=abs(Z_old[1]).min())\n",
    "\n",
    "figSy, axSy = plt.subplots(figsize=(15,15))\n",
    "imSy = axSy.imshow(Z_old[2], interpolation='bilinear', \n",
    "origin='lower', extent=[min(x), max(x), min(y), max(y)],cmap=cm.RdYlGn,\n",
    "              vmax=abs(Z_old[2]).max(), vmin=abs(Z_old[2]).min())\n",
    "\n",
    "figD.savefig(\"D/D\" + str(0) + \".png\")\n",
    "figSx.savefig(\"Sx/Sx\" + str(0) + \".png\")\n",
    "figSy.savefig(\"Sy/Sy\" + str(0) + \".png\")\n",
    "\n",
    "for i in range(1,Nt+1):\n",
    "    Z_new = rk4_step(RHS,Z_old,dt)\n",
    "    t.append(t[0]+i*dt)\n",
    "    \n",
    "    print(i)\n",
    "    \n",
    "    # Save D figure\n",
    "    imD = axD.imshow(Z_old[0], interpolation='bilinear', origin='lower', extent=[min(x), max(x), min(y), max(y)],cmap=cm.RdYlGn,vmax=abs(Z_old[0]).max(), vmin=abs(Z_old[0]).min())\n",
    "    figD.savefig(\"D/D\" + str(int(i/10)) + \".png\")\n",
    "    \n",
    "    # Save Sx figure\n",
    "    imSx = axSx.imshow(Z_old[1], interpolation='bilinear', \n",
    "               origin='lower', extent=[min(x), max(x), min(y), max(y)],cmap=cm.RdYlGn,\n",
    "               vmax=abs(Z_old[1]).max(), vmin=abs(Z_old[1]).min())\n",
    "    figSx.savefig(\"Sx\" + str(i) + \".png\")\n",
    "        \n",
    "    # Save Sy figure\n",
    "    imSy = axSy.imshow(Z_old[2], interpolation='bilinear', \n",
    "               origin='lower', extent=[min(x), max(x), min(y), max(y)],cmap=cm.RdYlGn,\n",
    "               vmax=abs(Z_old[2]).max(), vmin=abs(Z_old[2]).min())\n",
    "    figSy.savefig(\"Sy\" + str(i) + \".png\")\n",
    "        \n",
    "    Z_old = Z_new"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}