{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# **Multi-GPU Particle Mesh Simulation with Advanced Solvers**\n",
    "\n",
    "<a href=\"https://colab.research.google.com/github/DifferentiableUniverseInitiative/JaxPM/blob/main/notebooks/04-MultiGPU_PM_Solvers.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": [
    "import os\n",
    "os.environ[\"EQX_ON_ERROR\"] = \"nan\"\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import jax_cosmo as jc\n",
    "\n",
    "from jaxpm.kernels import interpolate_power_spectrum\n",
    "from jaxpm.painting import cic_paint_dx\n",
    "from jaxpm.pm import linear_field, lpt, make_diffrax_ode\n",
    "from functools import partial\n",
    "from diffrax import ConstantStepSize, LeapfrogMidpoint,Dopri5 , PIDController , ODETerm, SaveAt, diffeqsolve"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Note**: This notebook requires 8 devices (GPU or TPU).\\\n",
    "> If you're running on CPU or don't have access to 8 devices,\\\n",
    "> you can simulate multiple devices by adding the following code at the start **BEFORE IMPORTING JAX**:\n",
    "\n",
    "```python\n",
    "import os\n",
    "os.environ[\"JAX_PLATFORM_NAME\"] = \"cpu\"\n",
    "os.environ[\"XLA_FLAGS\"] = \"--xla_force_host_platform_device_count=8\"\n",
    "```\n",
    "\n",
    "**Recommended only for debugging**. If used, you must probably lower the resolution of the mesh."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert jax.device_count() >= 8, \"This notebook requires a TPU or GPU runtime with 8 devices\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setting Up Device Mesh and Sharding for Multi-GPU Simulation\n",
    "\n",
    "This cell configures a **2x4 device mesh** across 8 devices and sets up named sharding to distribute data efficiently.\n",
    "\n",
    "- **Device Mesh**: `pdims = (2, 4)` arranges devices in a 2x4 grid. `create_device_mesh(pdims)` initializes this layout across available GPUs.\n",
    "- **Sharding with Mesh**: `Mesh(devices, axis_names=('x', 'y'))` assigns the mesh grid axes, which allows flexible mapping of array data across devices.\n",
    "- **PartitionSpec and NamedSharding**: `PartitionSpec` defines data partitioning across mesh axes `('x', 'y')`, and `NamedSharding(mesh, P('x', 'y'))` specifies this sharding scheme for arrays in the simulation.\n",
    "\n",
    "More info on Sharding in general in [Distributed arrays and automatic parallelization](https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from jax.experimental.mesh_utils import create_device_mesh\n",
    "from jax.experimental.multihost_utils import process_allgather\n",
    "from jax.sharding import Mesh, NamedSharding\n",
    "from jax.sharding import PartitionSpec as P\n",
    "\n",
    "all_gather = partial(process_allgather, tiled=False)\n",
    "\n",
    "pdims = (2, 4)\n",
    "devices = create_device_mesh(pdims)\n",
    "mesh = Mesh(devices, axis_names=('x', 'y'))\n",
    "sharding = NamedSharding(mesh, P('x', 'y'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@partial(jax.jit , static_argnums=(2,3,4,5))\n",
    "def run_simulation_with_fields(omega_c, sigma8,mesh_shape,box_size,halo_size , snapshots):\n",
    "    mesh_shape = (mesh_shape,) * 3\n",
    "    box_size = (box_size,) * 3\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(\n",
    "        jc.Planck15(Omega_c=omega_c, sigma8=sigma8), k)\n",
    "    pk_fn = lambda x: interpolate_power_spectrum(x, k, pk, sharding)\n",
    "\n",
    "    # Create initial conditions\n",
    "    initial_conditions = linear_field(mesh_shape,\n",
    "                                      box_size,\n",
    "                                      pk_fn,\n",
    "                                      seed=jax.random.PRNGKey(0),\n",
    "                                      sharding=sharding)\n",
    "\n",
    "\n",
    "    cosmo = jc.Planck15(Omega_c=omega_c, sigma8=sigma8)\n",
    "\n",
    "    # Initial displacement\n",
    "    dx, p, f = lpt(cosmo,\n",
    "                   initial_conditions,\n",
    "                   a=0.1,\n",
    "                   order=2,\n",
    "                   halo_size=halo_size,\n",
    "                   sharding=sharding)\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",
    "    ode_fields = [cic_paint_dx(sol[0], halo_size=halo_size, sharding=sharding) for sol in res.ys]\n",
    "    lpt_field = cic_paint_dx(dx , halo_size=halo_size, sharding=sharding)\n",
    "    return initial_conditions, lpt_field, ode_fields, res.stats"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Large-Scale Simulation Across Multiple Devices\n",
    "\n",
    "In this cell, we run a large simulation that would not be feasible on a single device. By distributing data across multiple devices, we achieve a higher resolution (`mesh_shape = 1024` and `box_size = 1000.`) with effective boundary handling using a `halo_size` of 128.\n",
    "\n",
    "We gather initial conditions and computed fields from all devices for visualization.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "45.6 s ± 11.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 1965172 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",
    "mesh_shape = 1024\n",
    "box_size = 1000.\n",
    "halo_size = 128\n",
    "snapshots = (0.5 , 1.0)\n",
    "\n",
    "initial_conditions , lpt_field , ode_fields , solver_stats = run_simulation_with_fields(0.25, 0.8 , mesh_shape, box_size , halo_size , snapshots)\n",
    "ode_fields[-1].block_until_ready()\n",
    "%timeit initial_conditions , lpt_field , ode_fields , solver_stats = run_simulation_with_fields(0.25, 0.8 , mesh_shape, box_size , halo_size , snapshots);ode_fields[-1].block_until_ready()\n",
    "\n",
    "initial_conditions_g = all_gather(initial_conditions)\n",
    "lpt_field_g = all_gather(lpt_field)\n",
    "ode_fields_g = [all_gather(p) for p in ode_fields]\n",
    "\n",
    "fields = {\"Initial Conditions\" : initial_conditions_g , \"LPT Field\" : lpt_field_g}\n",
    "for i , field in enumerate(ode_fields_g):\n",
    "    fields[f\"field_{i}\"] = field\n",
    "plot_fields_single_projection(fields,project_axis=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This simulation runs in **45 seconds** (less than half a second per step), which is impressive for a setup with **over one billion particles** (since \\( 1024^3 \\approx 1.07 \\) billion). This performance demonstrates the efficiency of distributing data and computation across multiple devices.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Comparing ODE Solvers: Leapfrog vs. Dopri5\n",
    "\n",
    "Next, we compare the **Leapfrog** solver with **Dopri5** (an adaptive Runge-Kutta method) to observe differences in accuracy and performance for particle evolution.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "5.04 s ± 9.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "mesh_shape = 512\n",
    "box_size = 512.\n",
    "halo_size = 64\n",
    "snapshots = (0.5 , 1.0)\n",
    "\n",
    "initial_conditions , lpt_field , ode_fields , solver_stats = run_simulation_with_fields(0.25, 0.8 , mesh_shape, box_size , halo_size , snapshots)\n",
    "ode_fields[-1].block_until_ready()\n",
    "%timeit initial_conditions , lpt_field , ode_fields , solver_stats = run_simulation_with_fields(0.25, 0.8 , mesh_shape, box_size , halo_size , snapshots);ode_fields[-1].block_until_ready()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4.44 s ± 8.12 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(12, dtype=int32, weak_type=True), 'num_rejected_steps': Array(0, dtype=int32, weak_type=True), 'num_steps': Array(12, dtype=int32, weak_type=True)}\n"
     ]
    }
   ],
   "source": [
    "mesh_shape = 512\n",
    "box_size = 512.\n",
    "halo_size = 64\n",
    "snapshots = (0.5, 1.0)\n",
    "\n",
    "@partial(jax.jit , static_argnums=(2,3,4,5))\n",
    "def run_simulation_with_dopri(omega_c, sigma8,mesh_shape,box_size,halo_size , snapshots):\n",
    "    mesh_shape = (mesh_shape,) * 3\n",
    "    box_size = (box_size,) * 3\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(\n",
    "        jc.Planck15(Omega_c=omega_c, sigma8=sigma8), k)\n",
    "    pk_fn = lambda x: interpolate_power_spectrum(x, k, pk, sharding)\n",
    "\n",
    "    # Create initial conditions\n",
    "    initial_conditions = linear_field(mesh_shape,\n",
    "                                      box_size,\n",
    "                                      pk_fn,\n",
    "                                      seed=jax.random.PRNGKey(0),\n",
    "                                      sharding=sharding)\n",
    "\n",
    "\n",
    "    cosmo = jc.Planck15(Omega_c=omega_c, sigma8=sigma8)\n",
    "\n",
    "    # Initial displacement\n",
    "    dx, p, f = lpt(cosmo,\n",
    "                   initial_conditions,\n",
    "                   a=0.1,\n",
    "                   order=2,\n",
    "                   halo_size=halo_size,\n",
    "                   sharding=sharding)\n",
    "\n",
    "    # Evolve the simulation forward\n",
    "    ode_fn = ODETerm(\n",
    "        make_diffrax_ode(cosmo, mesh_shape, paint_absolute_pos=False))\n",
    "    solver = Dopri5()\n",
    "\n",
    "    stepsize_controller = PIDController(rtol=1e-5,atol=1e-5)\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",
    "    ode_fields = [cic_paint_dx(sol[0], halo_size=halo_size, sharding=sharding) for sol in res.ys]\n",
    "    lpt_field = cic_paint_dx(dx , halo_size=halo_size, sharding=sharding)\n",
    "    return initial_conditions, lpt_field, ode_fields, res.stats\n",
    "\n",
    "initial_conditions , lpt_field , ode_fields , solver_stats = run_simulation_with_dopri(0.25, 0.8 , mesh_shape, box_size , halo_size , snapshots)\n",
    "ode_fields[-1].block_until_ready()\n",
    "%timeit initial_conditions , lpt_field , ode_fields , solver_stats = run_simulation_with_dopri(0.25, 0.8 , mesh_shape, box_size , halo_size , snapshots);ode_fields[-1].block_until_ready()\n",
    "\n",
    "print(f\"Solver Stats : {solver_stats}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 1902496 characters]",
      "text/plain": [
       "<Figure size 2000x500 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "initial_conditions_g = all_gather(initial_conditions)\n",
    "lpt_field_g = all_gather(lpt_field)\n",
    "ode_fields_g = [all_gather(p) for p in ode_fields]\n",
    "\n",
    "fields = {\"Initial Conditions\" : initial_conditions_g , \"LPT Field\" : lpt_field_g}\n",
    "for i , field in enumerate(ode_fields_g):\n",
    "    fields[f\"field_{i}\"] = field\n",
    "plot_fields_single_projection(fields,project_axis=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see how **easily we can switch solvers** here. Although Dopri5 offers adaptive stepping, it didn’t yield a significant performance boost over Leapfrog in this case.\n",
    "\n",
    "> **Note**: Dopri5 uses a **PIDController** for adaptive stepping, which might face challenges in distributed setups. In my experience, it works well without triggering all-gathers, but make sure to set:\n",
    "> ```python\n",
    "> os.environ[\"EQX_ON_ERROR\"] = \"nan\"\n",
    "> ```\n",
    "> before importing `diffrax` to handle any errors gracefully.\n",
    "\n",
    "However, **Dopri5 requires more memory** than Leapfrog, making a $1024^3$ mesh simulation unfeasible on eight A100 GPUs with 80GB memory each!!. For larger setups, we’ll need more compute resources—this is covered in the final notebook, **05-MultiHost_PM.ipynb**.\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "a100",
   "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
}
