guilhem_lavaux/JaxPM
0
0
Fork 1
mirror of https://github.com/DifferentiableUniverseInitiative/JaxPM.git synced 2025-05-15 04:21:12 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

apply formatting

Browse source
This commit is contained in:
Wassim KABALAN 2024-10-27 00:52:14 +02:00
parent 11f7e90066
commit 4342279817
5 changed files with 26 additions and 22 deletions
Download patch file Download diff file Expand all files Collapse all files

28
notebooks/03-MultiHost_PM.py
View file

@ -7,21 +7,21 @@ jax.distributed.initialize()
rank = jax.process_index()
size = jax.process_count()
from functools import partial
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
import numpy as np
from diffrax import (ConstantStepSize, LeapfrogMidpoint, ODETerm, SaveAt,
diffeqsolve)
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
from jaxpm.kernels import interpolate_power_spectrum
from jaxpm.painting import cic_paint_dx
from jaxpm.pm import linear_field, lpt, make_ode_fn
all_gather = partial(process_allgather, tiled=True)
@ -84,7 +84,8 @@ def run_simulation(omega_c, sigma8):
return initial_conditions, dx, res.ys, res.stats
initial_conditions , lpt_displacements , ode_solutions , solver_stats = run_simulation(0.25, 0.8)
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}")
@ -95,10 +96,9 @@ 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)
initial_conditions=initial_conditions,
lpt_displacements=lpt_displacements,
ode_solutions=ode_solutions,
solver_stats=solver_stats)
print(f"[{rank}] Simulation results saved")

12
notebooks/visualize.py
View file

@ -1,6 +1,7 @@
import numpy as np
import jax.numpy as jnp
import matplotlib.pyplot as plt
import numpy as np
def plot_fields(fields_dict, sum_over=None):
"""
@ -23,8 +24,10 @@ def plot_fields(fields_dict, sum_over=None):
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].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)
@ -32,7 +35,8 @@ def plot_fields(fields_dict, sum_over=None):
# 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}")
plot_subplots(proj_axis, field, i,
f"{name} projection {proj_axis}")
plt.tight_layout()
plt.show()