{
 "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 *\n",
    "from jaxpm.growth 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": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2022-02-14 00:21:13.509002: E tensorflow/stream_executor/cuda/cuda_blas.cc:232] failed to create cublas handle: CUBLAS_STATUS_NOT_INITIALIZED\n",
      "2022-02-14 00:21:13.509025: E tensorflow/stream_executor/cuda/cuda_blas.cc:234] Failure to initialize cublas may be due to OOM (cublas needs some free memory when you initialize it, and your deep-learning framework may have preallocated more than its fair share), or may be because this binary was not built with support for the GPU in your machine.\n",
      "2022-02-14 00:21:13.509034: W tensorflow/stream_executor/stream.cc:1260] attempting to perform BLAS operation using StreamExecutor without BLAS support\n"
     ]
    },
    {
     "ename": "InternalError",
     "evalue": "BlasScal failed : in.shape=[1,100,100,100] [Op:IFFT3D]",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mInternalError\u001b[0m                             Traceback (most recent call last)",
      "Input \u001b[0;32mIn [6]\u001b[0m, in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      6\u001b[0m pk \u001b[38;5;241m=\u001b[39m linear_matter_power(cosmology, k)\n\u001b[1;32m      7\u001b[0m pk_fun \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mlambda\u001b[39;00m x: tf\u001b[38;5;241m.\u001b[39mcast(tf\u001b[38;5;241m.\u001b[39mreshape(interpolate\u001b[38;5;241m.\u001b[39minterp_tf(tf\u001b[38;5;241m.\u001b[39mreshape(tf\u001b[38;5;241m.\u001b[39mcast(x, tf\u001b[38;5;241m.\u001b[39mfloat32), [\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m]), k, pk), x\u001b[38;5;241m.\u001b[39mshape), tf\u001b[38;5;241m.\u001b[39mcomplex64)\n\u001b[0;32m----> 8\u001b[0m initial_conditions \u001b[38;5;241m=\u001b[39m \u001b[43mflowpm\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mlinear_field\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnc\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m      9\u001b[0m \u001b[43m                                      \u001b[49m\u001b[43mbox_size\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\n\u001b[1;32m     10\u001b[0m \u001b[43m                                      \u001b[49m\u001b[43mpk_fun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m         \u001b[49m\n\u001b[1;32m     11\u001b[0m \u001b[43m                                      \u001b[49m\u001b[43mbatch_size\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mbatch_size\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     13\u001b[0m initial_state \u001b[38;5;241m=\u001b[39m flowpm\u001b[38;5;241m.\u001b[39mlpt_init(cosmology, initial_conditions, \u001b[38;5;241m0.1\u001b[39m, order\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1\u001b[39m)\n",
      "File \u001b[0;32m~/repo/flowpm/flowpm/tfpm.py:62\u001b[0m, in \u001b[0;36mlinear_field\u001b[0;34m(nc, boxsize, pk, kvec, batch_size, seed, dtype, name)\u001b[0m\n\u001b[1;32m     59\u001b[0m whitec \u001b[38;5;241m=\u001b[39m white_noise(nc, batch_size\u001b[38;5;241m=\u001b[39mbatch_size, seed\u001b[38;5;241m=\u001b[39mseed, \u001b[38;5;28mtype\u001b[39m\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mcomplex\u001b[39m\u001b[38;5;124m'\u001b[39m)\n\u001b[1;32m     60\u001b[0m lineark \u001b[38;5;241m=\u001b[39m tf\u001b[38;5;241m.\u001b[39mmultiply(whitec, (pkmesh \u001b[38;5;241m/\u001b[39m\n\u001b[1;32m     61\u001b[0m                                (boxsize[\u001b[38;5;241m0\u001b[39m] \u001b[38;5;241m*\u001b[39m boxsize[\u001b[38;5;241m1\u001b[39m] \u001b[38;5;241m*\u001b[39m boxsize[\u001b[38;5;241m2\u001b[39m]))\u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39m\u001b[38;5;241m0.5\u001b[39m)\n\u001b[0;32m---> 62\u001b[0m linear \u001b[38;5;241m=\u001b[39m \u001b[43mc2r3d\u001b[49m\u001b[43m(\u001b[49m\u001b[43mlineark\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mnorm\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mnc\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43m \u001b[49m\u001b[43mnc\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43m \u001b[49m\u001b[43mnc\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;241;43m2\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mname\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mname\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdtype\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdtype\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     63\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m linear\n",
      "File \u001b[0;32m~/repo/flowpm/flowpm/utils.py:271\u001b[0m, in \u001b[0;36mc2r3d\u001b[0;34m(cfield, norm, dtype, name)\u001b[0m\n\u001b[1;32m    268\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    269\u001b[0m   norm \u001b[38;5;241m=\u001b[39m tf\u001b[38;5;241m.\u001b[39mcast(norm, dtype)\n\u001b[1;32m    270\u001b[0m rfield \u001b[38;5;241m=\u001b[39m tf\u001b[38;5;241m.\u001b[39mmultiply(\n\u001b[0;32m--> 271\u001b[0m     tf\u001b[38;5;241m.\u001b[39mcast(\u001b[43mtf\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msignal\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mifft3d\u001b[49m\u001b[43m(\u001b[49m\u001b[43mcfield\u001b[49m\u001b[43m)\u001b[49m, dtype), norm, name\u001b[38;5;241m=\u001b[39mname)\n\u001b[1;32m    272\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m rfield\n",
      "File \u001b[0;32m/usr/lib/python3.10/site-packages/tensorflow/python/ops/gen_spectral_ops.py:809\u001b[0m, in \u001b[0;36mifft3d\u001b[0;34m(input, name)\u001b[0m\n\u001b[1;32m    807\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m _result\n\u001b[1;32m    808\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m _core\u001b[38;5;241m.\u001b[39m_NotOkStatusException \u001b[38;5;28;01mas\u001b[39;00m e:\n\u001b[0;32m--> 809\u001b[0m   \u001b[43m_ops\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mraise_from_not_ok_status\u001b[49m\u001b[43m(\u001b[49m\u001b[43me\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mname\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    810\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m _core\u001b[38;5;241m.\u001b[39m_FallbackException:\n\u001b[1;32m    811\u001b[0m   \u001b[38;5;28;01mpass\u001b[39;00m\n",
      "File \u001b[0;32m/usr/lib/python3.10/site-packages/tensorflow/python/framework/ops.py:7186\u001b[0m, in \u001b[0;36mraise_from_not_ok_status\u001b[0;34m(e, name)\u001b[0m\n\u001b[1;32m   7184\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mraise_from_not_ok_status\u001b[39m(e, name):\n\u001b[1;32m   7185\u001b[0m   e\u001b[38;5;241m.\u001b[39mmessage \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m (\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m name: \u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;241m+\u001b[39m name \u001b[38;5;28;01mif\u001b[39;00m name \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m-> 7186\u001b[0m   \u001b[38;5;28;01mraise\u001b[39;00m core\u001b[38;5;241m.\u001b[39m_status_to_exception(e) \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;28mNone\u001b[39m\n",
      "\u001b[0;31mInternalError\u001b[0m: BlasScal failed : in.shape=[1,100,100,100] [Op:IFFT3D]"
     ]
    }
   ],
   "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, order=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "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": 8,
   "id": "7d4780a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "final_state = solve_tf(initial_state)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "6a94dae9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.3 s ± 2.59 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": 9,
   "id": "e5148fed",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x7fd56040fdc0>"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 18616 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": "cf61ff87",
   "metadata": {},
   "outputs": [],
   "source": [
    "mesh_shape = nc\n",
    "kvec = fftk(mesh_shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "bc83bc0c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# We need the vel and forces from lpt\n",
    "def lpt1(delta_k, pos):   \n",
    "    pot_k = delta_k * laplace_kernel(kvec)\n",
    "    return jnp.stack([cic_read(jnp.fft.irfftn(gradient_kernel(kvec, i)*pot_k), pos) for i in range(3)],axis=-1)\n",
    "    \n",
    "def lpt(cosmo, linear, pos, a):\n",
    "    delta_k = jnp.fft.rfftn(linear)\n",
    "    a = jnp.atleast_1d(a)\n",
    "    lpt_dx = lpt1(delta_k, pos)\n",
    "    dx = growth_factor(cosmo, a) * lpt_dx\n",
    "    p = a**2 * growth_rate(cosmo, a) * jnp.sqrt(jc.background.Esqr(cosmo, a)) * dx\n",
    "    f = a**2 * jnp.sqrt(jc.background.Esqr(cosmo, a)) * dGfa(cosmo, a) / growth_factor(cosmo, a) * dx\n",
    "    return lpt_dx, dx, p, f"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "ab85a109",
   "metadata": {},
   "outputs": [],
   "source": [
    "pos0 = jnp.stack(jnp.meshgrid(jnp.arange(100), jnp.arange(100), jnp.arange(100)), axis=-1).reshape([-1,3])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "02141439",
   "metadata": {},
   "outputs": [],
   "source": [
    "lpt_dx, x, p,f = lpt(jc.Planck15(), initial_conditions[0].numpy(), pos0, 1.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "4225880e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.colorbar.Colorbar at 0x7fce887a5c30>"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 69432 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), pos0+x).sum(axis=0)); colorbar()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "167d193f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the ODE\n",
    "def f(state, a, cosmo, lpt_dx):\n",
    "    # Extracts positions and velocity at a given point in\n",
    "    # the simulation\n",
    "    res_pos, res_vel = state\n",
    "    a =  jnp.atleast_1d(a)\n",
    "    pos = pos0 + growth_factor(cosmo, a) * lpt_dx + res_pos\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))) * res_vel\n",
    "    \n",
    "    # Computes the update of velocity (kick)\n",
    "    dvel = 1. / (a**2 * jnp.sqrt(jc.background.Esqr(cosmo, a))) * forces - dGfa(cosmo, a) * lpt_dx\n",
    "    \n",
    "    return dpos, dvel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "ce83d3b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import diffrax\n",
    "from diffrax import diffeqsolve, ODETerm, Dopri5, LeapfrogMidpoint, PIDController"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "d4e7b055",
   "metadata": {},
   "outputs": [],
   "source": [
    "cosmo = jc.Planck15()\n",
    "\n",
    "def f_diffrax(t, state, args):\n",
    "    return jnp.stack(f(state, t, cosmo, lpt_dx),axis=0)\n",
    "\n",
    "\n",
    "term = ODETerm(f_diffrax)\n",
    "solver = Dopri5()\n",
    "stepsize_controller = PIDController(rtol=1e-5, atol=1e-5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "37285eb4",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/francois/.local/lib/python3.10/site-packages/jax/_src/numpy/lax_numpy.py:6655: UserWarning: Explicitly requested dtype <class 'jax._src.numpy.lax_numpy.int64'> requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/google/jax#current-gotchas for more.\n",
      "  lax._check_user_dtype_supported(dtype, \"astype\")\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(1000000, 8)\n",
      "(1000000, 8)\n",
      "(1000000, 8)\n",
      "(1000000, 8)\n",
      "(1000000, 8)\n",
      "(1000000, 8)\n",
      "(1000000, 8)\n",
      "(1000000, 8)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2022-02-14 00:10:24.488749: E external/org_tensorflow/tensorflow/compiler/xla/pjrt/pjrt_stream_executor_client.cc:2124] Execution of replica 0 failed: INTERNAL: Failed to load in-memory CUBIN: CUDA_ERROR_OUT_OF_MEMORY: out of memory\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)",
      "Input \u001b[0;32mIn [31]\u001b[0m, in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      1\u001b[0m init_state \u001b[38;5;241m=\u001b[39m jnp\u001b[38;5;241m.\u001b[39mstack([jnp\u001b[38;5;241m.\u001b[39mzeros((\u001b[38;5;241m100\u001b[39m,\u001b[38;5;241m100\u001b[39m,\u001b[38;5;241m100\u001b[39m,\u001b[38;5;241m3\u001b[39m))\u001b[38;5;241m.\u001b[39mreshape((\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m,\u001b[38;5;241m3\u001b[39m)),\n\u001b[1;32m      2\u001b[0m               jnp\u001b[38;5;241m.\u001b[39mzeros((\u001b[38;5;241m100\u001b[39m,\u001b[38;5;241m100\u001b[39m,\u001b[38;5;241m100\u001b[39m,\u001b[38;5;241m3\u001b[39m))\u001b[38;5;241m.\u001b[39mreshape((\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m,\u001b[38;5;241m3\u001b[39m))], axis\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m0\u001b[39m)\n\u001b[0;32m----> 4\u001b[0m solution \u001b[38;5;241m=\u001b[39m \u001b[43mdiffeqsolve\u001b[49m\u001b[43m(\u001b[49m\u001b[43mterm\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43msolver\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mt0\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m0.1\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mt1\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m1.\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdt0\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m0.1\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m      5\u001b[0m \u001b[43m                       \u001b[49m\u001b[43my0\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43minit_state\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mstepsize_controller\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mstepsize_controller\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/.local/lib/python3.10/site-packages/equinox/jit.py:90\u001b[0m, in \u001b[0;36mfilter_jit.<locals>.fun_wrapper\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m     88\u001b[0m static_leaves, static_treedef \u001b[38;5;241m=\u001b[39m jax\u001b[38;5;241m.\u001b[39mtree_flatten((static_args, static_kwargs))\n\u001b[1;32m     89\u001b[0m static_leaves \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mtuple\u001b[39m(static_leaves)\n\u001b[0;32m---> 90\u001b[0m dynamic_out, static_out \u001b[38;5;241m=\u001b[39m \u001b[43m_filter_jit_cache\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mjitkwargs\u001b[49m\u001b[43m)\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m     91\u001b[0m \u001b[43m    \u001b[49m\u001b[43mstatic_leaves\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     92\u001b[0m \u001b[43m    \u001b[49m\u001b[43mstatic_treedef\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     93\u001b[0m \u001b[43m    \u001b[49m\u001b[43mdynamic_args\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     94\u001b[0m \u001b[43m    \u001b[49m\u001b[43mdynamic_kwargs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     95\u001b[0m \u001b[43m    \u001b[49m\u001b[43mfilter_spec_return\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     96\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     97\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m combine(dynamic_out, static_out\u001b[38;5;241m.\u001b[39mvalue)\n",
      "    \u001b[0;31m[... skipping hidden 6 frame]\u001b[0m\n",
      "File \u001b[0;32m~/.local/lib/python3.10/site-packages/jax/_src/dispatch.py:444\u001b[0m, in \u001b[0;36m_execute_compiled\u001b[0;34m(name, compiled, output_buffer_counts, result_handlers, kept_var_idx, *args)\u001b[0m\n\u001b[1;32m    441\u001b[0m device, \u001b[38;5;241m=\u001b[39m compiled\u001b[38;5;241m.\u001b[39mlocal_devices()\n\u001b[1;32m    442\u001b[0m input_bufs \u001b[38;5;241m=\u001b[39m util\u001b[38;5;241m.\u001b[39mflatten(\n\u001b[1;32m    443\u001b[0m     device_put(x, device) \u001b[38;5;28;01mfor\u001b[39;00m i, x \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28menumerate\u001b[39m(args) \u001b[38;5;28;01mif\u001b[39;00m i \u001b[38;5;129;01min\u001b[39;00m kept_var_idx)\n\u001b[0;32m--> 444\u001b[0m out_bufs \u001b[38;5;241m=\u001b[39m \u001b[43mcompiled\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mexecute\u001b[49m\u001b[43m(\u001b[49m\u001b[43minput_bufs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    445\u001b[0m check_special(name, out_bufs)\n\u001b[1;32m    446\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m output_buffer_counts \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n",
      "\u001b[0;31mRuntimeError\u001b[0m: INTERNAL: Failed to load in-memory CUBIN: CUDA_ERROR_OUT_OF_MEMORY: out of memory"
     ]
    }
   ],
   "source": [
    "init_state = jnp.stack([jnp.zeros((100,100,100,3)).reshape((-1,3)),\n",
    "              jnp.zeros((100,100,100,3)).reshape((-1,3))], axis=0)\n",
    "\n",
    "solution = diffeqsolve(term, solver, t0=0.1, t1=1., dt0=0.1,\n",
    "                       y0=init_state, stepsize_controller=stepsize_controller)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "9fc54299",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "421 ms ± 2.28 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%timeit solution = diffeqsolve(term, solver, t0=0.1, t1=1., dt0=0.1, y0=init_state+0, stepsize_controller=stepsize_controller)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "325ee24c",
   "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([[[[-0.05698265, -0.4233138 , -0.23527317],\n",
       "               [-0.4048839 , -0.18278353,  0.10613398],\n",
       "               [-0.39811632,  0.36003193,  0.5207157 ],\n",
       "               ...,\n",
       "               [-0.12596521,  0.04650298,  0.12780993],\n",
       "               [ 0.10532238,  0.22021866, -0.00881826],\n",
       "               [ 0.21299808,  0.42117628, -0.05347536]],\n",
       "\n",
       "              [[-0.07113229, -0.40208206, -0.18125573],\n",
       "               [-0.3594675 , -0.14042078,  0.1074788 ],\n",
       "               [-0.29909593,  0.34923548,  0.42418486],\n",
       "               ...,\n",
       "               [-0.2394397 ,  0.07868364,  0.09896102],\n",
       "               [ 0.06142797,  0.22736594,  0.05615562],\n",
       "               [ 0.18805356,  0.38284227, -0.00324537]]]],            dtype=float32, weak_type=True), interpolation=None, stats={'compiled_num_steps': None, 'max_steps': 4096, 'num_accepted_steps': DeviceArray(22, dtype=int32, weak_type=True), 'num_rejected_steps': DeviceArray(6, dtype=int32, weak_type=True), 'num_steps': DeviceArray(28, dtype=int32, weak_type=True)}, result=DeviceArray(0, dtype=int32, weak_type=True), solver_state=None, controller_state=None, made_jump=None)"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "f131301b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1, 2, 1000000, 3)"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "solution.ys.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "a5f4d029",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.colorbar.Colorbar at 0x7f22ef156230>"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 59320 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "resx = solution.ys[0,0]+pos0+x\n",
    "imshow(cic_paint(jnp.zeros(mesh_shape), resx).sum(axis=0),cmap='gist_stern'); colorbar()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "654d9c69",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.colorbar.Colorbar at 0x7ff54ab2b1c0>"
      ]
     },
     "execution_count": 57,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 17836 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "imshow((flowpm.cic_paint(tf.zeros([1]+mesh_shape), final_state[0]).numpy()[0]).sum(axis=0)); colorbar()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "id": "80da1b98",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.colorbar.Colorbar at 0x7ff546f8ceb0>"
      ]
     },
     "execution_count": 58,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 12648 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) - \n",
    "        flowpm.cic_paint(tf.zeros([1]+mesh_shape), final_state[0]).numpy()[0]).sum(axis=0)); colorbar()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32357464",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "547af8a9",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f5cac4d",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "aa14a004",
   "metadata": {},
   "outputs": [],
   "source": [
    "from jax.experimental.ode import odeint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "973ca031",
   "metadata": {},
   "outputs": [],
   "source": [
    "init_state = [jnp.zeros((64,64,64,3)).reshape((-1,3)),\n",
    "              jnp.zeros((64,64,64,3)).reshape((-1,3))]\n",
    "\n",
    "@jax.jit\n",
    "def solve_ode(init_state):\n",
    "    cosmo = jc.Planck15()\n",
    "    return odeint(f, init_state, \n",
    "             jnp.linspace(0.1,1.0,10), \n",
    "             cosmo,\n",
    "             lpt_dx, rtol=1e-5, atol=1e-5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "e801c3eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "res = solve_ode(init_state)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "de770afb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "765 ms ± 7.45 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": 22,
   "id": "4bdfe23f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.colorbar.Colorbar at 0x7fd5603ba4d0>"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 21884 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "resx = res[0][-1]+pos0+x\n",
    "imshow(cic_paint(jnp.zeros(mesh_shape), resx).sum(axis=0)); colorbar()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "9abade92",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.colorbar.Colorbar at 0x7fd53e682e00>"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 21000 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "imshow((flowpm.cic_paint(tf.zeros([1]+mesh_shape), final_state[0]).numpy()[0]).sum(axis=0),cmap='gist_stern'); colorbar()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "a7f8b5d5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.colorbar.Colorbar at 0x7fd5382000d0>"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 19788 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) - \n",
    "        flowpm.cic_paint(tf.zeros([1]+mesh_shape), final_state[0]).numpy()[0]).sum(axis=0)); colorbar()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "68b72b68",
   "metadata": {},
   "outputs": [],
   "source": [
    "from DifferentiableHOS.pk import power_spectrum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "785818e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "k, pk = power_spectrum(flowpm.cic_paint(tf.zeros([1]+mesh_shape), final_state[0]),\n",
    "                       boxsize=array([64.,64.,64.]),                      \n",
    "                       kmin=0.1,dk=2*np.pi/64.)\n",
    "\n",
    "k, pk_jax = power_spectrum(tf.convert_to_tensor(cic_paint(jnp.zeros(mesh_shape), resx).reshape([1,64,64,64])),\n",
    "                       boxsize=array([64.,64.,64.]),                      \n",
    "                       kmin=0.1,dk=2*np.pi/64.)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "18e3d30b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.legend.Legend at 0x7fd53838db40>"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 12196 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": 28,
   "id": "c2c1ecdb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x7fd530740190>]"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 15160 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": "703c8195",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6921f17",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3083d29",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65116b88",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6bd8cfd",
   "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
}
