JaxPM/dev/JaxPM_ODE.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 6,
   "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": 7,
   "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": 8,
   "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": 9,
   "id": "b588c18c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from jaxpm.kernels import *\n",
    "from jaxpm.painting import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "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": 11,
   "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": 43,
   "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.,10), nc)\n",
    "    return final_state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "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.92 s ± 9.63 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": 21,
   "id": "e5148fed",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x7f932c2f1150>"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAPsAAAD7CAYAAACscuKmAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/YYfK9AAAACXBIWXMAAAsTAAALEwEAmpwYAAC+20lEQVR4nOz9W8wtyZIehn2RWbXW+v996dN9zpyZES8akhrKoEQJFGRblgCBAG3YsAXxiQRtyKBoGuMXy/IN5lAvejIwgA3DfB3YMmhYgEnLAugHgpIw8MAQDBMUZYEySZMeUsPhzJw51+7e+7+sVVWZ4YeIyIzMqlr/v7sP926bnUD33nutWlVZeYnLF19EEjPj6/Z1+7r9/38LH7oDX7ev29ft/bSvN/vX7ev2D0n7erN/3b5u/5C0rzf71+3r9g9J+3qzf92+bv+QtK83+9ft6/YPSftSm52I/mtE9LeI6FeI6Od/XJ36un3dvm4//kZfNM5ORBHA3wbwXwHw6wD+CoD/JjP/jR9f975uX7ev24+rDV/it/8FAL/CzH8XAIjo/wjgDwPY3ewHOvKJXgCgjW97obN1jfsmBCBQd53eg/TzQOBQ/6TEoJSBzEBK3U31PoHk75mBnPW2LI+x+/q+ct9V0muDfLkkiEDl7dd8qpH8j6x/5J6x9e6Z8VwBTvauW63cw40LSMZEv+ufI13an7dyXyKZP3LX233ZXYd2uMr9Sd/frmeu00BU/o3VHdz7bH3a9725z8avfV9sbTC7sQMQA3iwbcYgBjDPsr7st8MADgFkv2UAnOVPW4M2V+U5uYwZ5/q8M99j4vPmS36Zzf7bAPx99+9fB/Bf7C8iop8D8HMAcMIt/rnhv6ob4R1asAXBoCi/pdMRdHOji0a/t4U0RIAIfHtCPg3IxwHpNGB4mBF/dA+aZvDbOxnwMogiPOhwAB9H0HkC3z8AnMEpy2Y7jHId1wXf3AMAYgSFINemjPzpZ+BlKcLFTww4PzkWFEjuGeWdKEZ5ni0g/3vO4MsEnqb6mT1Px5CIyiYNxyNg9/PNC8MgY0PHAxAj+OFR7s8M7gQmxVjGoL2fE5o5A+MIOp1kLo8H+eruHpgXuaddBxUo+g4UAxAC6OYk4z3Ncn0SgUomRHKuv+PcD2j9e/fe5ffW/3muAk3nmbxAiBE0DnV8iMDzDCxL+Q19/BHSt14XIUSXBPrO98Dni/w2RuCbHyO/OoGWDJoWYEmgaRZFcT5L3w4HEQqnAzBE0MMZ/HAG5gn58VzG6P8x/yXstS+z2Z+jnsHMvwjgFwHgdfgmI8b1pADtIPpmm5kZsDWkmpfnWQb8IAOQX9wAkUBLlXrhvIAuCfHuArosoPtHucVBF+9JJokeLzJJ41A2KmIAEkBkWpkBTqqZIpBT3eh+k2fR6LZhiAgco0x+0EWcuSw86jdbGeHQLjDbYBxAWHSxxTKW4CAbwN8i7gsT21jN2AeqAgJQQSfWAmUZExoH+Tzn9U3dvcgJQOm+m2vOMpQXFUzzIkLR96+3UGwjDoOM6bIAS3dNt9FNuJYx9gI2t7/lACCl1cImIiBSuxalg+Bplr8uSxWkXrg9nBE+lzWVXh6BEaBhAI0qKFMCPZwRl1bImeancazvHgh8e0S6PWDIDLpM4EXfZW8NufZlNvuvA/gd7t+/HcBvPuuXXhvDSdSt5rVEv7hSLpqJh4j88gAOhPC4gFKSTb+kYh6RLigi1YxDVG0tWgILwIHEpCranoGEuglU02nH1ZrQ66NuPGRgEYugfW8CslgGFAicWU3pnXfvNYkfByLAhJATnru/2Wq6uMqSDwG0sX9lo7M8j2TDUe5M3L35c1ZZ6StQ3tk25q7rYdq5H6PuHe2dy306jV7G+lqz+XXrbHMs3YY364ZsHG0c7B7LIlYiEYAjEFT4+77NswgY2+wh1D1i69/ebwjgIYgieu48a/sym/2vAPhZIvpdAH4DwB8D8N+69gNxbTqTl1nNVCdtbQN4EztXSYjMwDiIdr45IX3yEsurAz77x45II+HFdxPG+4zjdx8QP7uTAUxZfpsSmAIoivalx4vcd5pFW50n0JKAeanPNFM3OMFjpnxK8rlNypJkA82zPFMnn2wR2i2c77s5VmYtxFhN214opiQmNXkBFMQV6TWMPrPZVLawes3uW9LrlwXMTsgBtS/9v22ImIGkrtc4yILOLJr5dGwFwbKA98bEBENKYk1MgwhlJ0yfwinKRu/N+E67IzM4dJtcLQpRSmpVLkujhcvQ5K5P5lJdJgzzItbALNYAzDXz1oJt9EEXStY+MgMpI356j/D2DLp/BJ8vVXnZe14Zgy+82Zl5IaL/PoB/D7KE/y1m/utXf+QXq2kiIlkMw6ALmotP2ggGDm5Ck1x/PIBvjphfH3H5eMDd7wDygRFSQB4Jhx8F2Xy2kGzQkOtCL2ZYKj4vZmoBPBM6qjXLIsu81lgmAFIWDU5iKVSgpW6QRpr7ZgIiihAhqAYxiwNucU+zsyp0MYZBNLH6sgXI2doQbsOTFxpA1Y5ZBSXnskBXzfWt6V/OMlfFUgJoiMAoGxZEApqqP74Sgn5DZgZDN1AIFQfp30uv534jb7WtDd83G5uCb7i10bsC6hYVi2WRDY5pAi6X9r6D09w5u71BYMOmANkTaubz/YPM67KAVXg8t30ZzQ5m/osA/uIX+GH7cvaZbyYtPUDSfY8lAdOM4W4CD4TTDyLyCBzeZoz3uS6iHNYbam8TcG7mEkCzuLlfiKaqzSQ2oeI3Wvd+RCTYgG2sblM2m4lVG6sQYRWYRcsMChCpWce2Mb2hGLr3t82g15Be44HE3hwtWIFaXqz9pMMBdHtTgVFdkLQsOociyA3AQgzg4wHpoxtZ0EE2+zAvxfoqY9xHTGzM/brxPq6zSvxG33SVnuHj+sbMIM4o822Ky/CB7vnNM3StrwBNOgrOczqCDyP4NCLfjMjHiPnFAEqM8e2MMGeEzwRYppTWwJgfjyum/Zfa7F+qbQFz5p/ZmPCVCeGs/hAhvo0gBm6/OyCPwPGzhOFRffbeTO360INC7yIpZcKt77rRbZF3qHJpDk2ncdTf0FoTo5rByHl9H0Vy6XRq/DeaF13oLfjJZfOugStOqIsWaP1+v+lTkvv47w4j+OWtAKSnAbRkhMlQ7AQCgdQCE3M4IL844vKtEzgS8kAIM+PFj47AWTQfmVArc5JXYGXv95Z36xF2+3fo1tm7NJsTG+ctwbG34cszW/O+RFWiAMv5dsT8+oDpowHzbcDlGzIut9+PGB4zbi6z4E5qiVYr9fnt/W72fvHbIkwKAMGZq/1vN0AsQ4jpMiOEgNOnR+RIGO8WhFk2OwdxE2gh0WSZBVnt+1WAI51MD/D0msZAKnMz/IawsJVp+l6jK5BDlMFRnzEMZROShlcb68dZH8W/M8BsWQBW05qomtvdODXx2nQlAuDHlmgtZKwFsSbodMTy8S3yELC8HEFLxmnJgoU8nsU/tT4HEnCJGcNjEr87EsKcRUh1/S7msEUtzKqJQeeomxf3Ts37USihS6QERMJTPn4ZN6gfHoJGJsQ9M2upuFXvEE0ugjRJZIemGSES4jliOAYsN4T5JSFOhPQ5ISx0XWnljbW60d7rZme4RQQUwAwA8l7c2pr5sYPGJp1PRDmDphmny1QXvTXbjMMA8KzA2sb9zQ8/jGJaGUaQGUiq/dMifTuq+XU8gMdBAL2cQSkhG3DjQ0Dl/UUak5rbBFQTNyUNo3RxYwu3pQRQEPAo52o5BJINb76fwwtkPNduxDX032uorQ1hC19AtwP45S0ef/KE5RRw/pgQEpDH1xjvFozfewu6exCNnBlMBB5l7sZPJX7MRKCURKtbvwPJ7lFhzuCiBSlGGX8L32XDRWy69e9mHZpFoOvGgDXy7sJG89+xulAEqKCuPrxdRVtW3Gp8HaAZgmBJKYPuH0HzIpuRCOePIy6fMOIEHD4jUBYEvkStDLwGVuvsmgh7r5u9oPHa9oAVr31X0tV8pxwA0j/7tauDTuV++m+LdwNiPssD6u9YTVXfH64bp6DLywIOETTkEjffEk6Usyy
      "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": 12,
   "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": 13,
   "id": "aa14a004",
   "metadata": {},
   "outputs": [],
   "source": [
    "from jax.experimental.ode import odeint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "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())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "e801c3eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "res = solve_ode(init_state)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "de770afb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3.86 s ± 44.6 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": 17,
   "id": "73e7cc3c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x7f93d02f4490>"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAPsAAAD7CAYAAACscuKmAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/YYfK9AAAACXBIWXMAAAsTAAALEwEAmpwYAAEAAElEQVR4nOz9S6hl3ZYmhn1jzLnW2vucE/E/8t58VGWalFVpEBiDG7YbarhAGIwtXC0Jy1jIIMiWwUI2Vsl9mwKDsLsJNsgPsAQ2yA2BMQXVsDGmJFPYWIWsh0uVmXUz773/IyLO2Xuvteacw43xmHPtiPjvrcok/ovzXxBExDn7sdacc7y/8Q0SEfxw/XD9cP3//8Xf9w38cP1w/XB9musHYf/h+uH6C3L9IOw/XD9cf0GuH4T9h+uH6y/I9YOw/3D9cP0FuX4Q9h+uH66/INefSdiJ6L9KRP8+Ef2HRPTX/7xu6ofrh+uH68//on/UOjsRJQD/HwD/FQB/BOBvA/jnROTf+/O7vR+uH64frj+vK/8Z3vtfBPAfish/DABE9L8D8NcAfFTY8/lR5tdfQgCAAGr2CwFI/9KLgZYBIQBzA7OgNbJfEiAAp4YlF5zTjl9LL0gEZBAqBD/Zn3CtE7aaIJX1g5u+z78rvsw+VpIADBALiAREAEFQG0M2/Qyy9wjb/U8NKTVAgCYEJgFTf5gmhLInfRB7Vqr63CQfuAd/xOFn711Nf0UN8TyH1w//F3r/e4QAyfoM8Rn+PurP5n/H59o9y4fuy3/Gd4aDhxsSOn6X/Z/q8LnU/yZBX7PhI/yeJD6/fwVVgIv+v6+xHNZAn4v06/1Z0/G2Jd0953hm7teb+vf5Gh4+i+3zyR7E/iYC2M5aa6Qfa18qfl4xvM8v/7n0//saQoD9zdco15cPnp4/i7D/ZQB/OPz/jwD8l+5fRES/D+D3AWB69QX+yj/3L8dipg3gXfRmK+KQ1YVw+zWgngXlL6+YloKyJ0hhXQghPH5+xV/5tZ/j9179FP/8l/83fMkFP04L3rQN/5Of/pfxd7/9TfzRt5/h+rJAbgm0MagQeFPB46qLIwyAgf3zCjoX5KVgWYoKPICXywL85GTv1WcqjwJZBMuvX/DlqxdsJWMrCXOuOE97PPt1n/DzP3kNWhNoI1AhzG8J+WIHc8fhgLfJ/s3DYXMBZb1n3gEqwPQioAakTVSxZIIwkGw920Sok79HD1mdCG0Cbr9GaDPAmwqHK6CWgTbrn/1V032a9PP5yuAKtCx6eO2A+/pJFshiEmYHkh8KUq5ojSGV0AoDG+sB3Ri8A9M7BlVXQII662fxRkg3PchcTV9O+l11EQgDbRZIEvCuazu9Iyzf6JqM60NNdA0a0BZCy4RyIpSFUE+E8tjXWxjYXwtaBsACIejnV/079szOsCRdn+mFQMWeYxC1ehbsrwQyCeShglJDmho4NTyeV8y54rpN2EtCa4TWCLUktGsGWEBTA5mwi5AankpAJVAlpBfG9I7sWYH/6H/9r31UYP8swv4h7fFeTCAifwDgDwDg/Ju/I/UENNOkVPXgcgXy2t+aNv34uhAuecF+msAbgRrpBjBwnWa8e73gT9dX+Du338bn6YL/VP4GL3LGS1lQhMEsSLmhnQUyM1oltEKgwsALh5UUAEgCzgJpjG3L2G8Zcs2gnZCqClJ51IMujwU0NzydV7yaV7wDUCqjNsLzOsdz7FUfVFiASQ9xqQxhVRzZVlHMQrmgt+xKT8IKSxZQ1UOXVoIkO3gYPCQAQgQiORzeNnUrQQKktSs791iomqH0+wlrqh5Pfayo4+YLYYwAhQVI/mH6s/PDhqfTitoYpTKeX04oq21+FjQGSlNhkaz3TA2gRuBiik30jNBohd36FxPCjcA7qfEoEspLCFhfsb62+n135RhegnsM0j8bpOdMGBARgAjNXkAV4NWfm2KdZdLv5dqVKIhQzqYUoJ7jctqRuGHOFYkbEje0RCAisH0eTna/Sd2eakJOa9IzeSHkKyGtQL7oa1vq+/ah688i7H8E4HeG//82gH/wne8gvaHYMHeHBzeIbPPzVX+WL4RaWRevAZIIkoB9T9hrwqXM+Mn+BW4yY6aKTRJW0yZMAuIGBkNYLVVjNo1MIBCk2X2YSwUArTJkS0jPg4klgcyiVmepSLlhzgWZG5K9rzZGG9R6Kd0fFBaQEFoCaBJQoaPLPFzu+skEtMmsTBKgqcCDAF7177Z219XXOD7P3d3RHRd9vZhXI2P85O7voDz8c8BDrCUECW0AU8II91RMWOZccJ521MaomXBdJ5R4Rv28NpuwWxiFjewsUCghanaLoxsNv0/qnuEY2pjgtgnhiQAAm4KItR8EnoBDtHEIr1hATLqPlSJMEOmeF7ivnxuy7mJ3hTsKOgEWNgqYgdbUvW/pTmqbWnMUUmFfCemqhjGtppiIPhr9AX82Yf/bAH6PiP4xAH8M4L8J4L/1y7xRMtCSgE6kN9gEvOsh5KLaV+ModXl5p1AGvnl0TfjmcsatZFzKjIe84TdP78DUcKsTznnHnCv2nFB2Qi0cBxUsaCfbKdKDl84F87Jj3xNaY4BEDyKpgIMFdKrgJJjmgpwbamO87DOe1xmX24yyJ7SLLSlBPYdiHoQ9Q7qpZaam6+Bxa6wN2/pMgpZV0EkA7D1WEwbKg6AthDqTekYv6rJWcxdaAiQTqEqEHx7X8S6QOiibZO571vVts4TCoDZojDH2zWbJ/d+5YToVXbemb66Ncd2neDYiALnpwbXPdSFDMoVoAqU/1HXiva+LeCxPg2BZ2CBJQ0CQQFa13uVBQ5dy1venG8IjAoA6q8cWgs1AszCBN11bj9dDCboyHuXRvQ0Pz0Yvwe5TGiCNUCqDSTBnNUinqWBKDeuesbWEWhLqNXcBh52fSuCbeT1bD+nCQ3Nv7SPXP7Kwi0ghov8ugP+TbhX+lyLy//5l3tuSuacZaAK01Q+exleAPpCQbs69pZEG0EZYbxP2PeG6TZhSxVfnR8xccc47TmlHThUpNU2StcHtJEAm+9CkgjxNFXOuqFXdfY9DwQCWCkqC+bSDWZBzRSJBE8KtZKx7DkFPb1N8B4BI/njcl7Ye97WkLt99Eq0lMWG1+zNrFPfOgjYTmgh4QsSTGkNTfDYY6hoWOSS7uOhnuDWQZEKeVdDbhG714iDR4SCJrVskNXPDNBcQAaUIRAhNCFtJarXcsGWBVOhBHiynC7iwAEw97hV1zSURWjN3HojDHYk4mKBOKvwt6RrURQV6+0yVd36hULYawwvqSY5WftIP5Krh1n3iVDycGBWAP19z9334GXwdVcm1xqjS0IQwcQOTYGI1HuueIULAxupB7BTr74qPdxP4ap6w72v9DknHn82yQ0T+bQD/9j/s+7ha/CN9ge/9D9XY6vp+KOPMK2F/OwNZsC8q1ACQBnf8ss7Y94RaNUEUmtLiULAgnSqYGx5OK85TQRVNksgEE3oBWZa57FnjYSG0pJuVU8W6TmjPE9IzY/6GQ3iFNUEj7DEc2Wbp74gRQgh0zZy2HkBKhiYHS39+SdADCnX7qBJ4BVqiCHf8MPOulp0sRBIiVIvJqQEEtfJUAViCLqoNAqCQeQPUlZIAEEZjALmB5oaUKx5PGxI33LYJVTT8upYZtbDuQWHI2hN01PTZAECaxdamFEEqqOKJhHHvd4Kk7k7fWzMX+vizCNrSILOg7aotItkGdK3menrTZ+XVkrKuBNiUHInG/ON5tXVps77W8y7lQZVJWxr4rAnL07yDWS35bVPPpwlhL0mNRqHuUfk+mWJLq+5xvgL50kMp9QaHdfrA9WcS9n+UK5JBg7qMZNKoDZuAGsXBvY9r0wbgTVZ398RoWfAMtTK1MqT1z2+7ZTCbakpJAFJTa73smHLF69OKc95RGqM1xgaohgVCwOvWb6I1Ms+EUC8Z+W1CfiacvhY7xCqU65ekrrgdSt703sNyj/Gybai6gWapiwlaia/Ww/NZ6xvbND8gmSxMMOVSNDxKm0ReRDPe9lyetLKEUl0s2z5az2Lhx2r/HuJhyRqGpdSwLAVfni9gErwlQWm
      "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": 18,
   "id": "4bdfe23f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.colorbar.Colorbar at 0x7f936c1dc220>"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAATcAAAD7CAYAAAAPf9NJAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/YYfK9AAAACXBIWXMAAAsTAAALEwEAmpwYAADGzElEQVR4nOz9fbBtW3YXhv3GXGvtvc+5H+/1h7ol1MICW8KRlCAbBZOQODiyY8UmFnbFBKVsYVuxDAUGJ6QCIq7IZUpVJHw4Dq4oaQeCKYNAMcZWUcLio0JhVyThFlCAPjBCksVTt7r7ve737r3nnL33WmuO/DHGmHPMuebaZ597z7uvW++OqnvPOXuvj7k+5phj/MZvjEHMjFfySl7JK/mFJuG9HsAreSWv5JW8G/JKub2SV/JKfkHKK+X2Sl7JK/kFKa+U2yt5Ja/kF6S8Um6v5JW8kl+Q8kq5vZJX8kp+QcoLKTci+iYi+jtE9JNE9Lvua1Cv5JW8klfyokLPy3Mjog7AfwPgnwLwBoD/GsC3MPOP3d/wXskreSWv5Pmkf4F9fyWAn2TmnwIAIvqTAL4ZwKpy29CWd/TgOU5Fy08CARTKr0xPk34YAhAITAACAZFBcwSYgTnqDlSewo7JDESWbey4gcrzrA6VdAxyHllA+PR+p4RID+vHSsvbwjLucxcsIsrHrMWO4a+HIPeE5V99lnSk1jH1eGznDZSPCwAx5nMujp2fUzHmxvbk9qh/Wxlt+Ykfe+Mai+3tnvj7w/ac9WfXAX1XHBPjJD9t376Td9VdT3r3oj0Hcv+gzyECsXzee77CkfcrD/U8+af/iQf81ufms7b9kb95+AFm/qYXOd+7JS+i3L4cwN93f78B4B+rNyKibwfw7QCwwyV+Vf9P6xdBHs5tQpXnHAhEBNpsgIudbqLb2EPue4AIfLEFhh5x1yNuenQ3I8I7V8A4ga+u5fwUstIiAm1lHxyO4JsbeXnmGQgBtBlkO3vh/PjtOERA14GGAZhnxHeegMcJ4AiObqrU+566BUOvCiHI8UMA9b2cL+jEibOM9XgEH8f2gew6I8t93GxAXfvcPEe9JbrPdgvqAni/l+uZ5zyp7Hi2bddVB7PtRNFT34N2WyB06Z7y9Q34eNSJHctju+dEXSfPKJBc5zzL9nNMY2Dm9jMqbmpIC57dm8X45zndh8U2QPkcho38tDEBYGaEx48QP/xaug80zsDPvwk+HNK+9IHXER9dgKYITDNomtNx+OZGhrvbybi2G3Dfga734P0eOI6IN/s0pB8a/4v29d5B3vzcjB/+gY+dte3wZX/vwy98wndJXkS5tVaHxULHzB8H8HEAeBw+xOklJQIzpYmxehKb1FFfMlVkzAw6jkAXRJn1HfhyJ5N/mpN1gXFCmGaEcAQdRuBmLy//ZpBj77ayz/4AnmY5zmYQC6/rAMwAqwKY51IZRlXQXZeVWq+3dJryxAgExADqkJWcU2i0dv3+XIDcA7sX8wygAw2i8DiSfBbCyfuZj8UyeaKfvEtFlxQ7R/AsY6Kukwc9Tfn6Wsfwv+vYyZ+PY1LEPE1JKaxJstq6oPdvzIpQx5CUoio1v6As7rNXgDGAoUraKSd/beldtOuJURRyCGKN2bPxyny/R3gygIcefLkFd5AFxRRoZOBmj2DehL23JsOmGDLvNojbASFGYByBMJ28Z88njPkcw+MLXF5Eub0B4Cvc3x8D8Mlb97LVT1900xvFiujFXoKuc66SuSURQLZm4m4DdATaT/KCTrKqU8yKjqdJJmjfAZ0oMhCJosScJ0+y5gKAckIClWWi46MQknsh1oObrM5KoEBp0q0qNn+vvJiCs7F45X+LBbgQfyz7O527+tzcqEAA5FkkKwnrzy9tY2O0ewTo84zJWluO74RTeOq+6X5sVh8AUABHzve7nrxmHUZO72R9HS1h5kIp++2YWd7D4ygeK7ZpLMV20wQ6UHpvSBdKAAvLmkMAD6FaROj0vbqjyFP54s85fxHl9l8D+Coi+iUAfg7AbwDwvz5rT3swtpJ3nVgDbsVKD9WvcA5joL4X62u3RXz9IeaHWzz5JTvEnnD5mQn99YzNZ69AT64UY1MrZY5AYLHG5hl0cxAXZ5pEGe0PYrWNYzrnYiwmOh70TgmZVXUck/KhrpPPO0IKULdwMeciyT0QBWyuV3JLYZOKgf1haV1sNgsLIu3jh88s97+ytoqJZ7/Mc3MCpW3d2PL1yBShrpMxuvtIm03exxaeary18BzFXTiqxeK3W/MA1ixkvxCYouP2AkFmmeu9Ya+M/T0JlJSTucZimV4Bhx5BsTYej7KdWaCU8bak2NxCmZR/jAhPrxCugnggh4NABHru+5SI9jP4YpLnVm7MPBHRbwXwAxBT5o8w84+e2ofQWAFDSJOYgPyCmXvncSKO4grNDHQBNAzg7Qbzwy2Orw949uUBsQdo7rEdCMPbnViH5vKwKciQMZpR8SmdvDyOcg4uV335qUEMxbeSorXv9Dg8xxKT8u6MbkctxWNKzWE/1IWsYGq3CBC3CE6BBAIpwJ0moXfr42kFl5WVswaYRbGQjrFW8g4L9GOzfe1akujClAByXeDSM2qJuZkz5PkQLd1Yp+C4us6TFrLHf9csc1P6XQeaZ43dVO+Int9jfjzPwD4CYQQfjsvzdg5m8UGDvnOKjdNixVfXEqSyd9rda7LA0wsKgzG+z91SMPP3A/j+O+/orYUiOqWWUIiiQIgA//761U2tMRonhOsRQ0/YvdUjDsDmKqK/icCklhNR28iOjBQP8/iKB6QXF10BzPanuaI6QZNiqa0RXZVTxNAUpN+2GqOt2mlszorA0Gela+Obo9zDuVKupgzsGqJOrICEhbEd3+67jWmewUQAB3HzDUPre+BiJy5537sJOOd7sBkkwGJKcDMgPr4EdzquKaJTSzcFMmLEIlZZY5AtOeWeNYJTi+3XXHtTwGlfhVXOhQJa1uhmSFYsbwfwZkC8GBA3PebLHhQZ/bMRNM4pECZR/vJYq5DOcwoDmN/nbumLib0sHj9K3zFAjCI+0QCDeZpBxxHh2Q0oRlx+dkDsCZt3JnT7WSbsKboDkK2w2kI0sZe3Pr8GCWyR5zlmy1LB+pY7mKzXrpOXe56VnqIT2ysxdes8vaNwIwOJ0jC3CRDLExDF7oMwbnJlq0Y0GDHlSeqCIDV+BB07+4VpuwUePQB3HXjbg2YGTRP4cBQLxxTgbgvuxNXiiw2OH7pA7AkcCGGMuHhnBxyPik2lky7cvnRef10+oioXiIXcpoRORe/1nicr3csa3uU+Z7dApNOZBa9Rfb7cYnq0xfi4x3gZcHg9IIzA5Wc7dPuI7XESa3+sAkzvkrzfMbfnljRRDAtTiyAD1w3zekb7RVLrDfuA7dsjYhfQX8tqh2lWV81wjUY0zo7nj+tX6OJ7h4mZ6+CUAE8ZE1zDucTFm8VNspe9CwBTVizenTS3zbk8hlMKTSWWbrEFT/z1VFbhuQGNVb6cw0lpu8H86BI8BEwPNwhzxDDNoMMRfLOXQI1Z2natALr9jBBI4kETZ0xvzfIyEN5gjEAl3hRDsW1S3Am/DEkZNq/rhBuW7jmRc0OXQbFzQX1yz4rnKO/vISBsOoRDAF0EjJeE7siYtwE0Iy3Sa16IXdN9qCQGMJ/JlfxClpeq3BhIWBegeIQX93JwmJbRuCAh9PxyzGCOwA2DDkdsDsfSUrMoXWBQ34FH5/YWA3N/931SHOAo+L/hcSmSld0J9BosmGbBQgq3dInbcZzk+pnFvRt6OY6B6iEA45iVlt0zjjJfKYgbqYsCjUdwzJYbzAr0hFBgscqfBNcNP2o9w6DYTtcBQw++3OHwkQvMu4D9BzrQDDzsAvpnR/RvBoBucqSbCDzI7/3bwt9iEuoF7w9J2RORfE4kiiOGvEiZxUviHhfRSmd9p+szV3YYBL8cxQLilQCJYbL1QpruOVBGfCHWOLyS0zGknzWWqosWT5NatwSa5oT0HV/rcXwdCEfC9h0CzQTuXaDHPbslp+9+lNIXP+L2kpUbQVbQYuWsXzDnXvjt/OrJyruSl7Bxohq
      "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": 48,
   "id": "a7f8b5d5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.colorbar.Colorbar at 0x7f8f23f16800>"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAASwAAAD7CAYAAADQMalWAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/YYfK9AAAACXBIWXMAAAsTAAALEwEAmpwYAAB7g0lEQVR4nO29a8x1y1Ee+NRae7/f+92OfQ72sR37TGw0DoQhF4hFSBhlmDGZYRIU50dgSETkECJrpISQmwIhGkXzI1I0iUiQJhrpCBIxChpgCCPQCHGJE6TJj1iYy4TBjoPjENv44MPx5Vy+y/vuvVfNj+7qVV2rulev/e738pld0tbee61e3dW9uqufqq6uJmbGkY50pCM9CdRdNwNHOtKRjtRKR4F1pCMd6Ymho8A60pGO9MTQUWAd6UhHemLoKLCOdKQjPTF0FFhHOtKRnhi6kMAioq8noo8Q0UeJ6LsOxdSRjnSkI3lE+/phEVEP4N8D+KMAPgng5wH8aWb+0OHYO9KRjnSkkVYXeParAHyUmT8GAET0QwDeA6AosFa37/LJ/WcuUORCYoDiBwMDRBj6eI/iR6VN1518mumiz8+RzosA7sI3DQcuxxJNiq6XJ+1wSJ4oz3KWh6sk+949vqhwXa51AJt8yKRnrRMxUh8XOn/1s9g+euD1wmb67/7ru/yZz+6a0v7Cvz37aWb++ouUt4QuIrDeCuAT6v8nAfxBm4iI3gfgfQCwvvc03vlNf63+chs7OilkyF2eIQ0MMNBtgG4H9GeM1Rlje4uwuUvgDhjWlPMxhDy5o6xTEDu8WOHG02ezjsbm29bV3itcp4GDYEr1Bjb3COiA1UOg2zC4A5homk+B7xoRc0rDK4p5h//dDqAdh7LiQJR3wh1NBycX6tXCC43fTOOg7XbzdTgolYRQ/Ja2SX3G1tHUlXac0hIDu1uE7anKegD6c6S2YwJ2twBeAbSVdxB+S/6/9sPfc+FqvvTZHT7w029rSrt+y394w4ULXEAXEVhNWIKZnwfwPADcefY5zlLpF2k79EypbBOrv0wEAoNXwK4HhhVhexoHXJ/PYqmT0ZgnDY28SLkEDFpIeELO8Dgy0HgtPi+IivuAFnkV68CBbxbkWOO/MsiT0CFK7QjEwbULbckrgBnAMC1EBBiTGrgOmk3lNDc00sDutmqyIuf5Upa1PlaaMEr5OgKY7LVYb9bfcosRLgw83mOgP3P4kjozsHoU/u9uEYaTAp8XJsaOh/lk10AXEVifBPCc+v82AJ/aO7dav/WgtBZyIsRkkETiLqThDnlNbVl6QNmOa2dKzge1lDsrjGqDYCFKYFFtrQpxQbRBjj1T6tjtItqqtKMrPISvfVVEByWK+ssdwsRUKtfy0Eotk+ecUNTv1SIsSZahMxrrVsiXhiDQuh0DoiV4SP6CxACGG6Nr53QRgfXzAN5JRO8A8BsAvhnAn2l6cmammqTx/jvqBndxTG2C6sSd0vk9tY4B9Op3o9rCROPg5ootZQlKmyNGjnjioE0zMgf0mPG6jyXDIiJRdz2hbgR6SEcj2vDepRXQtfdaEHIaRU6EVQs62uf+HNk2UddpZ67rdpM6KPSs8+MutCXHtLtbwDAEbWGiEl+0DooGfIEhLGbeEtFfAvDTCMP+nzDzrx6Ms2LB6rd9QRFRSSdhO1Cg7o/SZ2LonC0H44wIAibvtoTYbJrWMjVqjANUBEJ/ru4tcVKZQ3ai4Q6c/Z/waR9jztS8JMRq5XmDudZ+CVE3CqurJAehSxtoRJWRFsB6AgUCimaA+9BPicb3kWyIB68CY/MFqBKCmX8SwE/un8EezziqoFC3idn2wM4iK0d4ZYNf593KZ4s9pERzaVruF8onsYskqWMTqG8jDJOKpdBknjcmK1eSlyx+JLugLDQ5gljzxqRYZPVdUK2Kdq9W1bN1jHtmiKV5ksOvQqTJPoVoH4xCa1jHdh6QbJPk9Pms7AOqhLsvQJVwf/LUs7kZH5h0XmvEJHnx6wiZN6EzZMjDzs5L3ktDR0+IYgktVd9qal9Ux8Itw0uDWlRFLYy0AqsFVEIHYlMTZBFVFkJEXQ7CSAZpuVWbHOzzVo3XyEQ/qleUNTJeSvu8p1J9GEmtp138jbDSPKyD6kdR5WfJy7wfz+Z4KPpCtGHtRy2CqUQGCZG9J+NHZiYeE2m7U3HZX7KfEzoWtakBUzU82zKF5w55m1RQYZHiTB6EA3zhU3hugmoaKHNnGOJAW4XnaVsQPLaOBLhuI7Vn5tI4qqvbDnN1rZW1L0LTFAW8nmBEeHWbIKjCJFxuo8tQB4Xt3Q0N7Hk9COsiQkuTqNlWxRP/FuR2nVm0MacCaCiPcdBm/jdW9W9Ak4IyqM1XL+MvK9sKqTm1yApGLlyT/xbhdnGAbYFdRAUAsN7IQOOUh9gMM/VIEBuQ+W5p432Rl7n66XSpf6DNwXafvtnyjOIjXeCAZigK737D6LZxMYUA7qfvM7ML7svvDN1MC9ZVCyzd6WtCaynsLr04q/41oJSWQe6tErqGe12mx5u6RrYztlDLbF+67wnSxjLTqpVcEoGrkZqT5+Q9VHjI2qEwcTSTRnP70KEEwwwfqQ9pdfsagA6DjzYsoczAmi7CV4O8dFZN0tA6lUFTNSumH3ofYmsnx6rAUEJLD0zyEIvz3OQakHxvxAibGb5rAsfr/CX1yFGViv8dHiUdic2FGBxf5nASbIaivpCo44jCzfO+15OJeh/au56Yw+pYZ8qPdjS7w2GCAkt1tKq8TeOo+un/0razKu2k3/HYh6Og4g7gVUjbbeD3A0G4cyr1HsQMbG6mvLp6gVUcYPuo43MIrdB5rYPpQWiJmiEd0NqP5KdV65YI8Quky1TdWnqLqOSys4Loupbo/3Noe44u8uwcldpgri1tnQv8ZUbz0kQpgl22fFmzgVah9xlDPmfYHS6zg9LVCqwB6M4Zu1uj45sYbCezmSfIJrN9nI09fazQUUbVq8JnAQlN/nsqjMcjnM6okeGQX8/5zf2aJrN/TTjPDOTkY2WeSYsVFsVQXIEVYaXrvwP6YcxDO7kWB1NJWFeQDwMBzhKSp39ml9LfM+i0pLqWynbta5aUgOEe0YjO+b20Q2JkkJhHhD0A2CKtGIIQFjQQbVvbUL7eogQgdxy+ADHCjqGbSNeHsEzDeoOgnAlyNUI9NOnwlzkDSxGlFRVV9l7uDosZwaSuRYEppAc2q/fDyGwqkm/afDzkabNv+duy0loj0SI1L4Yf978uoyT8KsjObTNvcprjveV1KxNDZhdl+UGjuhg/pPu/VbEbimyhI8KKlARTnI25o7AhuQewC6sk4TrgrhTJS+sAdwN0msXiT90xNZqx+S5BXLrIrJMJ76FOwojt4MQMHihHeeTUB2rQtKp9tedNXUW4J3vTEFb8wuwdZvXz14fv/iygBbFhzarjlicrPAqCNeNZPUsq72TDIhrfcaH++4zgmi9a62bt5GOlo2uUhFxCZZQQI20i2pKIDGeY1HNiw+O2bjJHjKPAGsmgIxFAQw/0u/EFB4Oi6sSglFZUg8n7V8LKpYK6eUibVkCJjCF2JhryTj4ao5Wqt6RvWFRQU330M5XrDIoqSRwsaeIIvlXDKhh/k6BqKdPjd8lzOgsb2cETUKZdPBQ0ubYA5S0mQUmt/Soi12RIB9L7AIJaSUMYJ+5iAw7XhxnAZtE+r6ujK3dryOJFyYsZ4oAYUPabcoKbldDRxCZkX6SB0W7sKpvW3rNQXo2ltEoJ1QFFEFvEwwDtRhvHrOPp3D2pr5fNTIeWiQM9sDuJYWTsFidBklxHIi45ap1Vpz3Xh9lyCu9lNp8Wsm1ZUj9LfVH1jWL+8ZuG2C1k5dCo4tCbo0t5HUBoMQi7G3rcw9ULrJ7GhpVxKshKbBXOwJ2EcNHaop1x40CYzESGlwnVhFXlf4qjJYNPIUEob+YM2Ylqq2bRWriUojNlrW7eICo8H4QsJzeE3UlU1U2/najAnpBdwqNCbK766wiIWTXZE9oFgVLkqUaaJ4eHohAuZcc8qpBxDCSTgiozTYY1gX0gGopqyvXS1auEM2igJGRIebV
      "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": 5,
   "id": "68b72b68",
   "metadata": {},
   "outputs": [],
   "source": [
    "from DifferentiableHOS.pk import power_spectrum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "785818e3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:AutoGraph could not transform <function power_spectrum.<locals>.fn at 0x7f8f98401bd0> and will run it as-is.\n",
      "Please report this to the TensorFlow team. When filing the bug, set the verbosity to 10 (on Linux, `export AUTOGRAPH_VERBOSITY=10`) and attach the full output.\n",
      "Cause: closure mismatch, requested ('Nsum', 'W', 'boxsize', 'dig', 'xsum'), but source function had ()\n",
      "To silence this warning, decorate the function with @tf.autograph.experimental.do_not_convert\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:AutoGraph could not transform <function power_spectrum.<locals>.fn at 0x7f8f98401bd0> and will run it as-is.\n",
      "Please report this to the TensorFlow team. When filing the bug, set the verbosity to 10 (on Linux, `export AUTOGRAPH_VERBOSITY=10`) and attach the full output.\n",
      "Cause: closure mismatch, requested ('Nsum', 'W', 'boxsize', 'dig', 'xsum'), but source function had ()\n",
      "To silence this warning, decorate the function with @tf.autograph.experimental.do_not_convert\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING: AutoGraph could not transform <function power_spectrum.<locals>.fn at 0x7f8f98401bd0> and will run it as-is.\n",
      "Please report this to the TensorFlow team. When filing the bug, set the verbosity to 10 (on Linux, `export AUTOGRAPH_VERBOSITY=10`) and attach the full output.\n",
      "Cause: closure mismatch, requested ('Nsum', 'W', 'boxsize', 'dig', 'xsum'), but source function had ()\n",
      "To silence this warning, decorate the function with @tf.autograph.experimental.do_not_convert\n",
      "WARNING:tensorflow:AutoGraph could not transform <function power_spectrum.<locals>.fn at 0x7f8f9852a830> and will run it as-is.\n",
      "Please report this to the TensorFlow team. When filing the bug, set the verbosity to 10 (on Linux, `export AUTOGRAPH_VERBOSITY=10`) and attach the full output.\n",
      "Cause: closure mismatch, requested ('Nsum', 'W', 'boxsize', 'dig', 'xsum'), but source function had ()\n",
      "To silence this warning, decorate the function with @tf.autograph.experimental.do_not_convert\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:AutoGraph could not transform <function power_spectrum.<locals>.fn at 0x7f8f9852a830> and will run it as-is.\n",
      "Please report this to the TensorFlow team. When filing the bug, set the verbosity to 10 (on Linux, `export AUTOGRAPH_VERBOSITY=10`) and attach the full output.\n",
      "Cause: closure mismatch, requested ('Nsum', 'W', 'boxsize', 'dig', 'xsum'), but source function had ()\n",
      "To silence this warning, decorate the function with @tf.autograph.experimental.do_not_convert\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING: AutoGraph could not transform <function power_spectrum.<locals>.fn at 0x7f8f9852a830> and will run it as-is.\n",
      "Please report this to the TensorFlow team. When filing the bug, set the verbosity to 10 (on Linux, `export AUTOGRAPH_VERBOSITY=10`) and attach the full output.\n",
      "Cause: closure mismatch, requested ('Nsum', 'W', 'boxsize', 'dig', 'xsum'), but source function had ()\n",
      "To silence this warning, decorate the function with @tf.autograph.experimental.do_not_convert\n"
     ]
    }
   ],
   "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": 46,
   "id": "18e3d30b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x7f8f28206200>]"
      ]
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAXYAAAD8CAYAAABjAo9vAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/YYfK9AAAACXBIWXMAAAsTAAALEwEAmpwYAAAe8UlEQVR4nO3deXRV1eH28e++NwlJGMIUIBISAkmAQICQhCHFiq1WbEWs4lRxqAOiaG1/zra1w9K2am21jkRxqOJUqhUsltYBZU5umDEJBCQQQEIYAoSETPv9Q/SlVDSQYd/h+ayVP+4m95xnAXk47LPvPsZai4iIBA+P6wAiItKyVOwiIkFGxS4iEmRU7CIiQUbFLiISZFTsIiJBJszlyY0xE4AJHTt2vC41NdVlFBGRgFNQUFBhrY09dtz4wzr2rKws6/P5XMcQEQkoxpgCa23WseOaihERCTIqdhGRIKNiFxEJMip2EZEg47TYjTETjDG5lZWVLmOIiAQVp8VurZ1jrZ0SExPjMoaISFAJ6KmY0qLlFC2b5zqGiIhfCehi3zX3fga+exErHjqHspI1ruOIiPiFgC72tOufZ3HCVFIP5tPzpdNY9uS17Nu1w3UsERGnArrYo9t3IufqB6iems/ybueQtXMWnidGsPSle6mprnIdT0TEiYAu9i90j0tg1E/+ypaL32dTVDqjNz7K3geHUfBOLraxwXU8EZE2FVTLHZPSMhl+579Z/d2XqPJ0ItN3OyW/G0Xhkndb5PgiIoEgaDcBa2howDfnafqufJie7GZldA7dzvs9fVKHt+h5RERcCblNwLxeL6POm0an21ezuO80kqtWEDfzdPKeuJq9u7a7jici0mqCtti/ENW+AzlX/Y7qG/PxdZ/IiPK3CH98BEv/+gtqDh10HU9EpMUFfbF/IbZnH0bf/AJll37AhujhjN70GPseHEbB7KdobNANVhEJHiFT7F/oOzCDjDv/xZozZnLA25nM5Xex6XfZfLL4HdfRRERaRMgV+xfSx55D/3vyyMt4gPYNlaT9+zJWPjieLcUrXEcTEWmWkC12AI/Xy8iJU+l8+yoWJ91M/6qVnPLKd8h7/Er27NzqOp6IyEkJqnXsJyuqfQdyrryPwzcWkB/7QzJ2zaHdk1kse/Fuag4dcJpNROREBe069uYoLV7J7rfvYcShRZTTla0Zt5Jxzg14vF7X0UREvhRy69ibI3HAcEbcMZc133uNfd5uZK74OZ/+Lot1C2e7jiYi8o1U7F8jPedsku9ZRl7mQ0Q3HGDwe5ez+oEz2VLkP/+7EBE5lor9G3i8XkZOmEKXO1exuN8t9D20lt6vnkH+Y5PZvXOL63giIv9Dxd5EkVHtybnit9TftJy8HpMYXjGXqCezWPb8HVQf3O86nojIl1TsJ6hrbBxjpj3LjsnzKewwklGl0zn4x2H43voLjfX1ruOJiKjYT1ZCylAyb3+HtePfYE9YLFmrfknp7zNZ9/FbrqOJSIhTsTfTkNFnkXL3UvKzHyaisZrBH1zFmj98l8JFc6ivPew6noiEIK1jb0E11YdYPushhpRMp5OpYj/tKek0CgacTfK3fkinzrGuI4pIEDneOnYVeys4eGAfxYtmU180l5R9i+jKfuqth/WR6RxMPIM+o88nrt8Q1zFFJMCp2B1pqK+nePl89q2cTdxn80lqLAVgq6c323uOo/PwCSRnnoE3LNxxUhEJNCp2P1H2aRFbl7xJ+9L/MLBmFRGmgX10YGNMDt6BZ5OScx7tY7q6jikiAcAvi90YMwGYkJycfN2GDRuc5XClct8e1i9+G1v0Lin7F9OFA9RZL+sjh1KVdCYJYy6gV+JA1zFFxE/5ZbF/IZSu2I+nvq6OYt/7VK6awyk759PXlgGw2ZPAzrjT6ZJxLsnDx+EJC3OcVET8hYo9wJSWrGXbsjfpWPo+Aw+vIdw0sJdObOycQ9ig75Oacy7RHbu4jikiDqnYA1jlnl2fT9kUv8uA/UuIMVXU2jDWRw3jUNL3iM/8AXH9BmM8+liCSChRsQeJ2tpaivP/w4HVc4gv/4gEux2AfXSkNHoINb2y6DzgWySmjyUyuqPjtCLSmlTsQchay5YNq9mx+gPM1jx67V9N4pG5+TrrZXN4f/Z2G05E3zHEDxtH91P6OU4sIi1JxR4iKsp3ULr6I2o2LiamYgX9a4uIMrUA7DTd2dZxKA2nZNNt0Kkkpo3CGx7hOLGInCwVe4g6fLiGTevy2Fu0gPDt+cQfXEMcFQBU2wg2txvI/h4jiO6fQ99hp9Gxay/HiUWkqVTsAnw+fbNjy0bK1nxEfelSuu1ZSb/6jYSbBuDzT8TujBlGdMYk0r59geO0IvJ1VOxyXAcPHuDT1QvZv34hUTsLSKpeSxcOsLzTd+g3+XE69+jtOqKIfAUVuzRZTU01vpm/YuSWGRwykWzK/DkZP5iq5ZQifuZ4xa6fVPkfkZFRjL3mQbZePI8dYX0YUXA3ax88k/Itxa6jiUgTqNjluPqnZZF69yIWp95Jv+q1dJhxKvmv3a9HAIr4ORW7fC2v10vOj+5h348XsD5qKNlFD1LywLfYUlTgOpqIHIeKXZqkd99Uht3xb5ZmPED3um30evVM8p+/jbrD1a6jicgxVOzSZMbjYfTEqTTcsJSVnU4nu/QZdjyQTUnB+66jichRVOxywmJ7xjPy1r+zfGwuEY3V9Jt9AflPXkv1gX2uo4kIrVDsxphBxpinjTGzjDE3tPTxxX+MOONion6az9LuF5C5cxaVf8pi3UezXMcSCXlNKnZjzHPGmHJjzNpjxscbY4qNMSXGmLsArLWF1tqpwEXA/6yvlOAS07krOTfP4JOz36CGdgz+8BqW/3kS+ys+cx1NJGQ19Yr9BWD80QPGGC/wBHA2kAZcaoxJO/Jr5wILAU2+hogho79HrzvyWdj7Wobs+4CGx7NZMedJGuvrXEcTCTlNKnZr7cfAnmOGRwIl1tpN1tpa4DVg4pHvn22tzQEua8mw4t8io6IZe93DlE76Fzu9cWQU3M2u+weS//K9HNi703U8kZDRnDn23sDWo16XAb2NMeOMMX8xxkwH5h7vzcaYKcYYnzHGt2vXrmbEEH+Tkj6S5LsWkT/6cXaGx5Nd8ijhjwzG99hkthbmuY4nEvSa82Rk8xVj1lo7H5j/TW+21uYCufD5XjHNyCF+KCw8nOzxl8P4y1m/Oo+KDx8jo+JfRL0+h8J2w6jPnsLgcRfjCQt3HVUk6DTnir0M6HPU63hge/PiSDBKHTqSnFte4tBNa1iYdAsxh7eTvnAa5fenkT/zVxzYW+46okhQafLujsaYvsA71tohR16HAeuB7wLbgHzgR9badU0+uTETgAnJycnXbdiw4QSjS6Cqra1jxXuvEL3iWdLrVlNtI1gXeza9zryF+AGZruOJBIxmbdtrjHkVGAd0B3YCv7LWzjDGfB94BPACz1lr7z+ZcNq2N3QVrVrK3g8fI2PvPCJNHYXthtMw8nrSTrsIT1hzZgpFgp/2Yxe/tqt8O0VznyR586vEUUEVkWxtl8LBrulEJGQRl5ZDbMJAMF91a0ckNPllsWsqRo5VW1vLivdepb5kPl0r19GvfhPtzOdr4ffTnrLIAVTFDiMyMZP4tG/RJS5JZS8hyy+L/Qu6Ypfjqa6u4dNCH3s3LMWzYyXd96+jb0Ppl89o3U1ntkcPpKbHMOJPnUxc/6GOE4u0HRW7BI0DBw+weW0elZuWEfbZSnocKKRv41Y8xrK6fQ5Rp91CSvZZupKXoKdil6C2Y9sWNv7zEdK2v0FXDrApLJkDI6Yy5Mwr8YZHuI4n0ir8stg1xy4t7eDBA6x652nii58n0W6j3HRja8oVDDrnZqI7dXMdT6RF+WWxf0FX7NLS6uvrWfH+32jne5KhdaupIpLCXhPp+4Nb6d5ngOt4Ii1CxS4h65PlC9j/waNkHvgAD42s7XQaUWOupn/2eLzh7VzHEzlpKnYJeWWlJXw69xGGfvYmMaaKg0RTEjMa76BzSM75IVGdurqOKHJCVOwiR+w/UEnhojk0Fv6T1MpFdKOSOutlQ9RQavq
      "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": 47,
   "id": "c2c1ecdb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x7f8f281433d0>]"
      ]
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAY8AAAD8CAYAAACPWyg8AAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/YYfK9AAAACXBIWXMAAAsTAAALEwEAmpwYAAA5/UlEQVR4nO3deXhbV5n48e8ryfu+L3ESZ3GSOnvjpE33LVubksLATAqlLQXKQMsAMyzpj2UYpoXODAUGaOl0KJCW0lJgaAJN0ybpRvc4+544uxPv8S7bsuzz+0NXjp3IlmTZlpf38zx+pHt1z9VRKHp1tveIMQallFIqGLZwV0AppdTIo8FDKaVU0DR4KKWUCpoGD6WUUkHT4KGUUipoGjyUUkoFzRHuCgyV9PR0k5+fH+5qKKXUiLJt27ZqY0zGhefHTPDIz8+nuLg43NVQSqkRRURO+jqv3VZKKaWCpsFDKaVU0EIKHiKSKiKbROSI9ZjSy3XLReSQiJSIyJpAyovIA9b1h0RkmXUuVkReFJGDIrJPRB4Opf5KKaX6J9SWxxpgizGmANhiHfcgInbgUWAFUAjcLiKFfZW3Xl8NzASWA49Z9wH4oTFmBjAfuFJEVoT4GZRSSgUp1OCxClhrPV8L3ObjmkVAiTHmmDHGBTxnleur/CrgOWNMmzHmOFACLDLGOI0xrwFY99oO5IX4GZRSSgUp1OCRZYwpA7AeM31cMw443e241DrXV/m+ygAgIsnArXhaLD6JyL0iUiwixVVVVYF+JqWUUn74naorIpuBbB8vfTPA9xAf5/zlge+zjIg4gGeBnxpjjvV2E2PME8ATAEVFRf3KPb/tZC2J0Q4KshL6U1wppUYlv8HDGHNTb6+JSIWI5BhjykQkB6j0cVkpML7bcR5w1nreW/m+yoAnIBwxxvzEX/1D0d7RyZee24EIvPCFK0mLjxrMt1NKqREj1G6r9cBd1vO7gHU+rtkKFIjIJBGJxDMQvt5P+fXAahGJEpFJQAHwAYCIPAgkAV8Ose5+Rdht/Pzjl1LZ0Mbnnt5Gm7tjsN9SKaVGhFCDx8PAEhE5AiyxjhGRXBHZAGCMcQP3Ay8DB4DnjTH7+ipvvf48sB/YCNxnjOkQkTw83WWFwHYR2SkinwnxM/Rp3vhkHvn7uRSfrOWBP+1Bd15USimQsfJlWFRUZEJJT/LTLUf40abDfG3ZdO67fuoA1kwppYYvEdlmjCm68PyYyW0Vqi/eMJWjVU3818uHmJwex4rZOeGuklJKhY2mJwmQiPAffzeHSyck85Xnd7K7tC7cVVJKqbDR4BGE6Ag7T9xZRFpcFJ9ZW0xZfUu4q6SUUmGhwSNI6fFR/OruhThdHXxmbTFOlzvcVVJKqSGnwaMfpmcn8LPb53OgrIEvP7eTzs6xMelAKaW8NHj00/UzMvn2ykJe2V/Bf758KNzVUUqpi9Q2uzhV4xyUH7gaPEJw9xX53HH5BB5/4yjPF5/2X0AppYbQs1tPcc1/vYaro3PA763BIwQiwr/eOpOrpqbzzT/v4b1jNeGuklJKdXG2dWATiHIM/Fe9Bo8QRdhtPPqJS5mQGss//nYbJ6qbw10lpZQCwOnqIC7SgYivXLOh0eAxAJJiIvjV3QsR4J61W6l3toe7SkophdPlJibS7v/CftDgMUAmpsXx+B0LOH3OyX2/2077IPQxKqVUMJpdHcRFDU4iEQ0eA+iyyWl8/8Ozeaukmn9dv0+TKCqlwsrZ5iZ2kFoemttqgH2saDxHq5p5/I2jTM2I556rJoW7SkqpMarZ5SYuUlseI8bXl01naWEWD764n1cPVoS7OkqpMcrp6iA2Ssc8RgybTfjJ6nlckpPIF3+3g4PlDeGuklJqDGpu05bHiBMb6eDJuxYSH+3g078ppqqxLdxVUkqNMU5Xx6CNeWjwGETZSdH88s6F1DS3ce/TxbS26za2Sqmh09zm1tlWI9XsvCR+8g/z2HGqjq//cbfOwFJKDQljjLY8Rrrls3L42rLprN91lp9uKQl3dZRSY4CroxN3pxm0lodO1R0iX7huCkermvjx5sNMyojjQ3Nzw10lpdQo5mzzdJNry2OEExF+8JHZLMxP4at/2MWOU7XhrpJSahRzWmOsOttqFIhy2PmfTxaRnRjNZ5/axpk63cZWKTU4nG2eXU51nccokRoXyZN3FdHW3sGnf7OVpjbdxlYpNfCaXdptNeoUZCXw809cypHKJr707A46dBtbpdQA62p5DMduKxFJFZFNInLEekzp5brlInJIREpEZE0g5UXkAev6QyKyzMc914vI3lDqH07XTsvgu7cWsuVgJT/YcCDc1VFKjTLelsdwHfNYA2wxxhQAW6zjHkTEDjwKrAAKgdtFpLCv8tbrq4GZwHLgMes+3nt+BGgKse5h98nF+dx9RT6/fOs4v3v/VLiro5QaRZyu4T3msQpYaz1fC9zm45pFQIkx5pgxxgU8Z5Xrq/wq4DljTJsx5jhQYt0HEYkH/hl4MMS6DwvfuuUSrp2WwXfW7eWdkupwV0cpNUo0tw3vlkeWMaYMwHrM9HHNOOB0t+NS61xf5fsq8+/AI4DTX+VE5F4RKRaR4qqqqsA+0RBz2G387OPzmZQexz/+dhvHqkZ8g0opNQyEveUhIptFZK+Pv1X+ynpv4eOcvxFin2VEZB4w1Rjz50De2BjzhDGmyBhTlJGREUiRsEiM9mxj67DbuOc3W6ltdoW7SkqpEc7b8oiNCFPwMMbcZIyZ5eNvHVAhIjkA1mOlj1uUAuO7HecBZ63nvZXvrcxiYIGInADeAqaJyOuBfdThbXxqLE98cgFn61r5/DPbcLl1G1ulVP85XW6iHDYc9sGZVBvqXdcDd1nP7wLW+bhmK1AgIpNEJBLPQPh6P+XXA6tFJEpEJgEFwAfGmF8YY3KNMfnAVcBhY8x1IX6GYaMoP5X/+Ohs3jt2jm+9sEeTKCql+q3ZNXgZdSH03FYPA8+LyKeBU8DHAEQkF/ilMeZmY4xbRO4HXgbswK+MMfv6Km+M2ScizwP7ATdwnzFmTOQz//D8PI5VNfOzV0uYmhnPvddMCXeVlFIjkLNt8DLqQojBwxhTA9zo4/xZ4OZuxxuADYGWt157CHioj/c+AcwKutIjwFdumsaxqmZ+8NJBJqXHs6QwK9xVUkqNMIO5fznoCvNhyWYTfvixucwel8SXntvBvrP14a6SUmqEGcz9y0GDx7AVE2nnl3cWkRQTwWfWFlPZ0BruKimlRhCnq0NbHmNVZmI0v7yriPqWdj77VDEtrjEx7KOUGgDNbe5BHfPQ4DHMzcz1bGO7+0w9//KHnXRqEkWlVAAGcwta0OAxIiydmc0DK2awYU85P958ONzVUUqNAE6Xm9hhPFVXDZHPXj2ZksomfvZqCVMy4rlt/jj/hZRSY1ZzWwdx2vJQIsKDt83mskmpfP2Puyk+cS7cVVJKDVMdnYaW9o5B28sDNHiMKJEOG4/fsYDc5Gg+9/Q2XtxdRmu7DqIrpXpq8e5frlN1lVdKXCRP3r2Q6Ag79/1uOwsf3MxX/7CLt45U646ESilg8HcRBB3zGJGmZMTz5tev592jNbyw8wwb95bzx22lZCZE8aG5udw2fxwzcxMR8ZWcWCk12nXtIjiILQ8NHiOU3SZcVZDOVQXpPHjbLDYfqOCFHWdZ++4JfvnWcaZkxHHbvHGsmjeOCWmx4a6uUmoINWvLQwUiOsLOyjm5rJyTS22ziw17y1i34yyPbDrMI5sOc+mEZG6bP45bZueQFh8V7uoqpQaZc5D3LwcNHqNOSlwkn7hsIp+4bCKltU7W7TzLup1n+M66fXzvL/u5ZloGq+blsqQwa1B/lSilwqd5kHcRBA0eo1peSiz3XT+VL1w3hQNljazbeYb1u87y6sFKYiPtLJuZzap5uVw1NX3QNoxRSg29Fm15qIEgIhTmJlKYm8g3ls/g/ePnWLfzDC/uKePPO86QHh/Jyjm5rJqXy7zxyTrQrtQId37MQ1seaoDYbMLiKWksnpLGv62ayWsHq1i38wy/++AUv3nnBPlpsayaN47b5o9jUnpcuKurlOq
      "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": 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
}