JaxPM/notebooks/01-Introduction.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7fb27b941602401d91542211134fc71a",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "# **Single-GPU Particle Mesh Simulation with JAXPM**\n",
    "\n",
    "In this notebook, we'll run a simple Particle Mesh (PM) simulation on a single GPU using `JAXPM`. This example provides a hands-on introduction to simulating the evolution of matter in the universe through basic PM techniques, allowing you to explore how cosmological structures form over time.\n",
    "\n",
    "<a href=\"https://colab.research.google.com/github/DifferentiableUniverseInitiative/JaxPM/blob/main/notebooks/01-Introduction.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9Jy5BL1XiK1s",
   "metadata": {
    "id": "9Jy5BL1XiK1s"
   },
   "source": [
    "!pip install --quiet git+https://github.com/DifferentiableUniverseInitiative/JaxPM.git"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "af3a08c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "jax.config.update(\"jax_enable_x64\", True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c5f42bbe",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "c5f42bbe",
    "outputId": "a7841b28-5f20-4856-bd1d-f8a3572095b5"
   },
   "outputs": [],
   "source": [
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import jax_cosmo as jc\n",
    "\n",
    "from jax.experimental.ode import odeint\n",
    "\n",
    "from jaxpm.painting import cic_paint , cic_paint_dx\n",
    "from jaxpm.pm import linear_field, lpt, make_ode_fn\n",
    "from jaxpm.distributed import uniform_particles"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "adc38ef4",
   "metadata": {},
   "source": [
    "### **Particle Mesh Simulation Setup**\n",
    "\n",
    "In this example, we initialize particles with uniform positions across the grid. This setup implicitly means that the Cloud-in-Cell (CIC) painting scheme will map absolute particle positions onto the grid.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "281b4d3b",
   "metadata": {
    "id": "281b4d3b"
   },
   "outputs": [],
   "source": [
    "mesh_shape = [128, 128, 128]\n",
    "box_size = [128., 128., 128.]\n",
    "snapshots = jnp.array([0.1, 0.5, 1.0])\n",
    "\n",
    "@jax.jit\n",
    "def run_simulation(omega_c, sigma8):\n",
    "    # Create a small function to generate the matter power spectrum\n",
    "    k = jnp.logspace(-4, 1, 128)\n",
    "    pk = jc.power.linear_matter_power(jc.Planck15(Omega_c=omega_c, sigma8=sigma8), k)\n",
    "    pk_fn = lambda x: jnp.interp(x.reshape([-1]), k, pk).reshape(x.shape)\n",
    "\n",
    "    # Create initial conditions\n",
    "    initial_conditions = linear_field(mesh_shape, box_size, pk_fn, seed=jax.random.PRNGKey(0))\n",
    "\n",
    "    particles = uniform_particles(mesh_shape)\n",
    "    # Create particles\n",
    "    cosmo = jc.Planck15(Omega_c=omega_c, sigma8=sigma8)\n",
    "    \n",
    "    # Initial displacement\n",
    "    dx, p, f = lpt(cosmo, initial_conditions, particles, a=0.1)\n",
    "    \n",
    "    # Evolve the simulation forward\n",
    "    res = odeint(make_ode_fn(mesh_shape), [particles + dx, p], snapshots, cosmo, rtol=1e-8, atol=1e-8)\n",
    "    \n",
    "    # Return the simulation volume at requested \n",
    "\n",
    "    return initial_conditions ,  particles + dx , res[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "826be667",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "826be667",
    "outputId": "dc43b5c4-a004-41bf-f2c8-128c17cf4de1"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/wassim/micromamba/envs/jax/lib/python3.10/site-packages/jax/_src/numpy/array_methods.py:118: 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",
      "/home/wassim/micromamba/envs/jax/lib/python3.10/site-packages/jax/_src/numpy/array_methods.py:118: 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",
      "/home/wassim/micromamba/envs/jax/lib/python3.10/site-packages/jax/_src/numpy/array_methods.py:118: 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": [
      "pm_forces\n",
      "pm_forces\n"
     ]
    }
   ],
   "source": [
    "initial_conditions , lpt_particles , ode_particles = run_simulation(0.25, 0.8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "f9c05d34",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAB8UAAAH/CAYAAADOlQwMAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjkuMywgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/GU6VOAAAACXBIWXMAAA9hAAAPYQGoP6dpAAEAAElEQVR4nOydeXxcZb3/v+fMviSZrE3StEm6r7TQUnbKJtAiiAhcF2RRRBEu6lWv21VBQa+K4lWviHpVBPWngCwqyiK7bC2UpbR0TbckzT5ZZp855/dHzXw/3yeZ0GKxFL9vXn3xZM4z5zzr93nOmZnPx3Jd1yVFURRFURRFURRFURRFURRFURRFURRFeQtiH+gCKIqiKIqiKIqiKIqiKIqiKIqiKIqiKMobhX4oriiKoiiKoiiKoiiKoiiKoiiKoiiKorxl0Q/FFUVRFEVRFEVRFEVRFEVRFEVRFEVRlLcs+qG4oiiKoiiKoiiKoiiKoiiKoiiKoiiK8pZFPxRXFEVRFEVRFEVRFEVRFEVRFEVRFEVR3rLoh+KKoiiKoiiKoiiKoiiKoiiKoiiKoijKWxb9UFxRFEVRFEVRFEVRFEVRFEVRFEVRFEV5y6IfiiuKoiiKoiiKoiiKoiiKoiiKoiiKoihvWfRDcUVRFEVRFEVRFEVRFEVRFEVRFEVRFOUti34orhyUWJZFV1999V7lbWlpoYsvvnifr7Ft2zayLIt+8Ytf7PN73ww88sgjZFkWPfLII8XXLr74Ymppadmr91999dVkWdYbU7iDiAM5DvalvxRFUZQ3H+OtxXvLL37xC7Isi7Zt2/aaeV/vXueN5ECVaV/aTVEURXlzsmrVKjr66KMpEomQZVn0wgsv/EP3p3t7X3WwPwN4Lf6Rfck/ygknnEAnnHDCP/26iqIoyj8PXb/fGPTZtKLsX/RDceWAMPrAcvXq1fvlfE8++SRdffXVFI/H98v5Xg9dXV30qU99iubMmUPhcJgikQgtWbKErr322gNarolIJpN09dVXH5CbYmUPHR0ddPXVV9MLL7xwoIsyhv/7v/+juXPnUjAYpJkzZ9L3v//9A10kRVGUvdpDjN40jv7zeDw0depUeuc731mMtxdffLHIU+rfRB/sjt7gj/fvRz/60X6uuYJ87Wtfo7vuuutAF2MMTz75JB177LEUDoepvr6errrqKhoZGTnQxVIURTloyOVydN5551F/fz/dcMMNdMstt1Bzc/OBLtZrovdOROvWraOrr776TffFNMdx6Jvf/Ca1trZSMBikQw45hH7zm98c6GIpiqK8pTgY1+8bb7yRzjvvPJo6depr3vu/ldFn08q/Gt4DXQBFeT2kUinyenn4Pvnkk3TNNdfQxRdfTLFYTOTdsGED2fYb+/2PVatW0cqVK2lkZIQuuOACWrJkCRERrV69mv77v/+bHnvsMbr//vvf0DLsDT/5yU/IcZzi38lkkq655hoiojHf2v6v//ov+uxnP/vPLN6bkubmZkqlUuTz+d6Q83d0dNA111xDLS0ttHjxYnHM7K9/JjfddBN95CMfoXe96130H//xH/T444/TVVddRclkkj7zmc8ckDIpiqLsK+95z3to5cqVVCgUaP369XTjjTfSn//8Z3r66afpwx/+MJ1yyinFvG1tbfSlL32JLrvsMjruuOOKr0+fPv01r3PjjTdSNBoVrx1xxBE0ffp0SqVS5Pf791+lDhLe6P3X1772NTr33HPp7LPPFq+///3vp3e/+90UCATesGuX4oUXXqCTTz6Z5s6dS9/5zndo165ddP3119OmTZvoz3/+8z+9PIqiKAcjW7Zsoe3bt9NPfvITuvTSS4uvv5nvTw+We6fjjz/+Dd2XrFu3jq655ho64YQTxvyq7EA+D/nCF75A//3f/00f+tCH6PDDD6e7776b3vve95JlWfTud7/7gJVLURTlrcTBuH5/4xvfoOHhYVq2bBl1dnYe6OKURJ9Nv7n3V8rBh34orhyUBIPBvc77Rj8Ujcfj9M53vpM8Hg+tWbOG5syZI45fd9119JOf/OQNLcPesi+Lp9frFV88eKvgui6l02kKhUJ7ld+yrH0ab/uTN2qz81qkUin6whe+QGeccQbdfvvtRET0oQ99iBzHoa9+9at02WWXUWVl5QEpm6Ioyr5w2GGH0QUXXFD8+5hjjqGzzjqLbrzxRrrpppvoqKOOKh5bvXo1felLX6KjjjpKvGdvOPfcc6mmpmbcYwdqDdnfJBIJikQie53/QHwoTUTk8XjI4/EckGt//vOfp8rKSnrkkUeovLyciPbIyH/oQx+i+++/n0499dQDUi5FUZSDie7ubiKiMV92f7Penx7IeyfHcSibze71XsO27QO2LzlQXxBsb2+nb3/723TFFVfQD37wAyIiuvTSS2n58uX06U9/ms4777wDtm9QFEV5K3Gwrd9ERI8++mjxV+Lml9zfSPTZ9Gujz6aVNxKVT1feNFx88cUUjUapvb2dzj77bIpGo1RbW0uf+tSnqFAoiLzoKX711VfTpz/9aSIiam1tLcqWjkp2mZ6W/f399KlPfYoWLlxI0WiUysvLacWKFfTiiy++rnLfdNNN1N7eTt/5znfGfCBORDRp0iT6r//6L/HaD3/4Q5o/fz4FAgFqbGykK664YozE+gknnEALFiygdevW0YknnkjhcJgmT55M3/zmN8dcY9euXXT22WdTJBKhuro6+sQnPkGZTGZMPvQB2bZtG9XW1hIR0TXXXFNsN2xX0/Mln8/TV7/6VZo+fToFAgFqaWmhz3/+82Ou1dLSQm9/+9vpiSeeoGXLllEwGKRp06bRL3/5S5Evl8vRNddcQzNnzqRgMEjV1dV07LHH0gMPPDC2oYFR6dzHHnuMPvzhD1N1dTWVl5fThRdeSAMDA+OW5b777qOlS5dSKBSim266iYiItm7dSueddx5VVVVROBymI488kv70pz+J95fybXn11Vfp3HPPpaqqKgoGg7R06VK65557xpQ1Ho/TJz7xCWppaaFAIEBNTU104YUXUm9vLz3yyCN0+OGHExHRJZdcUuyD0WuN59uSSCTok5/8JE2ZMoUCgQDNnj2brr/+enJdV+SzLIuuvPJKuuuuu2jBggUUCARo/vz59Je//GXCtiUievjhh6mvr48++tGPitevuOIKSiQSY9pIURTlYOGkk04ioj2/Cv9nUMq785lnnqHTTz+dKioqKBwO0/Lly+lvf/vba57PdV269tprqampicLhMJ144on0yiuv7FVZRtez66+/nm644QZqbm6mUChEy5cvp7Vr14q8o3uyLVu20MqVK6msrIze9773EdHer0PjeYrH43H6+Mc/XnzvjBkz6Bvf+MaYb547jkP/8z//QwsXLqRgMEi1tbV0+umnFyXzLcuiRCJBN9988xiZ+1Ke4vt772UyNDREDzzwAF1wwQXFD8SJiC688EKKRqP0u9/97jXPoSiK8q/OxRdfTMuXLyciovPOO48syyoqmpXyJL311ltpyZIlFAqFqKqqit797nfTzp07X/Na8XicLr74YqqoqKBYLEYXXXTR67I9+0fvnUbr9eqrr9L5559P5eXlVF1dTR/72MconU6LvKP3eL/61a+Ka9ro/d2aNWtoxYoVVF5eTtFolE4++WR6+umnxfv/0X1Je3s7ffCDH6TGxkYKBALU2tpKl19+OWWzWfrFL35B5513HhERnXjiicX1efRa43mKd3d30wc/+EGaNGkSBYNBWrRoEd18880iD+5ffvzjHxefRRx++OG0atWqCduWiOjuu++mXC4n+seyLLr88stp165d9NRTT73mORRFUZSJORjXb6I9v8B+vX7n+mxan00rBy9vzq/pKP+yFAoFOu200+iII46g66+/nh588EH69re/TdOnT6fLL7983Pecc845tHHjRvrNb35DN9xwQ/GXWqMf+Jps3bqV7rrrLjrvvPOotbWVurq66KabbqLly5fTunXrqLGxcZ/KfM8991AoFKJzzz13r/JfffXVdM0119App5xCl19+OW3YsIFuvPFGWrVqFf3tb38T38AaGBig008/nc4
      "text/plain": [
       "<Figure size 2000x500 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from jaxpm.plotting import plot_fields_single_projection\n",
    "\n",
    "fields = {\"Initial Conditions\" : initial_conditions , \"LPT Field\" : cic_paint(jnp.zeros(mesh_shape) ,lpt_particles)}\n",
    "for i , field in enumerate(ode_particles[1:]):\n",
    "    fields[f\"field_{i}\"] = cic_paint(jnp.zeros(mesh_shape) , field)\n",
    "plot_fields_single_projection(fields)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd1e9196",
   "metadata": {},
   "source": [
    "# **Particle Mesh Simulation Setup - Relative Position Painting**\n",
    "\n",
    "In the second example, we leave the initial particle positions as `None`, which applies a relative position (displacement-only) CIC painting scheme. This approach assumes uniform particle positions by default, saving memory and improving efficiency, though with the trade-off of assuming uniformity.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b71824ed",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/wassim/micromamba/envs/jax/lib/python3.10/site-packages/jax/_src/numpy/array_methods.py:118: 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",
      "/home/wassim/micromamba/envs/jax/lib/python3.10/site-packages/jax/_src/numpy/array_methods.py:118: 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",
      "/home/wassim/micromamba/envs/jax/lib/python3.10/site-packages/jax/_src/numpy/array_methods.py:118: 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": [
      "pm_forces\n",
      "pm_forces\n"
     ]
    }
   ],
   "source": [
    "mesh_shape = [128, 128, 128]\n",
    "box_size = [128., 128., 128.]\n",
    "snapshots = jnp.array([0.1, 0.5, 1.0])\n",
    "\n",
    "@jax.jit\n",
    "def run_simulation(omega_c, sigma8):\n",
    "    # Create a small function to generate the matter power spectrum\n",
    "    k = jnp.logspace(-4, 1, 128)\n",
    "    pk = jc.power.linear_matter_power(jc.Planck15(Omega_c=omega_c, sigma8=sigma8), k)\n",
    "    pk_fn = lambda x: jnp.interp(x.reshape([-1]), k, pk).reshape(x.shape)\n",
    "\n",
    "    # Create initial conditions\n",
    "    initial_conditions = linear_field(mesh_shape, box_size, pk_fn, seed=jax.random.PRNGKey(0))\n",
    "\n",
    "    # Create particles\n",
    "    cosmo = jc.Planck15(Omega_c=omega_c, sigma8=sigma8)\n",
    "    \n",
    "    # Initial displacement\n",
    "    dx, p, f = lpt(cosmo, initial_conditions, a=0.1)\n",
    "    \n",
    "    # Evolve the simulation forward\n",
    "    res = odeint(make_ode_fn(mesh_shape,paint_absolute_pos=False), [dx, p], snapshots, cosmo, rtol=1e-8, atol=1e-8)\n",
    "    \n",
    "    # Return the simulation volume at requested \n",
    "\n",
    "    return initial_conditions ,  dx , res[0]\n",
    "    \n",
    "initial_conditions , lpt_displacements , ode_displacements = run_simulation(0.25, 0.8)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "33b5e684",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAB8UAAAH/CAYAAADOlQwMAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjkuMywgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/GU6VOAAAACXBIWXMAAA9hAAAPYQGoP6dpAAEAAElEQVR4nOydeXxcZb3/v+fMviSZrE3StEm6r7TQUnbKJtAiiAhcF2RRRBEu6lWv21VBQa+K4lWviHpVBPWngCwqyiK7bC2UpbR0TbckzT5ZZp855/dHzXw/3yeZ0GKxFL9vXn3xZM4z5zzr93nOmZnPx3Jd1yVFURRFURRFURRFURRFURRFURRFURRFeQtiH+gCKIqiKIqiKIqiKIqiKIqiKIqiKIqiKMobhX4oriiKoiiKoiiKoiiKoiiKoiiKoiiKorxl0Q/FFUVRFEVRFEVRFEVRFEVRFEVRFEVRlLcs+qG4oiiKoiiKoiiKoiiKoiiKoiiKoiiK8pZFPxRXFEVRFEVRFEVRFEVRFEVRFEVRFEVR3rLoh+KKoiiKoiiKoiiKoiiKoiiKoiiKoijKWxb9UFxRFEVRFEVRFEVRFEVRFEVRFEVRFEV5y6IfiiuKoiiKoiiKoiiKoiiKoiiKoiiKoihvWfRDcUVRFEVRFEVRFEVRFEVRFEVRFEVRFOUti34orhyUWJZFV1999V7lbWlpoYsvvnifr7Ft2zayLIt+8Ytf7PN73ww88sgjZFkWPfLII8XXLr74Ymppadmr91999dVkWdYbU7iDiAM5DvalvxRFUZQ3H+OtxXvLL37xC7Isi7Zt2/aaeV/vXueN5ECVaV/aTVEURXlzsmrVKjr66KMpEomQZVn0wgsv/EP3p3t7X3WwPwN4Lf6Rfck/ygknnEAnnHDCP/26iqIoyj8PXb/fGPTZtKLsX/RDceWAMPrAcvXq1fvlfE8++SRdffXVFI/H98v5Xg9dXV30qU99iubMmUPhcJgikQgtWbKErr322gNarolIJpN09dVXH5CbYmUPHR0ddPXVV9MLL7xwoIsyhv/7v/+juXPnUjAYpJkzZ9L3v//9A10kRVGUvdpDjN40jv7zeDw0depUeuc731mMtxdffLHIU+rfRB/sjt7gj/fvRz/60X6uuYJ87Wtfo7vuuutAF2MMTz75JB177LEUDoepvr6errrqKhoZGTnQxVIURTloyOVydN5551F/fz/dcMMNdMstt1Bzc/OBLtZrovdOROvWraOrr776TffFNMdx6Jvf/Ca1trZSMBikQw45hH7zm98c6GIpiqK8pTgY1+8bb7yRzjvvPJo6depr3vu/ldFn08q/Gt4DXQBFeT2kUinyenn4Pvnkk3TNNdfQxRdfTLFYTOTdsGED2fYb+/2PVatW0cqVK2lkZIQuuOACWrJkCRERrV69mv77v/+bHnvsMbr//vvf0DLsDT/5yU/IcZzi38lkkq655hoiojHf2v6v//ov+uxnP/vPLN6bkubmZkqlUuTz+d6Q83d0dNA111xDLS0ttHjxYnHM7K9/JjfddBN95CMfoXe96130H//xH/T444/TVVddRclkkj7zmc8ckDIpiqLsK+95z3to5cqVVCgUaP369XTjjTfSn//8Z3r66afpwx/+MJ1yyinFvG1tbfSlL32JLrvsMjruuOOKr0+fPv01r3PjjTdSNBoVrx1xxBE0ffp0SqVS5Pf791+lDhLe6P3X1772NTr33HPp7LPPFq+///3vp3e/+90UCATesGuX4oUXXqCTTz6Z5s6dS9/5zndo165ddP3119OmTZvoz3/+8z+9PIqiKAcjW7Zsoe3bt9NPfvITuvTSS4uvv5nvTw+We6fjjz/+Dd2XrFu3jq655ho64YQTxvyq7EA+D/nCF75A//3f/00f+tCH6PDDD6e7776b3vve95JlWfTud7/7gJVLURTlrcTBuH5/4xvfoOHhYVq2bBl1dnYe6OKURJ9Nv7n3V8rBh34orhyUBIPBvc77Rj8Ujcfj9M53vpM8Hg+tWbOG5syZI45fd9119JOf/OQNLcPesi+Lp9frFV88eKvgui6l02kKhUJ7ld+yrH0ab/uTN2qz81qkUin6whe+QGeccQbdfvvtRET0oQ99iBzHoa9+9at02WWXUWVl5QEpm6Ioyr5w2GGH0QUXXFD8+5hjjqGzzjqLbrzxRrrpppvoqKOOKh5bvXo1felLX6KjjjpKvGdvOPfcc6mmpmbcYwdqDdnfJBIJikQie53/QHwoTUTk8XjI4/EckGt//vOfp8rKSnrkkUeovLyciPbIyH/oQx+i+++/n0499dQDUi5FUZSDie7ubiKiMV92f7Penx7IeyfHcSibze71XsO27QO2LzlQXxBsb2+nb3/723TFFVfQD37wAyIiuvTSS2n58uX06U9/ms4777wDtm9QFEV5K3Gwrd9ERI8++mjxV+Lml9zfSPTZ9Gujz6aVNxKVT1feNFx88cUUjUapvb2dzj77bIpGo1RbW0uf+tSnqFAoiLzoKX711VfTpz/9aSIiam1tLcqWjkp2mZ6W/f399KlPfYoWLlxI0WiUysvLacWKFfTiiy++rnLfdNNN1N7eTt/5znfGfCBORDRp0iT6r//6L/HaD3/4Q5o/fz4FAgFqbGykK664YozE+gknnEALFiygdevW0YknnkjhcJgmT55M3/zmN8dcY9euXXT22WdTJBKhuro6+sQnPkGZTGZMPvQB2bZtG9XW1hIR0TXXXFNsN2xX0/Mln8/TV7/6VZo+fToFAgFqaWmhz3/+82Ou1dLSQm9/+9vpiSeeoGXLllEwGKRp06bRL3/5S5Evl8vRNddcQzNnzqRgMEjV1dV07LHH0gMPPDC2oYFR6dzHHnuMPvzhD1N1dTWVl5fThRdeSAMDA+OW5b777qOlS5dSKBSim266iYiItm7dSueddx5VVVVROBymI488kv70pz+J95fybXn11Vfp3HPPpaqqKgoGg7R06VK65557xpQ1Ho/TJz7xCWppaaFAIEBNTU104YUXUm9vLz3yyCN0+OGHExHRJZdcUuyD0WuN59uSSCTok5/8JE2ZMoUCgQDNnj2brr/+enJdV+SzLIuuvPJKuuuuu2jBggUUCARo/vz59Je//GXCtiUievjhh6mvr48++tGPitevuOIKSiQSY9pIURTlYOGkk04ioj2/Cv9nUMq785lnnqHTTz+dKioqKBwO0/Lly+lvf/vba57PdV269tprqampicLhMJ144on0yiuv7FVZRtez66+/nm644QZqbm6mUChEy5cvp7Vr14q8o3uyLVu20MqVK6msrIze9773EdHer0PjeYrH43H6+Mc/XnzvjBkz6Bvf+MaYb547jkP/8z//QwsXLqRgMEi1tbV0+umnFyXzLcuiRCJBN9988xiZ+1Ke4vt772UyNDREDzzwAF1wwQXFD8SJiC688EKKRqP0u9/97jXPoSiK8q/OxRdfTMuXLyciovPOO48syyoqmpXyJL311ltpyZIlFAqFqKqqit797nfTzp07X/Na8XicLr74YqqoqKBYLEYXXXTR67I9+0fvnUbr9eqrr9L5559P5eXlVF1dTR/72MconU6LvKP3eL/61a+Ka9ro/d2aNWtoxYoVVF5eTtFolE4++WR6+umnxfv/0X1Je3s7ffCDH6TGxkYKBALU2tpKl19+OWWzWfrFL35B5513HhERnXjiicX1efRa43mKd3d30wc/+EGaNGkSBYNBWrRoEd18880iD+5ffvzjHxefRRx++OG0atWqCduWiOjuu++mXC4n+seyLLr88stp165d9NRTT73mORRFUZSJORjXb6I9v8B+vX7n+mxan00rBy9vzq/pKP+yFAoFOu200+iII46g66+/nh588EH69re/TdOnT6fLL7983Pecc845tHHjRvrNb35DN9xwQ/GXWqMf+Jps3bqV7rrrLjrvvPOotbWVurq66KabbqLly5fTunXrqLGxcZ/KfM8991AoFKJzzz13r/JfffXVdM0119App5xCl19+OW3YsIFuvPFGWrVqFf3tb38T38AaGBig008/nc4
      "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_displacements[1:]):\n",
    "    fields[f\"field_{i}\"] = cic_paint_dx(field)\n",
    "plot_fields_single_projection(fields)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02f18b52",
   "metadata": {},
   "source": [
    "We note that **painting only displacements is slower** than painting absolute particle positions.\n",
    "\n",
    "This slower performance occurs because **painting displacements requires processing in smaller chunks or batches**. Instead of creating a large array of particle positions with neighbors (e.g., `(NParticles, 3, 8)`), which consumes a significant amount of memory, we paint the particles in manageable batches.\n",
    "\n",
    "This trade-off allows for greater memory efficiency but comes at the expense of speed. By reducing the memory footprint, we avoid memory limitations, especially useful in large-scale or distributed PM simulations, even if it means slightly slower painting.\n",
    "\n",
    "We’ll see in later notebooks that retaining only the displacement is essential for distributed Particle Mesh (PM) simulations, where memory efficiency and computational speed are key.\n"
   ]
  }
 ],
 "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
}