mirror of
https://github.com/DifferentiableUniverseInitiative/JaxPM.git
synced 2025-04-08 04:40:53 +00:00
add notebook examples
This commit is contained in:
Wassim KABALAN
parent
0c96a4dc10
commit
49dd18a3f8
6 changed files with 861 additions and 192 deletions
320
notebooks/01-Introduction.ipynb
Normal file
320
notebooks/01-Introduction.ipynb
Normal file
|
|
@ -0,0 +1,320 @@
|
||||||
|
{
|
||||||
|
"cells": [
|
||||||
|
{
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"metadata": {
|
||||||
|
"colab_type": "text",
|
||||||
|
"id": "view-in-github"
|
||||||
|
},
|
||||||
|
"source": [
|
||||||
|
"<a href=\"https://colab.research.google.com/github/DifferentiableUniverseInitiative/JaxPM/blob/main/notebooks/Introduction.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 8,
|
||||||
|
"id": "9Jy5BL1XiK1s",
|
||||||
|
"metadata": {
|
||||||
|
"id": "9Jy5BL1XiK1s"
|
||||||
|
},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"!pip install --quiet git+https://github.com/DifferentiableUniverseInitiative/JaxPM.git"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 1,
|
||||||
|
"id": "c5f42bbe",
|
||||||
|
"metadata": {
|
||||||
|
"colab": {
|
||||||
|
"base_uri": "https://localhost:8080/"
|
||||||
|
},
|
||||||
|
"id": "c5f42bbe",
|
||||||
|
"outputId": "a7841b28-5f20-4856-bd1d-f8a3572095b5"
|
||||||
|
},
|
||||||
|
"outputs": [
|
||||||
|
{
|
||||||
|
"name": "stdout",
|
||||||
|
"output_type": "stream",
|
||||||
|
"text": [
|
||||||
|
"%pylab is deprecated, use %matplotlib inline and import the required libraries.\n",
|
||||||
|
"Populating the interactive namespace from numpy and matplotlib\n"
|
||||||
|
]
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"source": [
|
||||||
|
"import os\n",
|
||||||
|
"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\n",
|
||||||
|
"from jaxpm.pm import linear_field, lpt, make_ode_fn"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 2,
|
||||||
|
"id": "281b4d3b",
|
||||||
|
"metadata": {
|
||||||
|
"id": "281b4d3b"
|
||||||
|
},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"mesh_shape= [256, 256, 256]\n",
|
||||||
|
"box_size = [256.,256.,256.]\n",
|
||||||
|
"snapshots = jnp.linspace(0.1,1.,3)\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",
|
||||||
|
" particles = jnp.stack(jnp.meshgrid(*[jnp.arange(s) for s in mesh_shape]),axis=-1).reshape([-1 , 3])\n",
|
||||||
|
" cosmo = jc.Planck15(Omega_c=omega_c, sigma8=sigma8)\n",
|
||||||
|
" \n",
|
||||||
|
" # Initial displacement\n",
|
||||||
|
" dx, p, f = lpt(cosmo, initial_conditions, particles, 0.1)\n",
|
||||||
|
" \n",
|
||||||
|
" # Evolve the simulation forward\n",
|
||||||
|
" res = odeint(make_ode_fn(mesh_shape,particles), [particles + dx, p], snapshots, cosmo, rtol=1e-5, atol=1e-5)\n",
|
||||||
|
" \n",
|
||||||
|
" # Return the simulation volume at requested \n",
|
||||||
|
"\n",
|
||||||
|
" return initial_conditions , particles + dx , res[0]"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"id": "826be667",
|
||||||
|
"metadata": {
|
||||||
|
"colab": {
|
||||||
|
"base_uri": "https://localhost:8080/"
|
||||||
|
},
|
||||||
|
"id": "826be667",
|
||||||
|
"outputId": "dc43b5c4-a004-41bf-f2c8-128c17cf4de1"
|
||||||
|
},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"initial_conditions , lpt_particles , ode_particles = run_simulation(0.25, 0.8)\n",
|
||||||
|
"%timeit initial_conditions , lpt_particles , ode_particles = run_simulation(0.25, 0.8)"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"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": [
|
||||||
|
{
|
||||||
|
"name": "stdout",
|
||||||
|
"output_type": "stream",
|
||||||
|
"text": [
|
||||||
|
"shape of grid_mesh: (256, 256, 256)\n"
|
||||||
|
]
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"source": [
|
||||||
|
"from visualize import plot_fields\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):\n",
|
||||||
|
" fields[f\"field_{i}\"] = cic_paint(jnp.zeros(mesh_shape) , field)\n",
|
||||||
|
"plot_fields(fields)"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"id": "b71824ed",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"mesh_shape= [256, 256, 256]\n",
|
||||||
|
"box_size = [256.,256.,256.]\n",
|
||||||
|
"snapshots = jnp.linspace(0.1,1.,3)\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, 0.1)\n",
|
||||||
|
" \n",
|
||||||
|
" # Evolve the simulation forward\n",
|
||||||
|
" res = odeint(make_ode_fn(mesh_shape), [dx, p], snapshots, cosmo, rtol=1e-5, atol=1e-5)\n",
|
||||||
|
" \n",
|
||||||
|
" # Return the simulation volume at requested \n",
|
||||||
|
"\n",
|
||||||
|
" return initial_conditions , dx , res[0]"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"id": "e9c9fd56",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"initial_conditions , lpt_displacements , ode_displacements = run_simulation(0.25, 0.8)\n",
|
||||||
|
"%timeit initial_conditions , lpt_displacements , ode_displacements = run_simulation(0.25, 0.8)"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"id": "33b5e684",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"from visualize import plot_fields\n",
|
||||||
|
"\n",
|
||||||
|
"fields = {\"Initial Conditions\" : initial_conditions , \"LPT Field\" : cic_paint(jnp.zeros(mesh_shape) ,lpt_displacements)}\n",
|
||||||
|
"for i , field in enumerate(ode_displacements):\n",
|
||||||
|
" fields[f\"field_{i}\"] = cic_paint(jnp.zeros(mesh_shape) , field)\n",
|
||||||
|
"plot_fields(fields)"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"id": "4e050871",
|
||||||
|
"metadata": {
|
||||||
|
"id": "4e050871"
|
||||||
|
},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"!pip install diffrax\n",
|
||||||
|
"from diffrax import ConstantStepSize, LeapfrogMidpoint, ODETerm, SaveAt, diffeqsolve"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"id": "43504a1b",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"mesh_shape= [256, 256, 256]\n",
|
||||||
|
"box_size = [256.,256.,256.]\n",
|
||||||
|
"snapshots = jnp.linspace(0.1,1.,3)\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, 0.1)\n",
|
||||||
|
" \n",
|
||||||
|
" # Evolve the simulation forward\n",
|
||||||
|
" ode_fn = make_ode_fn(mesh_shape)\n",
|
||||||
|
" term = ODETerm(\n",
|
||||||
|
" lambda t, state, args: jnp.stack(ode_fn(state, t, args), axis=0))\n",
|
||||||
|
" solver = LeapfrogMidpoint()\n",
|
||||||
|
"\n",
|
||||||
|
" stepsize_controller = ConstantStepSize()\n",
|
||||||
|
" res = diffeqsolve(term,\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",
|
||||||
|
"\n",
|
||||||
|
"\n",
|
||||||
|
" return initial_conditions , dx , res.ys , res.stats"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"id": "19949ff1",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"initial_conditions , lpt_displacements , ode_solutions , solver_stats = run_simulation(0.25, 0.8)\n",
|
||||||
|
"%timeit initial_conditions , lpt_displacements , ode_solutions , solver_stats = run_simulation(0.25, 0.8)\n",
|
||||||
|
"print(f\"Solver Stats : {solver_stats}\")"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"id": "76a26e98",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"from visualize import plot_fields\n",
|
||||||
|
"\n",
|
||||||
|
"fields = {\"Initial Conditions\" : initial_conditions , \"LPT Field\" : cic_paint(jnp.zeros(mesh_shape) ,lpt_displacements)}\n",
|
||||||
|
"for i , field in enumerate(ode_solutions):\n",
|
||||||
|
" fields[f\"field_{i}\"] = cic_paint(jnp.zeros(mesh_shape) , field[0])\n",
|
||||||
|
"plot_fields(fields)"
|
||||||
|
]
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"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
|
||||||
|
}
|
||||||
242
notebooks/02-MultiGPU_PM.ipynb
Normal file
242
notebooks/02-MultiGPU_PM.ipynb
Normal file
|
|
@ -0,0 +1,242 @@
|
||||||
|
{
|
||||||
|
"cells": [
|
||||||
|
{
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"metadata": {
|
||||||
|
"colab_type": "text",
|
||||||
|
"id": "view-in-github"
|
||||||
|
},
|
||||||
|
"source": [
|
||||||
|
"<a href=\"https://colab.research.google.com/github/DifferentiableUniverseInitiative/JaxPM/blob/main/notebooks/Introduction.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 8,
|
||||||
|
"id": "9Jy5BL1XiK1s",
|
||||||
|
"metadata": {
|
||||||
|
"id": "9Jy5BL1XiK1s"
|
||||||
|
},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"!pip install --quiet git+https://github.com/DifferentiableUniverseInitiative/JaxPM.git\n",
|
||||||
|
"!pip install diffrax"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 1,
|
||||||
|
"id": "c5f42bbe",
|
||||||
|
"metadata": {
|
||||||
|
"colab": {
|
||||||
|
"base_uri": "https://localhost:8080/"
|
||||||
|
},
|
||||||
|
"id": "c5f42bbe",
|
||||||
|
"outputId": "a7841b28-5f20-4856-bd1d-f8a3572095b5"
|
||||||
|
},
|
||||||
|
"outputs": [
|
||||||
|
{
|
||||||
|
"name": "stdout",
|
||||||
|
"output_type": "stream",
|
||||||
|
"text": [
|
||||||
|
"%pylab is deprecated, use %matplotlib inline and import the required libraries.\n",
|
||||||
|
"Populating the interactive namespace from numpy and matplotlib\n"
|
||||||
|
]
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"source": [
|
||||||
|
"import os\n",
|
||||||
|
"import jax\n",
|
||||||
|
"import jax.numpy as jnp\n",
|
||||||
|
"import jax_cosmo as jc\n",
|
||||||
|
"\n",
|
||||||
|
"from jax.experimental.ode import odeint\n",
|
||||||
|
"from jaxpm.kernels import interpolate_power_spectrum\n",
|
||||||
|
"from jaxpm.painting import cic_paint\n",
|
||||||
|
"from jaxpm.pm import linear_field, lpt, make_ode_fn\n",
|
||||||
|
"from diffrax import ConstantStepSize, LeapfrogMidpoint, ODETerm, SaveAt, diffeqsolve"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"id": "38df34e3",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"assert jax.device_count() >= 8, \"This notebook requires a TPU or GPU runtime with 8 devices\""
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"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",
|
||||||
|
"from functools import partial\n",
|
||||||
|
"\n",
|
||||||
|
"all_gather = partial(process_allgather, tiled=True)\n",
|
||||||
|
"\n",
|
||||||
|
"pdims = (2, 4)\n",
|
||||||
|
"devices = create_device_mesh(pdims)\n",
|
||||||
|
"mesh = Mesh(devices, axis_names=('x', 'y'))\n",
|
||||||
|
"sharding = NamedSharding(mesh, P('x', 'y'))"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 2,
|
||||||
|
"id": "281b4d3b",
|
||||||
|
"metadata": {
|
||||||
|
"id": "281b4d3b"
|
||||||
|
},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"mesh_shape = [1024, 1024, 1024]\n",
|
||||||
|
"box_size = [1024., 1024., 1024.]\n",
|
||||||
|
"halo_size = 128\n",
|
||||||
|
"snapshots = jnp.linspace(0.1, 1., 3)\n",
|
||||||
|
"\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(\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",
|
||||||
|
" # Create particles\n",
|
||||||
|
" particles = jnp.stack(jnp.meshgrid(*[jnp.arange(s) for s in mesh_shape]),\n",
|
||||||
|
" axis=-1).reshape([-1, 3])\n",
|
||||||
|
" cosmo = jc.Planck15(Omega_c=omega_c, sigma8=sigma8)\n",
|
||||||
|
"\n",
|
||||||
|
" # Initial displacement\n",
|
||||||
|
" dx, p, f = lpt(cosmo,\n",
|
||||||
|
" initial_conditions,\n",
|
||||||
|
" particles,\n",
|
||||||
|
" 0.1,\n",
|
||||||
|
" halo_size=halo_size,\n",
|
||||||
|
" sharding=sharding)\n",
|
||||||
|
"\n",
|
||||||
|
" # Evolve the simulation forward\n",
|
||||||
|
" ode_fn = make_ode_fn(mesh_shape, halo_size=halo_size, sharding=sharding)\n",
|
||||||
|
" term = ODETerm(\n",
|
||||||
|
" lambda t, state, args: jnp.stack(ode_fn(state, t, args), axis=0))\n",
|
||||||
|
" solver = LeapfrogMidpoint()\n",
|
||||||
|
"\n",
|
||||||
|
" stepsize_controller = ConstantStepSize()\n",
|
||||||
|
" res = diffeqsolve(term,\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",
|
||||||
|
"\n",
|
||||||
|
" return initial_conditions, dx, res.ys, res.stats"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"id": "826be667",
|
||||||
|
"metadata": {
|
||||||
|
"colab": {
|
||||||
|
"base_uri": "https://localhost:8080/"
|
||||||
|
},
|
||||||
|
"id": "826be667",
|
||||||
|
"outputId": "dc43b5c4-a004-41bf-f2c8-128c17cf4de1"
|
||||||
|
},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"initial_conditions , lpt_displacements , ode_solutions , solver_stats = run_simulation(0.25, 0.8)\n",
|
||||||
|
"%timeit initial_conditions , lpt_displacements , ode_solutions , solver_stats = run_simulation(0.25, 0.8)\n",
|
||||||
|
"print(f\"Solver Stats : {solver_stats}\")"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"id": "042cc55c",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"initial_conditions = all_gather(initial_conditions)\n",
|
||||||
|
"lpt_particles = all_gather(lpt_particles)\n",
|
||||||
|
"ode_particles = [all_gather(p) for p in ode_particles]"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"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": [
|
||||||
|
{
|
||||||
|
"name": "stdout",
|
||||||
|
"output_type": "stream",
|
||||||
|
"text": [
|
||||||
|
"shape of grid_mesh: (256, 256, 256)\n"
|
||||||
|
]
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"source": [
|
||||||
|
"from visualize import plot_fields\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):\n",
|
||||||
|
" fields[f\"field_{i}\"] = cic_paint(jnp.zeros(mesh_shape) , field)\n",
|
||||||
|
"plot_fields(fields)"
|
||||||
|
]
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"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
|
||||||
|
}
|
||||||
157
notebooks/03-MultiHost_PM.ipynb
Normal file
157
notebooks/03-MultiHost_PM.ipynb
Normal file
|
|
@ -0,0 +1,157 @@
|
||||||
|
{
|
||||||
|
"cells": [
|
||||||
|
{
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"id": "22803ddc",
|
||||||
|
"metadata": {
|
||||||
|
"colab_type": "text",
|
||||||
|
"id": "view-in-github"
|
||||||
|
},
|
||||||
|
"source": [
|
||||||
|
"<a href=\"https://colab.research.google.com/github/DifferentiableUniverseInitiative/JaxPM/blob/main/notebooks/Introduction.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 8,
|
||||||
|
"id": "9Jy5BL1XiK1s",
|
||||||
|
"metadata": {
|
||||||
|
"id": "9Jy5BL1XiK1s"
|
||||||
|
},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"!pip install --quiet git+https://github.com/DifferentiableUniverseInitiative/JaxPM.git\n",
|
||||||
|
"!pip install diffrax"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 1,
|
||||||
|
"id": "c5f42bbe",
|
||||||
|
"metadata": {
|
||||||
|
"colab": {
|
||||||
|
"base_uri": "https://localhost:8080/"
|
||||||
|
},
|
||||||
|
"id": "c5f42bbe",
|
||||||
|
"outputId": "a7841b28-5f20-4856-bd1d-f8a3572095b5"
|
||||||
|
},
|
||||||
|
"outputs": [
|
||||||
|
{
|
||||||
|
"name": "stdout",
|
||||||
|
"output_type": "stream",
|
||||||
|
"text": [
|
||||||
|
"%pylab is deprecated, use %matplotlib inline and import the required libraries.\n",
|
||||||
|
"Populating the interactive namespace from numpy and matplotlib\n"
|
||||||
|
]
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"source": [
|
||||||
|
"!salloc --account=tkc@a100 -C a100 --gres=gpu:8 --ntasks-per-node=8 --time=00:30:00 --cpus-per-task=8 --hint=nomultithread --qos=qos_gpu-dev --nodes=2"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"id": "7ebdfc00",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"!squeue -u $USER"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"id": "c014316c",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"export JOB_ID=123456"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"id": "b7eabac5",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"!srun --jobid=$JOB_ID -n 16 python 03-MultiHost_PM.py"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"id": "472dd4bf",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"import numpy as np\n",
|
||||||
|
"\n",
|
||||||
|
"data = np.load(\"multihost_pm.npz\")\n",
|
||||||
|
"initial_conditions = data['initial_conditions']\n",
|
||||||
|
"lpt_displacements = data['lpt_displacements']\n",
|
||||||
|
"ode_solutions = data['ode_solutions']\n",
|
||||||
|
"solver_stats = data['solver_stats']\n",
|
||||||
|
"print(f\"Solver stats: {solver_stats}\")"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"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": [
|
||||||
|
{
|
||||||
|
"name": "stdout",
|
||||||
|
"output_type": "stream",
|
||||||
|
"text": [
|
||||||
|
"shape of grid_mesh: (256, 256, 256)\n"
|
||||||
|
]
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"source": [
|
||||||
|
"from visualize import plot_fields\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):\n",
|
||||||
|
" fields[f\"field_{i}\"] = cic_paint(jnp.zeros(mesh_shape) , field)\n",
|
||||||
|
"plot_fields(fields)"
|
||||||
|
]
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"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
|
||||||
|
}
|
||||||
104
notebooks/03-MultiHost_PM.py
Normal file
104
notebooks/03-MultiHost_PM.py
Normal file
|
|
@ -0,0 +1,104 @@
|
||||||
|
import os
|
||||||
|
|
||||||
|
os.environ["EQX_ON_ERROR"] = "nan" # avoid an allgather caused by diffrax
|
||||||
|
import jax
|
||||||
|
|
||||||
|
jax.distributed.initialize()
|
||||||
|
rank = jax.process_index()
|
||||||
|
size = jax.process_count()
|
||||||
|
|
||||||
|
import jax.numpy as jnp
|
||||||
|
import jax_cosmo as jc
|
||||||
|
|
||||||
|
from jaxpm.kernels import interpolate_power_spectrum
|
||||||
|
from jaxpm.painting import cic_paint_dx
|
||||||
|
from jaxpm.pm import linear_field, lpt, make_ode_fn
|
||||||
|
|
||||||
|
from jax.experimental.mesh_utils import create_device_mesh
|
||||||
|
from jax.experimental.multihost_utils import process_allgather
|
||||||
|
from jax.sharding import Mesh, NamedSharding
|
||||||
|
from jax.sharding import PartitionSpec as P
|
||||||
|
from functools import partial
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
from diffrax import ConstantStepSize, LeapfrogMidpoint, ODETerm, SaveAt, diffeqsolve
|
||||||
|
|
||||||
|
all_gather = partial(process_allgather, tiled=True)
|
||||||
|
|
||||||
|
pdims = (2, 4)
|
||||||
|
devices = create_device_mesh(pdims)
|
||||||
|
mesh = Mesh(devices, axis_names=('x', 'y'))
|
||||||
|
sharding = NamedSharding(mesh, P('x', 'y'))
|
||||||
|
|
||||||
|
mesh_shape = [2024, 1024, 1024]
|
||||||
|
box_size = [1024., 1024., 1024.]
|
||||||
|
halo_size = 512
|
||||||
|
snapshots = jnp.linspace(0.1, 1., 2)
|
||||||
|
|
||||||
|
|
||||||
|
@jax.jit
|
||||||
|
def run_simulation(omega_c, sigma8):
|
||||||
|
# Create a small function to generate the matter power spectrum
|
||||||
|
k = jnp.logspace(-4, 1, 128)
|
||||||
|
pk = jc.power.linear_matter_power(
|
||||||
|
jc.Planck15(Omega_c=omega_c, sigma8=sigma8), k)
|
||||||
|
pk_fn = lambda x: interpolate_power_spectrum(x, k, pk, sharding)
|
||||||
|
|
||||||
|
# Create initial conditions
|
||||||
|
initial_conditions = linear_field(mesh_shape,
|
||||||
|
box_size,
|
||||||
|
pk_fn,
|
||||||
|
seed=jax.random.PRNGKey(0),
|
||||||
|
sharding=sharding)
|
||||||
|
|
||||||
|
# Create particles
|
||||||
|
particles = jnp.stack(jnp.meshgrid(*[jnp.arange(s) for s in mesh_shape]),
|
||||||
|
axis=-1).reshape([-1, 3])
|
||||||
|
cosmo = jc.Planck15(Omega_c=omega_c, sigma8=sigma8)
|
||||||
|
|
||||||
|
# Initial displacement
|
||||||
|
dx, p, _ = lpt(cosmo,
|
||||||
|
initial_conditions,
|
||||||
|
particles,
|
||||||
|
0.1,
|
||||||
|
halo_size=halo_size,
|
||||||
|
sharding=sharding)
|
||||||
|
|
||||||
|
# Evolve the simulation forward
|
||||||
|
ode_fn = make_ode_fn(mesh_shape, halo_size=halo_size, sharding=sharding)
|
||||||
|
term = ODETerm(
|
||||||
|
lambda t, state, args: jnp.stack(ode_fn(state, t, args), axis=0))
|
||||||
|
solver = LeapfrogMidpoint()
|
||||||
|
|
||||||
|
stepsize_controller = ConstantStepSize()
|
||||||
|
res = diffeqsolve(term,
|
||||||
|
solver,
|
||||||
|
t0=0.1,
|
||||||
|
t1=1.,
|
||||||
|
dt0=0.01,
|
||||||
|
y0=jnp.stack([dx, p], axis=0),
|
||||||
|
args=cosmo,
|
||||||
|
saveat=SaveAt(ts=snapshots),
|
||||||
|
stepsize_controller=stepsize_controller)
|
||||||
|
|
||||||
|
return initial_conditions, dx, res.ys, res.stats
|
||||||
|
|
||||||
|
|
||||||
|
initial_conditions , lpt_displacements , ode_solutions , solver_stats = run_simulation(0.25, 0.8)
|
||||||
|
print(f"[{rank}] Simulation completed")
|
||||||
|
print(f"[{rank}] Solver stats: {solver_stats}")
|
||||||
|
|
||||||
|
# Gather the results
|
||||||
|
initial_conditions = all_gather(initial_conditions)
|
||||||
|
lpt_displacements = all_gather(lpt_displacements)
|
||||||
|
ode_solutions = [all_gather(sol) for sol in ode_solutions]
|
||||||
|
|
||||||
|
if rank == 0:
|
||||||
|
np.savez("multihost_pm.npz",
|
||||||
|
initial_conditions=initial_conditions,
|
||||||
|
lpt_displacements=lpt_displacements,
|
||||||
|
ode_solutions=ode_solutions,
|
||||||
|
solver_stats=solver_stats)
|
||||||
|
|
||||||
|
print(f"[{rank}] Simulation results saved")
|
||||||
|
|
||||||
|
|
@ -1,192 +0,0 @@
|
||||||
{
|
|
||||||
"cells": [
|
|
||||||
{
|
|
||||||
"cell_type": "markdown",
|
|
||||||
"metadata": {
|
|
||||||
"colab_type": "text",
|
|
||||||
"id": "view-in-github"
|
|
||||||
},
|
|
||||||
"source": [
|
|
||||||
"<a href=\"https://colab.research.google.com/github/DifferentiableUniverseInitiative/JaxPM/blob/main/notebooks/Introduction.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
|
|
||||||
]
|
|
||||||
},
|
|
||||||
{
|
|
||||||
"cell_type": "code",
|
|
||||||
"execution_count": 8,
|
|
||||||
"id": "9Jy5BL1XiK1s",
|
|
||||||
"metadata": {
|
|
||||||
"id": "9Jy5BL1XiK1s"
|
|
||||||
},
|
|
||||||
"outputs": [],
|
|
||||||
"source": [
|
|
||||||
"!pip install --quiet git+https://github.com/DifferentiableUniverseInitiative/JaxPM.git@ASKabalan/jaxdecomp_proto"
|
|
||||||
]
|
|
||||||
},
|
|
||||||
{
|
|
||||||
"cell_type": "code",
|
|
||||||
"execution_count": 1,
|
|
||||||
"id": "c5f42bbe",
|
|
||||||
"metadata": {
|
|
||||||
"colab": {
|
|
||||||
"base_uri": "https://localhost:8080/"
|
|
||||||
},
|
|
||||||
"id": "c5f42bbe",
|
|
||||||
"outputId": "a7841b28-5f20-4856-bd1d-f8a3572095b5"
|
|
||||||
},
|
|
||||||
"outputs": [
|
|
||||||
{
|
|
||||||
"name": "stdout",
|
|
||||||
"output_type": "stream",
|
|
||||||
"text": [
|
|
||||||
"%pylab is deprecated, use %matplotlib inline and import the required libraries.\n",
|
|
||||||
"Populating the interactive namespace from numpy and matplotlib\n"
|
|
||||||
]
|
|
||||||
}
|
|
||||||
],
|
|
||||||
"source": [
|
|
||||||
"%pylab inline\n",
|
|
||||||
"import os\n",
|
|
||||||
"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\n",
|
|
||||||
"from jaxpm.pm import linear_field, lpt, make_ode_fn"
|
|
||||||
]
|
|
||||||
},
|
|
||||||
{
|
|
||||||
"cell_type": "code",
|
|
||||||
"execution_count": 2,
|
|
||||||
"id": "281b4d3b",
|
|
||||||
"metadata": {
|
|
||||||
"id": "281b4d3b"
|
|
||||||
},
|
|
||||||
"outputs": [],
|
|
||||||
"source": [
|
|
||||||
"mesh_shape= [256, 256, 256]\n",
|
|
||||||
"box_size = [256.,256.,256.]\n",
|
|
||||||
"snapshots = jnp.linspace(0.1,1.,2)\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: jc.scipy.interpolate.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",
|
|
||||||
" particles = jnp.stack(jnp.meshgrid(*[jnp.arange(s) for s in mesh_shape]),axis=-1).reshape([-1,3])\n",
|
|
||||||
"\n",
|
|
||||||
" cosmo = jc.Planck15(Omega_c=omega_c, sigma8=sigma8)\n",
|
|
||||||
" \n",
|
|
||||||
" # Initial displacement\n",
|
|
||||||
" dx, p, f = lpt(cosmo, initial_conditions, particles, 0.1)\n",
|
|
||||||
" \n",
|
|
||||||
" # Evolve the simulation forward\n",
|
|
||||||
" res = odeint(make_ode_fn(mesh_shape), [particles+dx, p], snapshots, cosmo, rtol=1e-5, atol=1e-5)\n",
|
|
||||||
" \n",
|
|
||||||
" # Return the simulation volume at requested \n",
|
|
||||||
" return res[0]"
|
|
||||||
]
|
|
||||||
},
|
|
||||||
{
|
|
||||||
"cell_type": "code",
|
|
||||||
"execution_count": 3,
|
|
||||||
"id": "826be667",
|
|
||||||
"metadata": {
|
|
||||||
"colab": {
|
|
||||||
"base_uri": "https://localhost:8080/"
|
|
||||||
},
|
|
||||||
"id": "826be667",
|
|
||||||
"outputId": "dc43b5c4-a004-41bf-f2c8-128c17cf4de1"
|
|
||||||
},
|
|
||||||
"outputs": [
|
|
||||||
{
|
|
||||||
"name": "stdout",
|
|
||||||
"output_type": "stream",
|
|
||||||
"text": [
|
|
||||||
"particles are Traced<ShapedArray(int32[16777216,3])>with<DynamicJaxprTrace(level=1/0)>\n",
|
|
||||||
"pm_forces particles are Traced<ShapedArray(int32[16777216,3])>with<DynamicJaxprTrace(level=1/0)>\n",
|
|
||||||
"shape of displacement: (256, 256, 256)\n",
|
|
||||||
"pm_forces particles are Traced<ShapedArray(float32[16777216,3])>with<DynamicJaxprTrace(level=2/0)>\n"
|
|
||||||
]
|
|
||||||
}
|
|
||||||
],
|
|
||||||
"source": [
|
|
||||||
"res = run_simulation(0.25, 0.8)\n",
|
|
||||||
"#%timeit res = run_simulation(0.25, 0.8)"
|
|
||||||
]
|
|
||||||
},
|
|
||||||
{
|
|
||||||
"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": [
|
|
||||||
{
|
|
||||||
"name": "stdout",
|
|
||||||
"output_type": "stream",
|
|
||||||
"text": [
|
|
||||||
"shape of grid_mesh: (256, 256, 256)\n"
|
|
||||||
]
|
|
||||||
}
|
|
||||||
],
|
|
||||||
"source": [
|
|
||||||
"figure(figsize=[10,5])\n",
|
|
||||||
"subplot(121)\n",
|
|
||||||
"imshow(cic_paint(jnp.zeros(mesh_shape), res[0]).sum(axis=0),cmap='magma')\n",
|
|
||||||
"subplot(122)\n",
|
|
||||||
"imshow(cic_paint(jnp.zeros(mesh_shape), res[1]).sum(axis=0),cmap='magma')"
|
|
||||||
]
|
|
||||||
},
|
|
||||||
{
|
|
||||||
"cell_type": "code",
|
|
||||||
"execution_count": null,
|
|
||||||
"id": "4e050871",
|
|
||||||
"metadata": {
|
|
||||||
"id": "4e050871"
|
|
||||||
},
|
|
||||||
"outputs": [],
|
|
||||||
"source": []
|
|
||||||
}
|
|
||||||
],
|
|
||||||
"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
|
|
||||||
}
|
|
||||||
38
notebooks/visualize.py
Normal file
38
notebooks/visualize.py
Normal file
|
|
@ -0,0 +1,38 @@
|
||||||
|
import numpy as np
|
||||||
|
import jax.numpy as jnp
|
||||||
|
import matplotlib.pyplot as plt
|
||||||
|
|
||||||
|
def plot_fields(fields_dict, sum_over=None):
|
||||||
|
"""
|
||||||
|
Plots sum projections of 3D fields along different axes,
|
||||||
|
slicing only the first `sum_over` elements along each axis.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
- fields: list of 3D arrays representing fields to plot
|
||||||
|
- names: list of names for each field, used in titles
|
||||||
|
- sum_over: number of slices to sum along each axis (default: fields[0].shape[0] // 8)
|
||||||
|
"""
|
||||||
|
sum_over = sum_over or list(fields_dict.values())[0].shape[0] // 8
|
||||||
|
nb_rows = len(fields_dict)
|
||||||
|
nb_cols = 3
|
||||||
|
fig, axes = plt.subplots(nb_rows, nb_cols, figsize=(15, 5 * nb_rows))
|
||||||
|
|
||||||
|
def plot_subplots(proj_axis, field, row, title):
|
||||||
|
slicing = [slice(None)] * field.ndim
|
||||||
|
slicing[proj_axis] = slice(None, sum_over)
|
||||||
|
slicing = tuple(slicing)
|
||||||
|
|
||||||
|
# Sum projection over the specified axis and plot
|
||||||
|
axes[row, proj_axis].imshow(field[slicing].sum(axis=proj_axis) + 1,
|
||||||
|
cmap='magma', extent=[0, field.shape[proj_axis], 0, field.shape[proj_axis]])
|
||||||
|
axes[row, proj_axis].set_xlabel('Mpc/h')
|
||||||
|
axes[row, proj_axis].set_ylabel('Mpc/h')
|
||||||
|
axes[row, proj_axis].set_title(title)
|
||||||
|
|
||||||
|
# Plot each field across the three axes
|
||||||
|
for i, (name, field) in enumerate(fields_dict.items()):
|
||||||
|
for proj_axis in range(3):
|
||||||
|
plot_subplots(proj_axis, field, i, f"{name} projection {proj_axis}")
|
||||||
|
|
||||||
|
plt.tight_layout()
|
||||||
|
plt.show()
|
||||||
Loading…
Add table
Reference in a new issue