{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "2bb083b7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "%pylab is deprecated, use %matplotlib inline and import the required libraries.\n",
      "Populating the interactive namespace from numpy and matplotlib\n"
     ]
    }
   ],
   "source": [
    "%pylab inline\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "458efc98",
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import jax.lax as lax\n",
    "import jax_cosmo as jc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "da00e618",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import flowpm\n",
    "from flowpm.tfpower import linear_matter_power\n",
    "import flowpm.scipy.interpolate as interpolate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "b588c18c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from jaxpm.kernels import *\n",
    "from jaxpm.painting import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "f0cba7a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Below are a few parameters\n",
    "box_size = [100., 100., 100.]   # Transverse comoving size of the simulation volume\n",
    "nc = [100, 100, 100]            # Number of transverse voxels in the simulation volume\n",
    "batch_size = 1                  # Number of simulations to run in parallel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "8dad416f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Instantiates a cosmology with desired parameters\n",
    "cosmology = flowpm.cosmology.Planck15()\n",
    "\n",
    "# Create some initial conditions\n",
    "k = tf.constant(np.logspace(-4, 1, 128), dtype=tf.float32)\n",
    "pk = linear_matter_power(cosmology, k)\n",
    "pk_fun = lambda x: tf.cast(tf.reshape(interpolate.interp_tf(tf.reshape(tf.cast(x, tf.float32), [-1]), k, pk), x.shape), tf.complex64)\n",
    "initial_conditions = flowpm.linear_field(nc,\n",
    "                                      box_size, \n",
    "                                      pk_fun,         \n",
    "                                      batch_size=batch_size)\n",
    "\n",
    "initial_state = flowpm.lpt_init(cosmology, initial_conditions, 0.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "id": "5ad12a1e",
   "metadata": {},
   "outputs": [],
   "source": [
    "@tf.function\n",
    "def solve_tf(init_state):\n",
    "    final_state = flowpm.nbody(cosmology, initial_state, linspace(0.1,1.,40), nc)\n",
    "    return final_state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "id": "7d4780a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "final_state = solve_tf(initial_state)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "6a94dae9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4.89 s ± 12.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%timeit final_state = solve_tf(initial_state)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "e5148fed",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x7f9774507880>"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 53140 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "imshow(flowpm.cic_paint(tf.zeros([1]+nc), final_state[0]).numpy()[0].sum(axis=0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "167d193f",
   "metadata": {},
   "outputs": [],
   "source": [
    "mesh_shape = nc\n",
    "kvec = fftk(mesh_shape)\n",
    "\n",
    "# Define the ODE\n",
    "def f(state, a, cosmo):\n",
    "    # Extracts positions and velocity at a given point in\n",
    "    # the simulation\n",
    "    pos, vel = state\n",
    "    \n",
    "    # Computes the potential given the current positions\n",
    "    delta_k = jnp.fft.rfftn(cic_paint(jnp.zeros(mesh_shape), pos))\n",
    "    pot_k = delta_k * laplace_kernel(kvec) * longrange_kernel(kvec,r_split=0)\n",
    "    forces = jnp.stack([cic_read(jnp.fft.irfftn(gradient_kernel(kvec, i)*pot_k), pos) \n",
    "                      for i in range(3)],axis=-1)\n",
    "    forces = forces * 1.5 * cosmo.Omega_m\n",
    "    \n",
    "    # Computes the update of position (drift)\n",
    "    dpos = 1. / (a**3 * jnp.sqrt(jc.background.Esqr(cosmo, a))) * vel\n",
    "    \n",
    "    # Computes the update of velocity (kick)\n",
    "    dvel = 1. / (a**2 * jnp.sqrt(jc.background.Esqr(cosmo, a))) * forces\n",
    "    \n",
    "    return dpos, dvel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "aa14a004",
   "metadata": {},
   "outputs": [],
   "source": [
    "from jax.experimental.ode import odeint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "973ca031",
   "metadata": {},
   "outputs": [],
   "source": [
    "init_state = [initial_state[0,0].numpy(),\n",
    "              initial_state[1,0].numpy()]\n",
    "\n",
    "@jax.jit\n",
    "def solve_ode(init_state):\n",
    "    return odeint(f, init_state, \n",
    "             jnp.linspace(0.1,1.0,10), \n",
    "             jc.Planck15(),rtol=1e-5, atol=1e-5 )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "e801c3eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "res = solve_ode(init_state)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "de770afb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "352 ms ± 1.01 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%timeit res = solve_ode(init_state)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "73e7cc3c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x7f977c4261a0>"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 90284 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "imshow(cic_paint(jnp.zeros(mesh_shape), initial_state[0,0].numpy()).sum(axis=0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "4bdfe23f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.colorbar.Colorbar at 0x7f94181834c0>"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 56276 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "imshow(cic_paint(jnp.zeros(mesh_shape), res[0][-1]).sum(axis=0)); colorbar()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "a7f8b5d5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.colorbar.Colorbar at 0x7f940073dd80>"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 41284 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "imshow((cic_paint(jnp.zeros(mesh_shape), res[0][-1]) - \n",
    "        flowpm.cic_paint(tf.zeros([1]+mesh_shape), final_state[0]).numpy()[0]).sum(axis=0)); colorbar()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "68b72b68",
   "metadata": {},
   "outputs": [],
   "source": [
    "from DifferentiableHOS.pk import power_spectrum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "785818e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "k, pk = power_spectrum(flowpm.cic_paint(tf.zeros([1]+mesh_shape), final_state[0]),\n",
    "                       boxsize=np.array([100.,100.,100.]),                      \n",
    "                       kmin=0.1,dk=2*np.pi/100.)\n",
    "\n",
    "k, pk_jax = power_spectrum(tf.convert_to_tensor(cic_paint(jnp.zeros(mesh_shape), res[0][-1]).reshape([1,100,100,100])),\n",
    "                       boxsize=np.array([100.,100.,100.]),                      \n",
    "                       kmin=0.1,dk=2*np.pi/100.)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "18e3d30b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x7f94005fc850>]"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 10656 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "loglog(k,pk[0])\n",
    "loglog(k,pk_jax[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "c2c1ecdb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x7f94005693f0>]"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 16800 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "semilogx(k,(pk[0] - pk_jax[0])/pk[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "id": "703c8195",
   "metadata": {},
   "outputs": [],
   "source": [
    "import diffrax\n",
    "from diffrax import diffeqsolve, ODETerm, Dopri5, LeapfrogMidpoint, PIDController\n",
    "\n",
    "cosmo = jc.Planck15()\n",
    "\n",
    "def f_diffrax(t, state, args):\n",
    "    return jnp.stack(f(state, t, cosmo),axis=0)\n",
    "\n",
    "term = ODETerm(f_diffrax)\n",
    "solver = Dopri5()\n",
    "stepsize_controller = PIDController(rtol=1e-5, atol=1e-5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "id": "d6921f17",
   "metadata": {},
   "outputs": [],
   "source": [
    "solution = diffeqsolve(term, solver, t0=0.1, t1=1., dt0=0.1,\n",
    "                       y0=jnp.stack(init_state,axis=0), \n",
    "                       stepsize_controller=stepsize_controller)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "id": "b3083d29",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Solution(t0=DeviceArray(0.1, dtype=float32), t1=DeviceArray(1., dtype=float32), ts=DeviceArray([1.], dtype=float32, weak_type=True), ys=DeviceArray([[[[ -2.628675  ,   5.347485  ,  -1.5904427 ],\n",
       "               [ -2.5886178 ,   5.4675574 ,  -1.8807749 ],\n",
       "               [ -2.6215632 ,   5.6189113 ,  -1.9770483 ],\n",
       "               ...,\n",
       "               [ 97.6788    , 105.22156   ,  96.64059   ],\n",
       "               [ 97.459785  , 105.74062   ,  97.79143   ],\n",
       "               [ 97.26484   , 105.98253   ,  98.37041   ]],\n",
       "\n",
       "              [[ -1.7246462 ,  -0.27205068,   0.6488323 ],\n",
       "               [ -1.6212564 ,  -0.29098457,   0.29932806],\n",
       "               [ -1.4803871 ,  -0.176678  ,  -0.5965874 ],\n",
       "               ...,\n",
       "               [ -2.21784   ,   2.4735427 ,   1.3182787 ],\n",
       "               [ -2.3406084 ,   2.282835  ,   1.8463919 ],\n",
       "               [ -2.175532  ,   0.7863351 ,   1.1335845 ]]]],            dtype=float32, weak_type=True), interpolation=None, stats={'compiled_num_steps': None, 'max_steps': 4096, 'num_accepted_steps': DeviceArray(19, dtype=int32, weak_type=True), 'num_rejected_steps': DeviceArray(4, dtype=int32, weak_type=True), 'num_steps': DeviceArray(23, dtype=int32, weak_type=True)}, result=DeviceArray(0, dtype=int32, weak_type=True), solver_state=None, controller_state=None, made_jump=None)"
      ]
     },
     "execution_count": 67,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "id": "65116b88",
   "metadata": {},
   "outputs": [],
   "source": [
    "resx = solution.ys[0,0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "id": "c6bd8cfd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.colorbar.Colorbar at 0x7f9368964cd0>"
      ]
     },
     "execution_count": 69,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 51524 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "imshow(cic_paint(jnp.zeros(mesh_shape), resx).sum(axis=0),cmap='gist_stern'); colorbar()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "id": "d7ce5401",
   "metadata": {},
   "outputs": [],
   "source": [
    "k, pk = power_spectrum(flowpm.cic_paint(tf.zeros([1]+mesh_shape), final_state[0]),\n",
    "                       boxsize=array([100.,100.,100.]),                      \n",
    "                       kmin=0.1,dk=2*np.pi/100.)\n",
    "\n",
    "k, pk_jax = power_spectrum(tf.convert_to_tensor(cic_paint(jnp.zeros(mesh_shape), resx).reshape([1,100,100,100])),\n",
    "                       boxsize=array([100.,100.,100.]),                      \n",
    "                       kmin=0.1,dk=2*np.pi/100.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "id": "28aafee9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.legend.Legend at 0x7f97746730d0>"
      ]
     },
     "execution_count": 76,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 12396 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "loglog(k,pk[0], label='flowpm')\n",
    "loglog(k,pk_jax[0], label='jax')\n",
    "legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "id": "dc6a66c0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x7f978c1fb6a0>]"
      ]
     },
     "execution_count": 77,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 15448 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "semilogx(k,(pk_jax[0] - pk[0])/pk[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2b58a54",
   "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.10.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
