{
 "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": 1,
   "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 PIDController, Tsit5, ODETerm, SaveAt, diffeqsolve"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Particle Mesh Simulation with Diffrax Leapfrog Solver\n",
    "\n",
    "In this setup, we use the `Tsit5` solver from the `diffrax` library to evolve particle displacements over time in our Particle Mesh simulation.\n",
    "\n",
    "- **Efficient Displacement Tracking**: We initialize only displacements (`dx`) rather than absolute positions, which uses a the `pmwd` cic_painting algorithm which is more memory efficient at the cost of being slightly slower\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/wassim/micromamba/envs/jax/lib/python3.10/site-packages/jax/_src/numpy/array_methods.py:122: UserWarning: Explicitly requested dtype <class 'jax.numpy.int64'> requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.\n",
      "  return lax_numpy.astype(self, dtype, copy=copy, device=device)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "35 s ± 4.73 s 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(156, dtype=int32, weak_type=True), 'num_rejected_steps': Array(139, dtype=int32, weak_type=True), 'num_steps': Array(295, 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, particles=None,a=0.1,order=1)\n",
    "    \n",
    "    # Evolve the simulation forward\n",
    "    ode_fn = ODETerm(\n",
    "        make_diffrax_ode(mesh_shape, paint_absolute_pos=False))\n",
    "    solver = Tsit5()\n",
    "\n",
    "    stepsize_controller = PIDController(rtol=1e-6 , atol=1e-6)\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": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 346520 characters]",
      "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\" : jnp.log10(cic_paint_dx(lpt_displacements) + 1)}\n",
    "for i , field in enumerate(ode_solutions):\n",
    "    fields[f\"field_{i}\"] = jnp.log10(cic_paint_dx(field) + 1)\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": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/wassim/micromamba/envs/jax/lib/python3.10/site-packages/jax/_src/numpy/array_methods.py:122: UserWarning: Explicitly requested dtype <class 'jax.numpy.int64'> requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.\n",
      "  return lax_numpy.astype(self, dtype, copy=copy, device=device)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "19 ms ± 69.6 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n",
      "28.6 ms ± 82 μ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": "[encoded data, 182396 characters]",
      "text/plain": [
       "<Figure size 2000x500 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "lpt_fields = {\"First Order\" : jnp.log10(cic_paint_dx(lpt_displacements_1) + 1) , \"Second Order\" : jnp.log10(cic_paint_dx(lpt_displacements_2) + 1)}\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": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Solver Stats : {'max_steps': Array(4096, dtype=int32, weak_type=True), 'num_accepted_steps': Array(67, dtype=int32, weak_type=True), 'num_rejected_steps': Array(8, dtype=int32, weak_type=True), 'num_steps': Array(75, 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(mesh_shape))\n",
    "    solver = Tsit5()\n",
    "\n",
    "    stepsize_controller = PIDController(rtol=1e-6 , atol=1e-6)\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": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 346188 characters]",
      "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\" : jnp.log10(cic_paint(jnp.zeros(mesh_shape) ,lpt_particles) + 1)}\n",
    "for i , field in enumerate(ode_particles[1:]):\n",
    "    fields[f\"field_{i}\"] = jnp.log10(cic_paint(jnp.zeros(mesh_shape) , field)+1)\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": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 173264 characters]",
      "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(10)\n",
    "\n",
    "weighted = jnp.log10(cic_paint_dx(ode_solutions[0], weight=weights) + 1)\n",
    "unweighted = jnp.log10(cic_paint_dx(ode_solutions[0] , weight=1.0) + 1)\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": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 168568 characters]",
      "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(5)\n",
    "\n",
    "weighted = jnp.log10(cic_paint(jnp.zeros(mesh_shape),ode_particles[0], weight=weights) + 1)\n",
    "unweighted = jnp.log10(cic_paint(jnp.zeros(mesh_shape),ode_particles[0] , weight=1.0) + 1)\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
}
