{
 "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": "markdown",
   "metadata": {},
   "source": [
    "## Loading CAMELS data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 34/34 [01:48<00:00,  3.19s/it]\n"
     ]
    }
   ],
   "source": [
    "scales = []\n",
    "poss = []\n",
    "vels = []\n",
    "\n",
    "snapshot='/data/CAMELS/Sims/IllustrisTNG_DM/CV_0/snap_000.hdf5'\n",
    "\n",
    "header   = readgadget.header(snapshot)\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",
    "# 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",
    "    \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 * 32).reshape([256,256,256,3])[2::8,2::8,2::8,:].reshape([-1,3])\n",
    "    vel = (vel / 100 * (1./(1+redshift)) / BoxSize*32).reshape([256,256,256,3])[2::8,2::8,2::8,:].reshape([-1,3])\n",
    "    \n",
    "    scales.append((1./(1+redshift)))\n",
    "    poss.append(pos)\n",
    "    vels.append(vel)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Runing PM simulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "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, compensate_cic\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\n",
    "from jaxpm.utils import power_spectrum\n",
    "\n",
    "rng_seq = hk.PRNGSequence(1)\n",
    "\n",
    "mesh_shape= [32, 32, 32]\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": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run the simulation with Particle-Mesh scheme\n",
    "resi = odeint(make_ode_fn(mesh_shape), [poss[0], vels[0]], jnp.array(scales), cosmo, rtol=1e-5, atol=1e-5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 88848 characters]\n",
      "text/plain": [
       "<Figure size 720x720 with 16 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# High res simulation\n",
    "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": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 86112 characters]\n",
      "text/plain": [
       "<Figure size 720x720 with 16 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# JaxPM without correction\n",
    "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][::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": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/local/home/flanusse/.local/lib/python3.8/site-packages/jax/_src/lax/lax.py:488: ComplexWarning: Casting complex values to real discards the imaginary part\n",
      "  return _convert_element_type(operand, new_dtype, weak_type=False)\n"
     ]
    }
   ],
   "source": [
    "k, pk_ref = power_spectrum(\n",
    "      compensate_cic(cic_paint(jnp.zeros(mesh_shape), poss[-1])),\n",
    "      boxsize=np.array([25.] * 3),\n",
    "      kmin=np.pi / 25.,\n",
    "      dk=2 * np.pi / 25.)\n",
    "\n",
    "k, pk_i = power_spectrum(\n",
    "      compensate_cic(cic_paint(jnp.zeros(mesh_shape), resi[0][-1])),\n",
    "      boxsize=np.array([25.] * 3),\n",
    "      kmin=np.pi / 25.,\n",
    "      dk=2 * np.pi / 25.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0, 0.5, '$P(k)$')"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 17320 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "loglog(k,pk_ref, label='N-body')\n",
    "loglog(k,pk_i, label='JaxPM without correction')\n",
    "legend()\n",
    "plt.xlabel(r\"$k$ [$h \\ \\mathrm{Mpc}^{-1}$]\")\n",
    "plt.ylabel(r\"$P(k)$\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Implementation of a plain Hamiltonian Graph Neural Network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We start by computing a neighboorhood to compute PP interactions \n",
    "# Could we see if we could get a graph on this\n",
    "from sklearn.neighbors import radius_neighbors_graph\n",
    "\n",
    "A = radius_neighbors_graph(poss[-1], 1.5, mode='distance').tocoo()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "import jraph\n",
    "from frozendict import frozendict\n",
    "\n",
    "# Initialize graph structure\n",
    "static_graph = jraph.GraphsTuple(\n",
    "    n_node=np.asarray([A.shape[0]]),\n",
    "    n_edge=np.asarray([len(A.data)]),\n",
    "    nodes={\n",
    "        'position': poss[0],\n",
    "        'momentum': vels[0],\n",
    "    },\n",
    "    senders=A.row,\n",
    "    receivers=A.col,\n",
    "    edges={'grav_potential':jnp.zeros(len(A.data))},\n",
    "    globals={}\n",
    ")\n",
    "\n",
    "# Tell tree_util how to navigate frozendicts.\n",
    "jax.tree_util.register_pytree_node(\n",
    "    frozendict,\n",
    "    flatten_func=lambda s: (tuple(s.values()), tuple(s.keys())),\n",
    "    unflatten_func=lambda k, xs: frozendict(zip(k, xs)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<32768x32768 sparse matrix of type '<class 'numpy.float64'>'\n",
       "\twith 8060958 stored elements in COOrdinate format>"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "A"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 6756 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "hist(A.data,100);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Note**: we are not using a fully connected graph here, so we won't get the long range interactions, only the small scales"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "def update_edge_fn(edges, senders, receivers, globals_):\n",
    "    del globals_\n",
    "    # Models particle-particle interactions contribution to grav potential\n",
    "    distance = jnp.linalg.norm(senders[\"position\"] - receivers[\"position\"], axis=1)\n",
    "    grav_potential_per_edge = - 1. /(distance + 0.1) # The 0.1 is to soften the forces on small scales\n",
    "    return frozendict({\"grav_potential\": grav_potential_per_edge})\n",
    "\n",
    "def update_node_fn(nodes, sent_edges, received_edges, globals_):\n",
    "    del sent_edges, received_edges, globals_\n",
    "    \n",
    "    # Computes momentum\n",
    "    momentum_norm = jnp.linalg.norm(nodes[\"momentum\"],axis=1)\n",
    "    kinetic_energy = 0.5 * momentum_norm ** 2 \n",
    "    return frozendict({\"kinetic_energy\": kinetic_energy})\n",
    "\n",
    "def update_global_fn(nodes, edges, globals_):\n",
    "    del globals_\n",
    "    # At this point we will receive node and edge features aggregated (summed)\n",
    "    # for all nodes and edges in each graph.\n",
    "    hamiltonian_per_graph =  nodes[\"kinetic_energy\"] + edges[\"grav_potential\"]\n",
    "    return frozendict({\"hamiltonian\": hamiltonian_per_graph})\n",
    "\n",
    "# Create the Hamiltonian Graph\n",
    "hamiltonian_gnn = jraph.GraphNetwork(\n",
    "                      update_edge_fn=update_edge_fn,\n",
    "                      update_node_fn=update_node_fn,\n",
    "                      update_global_fn=update_global_fn)\n",
    "\n",
    "# Function that computes the hamiltonian for input position and momentum\n",
    "def hamiltonian_from_state_fn(position, momentum):\n",
    "    # Update variables in graph\n",
    "    graph = static_graph._replace(nodes={\n",
    "        'position': position,\n",
    "        'momentum': momentum})\n",
    "    output_graph = hamiltonian_gnn(graph)\n",
    "    return output_graph.globals[\"hamiltonian\"].sum()\n",
    "\n",
    "# Computes the derivatives of the Hamiltonian\n",
    "hamiltonian_gradients_fn = jax.grad(hamiltonian_from_state_fn, argnums=[0, 1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "@jax.jit\n",
    "def gnn_nbody_ode(state, a, cosmo):\n",
    "    \"\"\"\n",
    "    state is a tuple (position, velocities)\n",
    "    \"\"\"\n",
    "    pos, vel = state\n",
    "    \n",
    "    # Take the derivatives against position and momentum of the hamiltonian of the system\n",
    "    dh_dposition, dh_dmomentum = hamiltonian_gradients_fn(pos, vel)\n",
    "    \n",
    "    # Hamilton equations\n",
    "    dpos_da = 1. / (a**3 * jnp.sqrt(jc.background.Esqr(cosmo, a))) * dh_dmomentum\n",
    "    dvel_da = - 1.5 * cosmo.Omega_m / (a**2 * jnp.sqrt(jc.background.Esqr(cosmo, a))) * dh_dposition\n",
    "\n",
    "    return dpos_da, dvel_da"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "res = odeint(gnn_nbody_ode, [poss[0], vels[0]], jnp.array(scales), cosmo, rtol=1e-4, atol=1e-4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 97372 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][::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": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/local/home/flanusse/.local/lib/python3.8/site-packages/jax/_src/lax/lax.py:488: ComplexWarning: Casting complex values to real discards the imaginary part\n",
      "  return _convert_element_type(operand, new_dtype, weak_type=False)\n"
     ]
    }
   ],
   "source": [
    "k, pk_c = power_spectrum(\n",
    "      compensate_cic(cic_paint(jnp.zeros(mesh_shape), res[0][-1])),\n",
    "      boxsize=np.array([25.] * 3),\n",
    "      kmin=np.pi / 25.,\n",
    "      dk=2 * np.pi / 25.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0, 0.5, '$P(k)$')"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 22024 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "loglog(k, pk_ref, label='N-body')\n",
    "loglog(k, pk_i, label='JaxPM Pure PM')\n",
    "loglog(k, pk_c, '--', label='JaxPM Hamiltonian GNN')\n",
    "legend()\n",
    "plt.xlabel(r\"$k$ [$h \\ \\mathrm{Mpc}^{-1}$]\")\n",
    "plt.ylabel(r\"$P(k)$\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Try to combine large scale PM forces with PP interactions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "from jaxpm.nn import NeuralSplineFourierFilter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Instantiate the neural network\n",
    "model = hk.without_apply_rng(hk.transform(lambda x,a : NeuralSplineFourierFilter(n_knots=16, latent_size=32)(x,a)))\n",
    "params = model.init(next(rng_seq), jnp.zeros([64]), jnp.ones([1]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "metadata": {},
   "outputs": [],
   "source": [
    "def hamiltonian_from_state_fn(position, momentum, params):\n",
    "    def update_edge_fn(edges, senders, receivers, globals_):\n",
    "        del globals_\n",
    "        # Models particle-particle interactions contribution to grav potential\n",
    "        distance = jnp.linalg.norm(senders[\"position\"] - receivers[\"position\"], axis=1)\n",
    "        grav_potential_per_edge = - 1. / (distance + 0.5)\n",
    "        return frozendict({\"grav_potential\": grav_potential_per_edge})\n",
    "\n",
    "    def update_node_fn(nodes, sent_edges, received_edges, globals_):\n",
    "        del sent_edges, received_edges, globals_\n",
    "\n",
    "        # Compute gravitational potential by FFT\n",
    "        kvec = fftk(mesh_shape)\n",
    "        delta_k = jnp.fft.rfftn(cic_paint(jnp.zeros(mesh_shape), nodes[\"position\"]))\n",
    "        pot_k = - delta_k * laplace_kernel(kvec) * longrange_kernel(kvec, r_split=0)\n",
    "        \n",
    "        # Apply a correction filter to the potential, will take care of merging PM forces with PP\n",
    "        kk = jnp.sqrt(sum((ki/pi)**2 for ki in kvec))/2\n",
    "        pot_k = pot_k * model.apply(params, kk, jnp.ones([1]))\n",
    "    \n",
    "        grav_potential = 0.5 * (1 + cic_read(jnp.fft.irfftn(pot_k), nodes[\"position\"]))\n",
    " \n",
    "        # Computes momentum\n",
    "        momentum_norm = jnp.linalg.norm(nodes[\"momentum\"],axis=1)\n",
    "        kinetic_energy = 0.5 * momentum_norm ** 2 \n",
    "        return frozendict({\"kinetic_energy\": kinetic_energy, \"grav_potential\":grav_potential})\n",
    "\n",
    "    def update_global_fn(nodes, edges, globals_):\n",
    "        del globals_\n",
    "        # At this point we will receive node and edge features aggregated (summed)\n",
    "        # for all nodes and edges in each graph.\n",
    "        hamiltonian_per_graph = nodes['grav_potential'] + nodes[\"kinetic_energy\"] + edges[\"grav_potential\"]\n",
    "        return frozendict({\"hamiltonian\": hamiltonian_per_graph})\n",
    "\n",
    "    hamiltonian_gnn = jraph.GraphNetwork(\n",
    "                          update_edge_fn=update_edge_fn,\n",
    "                          update_node_fn=update_node_fn,\n",
    "                          update_global_fn=update_global_fn)\n",
    "\n",
    "\n",
    "    graph = static_graph._replace(nodes={\n",
    "            'position': position,\n",
    "            'momentum': momentum})\n",
    "    output_graph = hamiltonian_gnn(graph)\n",
    "    return output_graph.globals[\"hamiltonian\"].sum()\n",
    "\n",
    "hamiltonian_gradients_fn = jax.grad(hamiltonian_from_state_fn, argnums=[0, 1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "metadata": {},
   "outputs": [],
   "source": [
    "@jax.jit\n",
    "def gnn_nbody_ode(state, a, cosmo, params):\n",
    "    \"\"\"\n",
    "    state is a tuple (position, velocities)\n",
    "    \"\"\"\n",
    "    pos, vel = state\n",
    "            \n",
    "    # Take the derivatives against position and momentum of the hamiltonian of the system\n",
    "    dh_dposition, dh_dmomentum = hamiltonian_gradients_fn(pos, vel, params)\n",
    "    \n",
    "    # Hamilton equations\n",
    "    dpos_da = 1. / (a**3 * jnp.sqrt(jc.background.Esqr(cosmo, a))) * dh_dmomentum\n",
    "    dvel_da = - 1.5 * cosmo.Omega_m / (a**2 * jnp.sqrt(jc.background.Esqr(cosmo, a))) * dh_dposition\n",
    "\n",
    "    return dpos_da, dvel_da"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Precomputing a few data stuff\n",
    "ref_pos = jnp.stack(poss, axis=0)\n",
    "ref_vel = jnp.stack(vels, axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 95,
   "metadata": {},
   "outputs": [],
   "source": [
    "@jax.jit\n",
    "def loss_fn(params):\n",
    "  res = odeint(gnn_nbody_ode, [poss[0], vels[0]], jnp.array(scales), cosmo, params, rtol=1e-4, atol=1e-4) \n",
    "  distance = jnp.sum((res[0] - ref_pos)**2, axis=-1)\n",
    "  w = jnp.where(jnp.sqrt(distance) < 16, distance, 0.)\n",
    "\n",
    "  # Optional lines to include velocity in the loss\n",
    "#   vel = jnp.sum((res[1] - ref_vel)**2, axis=-1)\n",
    "#   wv = jnp.where(jnp.sqrt(distance) < 8, vel, 0.)\n",
    "    \n",
    "  return jnp.mean(w) #+ 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": 96,
   "metadata": {},
   "outputs": [],
   "source": [
    "import optax\n",
    "learning_rate=0.002\n",
    "optimizer = optax.adam(learning_rate)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 97,
   "metadata": {},
   "outputs": [],
   "source": [
    "losses = []\n",
    "opt_state = optimizer.init(params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 99,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/200 [00:00<?, ?it/s]\n"
     ]
    },
    {
     "ename": "RuntimeError",
     "evalue": "INTERNAL: Failed to load in-memory CUBIN: CUDA_ERROR_OUT_OF_MEMORY: out of memory",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-99-768b1d51abff>\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;36m200\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;31m[... skipping hidden 7 frame]\u001b[0m\n",
      "\u001b[0;31mRuntimeError\u001b[0m: INTERNAL: Failed to load in-memory CUBIN: CUDA_ERROR_OUT_OF_MEMORY: out of memory"
     ]
    }
   ],
   "source": [
    "for step in tqdm(range(200)):\n",
    "  l, params, opt_state = update(params, opt_state)\n",
    "  losses.append(l)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x7fb8d44bdd30>]"
      ]
     },
     "execution_count": 84,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 12732 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot(losses)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "metadata": {},
   "outputs": [],
   "source": [
    "res = odeint(gnn_nbody_ode, [poss[0], vels[0]], jnp.array(scales), cosmo, params, rtol=1e-4, atol=1e-4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 97548 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][::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": 87,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/local/home/flanusse/.local/lib/python3.8/site-packages/jax/_src/lax/lax.py:488: ComplexWarning: Casting complex values to real discards the imaginary part\n",
      "  return _convert_element_type(operand, new_dtype, weak_type=False)\n"
     ]
    }
   ],
   "source": [
    "k, pk_ref = power_spectrum(\n",
    "      compensate_cic(cic_paint(jnp.zeros(mesh_shape), poss[-1])),\n",
    "      boxsize=np.array([25.] * 3),\n",
    "      kmin=np.pi / 25.,\n",
    "      dk=2 * np.pi / 25.)\n",
    "\n",
    "k, pk_c = power_spectrum(\n",
    "      compensate_cic(cic_paint(jnp.zeros(mesh_shape), res[0][-1])),\n",
    "      boxsize=np.array([25.] * 3),\n",
    "      kmin=np.pi / 25.,\n",
    "      dk=2 * np.pi / 25.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0, 0.5, '$P(k)$')"
      ]
     },
     "execution_count": 88,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 22280 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "loglog(k, pk_ref, label='N-body')\n",
    "loglog(k, pk_i, label='JaxPM Pure PM')\n",
    "loglog(k, pk_c, '--', label='JaxPM Hamiltonian GNN')\n",
    "legend()\n",
    "plt.xlabel(r\"$k$ [$h \\ \\mathrm{Mpc}^{-1}$]\")\n",
    "plt.ylabel(r\"$P(k)$\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 20228 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "col = cm.viridis(np.linspace(0.,1.0,len(scales[::2])))   \n",
    "\n",
    "kvals = jnp.logspace(-2.,0,100)*sqrt(3)\n",
    "\n",
    "for i, a in enumerate(scales[::2]): \n",
    "    semilogx(kvals, model.apply(params, kvals , jnp.atleast_1d(1)), color=col[i])\n",
    "\n",
    "sm = plt.cm.ScalarMappable(cmap=\"viridis\", norm=plt.Normalize(vmin=0., vmax=1))\n",
    "plt.colorbar(sm, label='a')\n",
    "\n",
    "#xlim(kvals[0], 1)\n",
    "#ylim(0,5)\n",
    "grid()"
   ]
  },
  {
   "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
}
