{
 "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": null,
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "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 **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,
   "id": "9edd2246",
   "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": "markdown",
   "id": "74afa04a",
   "metadata": {},
   "source": [
    "### Multi-GPU Particle Mesh Simulation with Sharding\n",
    "\n",
    "This function is very similar to the single-GPU implementation, with the key difference being that `linear_field`, `lpt`, and `make_ode_fn` now take a `sharding` argument. This allows each stage of the simulation—initial conditions, displacements, and ODE evolution—to be distributed across the configured 2x4 device mesh, ensuring efficient parallel execution.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "281b4d3b",
   "metadata": {
    "id": "281b4d3b"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "778 ms ± 1.62 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 = 256\n",
    "box_size = 256.\n",
    "halo_size = 32\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(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_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",
    "%timeit 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": 5,
   "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\">  GPU 0  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #d6616b\">  GPU 1  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8ca252\">  GPU 2  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">  GPU 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\">  GPU 4  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #6b6ecf\">  GPU 5  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #a55194\">  GPU 6  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8c6d31\">  GPU 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;121mGPU 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;107mGPU 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;82mGPU 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;214mGPU 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;148mGPU 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;207mGPU 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;148mGPU 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;49mGPU 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": 6,
   "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": null,
   "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, 1555544 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\" : 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,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": [],
   "source": [
    "mesh_shape = 256\n",
    "box_size = 256.\n",
    "halo_size = 0\n",
    "snapshots = (0.5, 1.0)\n",
    "\n",
    "initial_conditions , lpt_displacements , ode_solutions , 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 = all_gather(lpt_displacements)\n",
    "ode_solutions_g = [all_gather(p) for p in ode_solutions]\n",
    "\n",
    "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,project_axis=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88f22b07",
   "metadata": {},
   "source": [
    "We can see very faint lines in the visualization, but the boundary conditions were handled successfully even without halo exchange. \n",
    "\n",
    "**Important**: Here, we painted the **all-gathered field** for visualization, which in a typical forward model should be avoided to maintain data distribution.\n",
    "\n",
    "Next, we’ll modify the simulation function to return density fields directly, allowing us to observe the impact of the `halo_size` on boundary accuracy.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c647b13",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "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 = 256\n",
    "box_size = 256.\n",
    "halo_size = 0\n",
    "snapshots = (0.5, 1.0)\n",
    "\n",
    "@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\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",
    "print(f\"Solver Stats : {solver_stats}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "id": "7b552794",
   "metadata": {},
   "outputs": [],
   "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]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0acb5253",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 1539972 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_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",
   "id": "ac9c8818",
   "metadata": {},
   "source": [
    "Now we can see that there are very apparent lines between the subdomains of the simulation. These lines highlight the **artifacts** that arise when running the simulation without a halo exchange, as boundary conditions are not accurately handled across device edges.\n"
   ]
  },
  {
   "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": 81,
   "id": "9395deea",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 2688196 characters]",
      "text/plain": [
       "<Figure size 2000x1000 with 8 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "mesh_shape = 256\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_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",
   "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": 82,
   "id": "57655904",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 3239152 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_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",
   "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": null,
   "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\">  GPU 0  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #d6616b\">  GPU 1  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8ca252\">  GPU 2  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">  GPU 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\">  GPU 4  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #6b6ecf\">  GPU 5  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #a55194\">  GPU 6  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8c6d31\">  GPU 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;121mGPU 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;107mGPU 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;82mGPU 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;214mGPU 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;148mGPU 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;207mGPU 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;148mGPU 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;49mGPU 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, 618444 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=2.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
}
