From 9bab01f8b1393a059aebe7d26817f3b767c10c31 Mon Sep 17 00:00:00 2001 From: Francois Lanusse Date: Sat, 21 Dec 2024 09:58:18 -0500 Subject: [PATCH] remove massive file Former-commit-id: 69f49ba2ed23c16bcea70024ea3bcff4d71b8a5b --- notebooks/06-Animating_PM_Fields.ipynb | 320 ++++++++++++++++++ ...6-Animating_PM_Fields.ipynb.REMOVED.git-id | 1 - 2 files changed, 320 insertions(+), 1 deletion(-) create mode 100644 notebooks/06-Animating_PM_Fields.ipynb delete mode 100644 notebooks/06-Animating_PM_Fields.ipynb.REMOVED.git-id diff --git a/notebooks/06-Animating_PM_Fields.ipynb b/notebooks/06-Animating_PM_Fields.ipynb new file mode 100644 index 0000000..850a81e --- /dev/null +++ b/notebooks/06-Animating_PM_Fields.ipynb @@ -0,0 +1,320 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# **Animating Particle Mesh density fields**\n", + "\n", + "In this tutorial, we will animate the density field of a particle mesh simulation. We will use the `manim` library to create the animation. \n", + "\n", + "The density fields are created exactly like in the notebook [**05-MultiHost_PM.ipynb**](05-MultiHost_PM.ipynb) using the same script [**05-MultiHost_PM.py**](05-MultiHost_PM.py)." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "To run a multi-host simulation, you first need to **allocate a job** with `salloc`. This command requests resources on an HPC cluster.\n", + "\n", + "just like in notebook [**05-MultiHost_PM.ipynb**]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "!salloc --account=XXX@a100 -C a100 --gres=gpu:8 --ntasks-per-node=8 --time=00:40:00 --cpus-per-task=8 --hint=nomultithread --qos=qos_gpu-dev --nodes=4 & " + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "**A few hours later**\n", + "\n", + "Use `!squeue -u $USER -o \"%i %D %b\"` to **check the JOB ID** and verify your resource allocation.\n", + "\n", + "In this example, we’ve been allocated **32 GPUs split across 4 nodes**.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "!squeue -u $USER -o \"%i %D %b\"" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Unset the following environment variables, as they can cause issues when using JAX in a distributed setting:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "del os.environ['VSCODE_PROXY_URI']\n", + "del os.environ['NO_PROXY']\n", + "del os.environ['no_proxy']" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Checking Available Compute Resources\n", + "\n", + "Run the following command to initialize JAX distributed computing and display the devices available for this job:\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "!srun --jobid=467745 -n 32 python -c \"import jax; jax.distributed.initialize(); print(jax.devices()) if jax.process_index() == 0 else None\"" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Multi-Host Simulation Script with Arguments (reminder)\n", + "\n", + "This script is nearly identical to the single-host version, with the main addition being the call to `jax.distributed.initialize()` at the start, enabling multi-host parallelism. Here’s a breakdown of the key arguments:\n", + "\n", + "- **`--pdims`** (`-p`): Specifies processor grid dimensions as two integers, like `16 2` for 16 x 2 device mesh (default is `[1, jax.devices()]`).\n", + "- **`--mesh_shape`** (`-m`): Defines the simulation mesh shape as three integers (default is `[512, 512, 512]`).\n", + "- **`--box_size`** (`-b`): Sets the physical box size of the simulation as three floating-point values, e.g., `1000. 1000. 1000.` (default is `[500.0, 500.0, 500.0]`).\n", + "- **`--halo_size`** (`-H`): Specifies the halo size for boundary overlap across nodes (default is `64`).\n", + "- **`--solver`** (`-s`): Chooses the ODE solver (`leapfrog` or `dopri8`). The `leapfrog` solver uses a fixed step size, while `dopri8` is an adaptive Runge-Kutta solver with a PID controller (default is `leapfrog`).\n", + "- **`--snapthots`** (`-st`) : Number of snapshots to save (warning, increases memory usage)\n", + "\n", + "### Running the Multi-Host Simulation Script\n", + "\n", + "To create a smooth animation, we need a series of closely spaced snapshots to capture the evolution of the density field over time. In this example, we set the number of snapshots to **10** to ensure smooth transitions in the animation.\n", + "\n", + "Using a larger number of GPUs helps process these snapshots efficiently, especially with a large simulation mesh or high-resolution data. This allows us to achieve both the desired snapshot frequency and the necessary simulation detail without excessive runtime.\n", + "\n", + "The command to run the multi-host simulation with these settings will look something like this:\n", + "\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import subprocess\n", + "\n", + "# Define parameters as variables\n", + "jobid = \"467745\"\n", + "num_processes = 32\n", + "script_name = \"05-MultiHost_PM.py\"\n", + "mesh_shape = (1024, 1024, 1024)\n", + "box_size = (1000., 1000., 1000.)\n", + "halo_size = 128\n", + "solver = \"leapfrog\"\n", + "pdims = (16, 2)\n", + "snapshots = 8\n", + "\n", + "# Build the command as a list, incorporating variables\n", + "command = [\n", + " \"srun\",\n", + " f\"--jobid={jobid}\",\n", + " \"-n\", str(num_processes),\n", + " \"python\", script_name,\n", + " \"--mesh_shape\", str(mesh_shape[0]), str(mesh_shape[1]), str(mesh_shape[2]),\n", + " \"--box_size\", str(box_size[0]), str(box_size[1]), str(box_size[2]),\n", + " \"--halo_size\", str(halo_size),\n", + " \"-s\", solver,\n", + " \"--pdims\", str(pdims[0]), str(pdims[1]),\n", + " \"--snapshots\", str(snapshots)\n", + "]\n", + "\n", + "# Execute the command as a subprocess\n", + "subprocess.run(command)\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Projecting the 3D Density Fields to 2D\n", + "\n", + "To visualize the 3D density fields in 2D, we need to create a projection:\n", + "\n", + "- **`project_to_2d` Function**: This function reduces the 3D array to 2D by summing over a portion of one axis.\n", + " - We sum the top one-eighth of the data along the first axis to capture a slice of the density field.\n", + "\n", + "- **Creating 2D Projections**: Apply `project_to_2d` to each 3D field (`initial_conditions`, `lpt_displacements`, `ode_solution_0`, and `ode_solution_1`) to get 2D arrays that represent the density fields.\n", + "\n", + "### Applying the Magma Colormap\n", + "\n", + "To improve visualization, apply the \"magma\" colormap to each 2D projection:\n", + "\n", + "- **`apply_colormap` Function**: This function maps values in the 2D array to colors using the \"magma\" colormap.\n", + " - First, normalize the array to the `[0, 1]` range.\n", + " - Apply the colormap to create RGB images, which will be used for the animation.\n" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [], + "source": [ + "from matplotlib import colormaps\n", + "\n", + "# Define a function to project the 3D field to 2D\n", + "def project_to_2d(field):\n", + " sum_over = field.shape[0] // 8\n", + " slicing = [slice(None)] * field.ndim\n", + " slicing[0] = slice(None, sum_over)\n", + " slicing = tuple(slicing)\n", + "\n", + " return field[slicing].sum(axis=0)\n", + "\n", + "\n", + "def apply_colormap(array, cmap_name=\"magma\"):\n", + " cmap = colormaps[cmap_name]\n", + " normalized_array = (array - array.min()) / (array.max() - array.min())\n", + " colored_image = cmap(normalized_array)[:, :, :3] # Drop alpha channel for RGB\n", + " return (colored_image * 255).astype(np.uint8)\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Loading and Visualizing Results\n", + "\n", + "After running the multi-host simulation, we load the saved results from disk:\n", + "\n", + "- **`initial_conditions.npy`**: Initial conditions for the simulation.\n", + "- **`lpt_displacements.npy`**: Linear perturbation displacements.\n", + "- **`ode_solution_*.npy`** : Solutions from the ODE solver at each snapshot.\n", + "\n", + "We will now project the fields to 2D maps and apply the color map\n" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "\n", + "initial_conditions = apply_colormap(project_to_2d(np.load('fields/initial_conditions.npy')))\n", + "lpt_displacements = apply_colormap(project_to_2d(np.load('fields/lpt_displacements.npy')))\n", + "ode_solutions = []\n", + "for i in range(8):\n", + " ode_solutions.append(apply_colormap(project_to_2d(np.load(f'fields/ode_solution_{i}.npy'))))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Animating with Manim\n", + "\n", + "To create animations with `manim` in a Jupyter notebook, we start by configuring some settings to ensure the output displays correctly and without a background.\n" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [], + "source": [ + "from manim import *\n", + "config.media_width = \"100%\"\n", + "config.verbosity = \"WARNING\"\n", + "config.background_color = \"#00000000\" # Transparent background" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Defining the Animation in Manim\n", + "\n", + "This animation class, `FieldTransition`, smoothly transitions through the stages of the particle mesh density field evolution.\n", + "\n", + "- **Setup**: Each density field snapshot is loaded as an image and aligned for smooth transitions.\n", + "- **Animation Sequence**:\n", + " - The animation begins with a fade-in of the initial conditions.\n", + " - It then transitions through the stages in sequence, showing each snapshot of the density field evolution with brief pauses in between.\n", + "\n", + "To run the animation, execute `%manim -v WARNING -qm FieldTransition` to render it in the Jupyter Notebook.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Define the animation in Manim\n", + "class FieldTransition(Scene):\n", + " def construct(self):\n", + " init_conditions_img = ImageMobject(initial_conditions).scale(4)\n", + " lpt_img = ImageMobject(lpt_displacements).scale(4)\n", + " snapshots_imgs = [ImageMobject(sol).scale(4) for sol in ode_solutions]\n", + "\n", + "\n", + " # Place the images on top of each other initially\n", + " lpt_img.move_to(init_conditions_img)\n", + " for img in snapshots_imgs:\n", + " img.move_to(init_conditions_img)\n", + "\n", + " # Show initial field and then transform between fields\n", + " self.play(FadeIn(init_conditions_img))\n", + " self.wait(0.2)\n", + " self.play(Transform(init_conditions_img, lpt_img))\n", + " self.wait(0.2)\n", + " self.play(Transform(lpt_img, snapshots_imgs[0]))\n", + " self.wait(0.2)\n", + " for img1, img2 in zip(snapshots_imgs, snapshots_imgs[1:]):\n", + " self.play(Transform(img1, img2))\n", + " self.wait(0.2)\n", + "\n", + "%manim -v WARNING -qm -o anim.gif --format=gif FieldTransition " + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.4" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/notebooks/06-Animating_PM_Fields.ipynb.REMOVED.git-id b/notebooks/06-Animating_PM_Fields.ipynb.REMOVED.git-id deleted file mode 100644 index aa7a5e1..0000000 --- a/notebooks/06-Animating_PM_Fields.ipynb.REMOVED.git-id +++ /dev/null @@ -1 +0,0 @@ -c4a44973e4f11841a8c14f4d200e7e87887419aa \ No newline at end of file