{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7fb27b941602401d91542211134fc71a",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "# **Multi-GPU Particle Mesh Simulation with Halo Exchange**\n",
    "\n",
    "<a href=\"https://colab.research.google.com/github/DifferentiableUniverseInitiative/JaxPM/blob/main/notebooks/02-MultiGPU_PM.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,
   "id": "9Jy5BL1XiK1s",
   "metadata": {
    "id": "9Jy5BL1XiK1s"
   },
   "outputs": [],
   "source": [
    "!pip install --quiet git+https://github.com/DifferentiableUniverseInitiative/JaxPM.git\n",
    "!pip install diffrax"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e84120d2",
   "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,
   "id": "c5f42bbe",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "c5f42bbe",
    "outputId": "a7841b28-5f20-4856-bd1d-f8a3572095b5"
   },
   "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",
    "from jax.debug import visualize_array_sharding\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, ODETerm, SaveAt, diffeqsolve, Tsit5, PIDController"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "38df34e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert jax.device_count() >= 8, \"This notebook requires a TPU or GPU runtime with 8 devices\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0fe6876",
   "metadata": {},
   "source": [
    "### Setting Up Device Mesh and Sharding for Multi-GPU Simulation\n",
    "\n",
    "This cell configures a **2×4 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 2×4 grid. This logical mesh layout defines how devices are organized for parallelism.\n",
    "  For more information on grid mesh layouts, see the [jaxDecomp GitHub repository](https://github.com/DifferentiableUniverseInitiative/jaxDecomp).\n",
    "* **Sharding with Mesh**: `Mesh(devices, axis_names=('x', 'y'))` assigns names to the mesh axes, enabling flexible and explicit data mapping across devices.\n",
    "* **PartitionSpec and NamedSharding**: `PartitionSpec('x', 'y')` specifies how to partition data along the mesh axes. `NamedSharding(mesh, P('x', 'y'))` applies this scheme, ensuring that arrays are correctly distributed across the 2D device mesh.\n",
    "\n",
    "More details on sharding and distributed computation are available in the JAX documentation: [Distributed arrays and automatic parallelization](https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "9edd2246",
   "metadata": {},
   "outputs": [],
   "source": [
    "from jax.experimental.multihost_utils import process_allgather\n",
    "from jax.sharding import PartitionSpec as P, NamedSharding\n",
    "\n",
    "all_gather = partial(process_allgather, tiled=True)\n",
    "\n",
    "pdims = (2, 4)\n",
    "mesh = jax.make_mesh(pdims, axis_names=('x', 'y'))\n",
    "sharding = NamedSharding(mesh, P('x', 'y'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74afa04a",
   "metadata": {},
   "source": [
    "### Multi-GPU Particle Mesh Simulation with Sharding\n",
    "\n",
    "This function closely follows the single-GPU simulation logic but extends it to support distributed computation across multiple devices. The key change is the use of the `sharding` argument in `linear_field`, `lpt`, and `make_diffrax_ode`.\n",
    "\n",
    "By passing `sharding`, each stage of the simulation—initial condition generation, LPT displacement calculation, and time integration via the ODE solver—is parallelized across the 2×4 device mesh. This enables efficient scaling and memory management in large-scale simulations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "281b4d3b",
   "metadata": {
    "id": "281b4d3b"
   },
   "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": [
      "CPU times: user 5min 42s, sys: 3.77 s, total: 5min 46s\n",
      "Wall time: 23.6 s\n",
      "Solver Stats : {'max_steps': Array(4096, dtype=int32, weak_type=True), 'num_accepted_steps': Array(8, dtype=int32, weak_type=True), 'num_rejected_steps': Array(1, dtype=int32, weak_type=True), 'num_steps': Array(9, dtype=int32, weak_type=True)}\n"
     ]
    }
   ],
   "source": [
    "mesh_shape = 128\n",
    "box_size = 256.\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(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(mesh_shape, paint_absolute_pos=False , sharding=sharding , halo_size=halo_size),)\n",
    "    solver = Tsit5()\n",
    "\n",
    "    stepsize_controller = PIDController(rtol=1e-3 , atol=1e-3)\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_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 , mesh_shape, box_size , halo_size , snapshots)\n",
    "ode_solutions[-1].block_until_ready()\n",
    "%time initial_conditions , lpt_displacements , ode_solutions , solver_stats = run_simulation(0.25, 0.8, mesh_shape, box_size , halo_size , snapshots);ode_solutions[-1].block_until_ready()\n",
    "print(f\"Solver Stats : {solver_stats}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "481bb668",
   "metadata": {},
   "source": [
    "\n",
    "All fields and particle grids remain distributed at all times (as seen below). `jaxPM` ensures they are **never gathered on a single device**. In a forward model scenario, it’s the **user's responsibility to maintain distributed data** to avoid memory bottlenecks.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "ca188e9a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #d6616b\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8ca252\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #d6616b\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8ca252\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">  CPU 0  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #d6616b\">  CPU 1  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8ca252\">  CPU 2  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">  CPU 3  </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #d6616b\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8ca252\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #d6616b\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8ca252\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #d6616b\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8ca252\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
       "<span style=\"color: #000000; text-decoration-color: #000000; background-color: #e7cb94\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #6b6ecf\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #a55194\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8c6d31\">         </span>\n",
       "<span style=\"color: #000000; text-decoration-color: #000000; background-color: #e7cb94\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #6b6ecf\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #a55194\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8c6d31\">         </span>\n",
       "<span style=\"color: #000000; text-decoration-color: #000000; background-color: #e7cb94\">  CPU 4  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #6b6ecf\">  CPU 5  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #a55194\">  CPU 6  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8c6d31\">  CPU 7  </span>\n",
       "<span style=\"color: #000000; text-decoration-color: #000000; background-color: #e7cb94\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #6b6ecf\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #a55194\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8c6d31\">         </span>\n",
       "<span style=\"color: #000000; text-decoration-color: #000000; background-color: #e7cb94\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #6b6ecf\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #a55194\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8c6d31\">         </span>\n",
       "<span style=\"color: #000000; text-decoration-color: #000000; background-color: #e7cb94\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #6b6ecf\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #a55194\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8c6d31\">         </span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107m         \u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107m         \u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mCPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107m  \u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107mCPU 1\u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107m  \u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82m  \u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82mCPU 2\u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mCPU 3\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107m         \u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107m         \u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107m         \u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
       "\u001b[38;2;0;0;0;48;2;231;203;148m         \u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207m         \u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148m         \u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49m         \u001b[0m\n",
       "\u001b[38;2;0;0;0;48;2;231;203;148m         \u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207m         \u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148m         \u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49m         \u001b[0m\n",
       "\u001b[38;2;0;0;0;48;2;231;203;148m  \u001b[0m\u001b[38;2;0;0;0;48;2;231;203;148mCPU 4\u001b[0m\u001b[38;2;0;0;0;48;2;231;203;148m  \u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207m  \u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207mCPU 5\u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207m  \u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148m  \u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148mCPU 6\u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148m  \u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49m  \u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49mCPU 7\u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49m  \u001b[0m\n",
       "\u001b[38;2;0;0;0;48;2;231;203;148m         \u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207m         \u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148m         \u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49m         \u001b[0m\n",
       "\u001b[38;2;0;0;0;48;2;231;203;148m         \u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207m         \u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148m         \u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49m         \u001b[0m\n",
       "\u001b[38;2;0;0;0;48;2;231;203;148m         \u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207m         \u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148m         \u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49m         \u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "visualize_array_sharding(ode_solutions[-1][:,:,0,0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b512d8a",
   "metadata": {},
   "source": [
    "> ⚠️ **Warning**: One caveat is that particle arrays usually have a shape of `(NPart, 3)`,\\\n",
    "> where `NPart = Nx * Nx * Nx`. However, this shape is **not shardable** in a distributed setup.\\\n",
    "> Instead, particle arrays will always have a shape of `(Nx, Ny, Nz, 3)` to ensure they remain distributed across devices.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "042cc55c",
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_conditions_g = all_gather(initial_conditions)\n",
    "lpt_displacements_g = all_gather(lpt_displacements)\n",
    "ode_solutions_g = [all_gather(p) for p in ode_solutions]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "4e012ce8",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 323
    },
    "id": "4e012ce8",
    "outputId": "75390318-8072-481f-ffb9-ec09cd71cb1d"
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 352088 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.log(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,project_axis=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0279167c",
   "metadata": {},
   "source": [
    "# **Halo Exchange**\n",
    "\n",
    "Let's start by running a simulation **without halo exchange**. Here, we set `halo_size = 0`, which means no overlapping regions between device boundaries. This configuration helps us observe the limitations of simulations without halo regions, especially for calculating forces near boundaries in multi-GPU setups.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02ba5519",
   "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"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 350620 characters]",
      "text/plain": [
       "<Figure size 2000x500 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "mesh_shape = 128\n",
    "box_size = 256.\n",
    "halo_size = 0\n",
    "snapshots = (0.5, 1.0)\n",
    "\n",
    "initial_conditions , lpt_displacements_bc , ode_solutions_bc , solver_stats = run_simulation(0.25, 0.8 , mesh_shape, box_size , halo_size , snapshots)\n",
    "\n",
    "initial_conditions_g = all_gather(initial_conditions)\n",
    "lpt_displacements_g_bc = all_gather(lpt_displacements_bc)\n",
    "ode_solutions_g_bc = [all_gather(p) for p in ode_solutions_bc]\n",
    "\n",
    "from jaxpm.plotting import plot_fields_single_projection\n",
    "\n",
    "fields = {\"Initial Conditions\" : initial_conditions_g , \"LPT Field\" : jnp.log(cic_paint_dx(lpt_displacements_g_bc) + 1)}\n",
    "for i , field in enumerate(ode_solutions_g_bc):\n",
    "    fields[f\"field_{i}\"] = jnp.log10(cic_paint_dx(field) + 1)\n",
    "plot_fields_single_projection(fields,project_axis=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88f22b07",
   "metadata": {},
   "source": [
    "We can clearly observe artifacts in the visualization—most notably, horizontal and vertical discontinuities appearing in the evolved density fields (`field_0`, `field_1`). These are a direct consequence of not using a halo exchange (`halo_size = 0`), which causes poor force computation across subdomain boundaries in a multi-device simulation.\n",
    "\n",
    "> 🔍 These artifacts highlight where the simulation fails to maintain physical continuity between neighboring partitions, especially as structures evolve and particles interact across device boundaries.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c707514",
   "metadata": {},
   "source": [
    "### Returning Density Fields for Forward Modeling\n",
    "\n",
    "In typical forward models, the goal is to produce **density fields**, not raw particle displacements. This function (`run_simulation_with_fields`) returns pre-painted density fields directly from each device, keeping data distributed and ready for downstream use (e.g. likelihood evaluation, comparison with observations).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "8c647b13",
   "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(mesh_shape, paint_absolute_pos=False , sharding=sharding , halo_size=halo_size))\n",
    "    solver = Tsit5()\n",
    "\n",
    "    stepsize_controller = PIDController(rtol=1e-3 , atol=1e-3)\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",
   "id": "039b197d",
   "metadata": {},
   "source": [
    "### Choosing the Right Halo Size\n",
    "\n",
    "In some cases, the halo size can be too small, leading to visible artifacts in the snapshots. Here, we see that boundaries are handled well in the first and second snapshots, but the lines become more pronounced with each successive step. This indicates that a larger halo size may be needed to fully capture interactions across device boundaries over time.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "9395deea",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 672208 characters]",
      "text/plain": [
       "<Figure size 2000x1000 with 8 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from jaxpm.plotting import plot_fields_single_projection\n",
    "\n",
    "mesh_shape = 128\n",
    "box_size = 128.\n",
    "halo_size = 4\n",
    "snapshots = (0.3 ,0.4, 0.5 , 0.6, 0.8, 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",
    "\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 , \"LPT Field\" : jnp.log(lpt_field + 1)}\n",
    "for i , field in enumerate(ode_fields):\n",
    "    fields[f\"field_{i}\"] = jnp.log10(field + 1)\n",
    "plot_fields_single_projection(fields,project_axis=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68663d3f",
   "metadata": {},
   "source": [
    "In other cases, if the **box size is too large**, particles must cover greater distances, resulting in smaller final displacements. This reduces the impact of insufficient halo size on boundary artifacts. \n",
    "\n",
    "### Explanation\n",
    "\n",
    "- **Large Box Sizes**: In larger simulation boxes, particles tend to have smaller relative displacements (or slower speeds). This reduces the frequency of interactions with particles in neighboring subdomains, making boundary artifacts less pronounced, even if the halo size is smaller.\n",
    "\n",
    "- **Smaller Box Sizes**: In smaller boxes, particles cover a greater relative distance, leading to more frequent interactions with boundary particles. Here, the halo size must be carefully chosen to capture these interactions accurately, reducing visible artifacts in the visualization.\n",
    "\n",
    "\n",
    "In this scenario, we can see that the insufficient halo size does not lead to severe artifacts, as particles are less affected by neighboring boundaries.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "57655904",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/wassim/micromamba/envs/jax/lib/python3.10/site-packages/jax/_src/interpreters/mlir.py:1132: UserWarning: A large amount of constants were captured during lowering (2.42GB total). If this is intentional, disable this warning by setting JAX_CAPTURED_CONSTANTS_WARN_BYTES=-1. To obtain a report of where these constants were encountered, set JAX_CAPTURED_CONSTANTS_REPORT_FRAMES=-1.\n",
      "  warnings.warn(message)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 3589804 characters]",
      "text/plain": [
       "<Figure size 2000x1000 with 8 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "mesh_shape = 256\n",
    "box_size = 1000.\n",
    "halo_size = 4\n",
    "snapshots = (0.3 ,0.4, 0.5 , 0.6, 0.8, 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",
    "\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 , \"LPT Field\" : jnp.log(lpt_field + 1)}\n",
    "for i , field in enumerate(ode_fields):\n",
    "    fields[f\"field_{i}\"] = jnp.log10(field + 1)\n",
    "plot_fields_single_projection(fields,project_axis=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dfb706ec",
   "metadata": {},
   "source": [
    "### General Guideline\n",
    "\n",
    "<span style=\"font-size: 1.1em;\">Start with a halo size that is **one-eighth of the box size**. Gradually reduce it until you begin to notice lines in the visualization, indicating an insufficient halo size.</span>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7586885e",
   "metadata": {},
   "source": [
    "# Applying Weights in a Distributed Setup\n",
    "\n",
    "We can apply weights just like before. In general, we want to apply weights on a **distributed particle grid**.\n",
    "\n",
    "> **Note**: When using weights in a distributed setting, ensure that the weights have the **same sharding** as the particle grid. If the sharding is not identical, JAX may perform an all-gather or other collective operations that could significantly impact performance.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "59cfba84",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #d6616b\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8ca252\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #d6616b\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8ca252\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">  CPU 0  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #d6616b\">  CPU 1  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8ca252\">  CPU 2  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">  CPU 3  </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #d6616b\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8ca252\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #d6616b\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8ca252\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #d6616b\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8ca252\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
       "<span style=\"color: #000000; text-decoration-color: #000000; background-color: #e7cb94\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #6b6ecf\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #a55194\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8c6d31\">         </span>\n",
       "<span style=\"color: #000000; text-decoration-color: #000000; background-color: #e7cb94\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #6b6ecf\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #a55194\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8c6d31\">         </span>\n",
       "<span style=\"color: #000000; text-decoration-color: #000000; background-color: #e7cb94\">  CPU 4  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #6b6ecf\">  CPU 5  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #a55194\">  CPU 6  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8c6d31\">  CPU 7  </span>\n",
       "<span style=\"color: #000000; text-decoration-color: #000000; background-color: #e7cb94\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #6b6ecf\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #a55194\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8c6d31\">         </span>\n",
       "<span style=\"color: #000000; text-decoration-color: #000000; background-color: #e7cb94\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #6b6ecf\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #a55194\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8c6d31\">         </span>\n",
       "<span style=\"color: #000000; text-decoration-color: #000000; background-color: #e7cb94\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #6b6ecf\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #a55194\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8c6d31\">         </span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107m         \u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107m         \u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mCPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107m  \u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107mCPU 1\u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107m  \u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82m  \u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82mCPU 2\u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mCPU 3\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107m         \u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107m         \u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107m         \u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
       "\u001b[38;2;0;0;0;48;2;231;203;148m         \u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207m         \u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148m         \u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49m         \u001b[0m\n",
       "\u001b[38;2;0;0;0;48;2;231;203;148m         \u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207m         \u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148m         \u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49m         \u001b[0m\n",
       "\u001b[38;2;0;0;0;48;2;231;203;148m  \u001b[0m\u001b[38;2;0;0;0;48;2;231;203;148mCPU 4\u001b[0m\u001b[38;2;0;0;0;48;2;231;203;148m  \u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207m  \u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207mCPU 5\u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207m  \u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148m  \u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148mCPU 6\u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148m  \u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49m  \u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49mCPU 7\u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49m  \u001b[0m\n",
       "\u001b[38;2;0;0;0;48;2;231;203;148m         \u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207m         \u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148m         \u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49m         \u001b[0m\n",
       "\u001b[38;2;0;0;0;48;2;231;203;148m         \u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207m         \u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148m         \u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49m         \u001b[0m\n",
       "\u001b[38;2;0;0;0;48;2;231;203;148m         \u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207m         \u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148m         \u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49m         \u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 140284 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",
    "field = ode_solutions[0]\n",
    "\n",
    "center = slice(field.shape[0] // 4, 3 * field.shape[0] // 4 )\n",
    "center3d = (slice(None) , center,center)  # All of X, Y=0, Z=0\n",
    "weights = jnp.ones_like(field[...,0])\n",
    "# Update weights for the down-left pencil by multiplying by 100\n",
    "weights = weights.at[center3d].multiply(3)\n",
    "visualize_array_sharding(weights[:,:,0])\n",
    "\n",
    "weighted = cic_paint_dx(field, weight=weights)\n",
    "unweighted = cic_paint_dx(field, weight=1.0)\n",
    "\n",
    "plot_fields_single_projection({\"Weighted\" : weighted , \"Unweighted\" : unweighted} , project_axis=0)"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "include_colab_link": true,
   "name": "Introduction.ipynb",
   "provenance": []
  },
  "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": 5
}
