JaxPM/notebooks/02-Advanced_usage.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# **Advanced Particle Mesh Simulation on a Single GPU**\n",
    "\n",
    "<a href=\"https://colab.research.google.com/github/DifferentiableUniverseInitiative/JaxPM/blob/main/notebooks/02-Advanced_usage.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install --quiet git+https://github.com/DifferentiableUniverseInitiative/JaxPM.git\n",
    "!pip install diffrax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import jax_cosmo as jc\n",
    "\n",
    "from jaxpm.painting import cic_paint , cic_paint_dx\n",
    "from jaxpm.pm import linear_field, lpt, make_diffrax_ode\n",
    "from jaxpm.distributed import uniform_particles\n",
    "from diffrax import ConstantStepSize, LeapfrogMidpoint, ODETerm, SaveAt, diffeqsolve"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Particle Mesh Simulation with Diffrax Leapfrog Solver\n",
    "\n",
    "In this setup, we use the `LeapfrogMidpoint` solver from the `diffrax` library to evolve particle displacements over time in our Particle Mesh simulation. The novelty here is the use of a **Leapfrog solver** from `diffrax` for efficient, memory-saving time integration.\n",
    "\n",
    "- **Leapfrog Integration**: This symplectic integrator is well-suited for simulations of gravitational dynamics, preserving energy over long timescales and allowing larger time steps without sacrificing accuracy.\n",
    "- **Efficient Displacement Tracking**: We initialize only displacements (`dx`) rather than absolute positions, which, combined with Leapfrog’s stability, enhances memory efficiency and speeds up computation.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4.05 s ± 1.54 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n",
      "Solver Stats : {'max_steps': Array(4096, dtype=int32, weak_type=True), 'num_accepted_steps': Array(90, dtype=int32, weak_type=True), 'num_rejected_steps': Array(0, dtype=int32, weak_type=True), 'num_steps': Array(90, dtype=int32, weak_type=True)}\n"
     ]
    }
   ],
   "source": [
    "mesh_shape = [128, 128, 128]\n",
    "box_size = [128., 128., 128.]\n",
    "snapshots = jnp.array([0.5, 1.0])\n",
    "\n",
    "@jax.jit\n",
    "def run_simulation(omega_c, sigma8):\n",
    "    # Create a small function to generate the matter power spectrum\n",
    "    k = jnp.logspace(-4, 1, 128)\n",
    "    pk = jc.power.linear_matter_power(jc.Planck15(Omega_c=omega_c, sigma8=sigma8), k)\n",
    "    pk_fn = lambda x: jnp.interp(x.reshape([-1]), k, pk).reshape(x.shape)\n",
    "\n",
    "    # Create initial conditions\n",
    "    initial_conditions = linear_field(mesh_shape, box_size, pk_fn, seed=jax.random.PRNGKey(0))\n",
    "\n",
    "    # Create particles\n",
    "    cosmo = jc.Planck15(Omega_c=omega_c, sigma8=sigma8)\n",
    "    \n",
    "    # Initial displacement\n",
    "    dx, p, f = lpt(cosmo, initial_conditions, a=0.1,order=1)\n",
    "    \n",
    "    # Evolve the simulation forward\n",
    "    ode_fn = ODETerm(\n",
    "        make_diffrax_ode(cosmo, mesh_shape, paint_absolute_pos=False))\n",
    "    solver = LeapfrogMidpoint()\n",
    "\n",
    "    stepsize_controller = ConstantStepSize()\n",
    "    res = diffeqsolve(ode_fn,\n",
    "                      solver,\n",
    "                      t0=0.1,\n",
    "                      t1=1.,\n",
    "                      dt0=0.01,\n",
    "                      y0=jnp.stack([dx, p], axis=0),\n",
    "                      args=cosmo,\n",
    "                      saveat=SaveAt(ts=snapshots),\n",
    "                      stepsize_controller=stepsize_controller)\n",
    "\n",
    "    ode_solutions = [sol[0] for sol in res.ys]\n",
    "    return initial_conditions ,  dx , ode_solutions , res.stats\n",
    "\n",
    "initial_conditions , lpt_displacements , ode_solutions , solver_stats = run_simulation(0.25, 0.8)\n",
    "ode_solutions[-1].block_until_ready()\n",
    "%timeit initial_conditions , lpt_displacements , ode_solutions , solver_stats = run_simulation(0.25, 0.8);ode_solutions[-1].block_until_ready()\n",
    "print(f\"Solver Stats : {solver_stats}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAB8QAAAH5CAYAAAD3DoQvAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjkuMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8hTgPZAAAACXBIWXMAAA9hAAAPYQGoP6dpAAEAAElEQVR4nOy9ebwlRXn//66qXs52t5m5MAs4DIuAqJiMYkRkQFEWlaACX2MUMCKKGNSoiUuMoKBxN2pE1LhEjd8IqOj35x5RMagBxX1DBJQBZrtz7z33LL1U1e+P6u5zzr0zMIPggKk3r8vc26dPd3V1dVX183mep4S11uLxeDwej8fj8Xg8Ho/H4/F4PB6Px+PxeDwez58Yck8XwOPxeDwej8fj8Xg8Ho/H4/F4PB6Px+PxeDyeewMviHs8Ho/H4/F4PB6Px+PxeDwej8fj8Xg8Ho/nTxIviHs8Ho/H4/F4PB6Px+PxeDwej8fj8Xg8Ho/nTxIviHs8Ho/H4/F4PB6Px+PxeDwej8fj8Xg8Ho/nTxIviHs8Ho/H4/F4PB6Px+PxeDwej8fj8Xg8Ho/nTxIviHs8Ho/H4/F4PB6Px+PxeDwej8fj8Xg8Ho/nTxIviHs8Ho/H4/F4PB6Px+PxeDwej8fj8Xg8Ho/nTxIviHs8Ho/H4/F4PB6Px+PxeDwej8fj8Xg8Ho/nTxIviHs8Ho/H4/F4PB6Px+PxeDwej8fj8Xg8Ho/nTxIviHvuNwghuOCCC3Zp3/3224+zzjprt89x8803I4TgIx/5yG5/977AN77xDYQQfOMb36i2nXXWWey333679P0LLrgAIcS9U7j7EXuyHezO/fJ4PB7PfZMdjce7ykc+8hGEENx88813ue/dne/cm+ypMu1OvXk8Ho/nvse1117LkUceSbPZRAjBD3/4wz/o/XRX36vu7zaAu+IPmZP8oRxzzDEcc8wxf/Tzejwej+ePhx+/7x28bdrjuXfwgrjnj0ZpqLzuuuvukeNdc801XHDBBczOzt4jx7s7bNq0iZe97GUccsghNBoNms0m69ev56KLLtqj5bozut0uF1xwwR55IfY4brvtNi644AJ++MMf7umiLOHf/u3fOPTQQ6nVahx00EG8+93v3tNF8ng8HmDX5hHlS2P5o5TiAQ94AE95ylOqPvess84a2WdnP3cm6pYv+Dv6ed/73ncPX7lnmDe84Q189rOf3dPFWMI111zDUUcdRaPRYOXKlZx//vksLCzs6WJ5PB7P/YIsyzjttNOYmZnhHe94Bx/72MdYu3btni7WXeLfneDnP/85F1xwwX3OIc0Yw5vf/GbWrVtHrVbjoQ99KJ/85Cf3dLE8Ho/nT4r74/h9ySWXcNppp/GABzzgLt/7/5TxtmnP/1aCPV0Aj2dX6fV6BMGgyV5zzTVceOGFnHXWWUxOTo7s+6tf/Qop711/j2uvvZaTTjqJhYUFnvnMZ7J+/XoArrvuOv75n/+Zb33rW3zlK1+5V8uwK3zgAx/AGFP93e12ufDCCwGWeGv/4z/+I694xSv+mMW7T7J27Vp6vR5hGN4rx7/tttu48MIL2W+//XjYwx428tni+/XH5NJLL+X5z38+T3va0/i7v/s7rr76as4//3y63S7/8A//sEfK5PF4PHeHv/qrv+Kkk05Ca80vfvELLrnkEr74xS/y3e9+l+c973kcd9xx1b433XQT//RP/8Q555zDYx7zmGr7AQcccJfnueSSS2i1WiPbHvnIR3LAAQfQ6/WIouieu6j7Cff2HOwNb3gDp556KqeccsrI9mc961k8/elPJ47je+3cO+OHP/whj3vc4zj00EN5+9vfzq233spb3/pWbrjhBr74xS/+0cvj8Xg89zduvPFGbrnlFj7wgQ9w9tlnV9vvy++n95d3p6OPPvpenZP8/Oc/58ILL+SYY45ZEk22J+0hr371q/nnf/5nnvvc5/KIRzyCK6+8kmc84xkIIXj605++x8rl8Xg8f0rcH8fvN73pTbTbbY444ghuv/32PV2cneJt0/ft+ZXn/osXxD33G2q12i7ve28bQ2dnZ3nKU56CUorrr7+eQw45ZOTziy++mA984AP3ahl2ld0ZOIMgGHE6+FPBWku/36der+/S/kKI3Wpv9yT31kTnruj1erz61a/miU98IpdffjkAz33uczHG8PrXv55zzjmHqampPVI2j8fj2V3+/M//nGc+85nV349+9KM5+eSTueSSS7j00kt51KMeVX123XXX8U//9E886lGPGvnOrnDqqaeyYsWKHX62p8aRe5pOp0Oz2dzl/feEIA2glEIptUfO/apXvYqpqSm+8Y1vMD4+DrjU8c997nP5yle+whOe8IQ9Ui6Px+O5v7B582aAJY7u99X30z357mSMIU3TXZ5nSCn32JxkTzkGbty4kbe97W2cd955vOc97wHg7LPPZsOGDbz85S/ntNNO22NzBo/H4/lT4v42fgN885vfrKLDFzu335t42/Rd423Tnj8GPmW6Z49y1lln0Wq12LhxI6eccgqtVovp6Wle9rKXobUe2Xd4DfELLriAl7/85QCsW7euSlNapulavH7lzMwML3vZy3jIQx5Cq9VifHycE088kR/96Ed3q9yXXnopGzdu5O1vf/sSMRxg77335h//8R9Htr33ve/lsMMOI45jVq9ezXnnnbckrfoxxxzDgx/8YH7+859z7LHH0mg0WLNmDW9+85uXnOPWW2/llFNOodlsstdee/GSl7yEJEmW7De87sfNN9/M9PQ0ABdeeGFVb8P1uniNlzzPef3rX88BBxxAHMfst99+vOpVr1pyrv32248nPelJfPvb3+aII46gVqux//778+///u8j+2VZxoUXXshBBx1ErVZj+fLlHHXUUXz1q19dWtFDlKlyv/Wtb/G85z2P5cuXMz4+zhlnnMH27dt3WJYvf/nLPPzhD6der3PppZcC8Nvf/pbTTjuNZcuW0Wg0+Iu/+Av+v//v/xv5/s7WafnlL3/JqaeeyrJly6jVajz84Q/nc5/73JKyzs7O8pKXvIT99tuPOI7ZZ599OOOMM9i6dSvf+MY3eMQjHgHAs5/97OoelOfa0TotnU6Hl770pey7777EcczBBx/MW9/6Vqy1I/sJIXjhC1/IZz/7WR784AcTxzGHHXYYX/rSl+60bgGuuuoqtm3bxgte8IKR7eeddx6dTmdJHXk8Hs/9icc+9rGAiwb/Y7Cz9Tq/973vccIJJzAxMUGj0WDDhg3893//910ez1rLRRddxD777EOj0eDYY4/lZz/72S6VpRzT3vrWt/KOd7yDtWvXUq/X2bBhAz/96U9H9i3nZTfeeCMnnXQSY2Nj/PVf/zWw62PRjtYQn52d5cUvfnH13QMPPJA3velNS7zOjTH8y7/8Cw95yEOo1WpMT09zwgknVGnyhRB0Oh0++tGPLkltv7M1xO/p+ddi5ufn+epXv8ozn/nMSgwHOOOMM2i1WnzqU5+6y2N4PB7P/2bOOussNmzYAMBpp52GEKLKZLazNUg//vGPs379eur1OsuWLePpT386v//97+/yXLOzs5x11llMTEwwOTnJmWeeebeWOvtD353K6/rlL3/J6aefzvj4OMuXL+dFL3oR/X5/ZN/yHe8Tn/hENZ6V73fXX389J554IuPj47RaLR73uMfx3e9+d+T7f+icZOPGjTznOc9h9erVxHHMunXrOPfcc0nTlI985COcdtppABx77LHV2Fyea0driG/evJnnPOc57L333tRqNQ4//HA++tGPjuwzPHd5//vfX9kiHvGIR3Dttdfead0CXHnllWRZNnJ/hBCce+653HrrrXznO9+5y2N4PB6P5865P47f4CKv7+765t427W3Tnvs/901XHc//KrTWHH/88TzykY/krW99K1/72td429vexgEHHMC55567w+889alP5de//jWf/OQnecc73lFFZ5Vi72J++9vf8tnPfpbTTjuNdevWsWnTJi699FI2bNjAz3/+c1avXr1bZf7c5z5HvV7n1FNP3aX9L7jgAi688EKOO+44zj33XH71q19xySWXcO211/L
      "text/plain": [
       "<Figure size 2000x500 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from jaxpm.plotting import plot_fields_single_projection\n",
    "\n",
    "fields = {\"Initial Conditions\" : initial_conditions , \"LPT Field\" : cic_paint_dx(lpt_displacements)}\n",
    "for i , field in enumerate(ode_solutions):\n",
    "    fields[f\"field_{i}\"] = cic_paint_dx(field)\n",
    "plot_fields_single_projection(fields)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### First and Second Order Lagrangian Perturbation Theory (LPT) Displacements\n",
    "\n",
    "This section introduces **first-order and second-order LPT** simulations, controlled by the `order` argument. First-order LPT captures linear displacements, while second-order LPT includes nonlinear corrections, allowing more accurate modeling of structure formation.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "32.3 ms ± 9.42 μs per loop (mean ± std. dev. of 7 runs, 10 loops each)\n",
      "42.8 ms ± 9.74 μs per loop (mean ± std. dev. of 7 runs, 10 loops each)\n"
     ]
    }
   ],
   "source": [
    "from functools import partial \n",
    "\n",
    "mesh_shape = [128, 128, 128]\n",
    "box_size = [128., 128., 128.]\n",
    "snapshots = jnp.array([0.5,1.])\n",
    "\n",
    "@partial(jax.jit , static_argnums=(2,))\n",
    "def lpt_simulation(omega_c, sigma8, order=1):\n",
    "    # Create a small function to generate the matter power spectrum\n",
    "    k = jnp.logspace(-4, 1, 128)\n",
    "    pk = jc.power.linear_matter_power(jc.Planck15(Omega_c=omega_c, sigma8=sigma8), k)\n",
    "    pk_fn = lambda x: jnp.interp(x.reshape([-1]), k, pk).reshape(x.shape)\n",
    "\n",
    "    # Create initial conditions\n",
    "    initial_conditions = linear_field(mesh_shape, box_size, pk_fn, seed=jax.random.PRNGKey(0))\n",
    "\n",
    "    # Create particles\n",
    "    cosmo = jc.Planck15(Omega_c=omega_c, sigma8=sigma8)\n",
    "    \n",
    "    # Initial displacement\n",
    "    dx, p, f = lpt(cosmo, initial_conditions, a=0.8,order=order)\n",
    "\n",
    "    return initial_conditions ,  dx\n",
    "\n",
    "initial_conditions_1 , lpt_displacements_1 = lpt_simulation(0.25, 0.8 , order=1)\n",
    "lpt_displacements_1.block_until_ready()\n",
    "initial_conditions_2 , lpt_displacements_2 = lpt_simulation(0.25, 0.8 , order=2)\n",
    "lpt_displacements_2.block_until_ready()\n",
    "%timeit initial_conditions_1 , lpt_displacements_1 = lpt_simulation(0.25, 0.8 , order=1);lpt_displacements_1.block_until_ready()\n",
    "%timeit initial_conditions_2 , lpt_displacements_2 = lpt_simulation(0.25, 0.8, order=2);lpt_displacements_2.block_until_ready()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAA+YAAAH7CAYAAABBt4VLAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjkuMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8hTgPZAAAACXBIWXMAAA9hAAAPYQGoP6dpAAEAAElEQVR4nOy9e7RtWV3f+fnNudbar/O477pVUBZlKfK0iYAJ8rQVGYAYjdXERyIUMRBEiNhqTA9sNGqqNYmNTTTiyBhqFLptDISEqDziWzQdBTSIIJSghfW+de89j733esw5+4/5WHPtc25VgWCVYf1q3FHn7LP2esw51/y9vr/vT5xzjlFGGWWUUUYZZZRRRhlllFFGGeVBEfVg38Aoo4wyyiijjDLKKKOMMsooo3w2y+iYjzLKKKOMMsooo4wyyiijjDLKgyijYz7KKKOMMsooo4wyyiijjDLKKA+ijI75KKOMMsooo4wyyiijjDLKKKM8iDI65qOMMsooo4wyyiijjDLKKKOM8iDK6JiPMsooo4wyyiijjDLKKKOMMsqDKKNjPsooo4wyyiijjDLKKKOMMsooD6KMjvkoo4wyyiijjDLKKKOMMsooozyIMjrmo4wyyiijjDLKKKOMMsooo4zyIMromI/yWSkf//jHERF++qd/+sG+lc+ovPjFL+YRj3jEg30bn3b53u/9XkTkQbm2iPC93/u9D8q1RxlllFFGeejJQ0kvjHr/0y8Ppfkd5X9sGR3zUf6HlJ/+6Z9GRI79993f/d2fkWv+83/+z/kP/+E/fFLfuXDhAt/5nd/JF3zBFzCdTjl16hTPec5zePvb3/4ZucdRHpj84i/+4kNSCV+6dImXvvSlnD17lsViwZd+6Zfy3ve+98G+rVFGGWWU+5T//t//OzfeeCPXXXcd0+mUhz3sYTz72c/m9a9//YN9a3/lMur9h6aMen+Uh4KIc8492Dcxyiifbvnpn/5pbrrpJv7ZP/tnXH/99YO/Pe5xj+N/+p/+J+q6pixLtNaflmtubW1x4403PuAs/Ic//GG+7Mu+jLvvvpubbrqJJz3pSVy6dIk3vvGNvP/97+c7vuM7+Bf/4l/8pe7pxS9+Mb/2a7/Gxz/+8b/UeR5q0nUdXdcxnU4/I+f/1m/9Vn7sx36M47bH9XpNURQURfEZufaVxFrL05/+dP7gD/6A7/zO7+TMmTP8+I//OLfeeiu///u/z+d//uf/ld7PKKOMMsoDkfe85z186Zd+KZ/zOZ/Di170Is6fP8+tt97K7/7u73LLLbfw0Y9+9MG+xb+0iAivfe1r79exG/X+py6j3h/1/meD/NWusFFG+SuW5z73uTzpSU869m8PZHM/PDxksVh8um+Ltm258cYbuXjxIr/xG7/B3/ybfzP97dWvfjXf+I3fyL/8l/+SJz3pSfzdv/t3r3ie9XpNVVUo9ZkHvzjnWK/XzGazT/u5P9nneDAUZJTPlFFwf/ILv/ALvOc97+HNb34zN954IwAvfOELeeQjH8lrX/ta3vSmNz0o9zXKKKOMcl/ygz/4g+zu7vLf/tt/48SJE4O/3XXXXQ/OTT0IMur9oYx6//5l1PuffTJC2Uf5rJTjasxf/OIXs7W1xS233MLznvc8tre3+cZv/EYAPvKRj/C1X/u1nD9/nul0ysMf/nC+7uu+jsuXLwM+Wn54eMjP/MzPJMj8i1/84ite/9//+3/PBz7wAb77u797oJwBtNa84Q1v4MSJE4Po+6/92q8hIvw//8//w2te8xoe9rCHMZ/P2dvbA+A//If/wOMe9zim0ymPe9zjeOtb33rsta21vO51r+Oxj30s0+mUq666ipe97GVcvHhxcNwjHvEIvvIrv5J3vOMdPOlJT2I2m/GGN7zhis/0rGc9i8c97nH8/u//Pl/yJV/CbDbj+uuv5yd+4icGx93fc7z5zW/miU98IrPZjDNnzvD3/t7f4y/+4i8G57hSrdnP/dzPpe+eOnWKr/u6r+PWW289ctx//a//lec973mcPHmSxWLBF37hF/KjP/qjgF8HP/ZjPwYwKIGIclyt2fve9z6e+9znsrOzw9bWFl/2ZV/G7/7u7w6OieUVv/3bv823f/u3J1ja13zN13D33XdfcVyj/MIv/AJXXXUVf+fv/J302dmzZ3nhC1/I2972Nuq6vt9zjDLKKKP8Vcstt9zCYx/72CNOOcC5c+eOfPbp2Mej/Mqv/ApPf/rTWSwWnDhxgr/9t/82f/zHfzw4JuqTj370o7z4xS/mxIkT7O7uctNNN7FcLgfH1nXNq1/9as6ePcv29jZf9VVfxSc+8YkHNA6j3h/1/qj3R7k/GTPmo/wPLZcvX+aee+4ZfHbmzJkrHt91Hc95znN42tOexr/8l/+S+XxO0zQ85znPoa5rXvnKV3L+/Hn+4i/+gre//e1cunSJ3d1dfvZnf5Zv/uZv5ou/+It56UtfCsANN9xwxev8p//0nwD4pm/6pmP/vru7y9/+23+bn/mZn+GjH/0on/d5n5f+9v3f//1UVcV3fMd3UNc1VVXxzne+k6/92q/lMY95DDfffDMXLlzgpptu4uEPf/iRc7/sZS9LUP9XvepVfOxjH+Nf/+t/zfve9z5++7d/m7Is07Ef/vCH+fqv/3pe9rKX8Q//4T/kC77gC674TAAXL17kec97Hi984Qv5+q//ev7f//f/5eUvfzlVVfGSl7xkcOxxzxHv68lPfjI333wzd955Jz/6oz/Kb//2b/O+973vWMMuyg/+4A/yPd/zPbzwhS/km7/5m7n77rt5/etfzzOe8YzBd9/1rnfxlV/5lVx99dX843/8jzl//jx//Md/zNvf/nb+8T/+x7zsZS/jtttu413vehc/+7M/e5/PC/BHf/RHPP3pT2dnZ4fv+q7voixL3vCGN/CsZz2LX//1Xz9igL3yla/k5MmTvPa1r+XjH/84r3vd6/jWb/1Wfv7nf/4+r/O+972PL/qiLzqSXfjiL/5ifvInf5I/+ZM/4fGPf/z93u8oo4wyyl+lXHfddfzO7/wOH/jAB3jc4x53n8d+uvZxgHe/+90897nP5XM/93P53u/9XlarFa9//et56lOfynvf+94jBGkvfOELuf7667n55pt573vfy7/9t/+Wc+fO8UM/9EPpmG/+5m/m537u5/iGb/gGvuRLvoRf+ZVf4fnPf/4DGodR7496f9T7o9yvuFFG+R9QfuqnfsoBx/5zzrmPfexjDnA/9VM/lb7zohe9yAHuu7/7uwfnet/73ucA9+Y3v/k+r7lYLNyLXvSiB3R/T3jCE9zu7u59HvMjP/IjDnD/8T/+R+ecc7/6q7/qAPe5n/u5brlcHjnf1Vdf7S5dupQ+e+c73+kAd91116XPfvM3f9MB7o1vfOPg+7/8y7985PPrrrvOAe6Xf/mXH9AzPfOZz3SA+1f/6l+lz+q6dk94whPcuXPnXNM09/kcTdO4c+fOucc97nFutVqlz9/+9rc7wP3v//v/nj577Wtf6/Lt6+Mf/7jTWrsf/MEfHNzTf//v/90VRZE+77rOXX/99e66665zFy9eHBxrrU0/v+IVr3BX2h4B99rXvjb9/tVf/dWuqip3yy23pM9uu+02t7297Z7xjGekz+Ka/PIv//LBtV796lc7rfVg7o6TxWLhXvKSlxz5/D//5//8Sc3TKKOMMspfpbzzne90WmuntXZPecpT3Hd913e5d7zjHUknRPl07+NR91y4cCF99gd/8AdOKeW+6Zu+KX0W9cnm/vo1X/M17vTp0+n397///Q5w3/It3zI47hu+4RuO6IXjZNT7o96PMur9Ua4kI5R9lP+h5cd+7Md417veNfh3f/Lyl7988Pvu7i4A73jHO47A2j5V2d/fZ3t7+z6PiX+PUK8oL3rRiwb1Xrfffjvvf//7edGLXpTuFeDZz342j3nMYwbfffOb38zu7i7Pfvazueeee9K/Jz7xiWxtbfGrv/qrg+Ovv/56nvOc5zzg5yqKgpe97GXp96qqeNnLXsZdd93F7//+79/nc/ze7/0ed911F9/yLd8yqOd6/vOfz6Me9Sj+83/+z1e
      "text/plain": [
       "<Figure size 2000x500 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "lpt_fields = {\"First Order\" : cic_paint_dx(lpt_displacements_1) , \"Second Order\" : cic_paint_dx(lpt_displacements_2)}\n",
    "plot_fields_single_projection(lpt_fields)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Custom ODE Solver with Absolute Positions\n",
    "\n",
    "Just like in the introduction notebook, this example uses **absolute particle positions** initialized on a uniform grid. We evolve these absolute positions forward using a Cloud-in-Cell (CIC) scheme, which enables clear tracking of particle movement across the simulation volume.\n",
    "\n",
    "Here, we integrate over multiple snapshots with `diffeqsolve` and a **Leapfrog solver**.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Solver Stats : {'max_steps': Array(4096, dtype=int32, weak_type=True), 'num_accepted_steps': Array(5, dtype=int32, weak_type=True), 'num_rejected_steps': Array(0, dtype=int32, weak_type=True), 'num_steps': Array(5, dtype=int32, weak_type=True)}\n"
     ]
    }
   ],
   "source": [
    "mesh_shape = [128, 128, 128]\n",
    "box_size = [128., 128., 128.]\n",
    "snapshots = jnp.array([0.1 ,0.5, 1.])\n",
    "\n",
    "@jax.jit\n",
    "def run_simulation(omega_c, sigma8):\n",
    "    # Create a small function to generate the matter power spectrum\n",
    "    k = jnp.logspace(-4, 1, 128)\n",
    "    pk = jc.power.linear_matter_power(jc.Planck15(Omega_c=omega_c, sigma8=sigma8), k)\n",
    "    pk_fn = lambda x: jnp.interp(x.reshape([-1]), k, pk).reshape(x.shape)\n",
    "\n",
    "    # Create initial conditions\n",
    "    initial_conditions = linear_field(mesh_shape, box_size, pk_fn, seed=jax.random.PRNGKey(0))\n",
    "\n",
    "    # Create particles\n",
    "    cosmo = jc.Planck15(Omega_c=omega_c, sigma8=sigma8)\n",
    "    \n",
    "    particles = uniform_particles(mesh_shape)\n",
    "    # Initial displacement\n",
    "    dx, p, f = lpt(cosmo, initial_conditions,particles=particles,a=0.1,order=2)\n",
    "    \n",
    "    # Evolve the simulation forward\n",
    "    ode_fn = ODETerm(\n",
    "        make_diffrax_ode(cosmo, mesh_shape))\n",
    "    solver = LeapfrogMidpoint()\n",
    "\n",
    "    stepsize_controller = ConstantStepSize()\n",
    "    res = diffeqsolve(ode_fn,\n",
    "                      solver,\n",
    "                      t0=0.1,\n",
    "                      t1=1.,\n",
    "                      dt0=0.2,\n",
    "                      y0=jnp.stack([particles + dx, p], axis=0),\n",
    "                      args=cosmo,\n",
    "                      saveat=SaveAt(ts=snapshots),\n",
    "                      stepsize_controller=stepsize_controller)\n",
    "\n",
    "    ode_particles = [sol[0] for sol in res.ys]\n",
    "    return initial_conditions ,  particles + dx , ode_particles , res.stats\n",
    "\n",
    "initial_conditions , lpt_particles , ode_particles , solver_stats = run_simulation(0.25, 0.8)\n",
    "print(f\"Solver Stats : {solver_stats}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAB8QAAAH5CAYAAAD3DoQvAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjkuMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8hTgPZAAAACXBIWXMAAA9hAAAPYQGoP6dpAAEAAElEQVR4nOy9ebwlRX2w/1RVb2e568zADNswLIKiYjKKEZFFURaVoAKvMQoYEUUMatTEJUZwjbtRI6LGJWp8I6Ci7889omJQA4r7hggIwzj7Xc7WS1X9/qjuPufMnYEZBAFTj5/rcPv2OV1dVV317e8qrLUWj8fj8Xg8Ho/H4/F4PB6Px+PxeDwej8fj8Xj+xJD3dAM8Ho/H4/F4PB6Px+PxeDwej8fj8Xg8Ho/H47k78AZxj8fj8Xg8Ho/H4/F4PB6Px+PxeDwej8fj8fxJ4g3iHo/H4/F4PB6Px+PxeDwej8fj8Xg8Ho/H4/mTxBvEPR6Px+PxeDwej8fj8Xg8Ho/H4/F4PB6Px/MniTeIezwej8fj8Xg8Ho/H4/F4PB6Px+PxeDwej+dPEm8Q93g8Ho/H4/F4PB6Px+PxeDwej8fj8Xg8Hs+fJN4g7vF4PB6Px+PxeDwej8fj8Xg8Ho/H4/F4PJ4/SbxB3OPxeDwej8fj8Xg8Ho/H4/F4PB6Px+PxeDx/kniDuMfj8Xg8Ho/H4/F4PB6Px+PxeDwej8fj8Xj+JPEGcc99BiEEF1544S6du//++3P22Wfv9jVuuukmhBB85CMf2e3P3hv4xje+gRCCb3zjG/Wxs88+m/3333+XPn/hhRcihLh7Gncf4p6cB7szXh6Px+O5d7Kj/XhX+chHPoIQgptuuukOz72z8s7dyT3Vpt3pN4/H4/Hc+7jmmms48sgjabVaCCH44Q9/+Ae9n+7qe9V9XQdwR/whMskfyrHHHsuxxx77R7+ux+PxeP54+P377sHrpj2euwdvEPf80agUlddee+1d8n1XX301F154IXNzc3fJ990ZNmzYwEte8hIOPfRQms0mrVaLtWvX8rrXve4ebdft0ev1uPDCC++RF2KP47bbbuPCCy/khz/84T3dlCX827/9G/e///1JkoSDDz6Yd7/73fd0kzwejwfYNTmiemmsfpRS7LfffjzpSU+q19yzzz577Jyd/dyeUbd6wd/Rz/ve9767+M49o7zhDW/gs5/97D3djCVcffXVHHXUUTSbTVauXMkFF1xAp9O5p5vl8Xg89wnyPOf0009n69atvOMd7+BjH/sYq1evvqebdYf4dyf4+c9/zoUXXnivc0gzxvDmN7+ZNWvWkCQJD37wg/nkJz95TzfL4/F4/qS4L+7fF198Maeffjr77bffHb73/ynjddOe/60E93QDPJ5dpd/vEwTDKXv11Vdz0UUXcfbZZzM9PT127q9+9SukvHv9Pa655hpOPvlkOp0OT3/601m7di0A1157Lf/8z//Mt771Lb7yla/crW3YFT7wgQ9gjKl/7/V6XHTRRQBLvLX/8R//kZe97GV/zObdK1m9ejX9fp8wDO+W77/tttu46KKL2H///XnIQx4y9rftx+uPySWXXMJzn/tcnvKUp/B3f/d3XHXVVVxwwQX0ej3+4R/+4R5pk8fj8dwZ/uqv/oqTTz4ZrTW/+MUvuPjii/niF7/Id7/7XZ7znOdw/PHH1+feeOON/NM//RPnnnsuj3rUo+rjBx544B1e5+KLL6bdbo8de/jDH86BBx5Iv98niqK77qbuI9zdMtgb3vAGTjvtNE499dSx4894xjN46lOfShzHd9u1d8YPf/hDHvOYx3D/+9+ft7/97dx666289a1v5frrr+eLX/ziH709Ho/Hc1/jhhtu4Oabb+YDH/gA55xzTn383vx+el95dzr66KPvVpnk5z//ORdddBHHHnvskmiye1If8spXvpJ//ud/5tnPfjYPe9jDuOKKK3ja056GEIKnPvWp91i7PB6P50+J++L+/aY3vYnFxUWOOOII1q9ff083Z6d43fS9W77y3HfxBnHPfYYkSXb53LtbGTo3N8eTnvQklFJcd911HHrooWN/f/3rX88HPvCBu7UNu8rubJxBEIw5HfypYK1lMBjQaDR26XwhxG7Nt7uSu0vQuSP6/T6vfOUrefzjH89ll10GwLOf/WyMMbz2ta/l3HPPZWZm5h5pm8fj8ewuf/7nf87Tn/70+vdHPvKRnHLKKVx88cVccsklPOIRj6j/du211/JP//RPPOIRjxj7zK5w2mmnsXz58h3+7Z7aR+5qut0urVZrl8+/JwzSAEoplFL3yLVf8YpXMDMzwze+8Q0mJycBlzr+2c9+Nl/5yld43OMed4+0y+PxeO4rbNy4EWCJo/u99f30nnx3MsaQZdkuyxlSyntMJrmnHAPXrVvH2972Ns4//3ze8573AHDOOedwzDHH8NKXvpTTTz/9HpMZPB6P50+J+9r+DfDNb36zjg7f3rn97sTrpu8Yr5v2/DHwKdM99yhnn3027XabdevWceqpp9Jut1mxYgUveclL0FqPnTtaQ/zCCy/kpS99KQBr1qyp05RWabq2r1+5detWXvKSl/CgBz2IdrvN5OQkJ510Ej/60Y/uVLsvueQS1q1bx9vf/vYlxnCAPffck3/8x38cO/be976Xww47jDiO2WuvvTj//POXpFU/9thjeeADH8jPf/5zjjvuOJrNJnvvvTdvfvObl1zj1ltv5dRTT6XVarHHHnvwohe9iDRNl5w3WvfjpptuYsWKFQBcdNFFdb+N9uv2NV6KouC1r30tBx54IHEcs//++/OKV7xiybX2339/nvCEJ/Dtb3+bI444giRJOOCAA/j3f//3sfPyPOeiiy7i4IMPJkkSli1bxlFHHcVXv/rVpR09QpUq91vf+hbPec5zWLZsGZOTk5x55pls27Zth2358pe/zEMf+lAajQaXXHIJAL/97W85/fTTmZ2dpdls8hd/8Rf8f//f/zf2+Z3VafnlL3/JaaedxuzsLEmS8NCHPpTPfe5zS9o6NzfHi170Ivbff3/iOGafffbhzDPPZPPmzXzjG9/gYQ97GADPfOYz6zGorrWjOi3dbpcXv/jF7LvvvsRxzCGHHMJb3/pWrLVj5wkheP7zn89nP/tZHvjABxLHMYcddhhf+tKXbrdvAa688kq2bNnC8573vLHj559/Pt1ud0kfeTwez32JRz/60YCLBv9jsLN6nd/73vc48cQTmZqaotlscswxx/Df//3fd/h91lpe97rXsc8++9BsNjnuuOP42c9+tkttqfa0t771rbzjHe9g9erVNBoNjjnmGH7605+OnVvJZTfccAMnn3wyExMT/PVf/zWw63vRjmqIz83N8cIXvrD+7EEHHcSb3vSmJV7nxhj+5V/+hQc96EEkScKKFSs48cQT6zT5Qgi63S4f/ehHl6S231kN8bta/tqehYUFvvrVr/L0pz+9NoYDnHnmmbTbbT71qU/d4Xd4PB7P/2bOPvtsjjnmGABOP/10hBB1JrOd1SD9+Mc/ztq1a2k0GszOzvLUpz6VW2655Q6vNTc3x9lnn83U1BTT09OcddZZd6rU2R/67lTd1y9/+UvOOOMMJicnWbZsGS94wQsYDAZj51bveJ/4xCfq/ax6v7vuuus46aSTmJycpN1u85jHPIbvfve7Y5//Q2WSdevW8axnPYu99tqLOI5Zs2YN5513HlmW8ZGPfITTTz8dgOOOO67em6tr7aiG+MaNG3nWs57FnnvuSZIkHH744Xz0ox8dO2dUdnn/+99f6yIe9rCHcc0119xu3wJcccUV5Hk+Nj5CCM477zxuvfVWvvOd79zhd3g8Ho/n9rkv7t/gIq/vbH1zr5v2umnPfZ97p6uO538VWmtOOOEEHv7wh/PWt76Vr33ta7ztbW/jwAMP5LzzztvhZ5785Cfz61//mk9+8pO84x3vqKOzKmPv9vz2t7/ls5/9LKeffjpr1qxhw4YNXHLJJRxzzDH8/Oc/Z6+99tqtNn/uc5+j0Whw2mmn7dL5F154IRdddBHHH3885513Hr/61a+4+OKLueaaa/jv//7
      "text/plain": [
       "<Figure size 2000x500 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from jaxpm.plotting import plot_fields_single_projection\n",
    "\n",
    "fields = {\"Initial Conditions\" : initial_conditions , \"LPT Field\" : cic_paint(jnp.zeros(mesh_shape) ,lpt_particles)}\n",
    "for i , field in enumerate(ode_particles[1:]):\n",
    "    fields[f\"field_{i}\"] = cic_paint(jnp.zeros(mesh_shape) , field)\n",
    "plot_fields_single_projection(fields)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Weighted Field Projection for Central Region\n",
    "\n",
    "In this cell, we apply custom weights to enhance density specifically in the **central 3D region** of the grid. By updating weights in this area, we multiply density by a factor of 3, emphasizing the structure in the center of the simulation volume.\n",
    "\n",
    "We compare:\n",
    "- **Weighted**: Density increased in the central region.\n",
    "- **Unweighted**: Standard CIC painting without additional weighting.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAA+YAAAH7CAYAAABBt4VLAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjkuMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8hTgPZAAAACXBIWXMAAA9hAAAPYQGoP6dpAAEAAElEQVR4nOy9d/wlVX3//zznzNz2KbvLsgsLGkSioqAhEkUREDWK2L5olJ8lKiYRY+yJJSYaRGNMNOarUb+2+FW/9h4TuyYWBEsK2LAriLSl7O6n3DYz5/z+OHXm3s/ugppFnbcP3M+dO3fOmdPe7/fr3YQxxtBSSy211FJLLbXUUksttdRSSy0dEJIHugMttdRSSy211FJLLbXUUksttfTrTK1i3lJLLbXUUksttdRSSy211FJLB5BaxbylllpqqaWWWmqppZZaaqmllg4gtYp5Sy211FJLLbXUUksttdRSSy0dQGoV85ZaaqmlllpqqaWWWmqppZZaOoDUKuYttdRSSy211FJLLbXUUksttXQAqVXMW2qppZZaaqmlllpqqaWWWmrpAFKrmLfUUksttdRSSy211FJLLbXU0gGkVjFvqaWWWmqppZZaaqmlllpqqaUDSK1i3lJLP0c666yzuMUtbnGjf7u4uPjz7dANpLe85S0IIbjkkksOaD9SOpB9usUtbsFZZ531P95uSy211FJLN2069dRTOfXUU2/0b4899tifb4duIL3gBS9ACHFA+9CkA9knIQQveMELDkjbLbXkqVXMW/qVp/e+970IIfjQhz40891v/dZvIYTgs5/97Mx3v/Ebv8GJJ574P9HFG0TD4ZAXvOAFfO5znzvQXfmVoAsuuIAXvOAF7N69+0B3pUaTyYTnPOc5HHbYYfT7fU444QQ+/elPH+hutdRSSy393MgrYtdee+3c74899tgbrfz+KtAVV1zBC17wAi666KID3ZVfCfrYxz52k1S+d+/ezdlnn822bdtYWFjgHve4B//93/99oLvV0gGgVjFv6VeeTjrpJAC++MUv1q6vrKzwzW9+kyzLOP/882vfXXbZZVx22WXht/tLb3zjG/nud7/7s3V4HzQcDjn33HN/bRTzRz/60YxGI4444ohfyPMvuOACzj333LmK+Xe/+13e+MY3/kLa3RedddZZ/MM//AOPetSjeOUrX4lSivvd734z67illlpqqaX/efrUpz7Fpz71qV9oG1dccQXnnnvur41i/rznPY/RaPQLe/7HPvYxzj333LnfjUYjnve85/3C2t6ItNbc//73553vfCdPfvKTeelLX8rOnTs59dRT+f73v/8/3p+WDixlB7oDLbX0i6bDDjuMI488ckah+dKXvoQxhoc97GEz3/nPN1Qxz/P8Z+vsrwGVZYnWmk6ns1/3K6VQSv2CezWfut3uAWn3q1/9Ku9+97t52ctexjOf+UwAHvOYx3Dsscfy7Gc/mwsuuOCA9KulllpqqSVL+8vDfp1pPB7T6XSQcv/sgFmWkWUHRjXp9XoHpN33v//9XHDBBbzvfe/joQ99KABnnnkmt771rTnnnHN45zvfeUD61dKBodZi3tKvBZ100klceOGFNST2/PPP55hjjuH000/ny1/+Mlrr2ndCCO52t7uFa29/+9s5/vjj6ff7HHTQQTz84Q/nsssuq7UzL8b8uuuu49GPfjTLy8ts3ryZxz72sXzta19DCMFb3vKWmb5efvnlnHHGGSwuLrJt2zae+cxnUlUVAJdccgnbtm0D4Nxzz0UIMRMX9Z3vfIeHPvShHHTQQfR6PX7nd36Hf/mXf5lp51vf+hb3vOc96ff73OxmN+Ov//qva2OwN/Lx8D/60Y847bTTWFhY4LDDDuOFL3whxphw3yWXXIIQgr//+7/nFa94BUcddRTdbpeLL74YgH//93/n5JNPZmFhgc2bN/O//tf/4tvf/natrY1izD/+8Y+H3y4tLXH/+9+fb33rWzN9/c53vsOZZ57Jtm3b6Pf73OY2t+Ev//IvAetG+axnPQuAI488Moynb2tejPmPfvQjHvawh3HQQQcxGAy4y13uwkc/+tHaPZ/73OcQQvDe976XF7/4xdzsZjej1+txr3vdix/84Af7HN/3v//9KKU4++yzw7Ver8cf/uEf8qUvfWlm3bXUUkst/TrQ/p6t//iP/4hSquYJ9fKXvxwhBH/6p38arlVVxdLSEs95znPCNa01r3jFKzjmmGPo9XoccsghPOEJT2DXrl21vsyLMb/00kt50IMexMLCAtu3b+cZz3gGn/zkJxFCzPVyu/jii7nHPe7BYDDg8MMP56UvfWntXe90pzsB8LjHPS7wp1Ru+MpXvsJ973tfNm3axGAw4O53v/uMByBYY8Od7nQner0eRx11FK9//ev3Os7N9zz22GP5r//6L0488UT6/T5HHnkkr3vd62r3+bl597vfzfOe9zwOP/xwBoMBKysrALzvfe8LMtTBBx/M7//+73P55ZfXnrFRjPn+yF9+PO53v/uxZcsWFhYWuMMd7sArX/lKwMotr3nNawDCWKZtzYsxv/DCCzn99NNZXl5mcXGRe93rXnz5y1+u3eNllPPPP58//dM/De7oD37wg7nmmmv2Ob7vf//7OeSQQ3jIQx4Srm3bto0zzzyTD3/4w0wmk30+o6VfHWot5i39WtBJJ53E2972Nr7yla8ERnr++edz4okncuKJJ7Jnzx6++c1vcoc73CF8d/TRR7N161YAXvziF/P85z+fM888kz/6oz/immuu4VWvehWnnHIKF154IZs3b57brtaaBz7wgXz1q1/liU98IkcffTQf/vCHeexjHzv3/qqqOO200zjhhBP4+7//ez7zmc/w8pe/nKOOOoonPvGJbNu2jde+9rU88YlP5MEPfnA4yH2/v/Wtb3G3u92Nww8/nD//8z9nYWGB9773vZxxxhl84AMf4MEPfjAAV111Ffe4xz0oyzLc94Y3vIF+v7/fY1pVFfe97325y13uwktf+lI+8YlPcM4551CWJS984Qtr9775zW9mPB5z9tln0+12Oeigg/jMZz7D6aefzi1veUte8IIXMBqNeNWrXsXd7nY3/vu//3uvSfTe9ra38djHPpbTTjuNv/u7v2M4HPLa1742ADD+t1//+tc5+eSTyfOcs88+m1vc4hb88Ic/5F//9V958YtfzEMe8hC+973v8a53vYv//b//NwcffDBAAD+adPXVV3PiiScyHA556lOfytatW3nrW9/Kgx70IN7//veH8fX0t3/7t0gpeeYzn8mePXt46UtfyqMe9Si+8pWv7HVsL7zwQm5961uzvLxcu37nO98ZgIsuuoib3/zme31GSy211NKvKu3rbD355JPRWvPFL36RBzzgAQCcd955SCk577zzwnMuvPBC1tbWOOWUU8K1JzzhCbzlLW/hcY97HE996lP58Y9/zKtf/WouvPBCzj///A0949bX17nnPe/JlVdeydOe9jQOPfRQ3vnOd87NYQOwa9cu7nvf+/KQhzyEM888k/e///085znP4fa3vz2nn346t73tbXnhC1/IX/3VX3H22Wdz8sknA4TcN//+7//O6aefzvHHH88555yDlJI3v/nN3POe9+S8884L/OIb3/gG97nPfdi2bRsveMELKMuSc845h0MOOWS/x3vXrl3c737348wzz+QRj3gE733ve3niE59Ip9PhD/7gD2r3vuhFL6LT6fDMZz6TyWRCp9MJ43mnO92Jl7zkJVx99dW88pWv5Pzzz9+rDAX7L399+tOf5gEPeAA7duwI4//tb3+bj3zkIzztaU/jCU94AldccQWf/vSnedvb3rbPd/7Wt77FySefzPLyMs9+9rPJ85zXv/71nHrqqXz+85/nhBNOqN3/lKc8hS1btnDOOedwySWX8IpXvIInP/nJvOc979lrOxdeeCF3vOMdZ7wK7nznO/OGN7yB733ve9z+9rffZ39b+hUh01JLvwb0rW99ywDmRS96kTHGmKIozMLCgnnrW99qjDHmkEMOMa95zWuMMcasrKwYpZR5/OMfb4wx5pJLLjFKKfPiF7+49sxvfOM
      "text/plain": [
       "<Figure size 2000x500 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from jaxpm.plotting import plot_fields_single_projection\n",
    "\n",
    "center = slice(mesh_shape[0] // 4, 3 * mesh_shape[0] // 4 )\n",
    "center3d = (slice(None) , center,center) \n",
    "weights = jnp.ones_like(initial_conditions)\n",
    "weights = weights.at[center3d].multiply(3)\n",
    "\n",
    "weighted = cic_paint_dx(ode_solutions[0], weight=weights)\n",
    "unweighted = cic_paint_dx(ode_solutions[0] , weight=1.0)\n",
    "\n",
    "plot_fields_single_projection({\"Weighted\" : weighted , \"Unweighted\" : unweighted} , project_axis=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Weighted Field Projection with Absolute Positions\n",
    "\n",
    "For simulations with absolute positions, we apply a weight factor of **1.3** to the central 3D region. Unlike previous cases using displacements, here the weight affects the absolute particle positions directly, impacting the overall density field differently.\n",
    "\n",
    "\n",
    "__Note:__ Since the weights apply to absolute positions (not displacements), the result differs, affecting the particle density distribution directly.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAA+YAAAH7CAYAAABBt4VLAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjkuMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8hTgPZAAAACXBIWXMAAA9hAAAPYQGoP6dpAAEAAElEQVR4nOy9edzu13jv/15rfYd7eJ5nT9k7IxGpOVSlhEgiKBHKQcnpRKJD/BQtrdL20AhVRfVQHEPrqKN0MLRa1NQaIlE6JFWCKkJkzp6e4b7v77DW9fvjWuv7vZ/snUhSuoP749VmP/f0ndZwDZ/rcxkRERZYYIEFFlhggQUWWGCBBRZYYIFDAnuoT2CBBRZYYIEFFlhggQUWWGCBBX6QsXDMF1hggQUWWGCBBRZYYIEFFljgEGLhmC+wwAILLLDAAgsssMACCyywwCHEwjFfYIEFFlhggQUWWGCBBRZYYIFDiIVjvsACCyywwAILLLDAAgsssMAChxALx3yBBRZYYIEFFlhggQUWWGCBBQ4hFo75AgsssMACCyywwAILLLDAAgscQiwc8wUWWGCBBRZYYIEFFlhggQUWOIRYOOYLLLDAAgsssMACCyywwAILLHAIsXDMF1jgO4hzzjmHO9zhDrf6u0tLS9/ZE7qF+JM/+ROMMVx22WWH9DzmcSjP6Q53uAPnnHPOf/txF1hggQUWuG3j9NNP5/TTT7/V3z3hhBO+syd0C/HCF74QY8whPYcb4lCekzGGF77whYfk2AsskLBwzBf4vsdf/uVfYozhr/7qrw5474d/+IcxxvCxj33sgPduf/vbc/LJJ/93nOItwmQy4YUvfCEf//jHD/WpfF/goosu4oUvfCH79u071KeyCVVV8bznPY+jjjqK4XDISSedxEc+8pFDfVoLLLDAAt8xJEfs+uuvP+j7J5xwwq12fr8fcOWVV/LCF76QSy655FCfyvcFPvCBD9wmne99+/Zx7rnnsnPnTsbjMQ9+8IP513/910N9WgscAiwc8wW+73HKKacA8KlPfWrT66urq3z+858nyzIuvPDCTe9dfvnlXH755d13by7+6I/+iC9/+cv/tRP+NphMJpx//vk/MI75k570JKbTKccee+x35fcvuugizj///IM65l/+8pf5oz/6o+/Kcb8dzjnnHP7gD/6An/mZn+HVr341zjke+chHHjCOF1hggQUW+O/Hhz/8YT784Q9/V49x5ZVXcv755//AOObPf/7zmU6n37Xf/8AHPsD5559/0Pem0ynPf/7zv2vHvjGEEHjUox7FO97xDp7xjGfw8pe/nGuvvZbTTz+dr3zlK//t57PAoUV2qE9ggQW+2zjqqKM47rjjDnBoPv3pTyMiPPGJTzzgvfT3LXXM8zz/r53sDwDatiWEQFEUN+vzzjmcc9/lszo4yrI8JMf97Gc/y5//+Z/zile8guc85zkAPPnJT+aEE07guc99LhdddNEhOa8FFlhggQUUN3cP+0HGbDajKAqsvXl5wCzLyLJD45oMBoNDctx3vetdXHTRRbzzne/kCU94AgBnnXUWd77znTnvvPN4xzvecUjOa4FDg0XGfIEfCJxyyilcfPHFmyKxF154Ife4xz0488wz+cd//EdCCJveM8bwwAc+sHvtT//0TznxxBMZDods376dn/zJn+Tyyy/fdJyD1Zjv3r2bJz3pSaysrLB161bOPvts/u3f/g1jDH/yJ39ywLleccUVPPaxj2VpaYmdO3fynOc8B+89AJdddhk7d+4E4Pzzz8cYc0Bd1Je+9CWe8IQnsH37dgaDAT/6oz/K3/zN3xxwnC984Qs85CEPYTgccswxx/A7v/M7m+7BTSHVw3/ta1/jjDPOYDwec9RRR/GiF70IEek+d9lll2GM4fd///d51atexfHHH09Zllx66aUA/MM//AOnnnoq4/GYrVu38j/+x//gi1/84qZj3ViN+d/93d91311eXuZRj3oUX/jCFw441y996UucddZZ7Ny5k+FwyF3uchf+1//6X4DSKH/9138dgOOOO667n+lYB6sx/9rXvsYTn/hEtm/fzmg04v73vz/vf//7N33m4x//OMYY/vIv/5KXvOQlHHPMMQwGAx760Ifyn//5n9/2/r7rXe/COce5557bvTYYDPj5n/95Pv3pTx8w7hZYYIEFfhBwc9fWP/zDP8Q5t4kJ9cpXvhJjDL/6q7/avea9Z3l5mec973ndayEEXvWqV3GPe9yDwWDA4YcfzlOf+lT27t276VwOVmP+jW98g8c85jGMx2N27drFs5/9bD70oQ9hjDkoy+3SSy/lwQ9+MKPRiKOPPpqXv/zlm671vve9LwBPecpTuv1p3m74zGc+wyMe8Qi2bNnCaDTiQQ960AEMQNBkw33ve18GgwHHH388b3zjG2/yPt/wOk844QT+5V/+hZNPPpnhcMhxxx3HG97whk2fS8/mz//8z3n+85/P0UcfzWg0YnV1FYB3vvOdnQ112GGH8bM/+7NcccUVm37jxmrMb479le7HIx/5SLZt28Z4POZe97oXr371qwG1W173utcBdPdy/lgHqzG/+OKLOfPMM1lZWWFpaYmHPvSh/OM//uOmzyQb5cILL+RXf/VXOzr64x73OK677rpve3/f9a53cfjhh/P4xz++e23nzp2cddZZvPe976Wqqm/7Gwt8/2CRMV/gBwKnnHIKb3vb2/jMZz7TbaQXXnghJ598MieffDL79+/n85//PPe617269+5617uyY8cOAF7ykpfwghe8gLPOOotf+IVf4LrrruM1r3kNp512GhdffDFbt2496HFDCDz60Y/ms5/9LE972tO4613vynvf+17OPvvsg37ee88ZZ5zBSSedxO///u/z0Y9+lFe+8pUcf/zxPO1pT2Pnzp28/vWv52lPexqPe9zjuoU8nfcXvvAFHvjAB3L00UfzG7/xG4zHY/7yL/+Sxz72sbz73e/mcY97HABXX301D37wg2nbtvvcm970JobD4c2+p957HvGIR3D/+9+fl7/85Xzwgx/kvPPOo21bXvSiF2367Fve8hZmsxnnnnsuZVmyfft2PvrRj3LmmWdyxzvekRe+8IVMp1Ne85rX8MAHPpB//dd/vUkRvbe97W2cffbZnHHGGbzsZS9jMpnw+te/vgvApO9+7nOf49RTTyXPc84991zucIc78NWvfpW//du/5SUveQmPf/zj+Y//+A/+7M/+jP/9v/83hx12GEAX/LghrrnmGk4++WQmkwm//Mu/zI4dO3jrW9/KYx7zGN71rnd19zfh937v97DW8pznPIf9+/fz8pe/nJ/5mZ/hM5/5zE3e24svvpg73/nOrKysbHr9fve7HwCXXHIJt7vd7W7yNxZYYIEFvl/x7dbWU089lRACn/rUp/jxH/9xAC644AKstVxwwQXd71x88cWsr69z2mmnda899alP5U/+5E94ylOewi//8i/z9a9/nde+9rVcfPHFXHjhhTfKjNvY2OAhD3kIV111Fb/yK7/CEUccwTve8Y6DatgA7N27l0c84hE8/vGP56yzzuJd73oXz3ve87jnPe/JmWeeyd3udjde9KIX8du//duce+65nHrqqQCd9s0//MM/cOaZZ3LiiSdy3nnnYa3lLW95Cw95yEO44IILuv3i3//933n4wx/Ozp07eeELX0jbtpx33nkcfvjhN/t+7927l0c+8pGcddZZ/NRP/RR/+Zd/ydOe9jSKouDnfu7nNn32xS9+MUVR8JznPIeqqiiKoruf973vfXnpS1/KNddcw6tf/WouvPDCm7Sh4ObbXx/5yEf48R//cY488sju/n/xi1/kfe97H7/yK7/CU5/6VK688ko+8pGP8La3ve3bXvMXvvAFTj31VFZWVnjuc59Lnue88Y1v5PTTT+cTn/gEJ5100qbPP/OZz2Tbtm2cd955XHbZZbzqVa/iGc94Bn/xF39xk8e5+OKLuc997nMAq+B+97sfb3rTm/iP//gP7nnPe37b813g+wSywAI/APjCF74ggLz4xS8WEZGmaWQ8Hstb3/pWERE5/PDD5XWve52IiKyurop
      "text/plain": [
       "<Figure size 2000x500 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from jaxpm.plotting import plot_fields_single_projection\n",
    "\n",
    "center = slice(mesh_shape[0] // 4, 3 * mesh_shape[0] // 4 )\n",
    "center3d = (slice(None) , center,center)  \n",
    "weights = jnp.ones_like(initial_conditions)\n",
    "weights = weights.at[center3d].multiply(1.3)\n",
    "\n",
    "weighted = cic_paint(jnp.zeros(mesh_shape),ode_particles[0], weight=weights)\n",
    "unweighted = cic_paint(jnp.zeros(mesh_shape),ode_particles[0] , weight=2.0)\n",
    "\n",
    "plot_fields_single_projection({\"Weighted\" : weighted , \"Unweighted\" : unweighted} , project_axis=0)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}