JaxPM/dev/JaxPM_ODE.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "2bb083b7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "%pylab is deprecated, use %matplotlib inline and import the required libraries.\n",
      "Populating the interactive namespace from numpy and matplotlib\n"
     ]
    }
   ],
   "source": [
    "%pylab inline\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "458efc98",
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import jax.lax as lax\n",
    "import jax_cosmo as jc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "da00e618",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import flowpm\n",
    "from flowpm.tfpower import linear_matter_power\n",
    "import flowpm.scipy.interpolate as interpolate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "b588c18c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from jaxpm.kernels import *\n",
    "from jaxpm.painting import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "f0cba7a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Below are a few parameters\n",
    "box_size = [100., 100., 100.]   # Transverse comoving size of the simulation volume\n",
    "nc = [100, 100, 100]            # Number of transverse voxels in the simulation volume\n",
    "batch_size = 1                  # Number of simulations to run in parallel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "8dad416f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Instantiates a cosmology with desired parameters\n",
    "cosmology = flowpm.cosmology.Planck15()\n",
    "\n",
    "# Create some initial conditions\n",
    "k = tf.constant(np.logspace(-4, 1, 128), dtype=tf.float32)\n",
    "pk = linear_matter_power(cosmology, k)\n",
    "pk_fun = lambda x: tf.cast(tf.reshape(interpolate.interp_tf(tf.reshape(tf.cast(x, tf.float32), [-1]), k, pk), x.shape), tf.complex64)\n",
    "initial_conditions = flowpm.linear_field(nc,\n",
    "                                      box_size, \n",
    "                                      pk_fun,         \n",
    "                                      batch_size=batch_size)\n",
    "\n",
    "initial_state = flowpm.lpt_init(cosmology, initial_conditions, 0.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "id": "5ad12a1e",
   "metadata": {},
   "outputs": [],
   "source": [
    "@tf.function\n",
    "def solve_tf(init_state):\n",
    "    final_state = flowpm.nbody(cosmology, initial_state, linspace(0.1,1.,40), nc)\n",
    "    return final_state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "id": "7d4780a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "final_state = solve_tf(initial_state)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "6a94dae9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4.89 s ± 12.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%timeit final_state = solve_tf(initial_state)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "e5148fed",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x7f9774507880>"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAPsAAAD7CAYAAACscuKmAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/YYfK9AAAACXBIWXMAAAsTAAALEwEAmpwYAACbG0lEQVR4nO39Xaxlz3Ifhv2q19r7nJn5///3gxSpa1I2ZYCwYwhQZAiRZAUJYcqBIwvmkwTZUUA5Cm4eElt2DFhk8mDkIQAfDMN6CALc2DGYWIilyIIpCIYl4yp8yIMJSbbgD9G0FJO4/LgieXk//jNzztl7rdWVh6rqru7VvfbaZ+bOmes7BQzm7L3XR69e3V1Vv/pVNTEzPsgH+SD//Zfw1A34IB/kg7wb+TDZP8gH+S6RD5P9g3yQ7xL5MNk/yAf5LpEPk/2DfJDvEvkw2T/IB/kukTea7ET0TxPRLxDR3yOin3hbjfogH+SDvH2hx8bZiWgA8N8C+KcA/AqAvwHgn2Pmv/P2mvdBPsgHeVsyvsG5/yMAf4+Z/zsAIKJ/H8CPAehO9iPd8C1eXHcXWv0BgKvP+olI/mAAzOB0bP7v+ntT9SX3r+XaStYWCo3ruAswA1Fbqgsv+eeoJT3XNY9QXk/us7NfqPwgbSMgUH42IrlOjHJYkO95DOBACFMElkXvxeXF7fr2G+s/ovx7fX1tAx9GYCBgiaDI8nuM6f3nY90teeP91TKEsh0xAkuUd/A2yGjFswOtodY7Nr0HG1vMYI544Nc486m+EoA3m+w/AOCX3edfAfD76oOI6IsAvggAt3iO3xf+UHXAtidBgfrHhfxMRAQ6HgEAvCzAsoCXCHBMx3C87gWRDWi7T2SAY/c6qa3DgHBzAwwD6PkzYBx0ApAM+nnRQcngeQafzkCM4HmWl3g4gIjygLJBzizPZm3Z8wxDAA6H/IXeB3qd9Cwcm31Mro9BAXQ8gIYBeHYLenYLhAAeByBG0Ot7udRHz4FxwPy551huRxy/9hrh6y/lneizpME66hCcZ/ntPEnbDiMQBmn/OMq5Dydp5+EIOh6w/PbPYfnoiOHTM8LrB9DpDH71Gojs+nLUSaHPNp1lXLTE+lYaCProBehwkEkfAnD/gPjyFbAsiOdpu+O5cY9L/bu6RPmOKVD6LhwPwCDjiojAywI+nfCfTn+1e703meytVq5GIDN/CcCXAOAT+jxfmtzru+w4PjJwGIDjQQbkdAaHAPAZ8O/Pddbl21J5b53ol46nIYCOR9CL5+AXz/DwD34eyzPRcABw87UTxm/ege5P4PsHOZcIHAJIBz6NIzAEkC5WDADTJJP/UvvrhWkBwGcZ+CHk3yiUz6PP2h18+jvPs0ykZQHuH7KG54i4RFl0AWAcMc4LxnEAvbwDv76T+0XRikmZUXk/XhbRvrYYBenT1O5A8lzzjPCtO4T7CThPoHnJC4VddxhkwRgH8ItnMjHuHoBJj0sWBGQhfYh5QY0MvHqt15BntGdH5O2xtDFO1odyW6E1ruHvx2ph6MNKn14YG28y2X8FwO9wn38QwK9dcwE/sDYbWg9Ae+m1DINMnDmAaKe5a526tajstQiCTqhBF56bI86fGTE9z6bkeDdieD2ATkEnpDwLRRnY8hxBLQqWxSq6F7+nvfXzLQBTAIUd57Su6wees244zI3zCTgeQMwgnUx8OqulpRMsxmy12GRLC5Hcy7s0vIhVRMOg7ZFFjx7OYiXpZ44xv6sgrgYN+X3wGEDzktu1xDyOvFa3Z7SFzRZKf33sUB71grpHAu0bb/peQQF7Hbs3mex/A8APE9HvBPCrAP44gH/+mgv4Vc3+X3VeoGL1Z+ZSe8lJ4GkG7u71JS551avuZ8dvyUq7dVyBYkUOMhhpHIFxBN8cwWNAmBjjCQgTI8yMw6dnhFcPwOkMzLNqX/U1l0UGlZqIrBMDy7LW6vUz7LKAzKduD6j0PP75e33mBltxTgiqHSOwnPOxh1Emq016e0YbqEssr6P35gBQtQCAZrnPPLuFUV2kIBqdDgd5Dy/EjVpeyPtAFNeG7k/A/ADMXFptsVrYrI3WPnN3dOEhqly7R7yXQru78yjk69aKMY/DZdMV8PLoyc7MMxH97wD8VQADgP87M//XV1+n9aAqZBPdaSQyreBekFwjgs/nco3zZs9bnujuhLwgUdYmPA6iSSIjTMD4esEwRYTXJ+D+ATxNYhIC2X935mv6Hthnvjufm4h2r/Yr7WSLry6w9XVWg8/uG7k0t8nArCgTQ/1ufaAm9kGBgNieHKu+AeSzaXzT+hSAoPc7jIjPb8CHgPhsFLDwdhTNfp6Su5B8+FpBNN59Uk56P8aiba7a1pCuQkMew8CQF2Od8Pm5ZEFL7yyN/31W3ptodjDzfwTgP3qTawBtDe/uISu7AVxA03/myLLKAk0TP3d0u2NWPnr+wf0ZE0iF2iS2Nho4dH9CmBfchgAeCOF+As0RdPcg5iZzAlZQg25D1YZqYeu2nyMQ1ay7wnwsJry6SEmLbgCSenJe7A4H0DiCPvlIJp8uarYYJF+8nrReGjgCsPbt8/GcLLnspxMwDuDbI+bP3IAHQfMpMmha1KrSdhiKb1YURPMz00Wc5tsmvr/VNUVaXNDvvwttfaPJfrVQf3Xb0vByrmh4nueMSNeSTLFQfvaXqTRZbYrL+WtXwq5LQwA9uwWNY54Qhm4r+ETzDL67B4WA4U5AOFZQKF25Blu82MJRPNq2mejNutWztcS0aGuAGCZi5mIHqU/30cFI4wi6OSJ+5gX4MGBQvz1ZJtHAr/4CsrIcIkt4rWh71ZZlEStpGPL4OozgmwPOnxnBA+HwakE4RdC0gB7O4POU22NgoFswCFAwrtNHRNkSCARiHVfeT2/028VITmVl0TCI9XIY9dpmwl+2JGp5t5PdyZZJowekP5lZ0N895uwlH4kCyDRnbyWsJknSKkPpUsBpw+R3swBFRIrAB72Z+ufFohIIhAEYysFGx6No12nur+LVc9Y+XXeih5A1Wut8Z5bX92qixsH1iVoEdJpBCwOzC38uMWn1VozaW3Z1ODBNQkAm4u1NOdmmWf33Ibd1moH7M26+PoFHwvB6Ak1RjjWegE1qBexIw1g46LQ4B4mCLNmqS+/OLAHz6xv91VqYaUAbcG6JuUNqMSJkywOARFrg59L22H/Hk72KW1cDuQVGZMJJFReur9wxw4trunvL4Amb5IiEF3gTUgc2q9+JRePWFhpblhSmAbJvl4Cfw0FeloE8xKDjsXRXnt1q+x6Ua8J5Ie+FyDYGmLSb1hqxd17LIrLnqABT6x/TQIgs8XYi8PkMxAWY5i73oWiHgWzw701dkyBtoGEAffwR+DBqaJLFirp/kP4cdGI8nECnM44PJ7n+vOT2BpnYPI4aylO84eZGnu3mKM/y6rX8r9ZkwgaWRRZi7atNheWjJ97d8TwKs3a8MrOFzkJsRJBOyItUfgc2phmIfWvu3Wv2mqQROtrVS8dfXQ34Grmf29cuBmvPZG+2w9DxmP3PyG20PF0/FgsExaiaRcEYa7cSvwDI9YYh+7l7Q3/NNq9NSYnr58VMTOsdYaIenuJJOxXRiV0cO7Vnp6wWfYu8HUYJpdlxs5rxjjxjYTU6nUsrYNQJO5glZuBqyL6xtbllJVZh0Iu8Dbc422JBx4MskMajiJ3FD1hbEl5qAJsiaIPr804nO/lYKRTJxLAdUuqhtvmD/B9IWGvmaxrtMpax4BXQ0yKW+NivXzxM63j2VEeT1lEEWwwSe+2gTDTznUOUlTtGiUsDakKuNWHLRK9xjKbvSxbeCtldGEfQ6SSRjMU9Txqk65BYsbCZP21mMVGmmfKovIc5E1V6ouBi8/n8gjUM4GdHxNsDlucH8EA4EiH4ibAs0odGSKKQQ3THA/jZEXRHgq2YRWFMQ3/v1Bc6qRIIm9/JHqIWWd/c3AhZ6qMXOcQ6zwAe1jhUFRUpqchZu9MwZLcjsrgeHXn3mt1NHhpHmZA2/2NMfkg29avTWxP9Cin8P7vGFgCVTCiHVlcL0MqX7ZnMtXR9NReK26M9UlvbSHZqU/q
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "imshow(flowpm.cic_paint(tf.zeros([1]+nc), final_state[0]).numpy()[0].sum(axis=0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "167d193f",
   "metadata": {},
   "outputs": [],
   "source": [
    "mesh_shape = nc\n",
    "kvec = fftk(mesh_shape)\n",
    "\n",
    "# Define the ODE\n",
    "def f(state, a, cosmo):\n",
    "    # Extracts positions and velocity at a given point in\n",
    "    # the simulation\n",
    "    pos, vel = state\n",
    "    \n",
    "    # Computes the potential given the current positions\n",
    "    delta_k = jnp.fft.rfftn(cic_paint(jnp.zeros(mesh_shape), pos))\n",
    "    pot_k = delta_k * laplace_kernel(kvec) * longrange_kernel(kvec,r_split=0)\n",
    "    forces = jnp.stack([cic_read(jnp.fft.irfftn(gradient_kernel(kvec, i)*pot_k), pos) \n",
    "                      for i in range(3)],axis=-1)\n",
    "    forces = forces * 1.5 * cosmo.Omega_m\n",
    "    \n",
    "    # Computes the update of position (drift)\n",
    "    dpos = 1. / (a**3 * jnp.sqrt(jc.background.Esqr(cosmo, a))) * vel\n",
    "    \n",
    "    # Computes the update of velocity (kick)\n",
    "    dvel = 1. / (a**2 * jnp.sqrt(jc.background.Esqr(cosmo, a))) * forces\n",
    "    \n",
    "    return dpos, dvel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "aa14a004",
   "metadata": {},
   "outputs": [],
   "source": [
    "from jax.experimental.ode import odeint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "973ca031",
   "metadata": {},
   "outputs": [],
   "source": [
    "init_state = [initial_state[0,0].numpy(),\n",
    "              initial_state[1,0].numpy()]\n",
    "\n",
    "@jax.jit\n",
    "def solve_ode(init_state):\n",
    "    return odeint(f, init_state, \n",
    "             jnp.linspace(0.1,1.0,10), \n",
    "             jc.Planck15(),rtol=1e-5, atol=1e-5 )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "e801c3eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "res = solve_ode(init_state)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "de770afb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "352 ms ± 1.01 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%timeit res = solve_ode(init_state)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "73e7cc3c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x7f977c4261a0>"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAPsAAAD7CAYAAACscuKmAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/YYfK9AAAACXBIWXMAAAsTAAALEwEAmpwYAAEAAElEQVR4nOz9Xaht25YmCH2t9T7GnHP97HPOjYibERkRZCjUW73Uiz6UQmIhiBbmi5WoKCUWxFOhomJG1bsQIIj14EugQomCWaBQPhSKJOSDL5KZIhQaVmVZRIYZERn33jg/e6+15pxj9N6aD+2n97n23ufcrBu1T2CcwV1377P2/Bmj997+vva11khV8cP1w/XD9f//F3/fN/DD9cP1w/Vprh+E/Yfrh+svyfWDsP9w/XD9Jbl+EPYfrh+uvyTXD8L+w/XD9Zfk+kHYf7h+uP6SXL+QsBPRf4GI/j0i+g+I6Hf+vG7qh+uH64frz/+i/6h5diIqAP59AP95AP8YwN8D8F9T1f/Xn9/t/XD9cP1w/Xld9Rd4738KwH+gqv8hABDR/w7A3wDwUWFf1ns9nr4AFCAASgAIUCYo+9/p9j0kAD6kj+jVa/DqdfPnKECzUlOAmgJQkP9amQAmSCVIifu6/TypH/rc6Xvn72f/TwbySwBAaLxHp89TgPurz+Lbl8zfEc8c9zp/Ttxvrin757fxJ4n9XgnQCmjxtWf/Rpmeb/5cnl5XFBCyH/hz0vh+YgX5BzApCgsW7mMplNC0ILZGQagsKCRgKAoJmjI2qRAliBBUCSo01k/H38nvU9nvb16TuMeiACtqESylQ5TQlSFinw8A6P6nf3aewfh1fH6sEwDqt9scbyeZ9nXam7iv+T067Vlc8bnU9YNyQKIgUd8TwuX6Nbbtmd5/5S8m7L8O4P87/fc/BvCffu9miH4bwG8DwOH0Of65/8x/B+Uq4C5oxwJZGO1EaEeGLEA/AEpkh0UU5WKLRaK54EqmHJTt38qmILEFodhshDIhcFPwbiulDJSr4vBnF6AJqNvh6w8HyKHg/MsLrp8x+kqQg29AAWQBrj8SyAI7MAD4zOB9HAbuBGr2+n5UKNufKAple085M3gn0A7wbs+pRcE7Yf0GoAawb6yspngoDhzse6gr6tm+Z3tDkGUIKzf7aUegn9TW9Oif/xWhXIHTzxRlU7SDvffyy4Ttc0VfFXISUCfwmUEd/nyU398PCq1Ae+ighwbdGPRcAVLo6qe3mkAtx4a6dNTaUVnw2emCv3L3DuyndpOCr6536MJowmBSfH4846FecSo77usVP7s+4I+fP8OlVbw7H9Bawfa0Ao1AjQEBeCNQcyUqtlf9od8q0jODBOifN5S7hh999ozfePwaz/sB31yPuO4VT89HSCfI8wLqBNrJ1uBKKBtBitpZWBX9XuxzLwzqhOXJ9j73yQW3noH1azufZVdXRnZjZfdz62e7HQntRPYd1c7t+lbBTbE8CXiXW+MCoFwayssOXQra/YK///f/5x8V2F9E2D+kPd7TPar6ewB+DwAeP/8NBcNuugkK2QZpZciiAPmDFkU/mblStsPN3RZ+tvQh2PsducUikKopC7/DFICF0jrXRbG8K2BM1psIJIp6UUi1TdgWAjHS8sjqB3oVECuk26GAa3Bhha648QqWd3YTsgJa7bC0RVBeTFGY5bD77UfX5M8AC8C7gtq0lsWEE9X/9OcJxacMwA0ndwCXcQDT8hDQF1ucfO/sUSmATsP6F1dGze6TdwKaQqmgwxVBB8AEFQWYQEsHVzv1rTG2rUKFcL6ueLoesNaGh2UDkeJQGrgq7uqGSoJT2VG5491+xE+vD3jeV7y7rrjuC85PB+jOoHMBNcpzoAzbl8lS8oVdMRJACjnavpa7htNpw/264VgamhaUXdIDUXXFsVN+dnsQNFfK6akRzFsg+29ZAGKCFlPyJsRA74R+JJRNwS8u2HmA/Tv8Y7gD5ep7SmqKxn9IbY2FxlmlrtCFIccKWQrafbn1Rl9dv4iw/2MAvzn9928A+ONve4O6VYYquAm0MqjZQ6WbIwDK5J6KLSJIQYRcgPxMNk0+3KxxCEZYQAAUynbAuXnoUAj0Km6gblZPFgJ3u53QHXD3lZcOYkCqHW7IFJJkaOL3urlVZIXAD0O1Q5MhRCiT4oip3a4J6PQcWiavJdbnxouZhNYF3D5vEozwVMTf+zpsUnK30X/x6vPJXWLuCtnZXqv2DKQEhYKLglmgwlAlyM7QnbEp8EyKrRYUUiyl43G5onLHZ8sFC3cwFEyCb0B43ldcWsXeC/a9QK/FFNFGoE5+T5oK6cal73bT3NwTLAIsglo7ltKxckf17zP5yVjC1i4UHbkiKfZDrBZGtNuFs31RC/XYz3H3tfYzFx4a5X7rjcAD5tUpz6GeewNTyBvnI/5bKkMWRl95GLoPXL+IsP89AP8MEf0nAPwRgP8qgP/6t70hYw+yG7TfKaQS+kroB6Ddm0D2o1vlagtUNspNiIUgUfSV0O79C3wx64tvti9Q99iOu4J325h+KtBCKJcGagLaO2gH1l2wvCWsx4LD24K+MrYHQjsSZC1oR0b7FUU5dugqaKTgC6NsZuHKTi5Qpv1lHeEDKVAuBL4U1DOhXOz+0lVvrsndFY+rr8Olj88JBUkNKYRSkLF4rDcE6UEAABhodzRidrYDGZZbd/s33qd4NTyIonkPsip0EYtvNcIqE4i6NNQqOD8XyLUAO4N2gjTGpTO2tYNJcVx3/MrpCY/1ir92+jPc8YYv2z3OfYEoYe8FL9cVz09HyLmiflltfa4mDP1oZyfWWhdxobf9hgCysynohx117TgddhyWBgHhaT/g0isUwN4L2qUCO4EIkEWhJwGqgCYh107ATuBzuVHmoRxJAXXlr4udd1kJDYTt3rzPevU4u9i6SWFoQbr0cRZIgXZikCjasYxzr8jQNHSyVEI/0H88ll1VGxH9qwD+zwAKgP+Vqv4/v/1NHo8yzCKG1XDrLCuhHQGtHusC5rK44KYHgFgYQj8C+xs1RRfWUAm84dbSElCuBN4NzJDFzCJfCUQE6mKCcW1AV5SloD5X9LsKpQXUgf7O3Lv2hWlPWsQ8hO5CogBvvtHVN/Jg3xcbxbu5luVqr4WDZzNw9jpcCYWgxV6vYm5+CqOvYezz7JK//hxlw0XiPfF76tOBleH+hxcVVk4qoMvwTtKjCteWFcti1vOlH4GNwVfDKUTUDJ4QrmsFs7ntb+oZP17e4g2fsWvBVexYdiVsrUBeKuhcsDx5eNHinsgUTNxDUdCxmyXs5sYHmLYeGg5rw2FpWNgW7dIX7FIc/GNg4wzLwAAfG+o63MjeCnpj82wuNIxQepXuphdA2JQjVTNeUMejClB2W3RTnIS+2pnnwHFg665kin72VrmH8QtlPQBlKfSepzZfv4hlh6r+OwD+nZ/39SSKcu6AmgD0hQdAd+cP7W6PLGYp+sF8mHK1hXgtCFphFj8PLOXBDVc34lXezUUvOyx8aGoovbo7VQBlkxo9FMihoB84ATBqtlH8XNCVDJFlO/Sy2I73wxBskIFeWsyikwDUyDd1WG++2toE2BfAjeYBorSwqhHbjUNuCC4l6p4WPYQiLH1YnAmlV3c5+TpwDor3xqEiAHcu8ItCVpnce187F3gqiofjFXfLjsu24FoWCC+Asltg27jWGK0WXHvFuaz4w+sv4cg7frI94l074mk/YO8F1+uC8ragnAnLk+8rI13l8EJUAFFO1z4PveMrqoTWGedtwcYF35yP7iUTemf0zqBTA5RARQACmDWPBpFCOplC2NiErhH4Oq0xAQxTavVqISLtw4WX1faoL8NdB9mZ0lSuCimURioViWMrbTEBZz/rcfEOLGcdAOEHrl9I2P+pL1GU8w45VGglyMroRxP2/d4PUwW0upBXdxUJkJcC3ibtF4AKzYfbBb2ZmyMUAJ/9d9ngPwLeFdzFwgogU28IzbwU9KMJe18pXV0AWJ7ItPy9DNdx9eyOW9NytQ2TowmGbsXc8x2uuOyeqNs9BVZgoKPHcZHmihgyhLkD5WLP3ReYkpq
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "imshow(cic_paint(jnp.zeros(mesh_shape), initial_state[0,0].numpy()).sum(axis=0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "4bdfe23f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.colorbar.Colorbar at 0x7f94181834c0>"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAATcAAAD7CAYAAAAPf9NJAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/YYfK9AAAACXBIWXMAAAsTAAALEwEAmpwYAACkS0lEQVR4nO39b7Bt21UfBv7GXHvvc8697z39QQYLSW2RapEEqLIdKEyHlIsYu1E7rgh3ld2iy0F2nChF4Rin3RWEvzgpl6r44JDY6TQVxRAgsY3VNi5ULjAGEjpxFQgLTCwEVqOACh4SCAlJ7957ztl7rzVHfxhzzDnmXHOutfY5+55z33t7VJ179157rTnnmmuuMcef3xiDmBknOtGJTvRKI3ffAzjRiU50oqdBJ+Z2ohOd6BVJJ+Z2ohOd6BVJJ+Z2ohOd6BVJJ+Z2ohOd6BVJJ+Z2ohOd6BVJt2JuRPR2IvooEX2MiN5zrEGd6EQnOtFtiW6KcyOiDsD/D8AfA/AigH8G4JuY+ZeON7wTnehEJ7oZrW5x7VcD+Bgz/yoAENEPAngHgCZz29AZn+PhYb1Q+YVrP8g3MseY5czbgJSp/MDZf+1rCKRjcSQHSX/k4noGPIchy/9E+b1lZy+6n3yOyvakjXIcDSqGQuTkGDnAufy+vJeTnANA4M4BDqDeA8NQGxqyuWXI82KWZ0nFvIV50uO8XgGOQEPo2zPAPrWj5+Y3377X8hrn8jF4D3gv83ff4Hcya0yXFXtc8xPseNteQAvoG/7th/yZ3x0Wnftz/2L7Y8z89tv097ToNsztTQB+w3x/EcAfKk8ioncDeDcAnOMB/pD7o/XWqK4hk6P8HPbj81140Ou1fB8GWYSDBw/LHlJG7AFyqW/tK/TNvr6w4/ldB1qtQOsV6OIC6DqgC8xg38v4mOUl6Xvw9Va+h7HSZgMikt+VYYTPcmzmxdI5MmNJvzF43wPs033U5rS8J/2+2Uib5+eg8zOgc+BVBxo8cHUtJ12cg1cd/AsX8GcrrD7zBPT5RzL2YQiMMTDeMDbuw7z0PXjwoM4Bq1Ua/zCAr0P76w1o1cF/0esxPFije7yFe3wN7Pbgqys5d98Dzsm1jhLD2u3Ha8I5mWd73DnQw4egzVqu7Rz48hr85AlY21/6DPIJnbyMHDXXV3beepWeq3Nyz7sdfmb/Y7PXztGnf3fAB3/szYvOXb/xf3/DrTt8SnQb5lbbHUZPhZnfB+B9APACvT6IJgea+uz5tWs9Ax2lhagLmBkY9LJi0ZSLr9XHIWMlJy/leg16cAG6OMf2rW9Af9EBDmAinH/qCt3vPgZt98D1Vq5z0odKn/pC0jAARHIvlrFNMKPsOMsLG4WlrmufH9osmVl5HitzDsyGiIAwtzwMwhSYQasVun0Pt+pAT66EMXnOGTYAdsU9xI0pbASAbA6W/ADuAff5J6CrNWifGKNeQ10nTK0LDO7iHNw50NUW2G7DZsGRyWIY4th48MDggctL8LVL97jvw7OYkv4aTE3nuXo438DJ+XkGNwxyTtjYmXkRU1xGjGHqPl4mdBvm9iKAt5jvbwbwidmrigc8u1MdwlxcWNCuB/h2juDJl3zqfCcMjlYr8GaN3WtW2D9wcTjrx2t0j1ZAPySGolKaVWWJwr0Ldx4xNiBKmCNylL+AKi3RAQu2kFazPocgve57YZx6756FufdrsGfQMMi97XZBWkyMrVSvreocJTw7fudEitPz2AO7PSiqoyySOrNR11yQmjvwZg2sOqAfQMrIhiGdBwDYZ/fMu30m9VWfwU2p9uz0u2Vwti9zPnsGufAsrEZzBGIAfpHN4tmm2zC3fwbgbUT0JQB+E8A7AfzfD21kdrdpMZlCpeIBgKot3sedWZnOSAUbD6TOKMz5tbFa1VWZGlYr4GwDrFdwO8aaPKgHyDNWn9+CLq+B7U4kDSvNhJeNd3thCip96gvcGncaTPETt5l0KSGoGq5tRImm0Zf97ItrPAPkwft0MW02ogoP4T6Dahyba21isc3wnIHElPpemJduBF6YKboOWK1Aqw783APwegX/YAN2hC4wP7reiirsGeC+Pr/sAe/AKNYOOdCaMkZXXcd2g2gxMnMtufw82fhd/Fyeb687ntQm5PEqltyYuSeivwDgxwB0AL6XmT+yvIHxg5TDFQZSNYijWFgDsNsdcgtNqjGEOcYmEmP4v+vAqw68cnB7D8BhdT2A9h7uegdsd+D9PrfbqA3OMwh7UddqqmhjXOSo+hJlLwDTaC6rZFSdqfuv9qtqpkpGnoF1YPh+AHkGE+U2PwDKRTMGCwjz16/K1PQ6VYODbQ7sjSraAWcb8IMz8LrDcL4CiOB2AzB4UWXDvPLgo7NgRMU49ZkLEyXZVM34G5NUPz5lXiBlasU57JO2E9f/ceGqDMb+Va6Wgpl/BMCP3Oji5ktI+QurFF6aaPOo7ZjlC2geUGy3pm5Z50FtfGFBxZ82G5EW9EVWSZFIjOF9D9ruAO9xRgReOdB1L2ra5XV4mTgzZC9xFCzanYO0Ud57/LlUB2tmAcNQpvosmXt0ojz3UFTAq2vwfi9zxT5JqXOexglpFMBImo/OiGASAABs1uDNGv3zZzL/Q+i396DdXqS+Uj3V+VEb35QKqs/uNrTE5KJzG8bEwd5ZpSMxJBEVXt1q6VOjqiqlzGQY2h7QmZcie5FrakJpq6q15Qh0diaSiJJ9UbwXW9TlNeAI7smVDG2/R4Sn2LEGW1E81lCNR/aXYmw1KUjveUSt+7RjKJ0vFVtp7D+8fLReAesN+IWH4HUHp4wsQjR8chg0+q9JScBYeo/MxUi3tF4lFTXY2fbPr8EdYf24h9sFiW23B/e9cSDk6ykyEgyAnzDwO5furzWXOkeHknrsg5c9en2vfVDhb4ACOIBe7Ta3w4nqtoP2+S5KEaQ71hLpZWR4NS/JlKOguoALKETwnIXG4//WPsbMUb1ilSTUvmPULiICd50YxQEwh3bXa/mt70HeLOYWFMbcV4s5gFxkELPSmzKzgpmOYDlhfkgxaUpq6O+HCHuB56j+RWY+M/44FnTJUeCcbC6dYSx9D6iEaLyftN1h87kduCN0V3vQfhDmpvPROTAQ29UNNEIsdruwmRp7WGEDHEl3Nc3AQousTXNmA4mka8ojMlTisT3uWOopAxjuG8d3BLpjyS23p9R2xMiMSrtLMG4sYogFZW2WzK3VnrWtWLhCARSGF9xYMk6rg2NI18fTBSKRYAodiD1AG/ktQD/oLGDJrq6DyjV+eVpqtM7piBE5cx+ll7BQu/WYvb41f9GA71wy6j++lCb2e2DwSW304bNt33gI7fjllHAf+oy6MGcvPAderwRbxyzq7yVEyglj4Mtr0HaH9dU23bOqnETC2HgFIg/2Xp7JZi3trwOk6JLE8dELBCTeaw0rWFLNg6neXt28jCMpbmwFw+MBERbEYROh1UqYsm54YY6oA+Bvhd+N9PK3uN2XWhp3P5cMpoHqO/c8U6PALDKasvEWBvPWOSOyhnJVRRfscskJInYnUphgkMqISBav6aPEhE13MJ7HjGEFD2ZkcCpxOTdtx1FaolpZsKxe48J9q/pof1tqIzJSTWQw6xV4EzagfgBWOWwD7MG9MD8KtlAAwZvdhTn2AHzmCMqk0Jbz5Qa2NtL5VlCyI5CuIUCkz/yCcZ/67EJ0Alt7XID50P7goY2IwSeb26FEQDL4AmK0tqJ0TYLAjKdSvshDtsh+ALzbAX6MJM+YFheSjrZVhuBoBEGQ1PTa+o26/HqjwjAGUTWxBq3qKgoHcK9KPOX9l3NUBSeTS46ZwPNFtfHRpkRBUsFuH+ZqbNcbPY9Ssug6ea5ECWeo0hEgD3nfw9dsU4XTZrJfFq8zAPDZGv58g+HhGrx22FAAPOtzHzywD57zPmDxFDKyWYMvzgSO0/fiWR0M01HV1jc2rbih1T2rI7VUv3edOKLU4UKUolUur/J2KxKyOj6iJuEcSLF/qkazB3a3V02Zgf3Ln7fdccoj3XXDnxqhqQuGUyt5TdjGDgXYRlIpsIjhnG1PF7nBynFgOvpnBpdfW9p
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "imshow(cic_paint(jnp.zeros(mesh_shape), res[0][-1]).sum(axis=0)); colorbar()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "a7f8b5d5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.colorbar.Colorbar at 0x7f940073dd80>"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAASwAAAD7CAYAAADQMalWAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/YYfK9AAAACXBIWXMAAAsTAAALEwEAmpwYAAB4X0lEQVR4nO29fawtyXEf9quec+597+0uuaQoUhsubdIwIVsQYMsgbMUKbCWUEcUWTAewFNmQQcsyiAC2TBsOLMpGICRAAAUxDAuwEWAhy2EswZYsK6AQCPoIbSUxEBNafcCKRCmmJYJacsUP8WOXu+/de85M5Y/+mOqa6p6ec+bce595Cni798z0dNf0dFdX/bqqmpgZZzrTmc70OJC7bQbOdKYznamVzgLrTGc602NDZ4F1pjOd6bGhs8A605nO9NjQWWCd6UxnemzoLLDOdKYzPTZ0lMAiom8kol8noo8Q0fvWYupMZzrTmSyiQ/2wiKgD8P8B+BMAXgDwcwD+HDP/6nrsnelMZzrTSJsjnv3DAD7CzL8BAET0zwC8C0BRYG0vnuB7D153RJN3m5jy38TyHgEUylC4NwDEDMRyJOshf+9IfugUfsEMsANABDByPhmggYGBQdc78L4HXVyAtx3gaNJHExL3Y9n4DrFfYhnq2bcFz0vs24xVihd40he6/slrxvuDfy8Y9ZeeO0m/N7R79ernsLt6pYHLMv3n/+kT/Duf7ZvK/vy/vfopZv7G0n0i+gEA3wTgU8z81eHa6wH8MIC3AvgogG9h5s+1tHeMwHozgN8Sv18A8EcMht8D4D0AcHn/afzBP/7e8ab1UVu7Wj6rn1lSry5LxnX5bGkgkp/E7LygoR5pMjMR+ksCd4RhAwwdodsxumsG9Yzuiv3kCM9z5yea2ythJoRDmoiV/mInJttKFAUEd4Rh6wWW20eJ4oXI9uUduqse7mOfxPDZz8P97jdj/6bXYtg69Pe6UNY/kwkjoiAIx37wAjCUCf0ybP1LX3xhj+5qCAuBvz5cuMRnKktRuGH8JqJ+yX9sO/YfALhrhtsN4I7AG0qLz9GkFirJw6RMjUT5X/o/v+9otj7z2R4f+qlnm8pun/n3b5gp8r8A+AcA/ldx7X0APsjM3xugpPcB+K6W9o4RWNYnm3QvMz8H4DkAeOrpZ/P7Sz76jKAAxoG9+IMX6k3PyoFVq5v8pEJP2Fz1YEe4em2HYTM+Rz28oBq8kKL4jKgzaWPOl/fv5ifgRFBV3m/SH3qC1Pqm9J7i2SQEZBkH9Pc68NZh85Y3gr789dg/dYn+sguTPZbzwo51vSSEseYvCJvYJ/sHHfr7Li0OSUN1XqByFIAA0Pm/5SLARGlBocEvHhNNjYJwJpfzpvvmEAE2tzjW6uXC36sQo+dhvlhLTcz/FxG9VV1+F4CvD3+/H8DP4gYE1gsA3iJ+PwvgE0fUV6ao9WvtIpIULjAE2JqsuEK9QdCwIz+R9wxsvDCKmgiF1dz1I7NpUmnTShAlbYSbVvdUbo5kv6k+LJWLkxywtTcmAm8c+o6BJy9A9zbgrRuFlRAEk/Yx/bayr2kQ/R8E/eBcEpxSmEUtKdadKS5aQEgznb3pKPngDsn8rdFkbNQE0QwVx5mi4pw4gjxacVKb9k3M/CIAMPOLRPTG1gePEVg/B+DtRPQ2AB8H8K0A/vwR9ZVJTabJhIyTqRt/TyaTVrV1E/HDw5sKBwk7AoZOTKw948End2AC9k90GDYE1/tJxd1oPuqxQcxg0GgKCU1Avq/VPuBNTmDhO5TqNLSyCR/yUzig7wg0BOE9uKS1UNDK2AmhqwRVEtoMTEHBKKxYmIrCvOuDOR7KsBvNx2h+ux2D9gOGrYOD/w79hecXPGTvN9Eejb6RAmO2PIz7+hUHzuuaq2pFQSVpQPPgeQMRPS9+PxesqpPQwQKLmfdE9FcB/BSADsAPMPOvHFyfMwZI47eIAmIy6CeNGHVq048K9+Z4YP8wdwD3BDgC7QdsvnAFGgbQ8AD9PZc0rYEJvNXPj3yQanh2YMaJ3xXuN2hPTeUrFPGnKDCdA6gfvwkPgAsLiZdHhXeKGrVgQi9QNCDhaf1FFI48Cvvey7sh9EfH3rym/eDNQscYgukYF5K0gEgeWrXVFRSSiSYd67WabzHrDyQGY9duEn6Gmd+xsIlPEtEzQbt6BsCnWh88RsMCM/8EgJ9Y9pD4W2rthdVJmnc8I8HiAF4ENBv8ZKtbnNDynhCoCWcK2kN35csMWwK7DtRf+FVzQ8n0GIhywUL5u0UB3Iwxld5HlzcwoWayyspJIzQl/+7+GjlkQmQIfWCa9fD9oIU1MBUcflyMwpD6oJ04SpsfAHDxiv+Y7tprXtFcjfepB7rdVMtN7VjCKuKODKDPTc/8YfX3jNyzrAb5bKuZeCwxgP60JuGPA3g3gO8N//9A64NHCayjKX7EQ/tGC7w5rKvGSjC7oomWYUVqNYv1D924ssfdpihs+gsXJtnGWzBB8A6bsd7Ep97Ni9rAgh3BrB9OTEWtI0yq/aXYgSMGccDwnP+X+CxovCwEUUaxvNYmedzpS7uMne/D7ct9MrMAoL/nMGxd0tBo8Lu09ouqtuOfbhwf8dsv0oLkdbVg6rEQ+zptHuj+OBGthWER0T+FB9jfQEQvAPgeeEH1I0T0HQA+BuCbW+u7XYEVqTAwTtqOQRmgHLSCcmFkkydrJuIsgHdViCukwCYihjXWNQLGVW1Gs7Fkg+GYfrYmnxTggxcQjtjv0DmvcTrxThlYXyIpGMT3kDuntXfQ2qkXTqJgbQwEMD8TigLkl22kBg7QUqOgtSwB3TcTPM2ilWEsBtCvlNiTmf9c4dY7D6nvVgRWNgi1ujsBIfXD83XXNCtrgpeeGZ0SMfILQOKR3vwASE5ogrlbJbfVuaOEr3geCA58sGZlkrF6+/dqfL7SbiYEGEFTCebX1eizNGwovdNkh7MFGhLPRP4lVpkEVHgvz9e4MPT3CODRpy37HhLzZHgci/y3SSafaDe2AQbcTo0NGN9KLmqhj0B+ZzNqZ8XNlJJ5WcGuvJ/a9PohdAOK+kF0KwLLFg7xJsoDufYxxArexEMBD5CrepGkGasHUrxW0LxS/Ty2T9E0iBPkFDs/rN63ZIo3LAjmM9GlQ92XO53HkhROAJIjbnRpGG/kQp/A6f+SJ03JebUbn0ta5KEzWGqhEvuMOBgMYWX9rl0XGvwa6haDT41hHUy3YxLqVYKnWo622UlpNakextSUEqvapOk46AXOoU0/KVS8ZmIAxEGLcnGVX4KXOa91dAFDGbUxX++wGeuaTBQlINO7t5h7csJIJ0oy+NdCV5uCungH9JcG0FJbgGq8GuWiRjVsfR/uHhD6LWHziNFdw2uocRHYh/fahO+3YbDQtNJ3FX0ybF3qf3YEt8MI6E/GZ8lGx2TsRZO5u/aaZw/f/70jYOtBf9rb1bUuwkPUCldY65iB3d2UV7cgsCyTYIXOIeaxvhMoKMColU2cJ2ObS98jaVYsbMoF9VlCZAEPs9rkjLCREzjrf6usBtvTg5jyPPceGqwu8FSjCGRbu76+QI57ze1Ql6jsDDyab1L7myXDbFxfIyf0p5pER9LtmIQxnCwEw7qex+3hMBgn6r+bXpNmTmnXygSJB5T9lSAGmaUwCKFltjd5oNJGNAFDfFpaIUUfVNso4Rxz94CEFw5bsp1sC89JE0u2Id0RNPbGndcApId/8Z3kNSkwhdCILgzbV4BtdCJVmlACtuOGQBwnMk4wCigF/DilXqS6Qru6P2rvknAyx+gvRuwKTOOrDWVhZTnUeizMzx3ekHd8jbjaCmA5A6iwdKt0a7uELGK2vGrOttlXoaWAfAvpFfFUnsRJS4jmmDA9SfbFnEnVSobWEidt8RHRF5k5VHJpsIBnUX6RJmEyhPEdOAe+J+PFgB2ysoX3lv2eFtb4vm76XWouHpMdT8dpo0n2wxJ8bJJxImz6IDjlrpUl4qxhAWmQ9Y5CZgI/SR1z8oAGEPCluJoaQGmr6VPBckhtt8tJVdLMWmgC5hdwoAhQS00lOb5qzEThVlUq4U4Cj626Bsh
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "imshow((cic_paint(jnp.zeros(mesh_shape), res[0][-1]) - \n",
    "        flowpm.cic_paint(tf.zeros([1]+mesh_shape), final_state[0]).numpy()[0]).sum(axis=0)); colorbar()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "68b72b68",
   "metadata": {},
   "outputs": [],
   "source": [
    "from DifferentiableHOS.pk import power_spectrum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "785818e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "k, pk = power_spectrum(flowpm.cic_paint(tf.zeros([1]+mesh_shape), final_state[0]),\n",
    "                       boxsize=np.array([100.,100.,100.]),                      \n",
    "                       kmin=0.1,dk=2*np.pi/100.)\n",
    "\n",
    "k, pk_jax = power_spectrum(tf.convert_to_tensor(cic_paint(jnp.zeros(mesh_shape), res[0][-1]).reshape([1,100,100,100])),\n",
    "                       boxsize=np.array([100.,100.,100.]),                      \n",
    "                       kmin=0.1,dk=2*np.pi/100.)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "18e3d30b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x7f94005fc850>]"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAXYAAAD8CAYAAABjAo9vAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/YYfK9AAAACXBIWXMAAAsTAAALEwEAmpwYAAAepUlEQVR4nO3dd3yVVaLu8d9Kg1ASIE1IIAkEAqElEIpSREUFFWSwF656HJFjn7mObfTgOOPMePTYK4o4jgLXsYIFFFQQAUnoJCQQWkgIhFBCSSjJXvcPy+UyIAkpa5fn+/nkj73I+77PB8jDYr1rv9tYaxEREf8R5DqAiIjULxW7iIifUbGLiPgZFbuIiJ9RsYuI+BkVu4iInwlxHQAgOjraJiUluY4hIuJTli5dWmatjTl+3CuKPSkpiezsbNcxRER8ijFmy4nGtRQjIuJnnBa7MWaUMWZSeXm5yxgiIn7FabFba2daa8dHRka6jCEi4le0FCMi4mdU7CIifkbFLiLiZ3y62DesWkTuwk85VLHfdRQREa/hFfvYT9eur5+n/97POTI7mPywzuyJziQ8ZTDJGecR0SbWdTwRESeMN3zQRmZmpj2dNyiV7y5l47KvObRhAa3Ksul0ZB1hphqATUGJ7GzTh+CkQbRPP5fYhE71HVtExCljzFJrbea/jftysR+v8uABNqycx768+bTYkUWnQzk0N4cAKDExFEekY9ufSVzPc2nfuTcmyKdXokQkwAVEsR+v6ugRNq75gV258wjbtpikg6uI4sc3Q+0hgi3Ne3K43UCi0oaR1GMgIaFh9Z5BRKShBGSxH896PBQWrGb76m8whYuI37eceLsDgArbhI1N0zgQ14+WXYbSqc8wmjZr2eCZREROl4r9JHYUb6Jw+VyqNy8kZvcykqs3E2QsR20wm0JT2B2TSXinwSRnDCciSjdkRcR7qNhrqHx3GZuWz6WiYAGtf7khWwXA5qAOlLbuQ1DSWbRPP5e49p0dpxWRQKZiP02VFQfZsGI++/Ln0WJHFh0rc2hhKgHYbmIobplOdYeBxPU4lw5d0nVDVkQajVcWuzFmFDAqJSXllvXr1zvLURtVR4+yKWcJZbnfEla8mMSDq4hmLwB7acnmZj053G4AiUOv5YwOXdyGFRG/5pXF/jNvnrGfivV42LphDSWrvsEULqTdvhUk2O0cscEsj7uMzpdNpE1cguuYIuKHVOyNqHhTPkUz/kzf3Z9xmDBWdxhHjyseokVEG9fRRMSPnKzYtSDcAOKTUxlw9zsUX/sNeS0GMHDr6xx9ujc/TP0zhw9VuI4nIn5Oxd6AElPT6fuHGeSPnsnWJikMWPcUe57oRfbHL1BdVeU6noj4KRV7I0jtM5SeD3zN6vPeZl9QKzJXPEzRX9NZ+dW7WI/HdTwR8TMq9kZijKHnkEtJeWgJSwc8h7Eeen9/G+v+diZ5iz93HU9E/IiKvZEFBQfRd+SNtH1oBT/0eJTWR0vpOusaVj8xnI2rF7mOJyJ+QMXuSGhoGAMu/x0t/7CaRZ3uoUNlLh0/GMHSp8eybUOO63gi4sNU7I6FN2/BmeP+BHetYmG7G+lWvoCYt4ew5MUb2VVS6DqeiPggFbuXiGwTzVnjn+PgrVksjR5Nxs4ZNHu1Lz+8fjf79pa5jiciPkTF7mVi2iUy8M632P6/vmNNxBAGFL+FfbYXP/xzIocqDriOJyI+QMXupdp36k6///0hBWO/YHPTNAZseJZ9/92T7A+eoeroEdfxRMSLqdi9XEqvs+j9wBzWXDCV3SFxZK5+lJK/9mb5rCnaAy8iJ6Ri9xE9zrqY1IcWsvSsl6kyIWQsvoeCv/YnZ8EnrqOJiJfRQ8B8UNXRoyz99DU6rHyGtpRREJLCrqRLSD77OmLb61HBIoFCT3f0Q4cqK1j+8XO0KfiA1Oofn2efH9qV8uRL6DjseqLbJTtOKCINScXu5woL1lC0YBqxhZ+T4tmIxxrym3Rnf6dRpAy7ljZxHVxHFJF6pmIPIJvzVlC8cCrtir4g2VNItTXkN+1FRefRdB52HZHRbV1HFJF60GjFbozpBtwNRANzrbWvnOoYFXvDsNayMSebHYumkrBtFh3sNqpsEHnhGRxKHUOXc64holWM65gicprqVOzGmDeBS4BSa22PY8ZHAM8BwcAb1tq/H/NrQcDr1tqbT3V+FXvDsx4P61cvomzxdBK3zybe7uCIDSavWSZV3caQevbVNI/UJzyJ+JK6FvtQ4ADw9s/FbowJBtYB5wNFQBZwjbU21xgzGngAeNFaO/VU51exNy7r8ZC3bD57sqbTccdXnEEZh20oeS364+n+G7oOvZLwFpGuY4rIKdR5KcYYkwR8ekyxnwk8aq298KfXDwJYa/92zDGfWWsvPtW5VezueKqrycv+hvLs6XTaOZdYdlNpw8hpcx7xl06kbVI31xFF5CROVuwhdThnPLD1mNdFwABjzDBgLNAEOOknSBhjxgPjATp00I4NV4KCg0kbMBwGDKe6upo1i2dzYOl00nd9TvCUL8mKvpjEMROJbZ/iOqqI1FBdit2cYMxaa78Fvj3VwdbaScAk+HHGXoccUk+Cg4PpMegiGHQR24s2semjx+hb9gm88TlL4sbQaexEos7QP8Ii3q4ujxQoAtof8zoB2Fa3OOItzkhI5sw7p1D2H4tZ1mYkGTs+ovkrfVjy6n+yt7TYdTwR+RV1KfYsoLMxJtkYEwZcDcyozQmMMaOMMZPKy8vrEEMaUrvELgy8+x1Kxn3Hyshz6VsyjbCXMvjhjbvZt7vUdTwROYGa7oqZBgzjx73pO4CJ1trJxpiLgGf5cbvjm9bax08nhG6e+o5Necsp+/RP9N3/LQdNU9YmjqP72Ae1VVLEAb3zVOpVwZof2Pf5Y/SpWEA5LcjveCM9x96nbZIijehkxa7H9sppSekxgD73fUb+pZ+ysWl3+m98kcqnerBk6mMcrtQnPYm45HTGbowZBYxKSUm5Zf369c5ySN3lLplL1Zw/0+vIcnbSms1pE+g9+m7Cmoa7jibit7QUIw3OWsuahV8Q9O3jdD+6hh0misIuN9Fj1F1aohFpACp2aTTW42Hld58QsuB/6HF09Y9r8O2vpMuoe2kVG+86nojfULGLE7lZc6n85mkyDn7PEUJYHXMJCRffR9vkNNfRRHyeVxa71tgDx6a8FeyY/RR9dn9BMNWsihhGq/PvJbnXYNfRRHyWVxb7zzRjDxw7ijezYeZT9Cx5n5amkpwmGZjB99Bt0GhMkDZpidSGil28Svme3eTMfJbOG98mhj1sDO7Ivr630fOCGwkOCXUdT8QnqNjFKx2qrGD5Z5Nol/M6ibaIEhNLcbebSRtxC80iolzHE/FqKnbxatXV1SyfM43mWS/QrSoPgBITS2mzFA63SaNpQi9iO/clLqkbJijYcVoR76BiF59grSV3yRz25n5N2K61xB5cT4KnmGDz49/TSpqwNTSZfRGpENedVsl9iE/NJDyitePkIo3PK4tdu2KkJg4c2M/W/GXs3bQcu301EeXraH9kA5Hm4C/fs83EURw3jLihN9MhbYDDtCKNxyuL/WeasUtteao9lBRtZPu6bA4VrSS8dAXdK5bQxFSxIaQTezpfQeoF/0HL1nGuo4o0GBW7+L2y0hLy50whtuB9Ons2cMSGkBsxiLDMcXQdNIYg7bYRP6Nil4BhrWXdqsWULXiTbju/oA37KaM1m+JHkXDOb2mb0tt1RJF6oWKXgHToUCWrvnmPkJVT6VW5hBDjYV1oVw50v4a082+iaXM9nEx8l4pdAt724kIK5k4mftMHJNutHKQpa6MvJPrsW0nqOch1PJFa88pi164YccFT7WFN1hwqFk6md/nXhJsjbAjpxN5u19Ltgv+gWUt9zJ/4Bq8s9p9pxi6u7Nm1k9wvJxO3fjopnk1U2CasjRpO66Hj6dj7bDDGdUSRk1Kxi/wK6/GQu3Qe+79/nZ575tDcHGZzcBK7ul5H94tu1Vq8eCUVu0gNle/ZTc6Xk4leN50u1QXspxl5Z1xKhxF3E5fUzXU8kV+o2EVqyXo85GR9TcV3L5G
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "loglog(k,pk[0])\n",
    "loglog(k,pk_jax[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "c2c1ecdb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x7f94005693f0>]"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAY8AAAD8CAYAAACPWyg8AAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/YYfK9AAAACXBIWXMAAAsTAAALEwEAmpwYAAAwpUlEQVR4nO3deXyV1Z348c839+ZmJ2QPEEISCLsiEEHBtYAs1uJSp6CjqLW2Vut0rL8pjmNXrbTTdjp16zBWxValOG6oKALiwiKryh4SEpaQkJUlC0lIcn5/3CcQ8Ibcm5vkLvm+X6+87n2ee77Pc24r+eYszzlijEEppZTyRIivK6CUUirwaPJQSinlMU0eSimlPKbJQymllMc0eSillPKYJg+llFIes/u6Aj0lMTHRZGRk+LoaSikVULZs2VJhjEk693yvSR4ZGRls3rzZ19VQSqmAIiIHXJ3XbiullFIe0+ShlFLKY5o8lFJKecyr5CEi8SKyQkTyrNe4dsrNEJFcEckXkfnuxIvIw1b5XBGZbp2LFJH3RGSPiOwUkQXe1F8ppVTneNvymA+sMsZkA6us47OIiA14GpgJjATmisjI88Vbn88BRgEzgGes6wD83hgzHBgLTBaRmV5+B6WUUh7yNnnMBhZZ7xcB17soMwHIN8YUGGMagcVW3PniZwOLjTENxphCIB+YYIypM8asBrCutRVI8/I7KKWU8pC3ySPFGFMCYL0muygzADjU5rjIOne++PPFACAifYHrcLZYXBKRe0Rks4hsLi8vd/c7neWzvHK2HDhKdf2pTsUrpVQw6vA5DxFZCaS6+OgRN+8hLs51tInIeWNExA68CvzZGFPQ3kWMMQuBhQA5OTmd2rjk0bd2sL+yDoABfSPITolmWEoM2SkxDEuJYUhyNBEOWwdXUUqp4NJh8jDGTG3vMxEpFZF+xpgSEekHlLkoVgQMbHOcBhRb79uLP18MOBNCnjHmTx3V31sv3TWR3NJq9lo/uUeqWZdfSWNzCwAikB4fydCUGIamRDM0JYZhqTFkJkYRZtekopQKTt4+Yb4UmAcssF7fdlFmE5AtIpnAYZwD4bd0EL8UeEVE/gj0B7KBjQAi8hgQC9ztZd3dkp4QSXpCJNNGppw+19Tcwv7KOvJKq9sklho+2lNGc4uzgWMLETITo84kFKu1kpEQid2mM6SVUoFNvNmGVkQSgCVAOnAQuNkYUyUi/YHnjDGzrHKzgD8BNuB5Y8zj54u3PnsEuAtoAn5sjHlfRNJwjoXsARqsajxljHmuo7rm5OSY7l6epKGpmcKKWnKPVJNXWnM6sRysqqP1f2aHLYSspCiGpcZYrRVnYkmLiyAkxFVvnVJK+Y6IbDHG5HztfG/Zw7wnkkd7TjY2k19Wc6brq9SZXA4fO3m6TESojeyUaLKTYxiWGn06sfSLDUdEk4pSyjfaSx69ZmFEX4pw2LggLZYL0mLPOl9df4q8shr2HjmTUD7NK+f1rUWny8SE2Z2D9G1aKkNTYkiMdmhSUUr5jLY8/NDR2kZnK6VNYtlbWs2xujPThROjHTw4bRhzJwzUJKKU6jba8gggcVEOJmYlMDEr4fQ5YwzlNQ3OsZQj1Xy46wj//uZ2Nu+v4rEbRhPp0P8rlVI9R3/jBAgRITkmnOSYcCYPSWTepAye+iifP63ay47i4zz7z+MZnBTt62oqpXoJnTMaoGwhwr9MzealuyZQUdPIt55cw7vbijsOVEqpLqDJI8Bdnp3Euz+6jGGpMdz/yhf8YulOGptafF0tpVSQ0+QRBPr3jWDxPZdy1+RMXly3n+8sXE9xm2nASinV1TR5BAmHPYSfXTeSZ24dR15pDdf++TM+2du5xSCVUqojmjyCzKwL+rH0/smk9Annjhc28l8r9p5eMkUppbqKJo8glJUUzZs/nMwNYwfw36vyuOOFjVTWNHQcqJRSbtLkEaQiHDb+cPMYFtx4ARsKq/jmk2vYcuCor6ullAoSmjyCmIgwZ0I6b9w7iVBbCN/5n/U8v6aQ3rKqgFKq+2jy6AVGD4jlnR9dxtXDk/nVu7u475WtujOiUsormjx6idiIUBbeNp6HZw5n+c5SZj+1lj1HTvi6WkqpAKXJoxcREb5/5WBeuXsi1Q1NXP/0Wl7fUtRxoFJKnUOTRy80MSuB9x64jIsG9uUnr33Fw29sp/5Us6+rpZQKIJo8eqnkmHD+/t2J/PCqwby68SA3PbuOg5V1vq6WUipAaPLoxey2EP5txnD+Oi+HQ1V1fPPJz1ixq9TX1VJKBQBNHoopI1J474HLGZQQxfde2syC9/fQ1KyLKyql2udV8hCReBFZISJ51mtcO+VmiEiuiOSLyHx34kXkYat8rohMd3HNpSKyw5v6qzMGxkfy2g8u5daJ6fzlk33c8twGyk7U+7paSik/5W3LYz6wyhiTDayyjs8iIjbgaWAmMBKYKyIjzxdvfT4HGAXMAJ6xrtN6zRuBGi/rrs4RHmrj8Rsu4I//NIbtRceZ9ec1rN9X6etqKaX8kLfJYzawyHq/CLjeRZkJQL4xpsAY0wgstuLOFz8bWGyMaTDGFAL51nUQkWjgQeAxL+uu2nHjuDTeum8yfSLs3Prc5zz78T5adHFFpVQb3iaPFGNMCYD1muyizADgUJvjIuvc+eLPF/Nr4A9Ah1ODROQeEdksIpvLy3V5ck8MS41h6f2XMfOCfvz2gz3c98pW6hqbfF0tpZSf6DB5iMhKEdnh4md2R7Gtl3BxrqM/Y13GiMhFwBBjzJvu3NgYs9AYk2OMyUlKSnInRLURHWbnqblj+Y9rR7B85xFu/otuMqWUcrJ3VMAYM7W9z0SkVET6GWNKRKQfUOaiWBEwsM1xGtC62XZ78e3FXAqMF5H9Vt2TReRjY8xVHX0P1Tkiwt2XZzE4KZoHXv2Cbz21loW3j2dcusu5EUqpXsLbbqulwDzr/TzgbRdlNgHZIpIpIg6cA+FLO4hfCswRkTARyQSygY3GmGeNMf2NMRnAZcBeTRw94+rhybzxw0lEhdmYs/Bz3tiqy5oo1Zt5mzwWANNEJA+YZh0jIv1FZBmAMaYJuB9YDuwGlhhjdp4v3vp8CbAL+AC4zxij62f4WHZKDG/9cDLj0+N4cMlXLHh/jw6kK9VLSW/Z2yEnJ8ds3rzZ19UICqeaW/jF0p28vOEgU0ck86c5Y4kO67AHVCkVgERkizEm59zz+oS58lioLYTHrh/Nr2aPYnVuOd9+dh2HqnRdLKV6E00eqlNEhNsvzWDRnRMoPnaS2U+vZWNhla+rpZTqIZo8lFcuy07krfsm0zcylFuf+5x/bDro6yoppXqAJg/ltaykaN784WQuyUrgp69v59fv7qJZB9KVCmqaPFSXiI0I5YU7LuaOSRn8dU0hd724iRO6T7pSQUuTh+oydlsIv/jWKH5zwwWsza/ghqfXsr+i1tfVUkp1A00eqsvdMjGdv989karaRmY/vZZ1+RW+rpJSqotp8lDd4pKsBN6+7zKSY8K47fmN/O3zA76uklKqC2nyUN0mPSGSN344iSuHJvHoWzv42ds7OKU7FCoVFDR5qG4VEx7K/96ew/evyOKl9Qe444WNHKtr9HW1lFJe0uShup0tRHh41gh+f/MYNhUe5fqn15JfphtBKhXINHmoHvPt8Wm8es9EahqauOGZtby68SCNTdqNpVQg0uShetT4QfG8dd9khiRH8/Ab27n69x/zt/X7qT+liyYrFUg0eagelxYXyRv3TuKFOy8mpU8Yj769kyv/czXPrynkZKMmEaUCgS7JrnzKGMO6fZX8eVUeGwqrSIx28L3Ls/jnSwYRpcu8K+Vz7S3JrslD+Y2NhVU8+VEen+VVEBcZyncvy+T2SRn0CQ/1ddWU6rU0eWjyCBhbDx7lyVV5rM4tp0+4nTsnZ3LX5ExiIzWJKNXTNHlo8gg424uO8+RHeXy4q5ToMDu3XzqIuy/PIj7K4euqKdVrdMtOgiISLyIrRCTPeo1rp9wMEckVkXwRme9OvIg8bJXPFZHpbc47RGShiOwVkT0icpM330H5rwvSYll4ew7v/8vlXDksiWc
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "semilogx(k,(pk[0] - pk_jax[0])/pk[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "id": "703c8195",
   "metadata": {},
   "outputs": [],
   "source": [
    "import diffrax\n",
    "from diffrax import diffeqsolve, ODETerm, Dopri5, LeapfrogMidpoint, PIDController\n",
    "\n",
    "cosmo = jc.Planck15()\n",
    "\n",
    "def f_diffrax(t, state, args):\n",
    "    return jnp.stack(f(state, t, cosmo),axis=0)\n",
    "\n",
    "term = ODETerm(f_diffrax)\n",
    "solver = Dopri5()\n",
    "stepsize_controller = PIDController(rtol=1e-5, atol=1e-5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "id": "d6921f17",
   "metadata": {},
   "outputs": [],
   "source": [
    "solution = diffeqsolve(term, solver, t0=0.1, t1=1., dt0=0.1,\n",
    "                       y0=jnp.stack(init_state,axis=0), \n",
    "                       stepsize_controller=stepsize_controller)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "id": "b3083d29",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Solution(t0=DeviceArray(0.1, dtype=float32), t1=DeviceArray(1., dtype=float32), ts=DeviceArray([1.], dtype=float32, weak_type=True), ys=DeviceArray([[[[ -2.628675  ,   5.347485  ,  -1.5904427 ],\n",
       "               [ -2.5886178 ,   5.4675574 ,  -1.8807749 ],\n",
       "               [ -2.6215632 ,   5.6189113 ,  -1.9770483 ],\n",
       "               ...,\n",
       "               [ 97.6788    , 105.22156   ,  96.64059   ],\n",
       "               [ 97.459785  , 105.74062   ,  97.79143   ],\n",
       "               [ 97.26484   , 105.98253   ,  98.37041   ]],\n",
       "\n",
       "              [[ -1.7246462 ,  -0.27205068,   0.6488323 ],\n",
       "               [ -1.6212564 ,  -0.29098457,   0.29932806],\n",
       "               [ -1.4803871 ,  -0.176678  ,  -0.5965874 ],\n",
       "               ...,\n",
       "               [ -2.21784   ,   2.4735427 ,   1.3182787 ],\n",
       "               [ -2.3406084 ,   2.282835  ,   1.8463919 ],\n",
       "               [ -2.175532  ,   0.7863351 ,   1.1335845 ]]]],            dtype=float32, weak_type=True), interpolation=None, stats={'compiled_num_steps': None, 'max_steps': 4096, 'num_accepted_steps': DeviceArray(19, dtype=int32, weak_type=True), 'num_rejected_steps': DeviceArray(4, dtype=int32, weak_type=True), 'num_steps': DeviceArray(23, dtype=int32, weak_type=True)}, result=DeviceArray(0, dtype=int32, weak_type=True), solver_state=None, controller_state=None, made_jump=None)"
      ]
     },
     "execution_count": 67,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "id": "65116b88",
   "metadata": {},
   "outputs": [],
   "source": [
    "resx = solution.ys[0,0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "id": "c6bd8cfd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.colorbar.Colorbar at 0x7f9368964cd0>"
      ]
     },
     "execution_count": 69,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAATcAAAD7CAYAAAAPf9NJAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/YYfK9AAAACXBIWXMAAAsTAAALEwEAmpwYAACWX0lEQVR4nO39f3Bk130fiH7O7duNnh5wQDRgipDIiUQAFJ9ErilaJKzIpVLC2MWhXSvtq1XWlmQpWWWVuKxI2U1qLbnqPXt3S6/0Ko52LXvLa27i2IpEy7Ls2HpeDv1DTy6tNzI4FM2EpkRpMJRqOCZICugZYDA9je7b9+wf93wvvv3F95x7u9EDDMn+VKG6cfve8+ue8z3f38dYazHBBBNM8EpDdNQNmGCCCSa4FpgQtwkmmOAViQlxm2CCCV6RmBC3CSaY4BWJCXGbYIIJXpGYELcJJpjgFYkDETdjzAPGmG8bY9aMMR8fV6MmmGCCCQ4KM6qfmzGmAuA7AH4UwAUAZwD8lLX2m+Nr3gQTTDDBaIgP8Ox9ANastc8CgDHmCwDeBcBL3CJjbIhV1BpTgRn4PVHujY1BDGAK2WcHFn0AHWuxi+GINy+f6o7Fb/1AmRXXhuPGoA6DEyZCxRikAFIYRLCsR0BqLVo2RQKLXXftuCsjEWXze5LApsT7EBuDGVdj4sq4Yi0S1g9tTHl/OGaMwXET4QYTYddUkBqDvutXPe0DAHajCvow2KlU0TURZvpd3NDvwcDmokIKwMIgNVn5FZsiAtC2KXatxXFjUDMRUgB9EyGyFrs2K79hIvRhsBlP4WoU43ia4FiaoGZTRGkfu66P9B6oH313vQOLWPQLQD5XKgBiGMxHFSSuDRYGU7aPF9IEuwCu2NQ7/nxMJQ6y4DiOmyifJxVjsGstNm2azX1r93duCDzwwAN2Y2Oj1L3f+MY3/tha+8BB6rtWOMhYvw7Ac+z/CwBW5E3GmA8D+DCQycAzRidvzRLX50yETTep5tj1ponQNBFWKjUAwNk0QcumWEsTnE1908yPlk3zMnldVHfLM7Hp/uUoxkqlhpVKDW+YOoF2FOd/80kHjTRB1aao2hS7/S4+32tj07UXAE7FdTRNhJZN87o23Xd+zYdN9ju1hf/2WL+bl8f7o70Dee1UXMdyFOM9tWmcr02jHcXYqtTQSBMs7m4DAM5NncBWpYbHjt+E87Vp/PjWebx95wVUbYpGmqBnIvRcuVuubTP9LhppgjP9Ls6mCZajGPdWagPlf6e7AwC4O65jq1LDbzeX8Ez9RtzT3sCdV1tY6LVxw+421tIEq/1uPieaJsKsq+/RpJOPM+9jy6YDc6VpIvx8/UZsxHX0TIR2FGNxdxv/n86lvPwQiuaID9SWItD8omfOpglOJx38x4J2lcHGxvdx5sxqqXujqDp/4AqvEQ5C3LTdYR87Ya19CMBDABAbY4HiFywxJwicBBGjU3EdsybKJzCfJHLSyAnUFMQyVF+onUuOmLx76gTOTZ3Ahxd+CM/Ub0S7UkXPRPjgxrfx4PZzWOi1cbK7g11gHyGlBbmWJjnR44RNI/DaeG2KBbsU6a+bj03o3cyZCI/1u/mmsdzdQdNEmDMRXrIp/kOaoGkivNum2KrUMJ90sFWp4Y7OJXynuzNAoHndHLQxnWUERLZ7J00wA+DtOy9gcXcbc0kHM/0uZvpdfJ09M2ciLDri+ITbZH5wdxsP9Nq46NpB9dNY06YIAL++u40m6+Ovu76HiE/ot7IbeFE5fJxo7pTZ9IaBtcMzBdcbDkLcLgC4lf1/C4Dnix6SL7hopxqGuExHMdpuIraGJKASwxJgur9pIiy7BbVebeCrJ16HtzXmkakogWd2XsRK+/uY6XdRZQSlZdO8r8RpNBmnKgkbkBEvbXzkmJ5Nk2xMhpj8klvldRIReMwRkiZbkEtRjFP9LuYZl1bptbHq4Rbl+PG+rrH207jSfVWbYqHXHuCCG2zBz7nxI475hWoDW5VaziFO2RSLjoj1TIQl0UcAOJ10BjYe7R2MCk5YCTTmc6I+3xhRW7lEMx5YWKcCeDnjIMTtDIBlY8wbAPwNgJ8E8N5hKytacD4iw18uACBN8LtObGmxBcgnZqg+bbINlO95lnNcxLXdW6lhtTaNzbiOhe4VfM3pnRppgn/U/j7e2LmEhV4bZ9iC52IRLSrOTfjazdsnCZ2vT/xezrERUeD9ksRU+y7HomVTzNoUadJBxf12Kq7nfTybJnnffG3XyqTxIc5qpd/FUhRj1om5O06cbZoIJys1bFVqWD3WxGZcz7hnU0HVpuiZKOOa+90gwZJqCM4xL0fxwHOhuaG9B41QyXHwSRPaXB4v12Zf3ZybtTYxxnwEwB8j08H+hrX26bLPyxcuiRCHnBhystEEPZ10huyFDo0gFBE2ThxIT0SiGQAsdi9jPungDkfY0qQzoLeRRIyLGiFugd+vcXIDO3+A2Ml+FYnx/F5eL5V/0Yl3mzbFSqWGk5Uamu5/4jS08jmBlfXS98fcJ3F0K5XawDOLUYwNxzl/2+nNzk2dQM9EWEiuopEmmE866LFNxbeByHZyDr1pIsAR2lHE0ZAqQL5HzklfCzF0ECnSdDxr6ShxIOONtfYRAI+M8qyPYPn0P1w34tsx5cveFJNS4+ToN5+IQOXw30/F9ZygUVm0aM+mCZb6XZzs7qCRJniPibBVqeHk7mXMpD28sXMJjTTBeTZJQ4tLG6MQqM+y774yNLVAWY5AEvdcyX2sie9Warh9dxtp0smU+SUItewDle2rl7Da7+ZtWHb6uY24jvVqA6vHb8JWVEXD9lG1KWaSXdzca2f6Oabf4/2lMkJtLav4D6HMRkP9Ir1jyEiWWd3Hg1e7WHrN4NvJSFTzvdwyi8K3W1L52oSVosL7qw2cZBbI887CR4t3td/F/bvbONndwR2dSwCA+aSDi67sHQxynyRm+XRoIdFUI9Yhhb0cB60uqZ8LiVV88a1Uarg7ruNLx2/CRlxH1aZYtCl6NtOJlSHkvrb7dLX095hNcZ8zxPTchvJCfAyrjR/AVjyFlZ0XsdDvYr6/iwVH3HLdnphPnJBo7eJtCI1liDMrAr2HJSdmE3f6uV57aP3p8HiVi6WjoAIT5CIk+Es8W4L958/xOkKLhKOMSMy5PDII0AIhIkQTs2kyax2AXL/GOT7iNPJF4D7JWkpiK41DiDj7OFI+Jj7RXy5Qn5Jas+pJXR0A3OwU/aRnJMskd9HxjbPGfXOumzaXxSjOy13td3Pujdox0+/i1t4V/Njlv8kttvNJB3NOTdBzY0+iLTdCkIvF6aSTzTvXXm0MJXfn22w0naY25hqkukJKGvx7kf9dGbzqdW6jVhbSCdH/mt6lLHELuYrMmf3iZ0iXRJ/aogawb3HxiU4LWBLzlUotd1MgwsbF0zkT4b3VBtpRDHR39vlT8Xb5+i0J+pxYENRnqcfhkAvWN35EoOmenTTBXVdbADJu9Zxz6SDCRv2RRFrTs2kEfTmK0WvM4y/iem4pfffuNrC7nbdjx/ncLfTaue8d3duO4txv7V63ibRsZuV9IK6jHcW5b9v7Opfyd/uYM14sR7HqKyihcXTLjAvjnOfZNMk3Nknw4IgvzY05E2GO+XNyiWPORHgpHYc4OSFuI2MY7qnIcZaXSboS+azv/jJlSpAY2WScSBkzPBdDm3ZQAU/Eha5dtClmS+rheLtkfZKw8mucAyjj7FzGJadpMheWDSB31OUEVXKVw/SNj9dWpYbNuI4XGJe24izM5KwLm+buNj0T5e4iZOgh1xEqkwxB3MF4lvV5GInDNzY0R0/F9bwcGne5iYWYACJoXCVAv3/DHNyJF5j4uQ2NujFeR1JA5yCAYnGRXjL37AcACIuk9qx0GZFcDk0i7kxL/l0+oiZFT+pDi5UBIBd/pAj4+V4bAAY4Hl/75fhIw4y2KOkdPBjXsRTFmZU56aiEOlTXpt1TvjdNlPsZEpHYqtRQjWI0mX7LV3ZRvZs2zX3RXqg2cK52A55qzOVc1o+mCSKbYgsZl/YkWc6ZU2/DRFivNvDs1Am8sXMJM/0uFp1ouuQ
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "imshow(cic_paint(jnp.zeros(mesh_shape), resx).sum(axis=0),cmap='gist_stern'); colorbar()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "id": "d7ce5401",
   "metadata": {},
   "outputs": [],
   "source": [
    "k, pk = power_spectrum(flowpm.cic_paint(tf.zeros([1]+mesh_shape), final_state[0]),\n",
    "                       boxsize=array([100.,100.,100.]),                      \n",
    "                       kmin=0.1,dk=2*np.pi/100.)\n",
    "\n",
    "k, pk_jax = power_spectrum(tf.convert_to_tensor(cic_paint(jnp.zeros(mesh_shape), resx).reshape([1,100,100,100])),\n",
    "                       boxsize=array([100.,100.,100.]),                      \n",
    "                       kmin=0.1,dk=2*np.pi/100.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "id": "28aafee9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.legend.Legend at 0x7f97746730d0>"
      ]
     },
     "execution_count": 76,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAXYAAAD8CAYAAABjAo9vAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/YYfK9AAAACXBIWXMAAAsTAAALEwEAmpwYAAAjvElEQVR4nO3dd3xW5cHG8d/9ZBBIQoAQVhIIkAAJhBn2VFkOwKptGVatCqWtWttaAV+rVF9bW2uLW3EPFK11gAo4mDITIMwAIQRCCCMDAgkEMu73D8dLKUhCxnnG9f18+OM5POecKyjX5+Y+9znHWGsRERHv4XI6gIiI1CwVu4iIl1Gxi4h4GRW7iIiXUbGLiHgZFbuIiJfxdzoAQNOmTW1MTIzTMUREPMr69evzrLUR5253i2KPiYkhJSXF6RgiIh7FGLPvfNs1FSMi4mUcLXZjzBhjzOzCwkInY4iIeBVHi91aO99aOyUsLMzJGCIiXsUt5thFxLeVlpaSnZ1NSUmJ01HcUlBQEFFRUQQEBFTq+yp2EXFcdnY2oaGhxMTEYIxxOo5bsdaSn59PdnY2bdu2rdQ+ungqIo4rKSkhPDxcpX4exhjCw8Or9K8Zjy72vVvXsHPNp5w5VeR0FBGpJpX6hVX1z8ajiz33yyfouHAiPNqaXY/0JXn2r9m2ZC7Fx3KdjiYiHujJJ58kPj6eyMhI7rjjDqfjXDKPnmNvd+Ms1m8cS8nulTTOT6HbgXcIzHkLlsE+vzbkNu5JQLsBRHcfTpNW7ZyOKyJu7tlnn2XBggUsW7bMo2+a9OhiD2/anPARE2DEBACKik6QtnEZJ3YtJ+RICp1yFxKS9yGsg0OmGTlh3aF1f1p2vZwW7bpiXB79DxYRqUFTp05lz549jB07lltvvfX77fv27ePWW28lNzeXiIgIXn31VSIjI4mLiyMjI4PCwkKaNGnC0qVLGTJkCIMHD+bVV1/lrbfeIiMjgwMHDrB//37uvfdeJk+ezNKlS3nwwQdp3rw5qampXHfddSQmJvLEE09w6tQpPvroI9q3b1+tn8Wji/1cISGhdBt8DQy+BoDTZ06zfctaCtKWEZSzlrbH1hJ+7HPY/CeOEUpWSFfORPalacIwWif0xxUQ6PBPICJ/mr+N7TnHa/SYCa0a8uCYzj/4neeff56FCxeyZMkSPvnkk++333HHHdx0003cfPPNvPLKK9x111189NFHdOjQge3bt5OZmUmvXr1YsWIFffv2JTs7m9jYWAA2b97MmjVrKC4upkePHlx99dUAbNq0ibS0NJo0aUK7du24/fbbWbduHU888QRPPfUUs2bNqtbP61XFfq56gfVI6DUEeg0BoKK8goz0zRzZshizfy2Rx1OJ3rkSdv6Dkx/WY29QAsUt+hDWcQgx3YYS2CDU4Z9ARJy2evVqPvjgAwB+9rOfce+99wIwePBgli9fTmZmJjNmzODFF19k6NCh9O7d+/t9x40bR/369alfvz6XXXYZ69ato1GjRvTu3ZuWLVsC0L59e0aOHAlAYmIiS5YsqXZmry72c7n8XLTv1J32nboD36wPPbA/k/2bFlOxdyXNjm6kU+ZsXHtfoHShH+mBsRRGJFE/dhBtegwnpHEzZ38AER9wsZG1075boTJ48GCef/55cnJyeOihh3jssce+n44597vnfq5Xr97321wu1/efXS4XZWVl1c7o05PMxhgiW7ej35jbGXDnq8Q+kMrRO3eRMvAF1racxBnrT+KBd+m8/JeEPBFH1sOJpDx9M5s+nU1BTobT8UWkDgwYMIC5c+cCMGfOHAYNGgRA3759WbVqFS6Xi6CgILp3784LL7zA4MGDv9/3448/pqSkhPz8fJYuXfofo/na5FMj9soIb9qM8BHjYcR4AIqKi9iRupwTO5cTcjiZjrmLCM37CJLhsGlKTsMe2Nb9aZl4OS1iu2Jcfs7+ACJSo5588kluvfVWHnvsse8vnsI3o+7o6Gj69esHfDOCf+edd0hMTPx+3z59+nD11VeTlZXFH//4R1q1asWuXbtqPbOx1tb6SS54cmPGAGNiY2Mnp6enO5ajKs6cKWX31jUUpC2j3oG1xJzcTATHACgkhH3BXTkT2Y82gyYQ0bqDs2FFPERaWhrx8fFOx6hRM2fOJCQkhHvuuadGjne+PyNjzHprbdK533V0xG6tnQ/MT0pKmuxkjqoIDAwgoedg6PnNP7cqyivI2L2Fw1uWYLLWEHU8lehdqziz8wlSmv+IuB//ibCIKIdTi4gv0VRMNbn8XLTv2I32HbsB31yQ3bdnJwfmP0Kfwx9w5pn5JLf+GZ1v+B8aNGzicFoRqSszZ8507Nw+ffG0NhhjaNO+EwPufpN9E5awLbgfvbNe4sw/urL+nYcpPX3S6Ygi4uVU7LWofafu9P7DPLZdM4+9gXH02vl3jj7ahY0fP0VFWanT8UTES6nY60DnpKF0m7GYjZe9ToErnB4b7+fAn7uz7au3sBUVTscTES+jYq8jxhh6DL2WuPvWsqb3E5RbS+cVvybjL/3YteZTp+OJiBdRsdcxPz8X/a6+hVYzUlnZeSYhpXl0WDiRbX+9gv1bVzodT8SnDRgwwOkINULF7pDAwEAG/vi3hP5hMyva3k2rkzuJfv8qUv9xLYcztzodT8QnrVq1yukINULF7rDg4BAG3/wn+E0qy1v+nLjCVYS/NpgNT9/E0UP7nI4n4lNCQkIoKiriiiuuoGfPniQmJvLxxx8DkJycTNeuXSkpKaG4uJjOnTuzdat7DsIcvfP0O0lJSdaTH2pfkw4eyCLj3zPpk/8RFbjYGj2R+BseILhRU6ejidSa/7ircsF0OLSlZk/QIhGufPSiXwsJCeHYsWOcPHmShg0bkpeXR79+/UhPT8cYw/33309JSQmnTp0iKiqKGTNm1GzOH1CVO081YnczLSNbM+iuVzhw4wo2hg6j5/43KJ/VlfVz/qh3u4rUAWst9913H127dmX48OEcOHCAw4cPA/DAAw/wxRdfkJKS8v3je92R7jx1U23jOtP2nvfZnrqakwseJCn9SfL+9ibZiXfSdcydeimIeK9KjKxr05w5c8jNzWX9+vUEBAQQExNDSUkJAAUFBRQVFVFaWkpJSQnBwcGOZr0QjdjdXEL3/vSavoiNI97lkF9Lum9+iIN/6crWRS9jK8qdjifidQoLC2nWrBkBAQEsWbKEffv+/1rXlClTePjhh5k0aRLTpk1zMOUPU7F7AGMMPQaOJmHGStb0f44SG0iX1b8j88+92bXyQ3CD6yQi3sAYw6RJk0hJSSEpKYk5c+bQqVMnAN544w38/f2ZOHEi06dPJzk5mcWLFzuc+Px08dQDlZaVsXbebNpu/ieRHGGvfzuOtr2G1kNuJDy6o9PxRKrMHR7bm5+fT8+ePf9jhO5OdPHUywX4+zPoul/ReNomlne4j2IbSI/0Jwl/uQ8Zj/Rm49yHOHZwj9MxRTxGTk4O/fv3r7FnpztNF089WIP6DRgycRowjYz0NLK/fpvm+z+jx47HYcfjpAcmUBQ7hvbDbqRhs9ZOxxVxW3X1ZqO6omL3Eu3j4mkf9zDWPsSuHZvJWfkOkQcW0GP7X6nY9jd2BXXhVIdxxF42ieAmrZyOKyK1qMaL3RgTD/wGaAp8Za19rqbPIRdmjKFDfDc6xHfD2r+QtnUDR1a/TfTBRXTb8r+Ub36EHQ26c6bjOOKGTaJ+o2ZORxYBvlk/boxxOoZbquq10EpdPDXGvAJcAxyx1nY5a/to4AnAD3jJWvvoWb/nAl601t52sePr4mntq6iwbN+0lry179D20Oe0IYcy6yI9uCdl8T8ibugEghqGOx1TfFRmZiahoaGEh4er3M9hrSU/P58TJ07Qtm3b//i9C108rWyxDwGKgDe+K3ZjjB+wCxgBZAPJwARr7XZjzFhgOvC0tfbtix1fxV63yssr2LLha44lv0vskc+J4gil1o/0kN7Q5UfEDv4pgSGNnY4pPqS0tJTs7OzvbwSS/xQUFERUVBQBAQH/sb1axf7tAWKAT84q9v7ATGvtqG8/zwCw1v7lrH0+tdZefbFjq9i
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "loglog(k,pk[0], label='flowpm')\n",
    "loglog(k,pk_jax[0], label='jax')\n",
    "legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "id": "dc6a66c0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x7f978c1fb6a0>]"
      ]
     },
     "execution_count": 77,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAYAAAAEACAYAAAC6d6FnAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/YYfK9AAAACXBIWXMAAAsTAAALEwEAmpwYAAAsrUlEQVR4nO3dd3hc1bXw4d9SL5Yl2ZJsS7LkJrnghpELzfSAaQ4QAiShJ4YAN5d0cpNLKrlJuEku5COASUwLJSEhYIKBUAJ2jDu4V1kuarb6qIykafv7Y0ZCFrI0kkaaMzPrfZ55sM6cPWcdbO01e59dxBiDUkqpyBMV7ACUUkoFhyYApZSKUJoAlFIqQmkCUEqpCKUJQCmlIpQmAKWUilB+JQARuURE9olIsYjc18P7IiIP+97fLiLz/CkrIv/he2+XiPxq8LejlFLKXzF9nSAi0cAjwEVAGbBJRFYaY3Z3OW0JUOB7LQQeBRb2VlZEzgOWArONMe0ikhXIG1NKKdU7f1oAC4BiY0yJMcYBvIi34u5qKfCM8VoPpInIuD7KfhX4hTGmHcAYUxWA+1FKKeUnfxJADlDa5ecy3zF/zumtbCFwtohsEJEPRGR+fwJXSik1OH12AQHSw7Hu60ec7JzeysYA6cAiYD7wFxGZZLqtTSEiy4BlAMnJyadNmzbNj5CVUkp12LJlS40xJrP7cX8SQBkwvsvPuUCFn+fE9VK2DHjZV+FvFBEPkAFUd/1gY8xyYDlAUVGR2bx5sx8hK6WU6iAiR3o67k8X0CagQEQmikgccD2wsts5K4GbfKOBFgE2Y0xlH2VfAc73BVeIN1nU9O+2lFJKDVSfLQBjjEtE7gHeAqKBFcaYXSJyp+/9x4BVwKVAMWAHbu2trO+jVwArRGQn4ABu7t79o5RSauhIKNW52gWklFL9JyJbjDFF3Y/rTGCllIpQmgCUUipCaQJQSqkIpQlAKaUsaleFjaF8TqsJQCmlLGh3RSOXPfxvVh8YutHxmgCUUsqCSmqaAThY1Txk19AEoJRSFlTZ0AZAeUPrkF1DE4BSSllQR8VfVm8fsmtoAlBKKQuq8CUAbQEopVSEqbD5EkC9JgCllIooHc8A6u1OWtpdQ3INTQBKKWUxbU43tS0OCseMAIauG0gTgFJKWUxH//+CiaOAoXsQrAlAKaUsptLm7f6ZP8GbAIbqOYAmAKWUspiOLp85uWnERUdRpl1ASikVGSoaWhGBcWkJZKclUKYtAKWUigwVDa1kjognPiaa3PQk7QJSSqlIUWlrY1xaIgA5aYnaAlBKqUhR3tBKTloCALnpidQ0t9PmdAf8OpoAlFLKQowxVDS0kp3qawGke/9bMQQPgjUBKKWUhdTbnbQ5PWR36QIChqQbSBOAUkpZSMc3/WxfF9C0sSP59bVzmDo2JeDXign4JyqllBowW6sTgPSkOABSk2K55rTcIbmWtgCUUspCOhZ+S44f+u/nmgCUUspC7A7vaJ+kuOghv5YmAKWUspAWh7YAlFIqItnbtQWglFIRqaMFkBSnLQCllIoodoebxNhooqNkyK+lCUAppSykpd1FcvzQd/+AJgCllLIUu8M9LN0/oAlAKaUspaXdNSwPgMHPBCAil4jIPhEpFpH7enhfRORh3/vbRWReX2VF5EciUi4iW32vSwNzS0opFbrsDvewDAEFPxKAiEQDjwBLgBnADSIyo9tpS4AC32sZ8KifZX9rjJnre60a7M0opVSoa3FYqwWwACg2xpQYYxzAi8DSbucsBZ4xXuuBNBEZ52dZpZRSPvZ2N8kWegaQA5R2+bnMd8yfc/oqe4+vy2iFiKT7HbVSSoWpFoeLJAuNAuppMKrx85zeyj4KTAbmApXAr3u8uMgyEdksIpurq6v9CFcppUKX3WGtFkAZML7Lz7lAhZ/nnLSsMea4McZtjPEAT+DtLvoUY8xyY0yRMaYoMzPTj3CVUip0tbRbqwWwCSgQkYkiEgdcD6zsds5K4CbfaKBFgM0YU9lbWd8zgg5XATsHeS9KKRXSXG4P7S7PsLUA+ryKMcYlIvcAbwHRwApjzC4RudP3/mPAKuBSoBiwA7f2Vtb30b8Skbl4u4QOA3cE8L6UUirk2J3DtxAc+LkjmG+I5qpuxx7r8mcD3O1vWd/xG/sVqVJKhbmOlUAtMw9AKaXU8PhkJVDrPANQSik1DDq3g7TQKCCllFLDoKVjMxgLjQJSSik1DOwObQEopVREanF0PATWFoBSSkUUe/vwbQcJmgCUUsoyOlsAmgCUUiqydLYAtAtIKaXCzwsbj/Lq1vIe32txuImLiSI2eniq5uFpZyillALgybWHiBJh6dzuq+p7RwElD9MkMNAEoJRSw8rW6qSm2UGb001C7ImVfUv78G0ID9oFpJTq4oHXd7O+pDbYYYS1BrsTt8ew/3hT5zHvcmq+FsAw9f+DJgCllE9FQytPrDnEL9/cG+xQwlab0027ywPAzvJGAFbtqGTeT9+mqc1Ji0NbAEqpIPjoaD0AHx9tYE9lY5CjCU+2Vmfnn3dW2ABYubWCeruTPZVN2Nu1BaBU2Hl1azk/fHUnh2tagh3KSW05Uk98TBRxMVG8sPFosMMJS10TwK5yG063h7XFNQDsO96kLQClwtFD7xzg6XVHOP/X7/P1P2+luKqp70LD7KMj9cwdn8alM8fy94/KO9elUYHTYPcmgOnjRrLnWBObDtXR5Bv7v/9Y07CPAtIEoNQQK62zU1LTwl3nTubLZ0/izZ3HuOi3q7n7uY8s09XS5nSzq6KRefnpfGFhPk3tLv6xrTLYYYWdjhbAmZNH43B5WLH2ENFRwtQxKew71uQdBTRMm8GAJgClhtyaA94m/tXzcvivS6ez9r7zuevcyXywv5olD63hK89sZntZQ1Bj3F5mw+UxnJaXzvwJ6UzJGsHz2g0UcJ0JYEoGAO/sqeK0vHROm5Du7QJq1xaAUmFlzYFqxqUmMDlzBACjkuP49sXTWPvd87n3wgI2lNRy5f9byxf/sJ7HPjjIhpLaYe9+2XLE+wB4Xn46IsINC/LYWtrA7gprtFDCRYPdAcDc8Wkk+uYALC7MYOqYFGytTlqdw/sMQCeCKTWEXL6HfEtmjkNETngvNSmWey8s5PazJvLs+iP8eVMpv3jDOwSzo1vg1Lw0Ts1L59S8NCaOTiYqSnq6zKBtOVLPpIxkRiXHAXDNvBx++eZent94hJ99dtaQXDMSNbY6EYHUxFhmZI9ky5F6zinMorn9k4Q/nKOANAEoNYS2ldlobHNxdmHGSc9JSYjlrnOncNe5U6htbmdraQNbSxv4+GgDK7dW8NwGb1fMBdOy+OMt8wMeozGGj47Wc/60rM5jaUlxXD5rHK98XMF/XTp9WL+VhrOGViepibFERQlnTB5NdVM7p2SPpKHL6CBtASgVJtYcqEYEzpx88gTQ1egR8VwwfQwXTB8DgMdjOFjdzO/fP8jKbRXYfBVIIB2ptVPX4mBeXvoJx29YmMfLH5fz2rYKrpufd8J7bo/haJ2diRnJg77++/uqKByTQnZa4qA/y+q6/v3de2Ehd583hagoYVRyHJkp8VQ3tes8AKXCxZoDNczOTSPd17XSX1FRQsGYFL64MA+3x/Chb8x4IHX0/5+Wf2ICKMpPpyBrBM9vOPFh8PayBq76/VrO+9/3eXv38UFdu67FwW1PbeJ37xUP6nNCRYPdSZovAURHyQlrAU0dkwIMbwtAE4BSQ8TW6mRraQOLC/z79t+buePTSEmI4YP91QGI7ERbjtaTEh9DQdaIE46LCF9YmMe2Mhs7y23YWp3c/+pOlj6ylkpbGzlpiTzw+m4cvqUNBuJfe6vwGNhR3jDIuwgNtlYnI0/Sgps61psAhmszGNAEoNSQWXewBrfHcHZB5qA/KyY6irOmZPDB/urOhcMC5aMj9Zyan97jA+arT80lPiaKH7+2iwt+/QF/Wn+Em0+fwLvfPIcHrprJ4Vo7z6w7POBrv7PH24LYd6yJdpd7wJ8TKhp76cLrbAFoF5BSoW/1gRpGxMdwal5aQD7vnMJMKm1tFFc1B+TzABrbnOw73sR
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "semilogx(k,(pk_jax[0] - pk[0])/pk[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2b58a54",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}