{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Populating the interactive namespace from numpy and matplotlib\n"
     ]
    }
   ],
   "source": [
    "%pylab inline\n",
    "import camels_library as cl\n",
    "import readgadget\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "init_cond = '/data/CAMELS/Sims/IllustrisTNG_DM/CV_0/ICs/ics'\n",
    "\n",
    "header   = readgadget.header(init_cond)\n",
    "BoxSize  = header.boxsize/1e3  #Mpc/h\n",
    "Nall     = header.nall         #Total number of particles\n",
    "Masses   = header.massarr*1e10 #Masses of the particles in Msun/h\n",
    "Omega_m  = header.omega_m      #value of Omega_m\n",
    "Omega_l  = header.omega_l      #value of Omega_l\n",
    "h        = header.hubble       #value of h\n",
    "redshift = header.redshift     #redshift of the snapshot\n",
    "Hubble   = 100.0*np.sqrt(Omega_m*(1.0+redshift)**3+Omega_l)#Value of H(z) in km/s/(Mpc/h)\n",
    "\n",
    "ptype = [1] #dark matter is particle type 1\n",
    "ids_i = np.argsort(readgadget.read_block(init_cond, \"ID  \", ptype)-1)  #IDs starting from 0\n",
    "pos_i = readgadget.read_block(init_cond, \"POS \", ptype)[ids_i]/1e3     #positions in Mpc/h\n",
    "vel_i = readgadget.read_block(init_cond, \"VEL \", ptype)[ids_i]         #peculiar velocities in km/s\n",
    "\n",
    "# Reordering data for simple reshaping\n",
    "pos_i = pos_i.reshape(4,4,4,64,64,64,3).transpose(0,3,1,4,2,5,6).reshape(-1,3)\n",
    "vel_i = vel_i.reshape(4,4,4,64,64,64,3).transpose(0,3,1,4,2,5,6).reshape(-1,3)\n",
    "pos_i = (pos_i/BoxSize*64).reshape([256,256,256,3])[::4,::4,::4,:].reshape([-1,3])\n",
    "vel_i = (vel_i / 100 * (1./(1+redshift)) / BoxSize*64).reshape([256,256,256,3])[::4,::4,::4,:].reshape([-1,3])\n",
    "a_i   = 1./(1+redshift)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 34/34 [01:54<00:00,  3.36s/it]\n"
     ]
    }
   ],
   "source": [
    "scales = []\n",
    "poss = []\n",
    "vels = []\n",
    "\n",
    "# Loading all the intermediate snapshots\n",
    "for i in tqdm(range(34)):\n",
    "    snapshot='/data/CAMELS/Sims/IllustrisTNG_DM/CV_0/snap_%03d.hdf5'%i\n",
    "    \n",
    "    header   = readgadget.header(snapshot)\n",
    "    \n",
    "    redshift = header.redshift     #redshift of the snapshot\n",
    "    h        = header.hubble       #value of h\n",
    "    \n",
    "    ptype = [1] #dark matter is particle type 1\n",
    "    ids = np.argsort(readgadget.read_block(snapshot, \"ID  \", ptype)-1)   #IDs starting from 0\n",
    "    pos = readgadget.read_block(snapshot, \"POS \", ptype)[ids]/1e3 #positions in Mpc/h\n",
    "    vel = readgadget.read_block(snapshot, \"VEL \", ptype)[ids]     #peculiar velocities in km/s\n",
    "\n",
    "    # Reordering data for simple reshaping\n",
    "    pos = pos.reshape(4,4,4,64,64,64,3).transpose(0,3,1,4,2,5,6).reshape(-1,3)\n",
    "    vel = vel.reshape(4,4,4,64,64,64,3).transpose(0,3,1,4,2,5,6).reshape(-1,3)\n",
    "    \n",
    "    pos = (pos / BoxSize * 64).reshape([256,256,256,3])[::4,::4,::4,:].reshape([-1,3])\n",
    "    vel = (vel / 100 * (1./(1+redshift)) / BoxSize*64).reshape([256,256,256,3])[::4,::4,::4,:].reshape([-1,3])\n",
    "    \n",
    "    scales.append((1./(1+redshift)))\n",
    "    poss.append(pos)\n",
    "    vels.append(vel)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x7f82862318b0>]"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 11256 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot(scales)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import jax_cosmo as jc\n",
    "import haiku as hk\n",
    "\n",
    "from jax.experimental.ode import odeint\n",
    "\n",
    "from jaxpm.painting import cic_paint, cic_read\n",
    "from jaxpm.pm import linear_field, lpt, make_ode_fn, pm_forces\n",
    "from jaxpm.kernels import fftk, gradient_kernel, laplace_kernel, longrange_kernel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "mesh_shape= [64,64, 64]\n",
    "box_size  = [25.,25.,25.]\n",
    "cosmo = jc.Planck15(Omega_c= 0.3 - 0.049, Omega_b=0.049, n_s=0.9624, h=0.671 , sigma8=0.8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 332876 characters]\n",
      "text/plain": [
       "<Figure size 720x720 with 16 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "figure(figsize=[10,10])\n",
    "for i in range(16):\n",
    "    subplot(4,4,i+1)\n",
    "    imshow(cic_paint(jnp.zeros(mesh_shape), poss[::2][i]).sum(axis=0), cmap='gist_stern', vmin=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Defining the force correction kernel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def deBoorVectorized(x, t, c, p):\n",
    "    \"\"\"\n",
    "    Evaluates S(x).\n",
    "\n",
    "    Args\n",
    "    ----\n",
    "    x: position\n",
    "    t: array of knot positions, needs to be padded as described above\n",
    "    c: array of control points\n",
    "    p: degree of B-spline\n",
    "    \"\"\"\n",
    "    k = jnp.digitize(x, t) -1\n",
    "    \n",
    "    d = [c[j + k - p] for j in range(0, p+1)]\n",
    "    for r in range(1, p+1):\n",
    "        for j in range(p, r-1, -1):\n",
    "            alpha = (x - t[j+k-p]) / (t[j+1+k-r] - t[j+k-p])\n",
    "            d[j] = (1.0 - alpha) * d[j-1] + alpha * d[j]\n",
    "    return d[p]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "class FreeFormFilter(hk.Module):\n",
    "  \"\"\"This is a NN that predicts the score.\"\"\"\n",
    "  def __call__(self, x, a):\n",
    "        \n",
    "    net = jax.nn.leaky_relu(hk.Linear(128)(a))\n",
    "    net = jax.nn.leaky_relu(hk.Linear(128)(net))\n",
    "    \n",
    "    w = hk.Linear(32)(net) # hk.get_parameter(\"w\", [32], init=jnp.ones)\n",
    "    k = hk.Linear(31)(net) # hk.get_parameter(\"k\", [31], init=jnp.ones)\n",
    "    \n",
    "    # make sure the knots sum to 1 and are in the interval 0,1\n",
    "    k = jnp.concatenate([jnp.zeros((1,)),\n",
    "                        jnp.cumsum(jax.nn.softmax(k))])\n",
    "\n",
    "    # Augment with repeating points\n",
    "    ak = jnp.concatenate([jnp.zeros((3,)), k, jnp.ones((3,))])\n",
    "\n",
    "    return deBoorVectorized(x, ak, w, 3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = hk.without_apply_rng(hk.transform(lambda x,a : FreeFormFilter()(x,a)))\n",
    "rng_seq = hk.PRNGSequence(12)\n",
    "params = model.init(next(rng_seq), jnp.zeros([64]), jnp.ones([1]))\n",
    "paramsi = model.init(next(rng_seq), jnp.zeros([64]), jnp.ones([1]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x7f82484566a0>]"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 22788 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot(jnp.linspace(0.01,0.99), model.apply(params, jnp.linspace(0.01,0.99), jnp.ones([1])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def neural_nbody_ode(state, a, cosmo, params):\n",
    "    \"\"\"\n",
    "    state is a tuple (position, velocities)\n",
    "    \"\"\"\n",
    "    pos, vel = state\n",
    "    \n",
    "    kvec = fftk(mesh_shape)\n",
    "\n",
    "    delta = cic_paint(jnp.zeros(mesh_shape), pos)\n",
    "    \n",
    "    delta_k = jnp.fft.rfftn(delta)\n",
    "    \n",
    "    # Computes gravitational potential\n",
    "    pot_k = delta_k * laplace_kernel(kvec) * longrange_kernel(kvec, r_split=0)\n",
    "    \n",
    "    # Apply a correction filter\n",
    "    kk = sum((ki/pi)**2 for ki in kvec)/4.\n",
    "    pot_k = pot_k *(1. + model.apply(params, kk, jnp.atleast_1d(a)))\n",
    "    \n",
    "    # Computes gravitational forces\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",
    "    \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": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "<ipython-input-12-b14f0c49fe59>:17: DeprecationWarning: Calling np.sum(generator) is deprecated, and in the future will give a different result. Use np.sum(np.fromiter(generator)) or the python sum builtin instead.\n",
      "  kk = sum((ki/pi)**2 for ki in kvec)/4.\n"
     ]
    }
   ],
   "source": [
    "# Run the simulation at the same steps\n",
    "resi = odeint(neural_nbody_ode, [pos_i, vel_i], jnp.array([a_i]+scales), cosmo, paramsi, rtol=1e-5, atol=1e-5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 288128 characters]\n",
      "text/plain": [
       "<Figure size 720x720 with 16 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "figure(figsize=[10,10])\n",
    "for i in range(16):\n",
    "    subplot(4,4,i+1)\n",
    "    imshow(cic_paint(jnp.zeros(mesh_shape), resi[0][1:][::2][i]).sum(axis=0), cmap='gist_stern',\n",
    "           vmax=cic_paint(jnp.zeros(mesh_shape), poss[::2][i]).sum(axis=0).max(),vmin=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.colorbar.Colorbar at 0x7f82486381f0>"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 71572 characters]\n",
      "text/plain": [
       "<Figure size 1080x360 with 5 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "figure(figsize=[15,5])\n",
    "subplot(131)\n",
    "imshow(cic_paint(jnp.zeros(mesh_shape), resi[0][0]).sum(axis=0))\n",
    "subplot(132)\n",
    "imshow(cic_paint(jnp.zeros(mesh_shape), resi[0][1]).sum(axis=0)); colorbar()\n",
    "subplot(133)\n",
    "imshow(cic_paint(jnp.zeros(mesh_shape), poss[0]).sum(axis=0)); colorbar()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "ref_pos = jnp.stack(poss, axis=0)\n",
    "ref_vel = jnp.stack(vels, axis=0)\n",
    "\n",
    "@jax.jit\n",
    "def loss_fn(params):\n",
    "  res = odeint(neural_nbody_ode, [pos_i, vel_i], jnp.array([a_i]+scales), cosmo, params, rtol=1e-5, atol=1e-5) \n",
    "  distance = jnp.sum((res[0][1:] - ref_pos)**2, axis=-1)\n",
    "  vel = jnp.sum((res[1][1:] - ref_vel)**2, axis=-1)\n",
    "  w = jnp.where(distance < 100, distance, 0.)\n",
    "  wv = jnp.where(distance < 100, vel, 0.)\n",
    "  return jnp.mean(w) #+ 0.001*jnp.mean(wv)\n",
    "\n",
    "@jax.jit\n",
    "def update(params, opt_state):\n",
    "    \"\"\"Single SGD update step.\"\"\"\n",
    "    loss, grads = jax.value_and_grad(loss_fn)(params)\n",
    "    updates, new_opt_state = optimizer.update(grads, opt_state)\n",
    "    new_params = optax.apply_updates(params, updates)\n",
    "    return loss, new_params, new_opt_state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "losses = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "import optax\n",
    "from tqdm import tqdm\n",
    "learning_rate=0.01\n",
    "optimizer = optax.adam(learning_rate)\n",
    "opt_state = optimizer.init(params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 23%|██▎       | 23/100 [03:31<11:48,  9.20s/it]\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-45-26234aa997fa>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mstep\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mtqdm\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m100\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m   \u001b[0ml\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mparams\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mopt_state\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mupdate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mparams\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mopt_state\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      3\u001b[0m   \u001b[0mlosses\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0ml\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.8/site-packages/haiku/_src/data_structures.py\u001b[0m in \u001b[0;36m<lambda>\u001b[0;34m(treedef, leaves)\u001b[0m\n\u001b[1;32m    260\u001b[0m     \u001b[0mFlatMapping\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    261\u001b[0m     \u001b[0;32mlambda\u001b[0m \u001b[0ms\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0ms\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_leaves\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0ms\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_structure\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m  \u001b[0;31m# pylint: disable=protected-access\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 262\u001b[0;31m     lambda treedef, leaves: FlatMapping(FlatComponents(leaves, treedef)))\n\u001b[0m\u001b[1;32m    263\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    264\u001b[0m \u001b[0;31m#      _                               _           _\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "for step in tqdm(range(100)):\n",
    "  l, params, opt_state = update(params, opt_state)\n",
    "  losses.append(l)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x7f7e846b6d00>]"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 12420 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot(losses[10:])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "pickle.dump( params, open( \"node.params\", \"wb\" ) )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [],
   "source": [
    "res = odeint(neural_nbody_ode, [pos_i, vel_i], jnp.array([a_i]+scales), cosmo, params, rtol=1e-5, atol=1e-5) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.colorbar.Colorbar at 0x7f7e845df820>"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 73488 characters]\n",
      "text/plain": [
       "<Figure size 1080x360 with 5 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "figure(figsize=[15,5])\n",
    "subplot(131)\n",
    "imshow(cic_paint(jnp.zeros(mesh_shape), res[0][0]).sum(axis=0))\n",
    "subplot(132)\n",
    "imshow(cic_paint(jnp.zeros(mesh_shape), res[0][1]).sum(axis=0));colorbar()\n",
    "subplot(133)\n",
    "imshow(cic_paint(jnp.zeros(mesh_shape), ref_pos[0]).sum(axis=0));colorbar()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.lines.Line2D at 0x7f8267c7c4c0>"
      ]
     },
     "execution_count": 74,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 93976 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "for i, a in enumerate(scales): \n",
    "    semilogx(jnp.linspace(0.00,0.99,100), model.apply(params, jnp.linspace(0.0,0.99,100), jnp.atleast_1d(a)))\n",
    "axvline(3./4.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import flowpm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [],
   "source": [
    "k, pk_ref = flowpm.power_spectrum(\n",
    "      tf.convert_to_tensor(cic_paint(jnp.zeros(mesh_shape), poss[0]).reshape([1,64,64,64])),\n",
    "      boxsize=np.array([25.] * 3),\n",
    "      kmin=np.pi / 25.,\n",
    "      dk=2 * np.pi / 25.)\n",
    "\n",
    "_, pki = flowpm.power_spectrum(\n",
    "      tf.convert_to_tensor(cic_paint(jnp.zeros(mesh_shape), resi[0][1]).reshape([1,64,64,64])),\n",
    "      boxsize=np.array([25.] * 3),\n",
    "      kmin=np.pi / 25.,\n",
    "      dk=2 * np.pi / 25.)\n",
    "\n",
    "_, pk = flowpm.power_spectrum(\n",
    "      tf.convert_to_tensor(cic_paint(jnp.zeros(mesh_shape), res[0][1]).reshape([1,64,64,64])),\n",
    "      boxsize=np.array([25.] * 3),\n",
    "      kmin=np.pi / 25.,\n",
    "      dk=2 * np.pi / 25.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x7f7e8436df10>]"
      ]
     },
     "execution_count": 54,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 15764 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "loglog(k,pk_ref[0])\n",
    "loglog(k,pki[0])\n",
    "loglog(k,pk[0],'--')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [],
   "source": [
    "resi = odeint(neural_nbody_ode, [pos_i, vel_i], jnp.array([a_i]+scales), cosmo, paramsi, rtol=1e-5, atol=1e-5) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.colorbar.Colorbar at 0x7f7ebc541bb0>"
      ]
     },
     "execution_count": 56,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 45292 characters]\n",
      "text/plain": [
       "<Figure size 1080x360 with 6 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "figure(figsize=[15,5])\n",
    "subplot(131)\n",
    "imshow(cic_paint(jnp.zeros(mesh_shape), resi[0][-1]).sum(axis=0));colorbar()\n",
    "subplot(132)\n",
    "imshow(cic_paint(jnp.zeros(mesh_shape), res[0][-1]).sum(axis=0));colorbar()\n",
    "subplot(133)\n",
    "imshow(cic_paint(jnp.zeros(mesh_shape), poss[-1]).sum(axis=0));colorbar()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [],
   "source": [
    "k, pk_ref = flowpm.power_spectrum(\n",
    "      tf.convert_to_tensor(cic_paint(jnp.zeros(mesh_shape), poss[-1]).reshape([1,64,64,64])),\n",
    "      boxsize=np.array([25.] * 3),\n",
    "      kmin=np.pi / 25.,\n",
    "      dk=2 * np.pi / 25.)\n",
    "\n",
    "_, pki = flowpm.power_spectrum(\n",
    "      tf.convert_to_tensor(cic_paint(jnp.zeros(mesh_shape), resi[0][-1]).reshape([1,64,64,64])),\n",
    "      boxsize=np.array([25.] * 3),\n",
    "      kmin=np.pi / 25.,\n",
    "      dk=2 * np.pi / 25.)\n",
    "\n",
    "_, pk = flowpm.power_spectrum(\n",
    "      tf.convert_to_tensor(cic_paint(jnp.zeros(mesh_shape), res[0][-1]).reshape([1,64,64,64])),\n",
    "      boxsize=np.array([25.] * 3),\n",
    "      kmin=np.pi / 25.,\n",
    "      dk=2 * np.pi / 25.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x7f7ee8151910>]"
      ]
     },
     "execution_count": 58,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 15244 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "loglog(k,pk_ref[0])\n",
    "loglog(k,pki[0])\n",
    "loglog(k,pk[0],'--')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 332876 characters]\n",
      "text/plain": [
       "<Figure size 720x720 with 16 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "figure(figsize=[10,10])\n",
    "for i in range(16):\n",
    "    subplot(4,4,i+1)\n",
    "    imshow(cic_paint(jnp.zeros(mesh_shape), poss[::2][i]).sum(axis=0), cmap='gist_stern',vmin=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 288132 characters]\n",
      "text/plain": [
       "<Figure size 720x720 with 16 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "figure(figsize=[10,10])\n",
    "for i in range(16):\n",
    "    subplot(4,4,i+1)\n",
    "    imshow(cic_paint(jnp.zeros(mesh_shape), resi[0][1:][::2][i]).sum(axis=0), cmap='gist_stern',\n",
    "           vmax=cic_paint(jnp.zeros(mesh_shape), poss[::2][i]).sum(axis=0).max(),vmin=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 311768 characters]\n",
      "text/plain": [
       "<Figure size 720x720 with 16 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "figure(figsize=[10,10])\n",
    "for i in range(16):\n",
    "    subplot(4,4,i+1)\n",
    "    imshow(cic_paint(jnp.zeros(mesh_shape), res[0][1:][::2][i]).sum(axis=0), cmap='gist_stern',\n",
    "           vmax=cic_paint(jnp.zeros(mesh_shape), poss[::2][i]).sum(axis=0).max(),vmin=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation parameter for omega_m 0.3\n"
     ]
    }
   ],
   "source": [
    "# So let's see what happens on a different simulation\n",
    "val_snap = '/data/CAMELS/Sims/IllustrisTNG_DM/1P_20/'\n",
    "init_cond_val = val_snap+'ICs/ics'\n",
    "\n",
    "header   = readgadget.header(init_cond_val)\n",
    "Omega_m_val  = header.omega_m      #value of Omega_m\n",
    "Omega_l_val  = header.omega_l      #value of Omega_l\n",
    "h_val        = header.hubble       #value of h    \n",
    "print('Validation parameter for omega_m', Omega_m_val)\n",
    "\n",
    "# For 1P_20, sigma8 = 0.96000\n",
    "cosmo_val = jc.Planck15(Omega_c= Omega_m_val - 0.049, Omega_b=0.049, n_s=0.9624, h=0.671 , sigma8=0.96)\n",
    "\n",
    "redshift = header.redshift     #redshift of the snapshot\n",
    "\n",
    "ptype = [1] #dark matter is particle type 1\n",
    "ids_val_i = np.argsort(readgadget.read_block(init_cond_val, \"ID  \", ptype)-1)  #IDs starting from 0\n",
    "pos_val_i = readgadget.read_block(init_cond_val, \"POS \", ptype)[ids_i]/1e3     #positions in Mpc/h\n",
    "vel_val_i = readgadget.read_block(init_cond_val, \"VEL \", ptype)[ids_i]         #peculiar velocities in km/s\n",
    "\n",
    "# Reordering data for simple reshaping\n",
    "pos_val_i = pos_val_i.reshape(4,4,4,64,64,64,3).transpose(0,3,1,4,2,5,6).reshape(-1,3)\n",
    "vel_val_i = vel_val_i.reshape(4,4,4,64,64,64,3).transpose(0,3,1,4,2,5,6).reshape(-1,3)\n",
    "pos_val_i = (pos_val_i/BoxSize*64).reshape([256,256,256,3])[::4,::4,::4,:].reshape([-1,3])\n",
    "vel_val_i = (vel_val_i / 100 * (1./(1+redshift)) / BoxSize*64).reshape([256,256,256,3])[::4,::4,::4,:].reshape([-1,3])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 34/34 [02:29<00:00,  4.39s/it]\n"
     ]
    }
   ],
   "source": [
    "poss_val = []\n",
    "vels_val = []\n",
    "\n",
    "# Loading all the intermediate snapshots\n",
    "for i in tqdm(range(34)):\n",
    "    snapshot= val_snap+ 'snap_%03d.hdf5'%i\n",
    "    \n",
    "    header   = readgadget.header(snapshot)\n",
    "    \n",
    "    redshift = header.redshift     #redshift of the snapshot\n",
    "    \n",
    "    ptype = [1] #dark matter is particle type 1\n",
    "    ids = np.argsort(readgadget.read_block(snapshot, \"ID  \", ptype)-1)   #IDs starting from 0\n",
    "    pos = readgadget.read_block(snapshot, \"POS \", ptype)[ids]/1e3 #positions in Mpc/h\n",
    "    vel = readgadget.read_block(snapshot, \"VEL \", ptype)[ids]     #peculiar velocities in km/s\n",
    "\n",
    "    # Reordering data for simple reshaping\n",
    "    pos = pos.reshape(4,4,4,64,64,64,3).transpose(0,3,1,4,2,5,6).reshape(-1,3)\n",
    "    vel = vel.reshape(4,4,4,64,64,64,3).transpose(0,3,1,4,2,5,6).reshape(-1,3)\n",
    "    \n",
    "    pos = (pos / BoxSize * 64).reshape([256,256,256,3])[::4,::4,::4,:].reshape([-1,3])\n",
    "    vel = (vel / 100 * (1./(1+redshift)) / BoxSize*64).reshape([256,256,256,3])[::4,::4,::4,:].reshape([-1,3])\n",
    "    \n",
    "    poss_val.append(pos)\n",
    "    vels_val.append(vel)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 281720 characters]\n",
      "text/plain": [
       "<Figure size 720x720 with 16 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "figure(figsize=[10,10])\n",
    "for i in range(16):\n",
    "    subplot(4,4,i+1)\n",
    "    imshow(cic_paint(jnp.zeros(mesh_shape), poss_val[::2][i]).sum(axis=0), cmap='gist_stern',vmin=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "<ipython-input-12-b14f0c49fe59>:17: DeprecationWarning: Calling np.sum(generator) is deprecated, and in the future will give a different result. Use np.sum(np.fromiter(generator)) or the python sum builtin instead.\n",
      "  kk = sum((ki/pi)**2 for ki in kvec)/4.\n"
     ]
    }
   ],
   "source": [
    "resi_val = odeint(neural_nbody_ode, [pos_val_i, vel_val_i], jnp.array([a_i]+scales), cosmo_val, paramsi, rtol=1e-5, atol=1e-5) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 242980 characters]\n",
      "text/plain": [
       "<Figure size 720x720 with 16 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "figure(figsize=[10,10])\n",
    "for i in range(16):\n",
    "    subplot(4,4,i+1)\n",
    "    imshow(cic_paint(jnp.zeros(mesh_shape), resi_val[0][1:][::2][i]).sum(axis=0), cmap='gist_stern',\n",
    "           vmax=cic_paint(jnp.zeros(mesh_shape), poss_val[::2][i]).sum(axis=0).max(),vmin=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [],
   "source": [
    "res_val = odeint(neural_nbody_ode, [pos_val_i, vel_val_i], jnp.array([a_i]+scales), cosmo_val, params, rtol=1e-5, atol=1e-5) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 264540 characters]\n",
      "text/plain": [
       "<Figure size 720x720 with 16 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "figure(figsize=[10,10])\n",
    "for i in range(16):\n",
    "    subplot(4,4,i+1)\n",
    "    imshow(cic_paint(jnp.zeros(mesh_shape), res_val[0][1:][::2][i]).sum(axis=0), cmap='gist_stern',\n",
    "           vmax=cic_paint(jnp.zeros(mesh_shape), poss_val[::2][i]).sum(axis=0).max(),vmin=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [],
   "source": [
    "k, pk_ref = flowpm.power_spectrum(\n",
    "      tf.convert_to_tensor(cic_paint(jnp.zeros(mesh_shape), poss_val[-1]).reshape([1,64,64,64])),\n",
    "      boxsize=np.array([25.] * 3),\n",
    "      kmin=np.pi / 25.,\n",
    "      dk=2 * np.pi / 25.)\n",
    "\n",
    "_, pki = flowpm.power_spectrum(\n",
    "      tf.convert_to_tensor(cic_paint(jnp.zeros(mesh_shape), resi_val[0][-1]).reshape([1,64,64,64])),\n",
    "      boxsize=np.array([25.] * 3),\n",
    "      kmin=np.pi / 25.,\n",
    "      dk=2 * np.pi / 25.)\n",
    "\n",
    "_, pk = flowpm.power_spectrum(\n",
    "      tf.convert_to_tensor(cic_paint(jnp.zeros(mesh_shape), res_val[0][-1]).reshape([1,64,64,64])),\n",
    "      boxsize=np.array([25.] * 3),\n",
    "      kmin=np.pi / 25.,\n",
    "      dk=2 * np.pi / 25.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x7f7e841d4550>]"
      ]
     },
     "execution_count": 70,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 15992 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "loglog(k,pk_ref[0])\n",
    "loglog(k,pki[0])\n",
    "loglog(k,pk[0],'--')"
   ]
  },
  {
   "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.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
