forked from guilhem_lavaux/JaxPM
df8602b318
* adding example of distributed solution * put back old functgion * update formatting * add halo exchange and slice pad * apply formatting * implement distributed optimized cic_paint * Use new cic_paint with halo * Fix seed for distributed normal * Wrap interpolation function to avoid all gather * Return normal order frequencies for single GPU * add example * format * add optimised bench script * times in ms * add lpt2 * update benchmark and add slurm * Visualize only final field * Update scripts/distributed_pm.py Co-authored-by: Francois Lanusse <EiffL@users.noreply.github.com> * Adjust pencil type for frequencies * fix painting issue with slabs * Shared operation in fourrier space now take inverted sharding axis for slabs * add assert to make pyright happy * adjust test for hpc-plotter * add PMWD test * bench * format * added github workflow * fix formatting from main * Update for jaxDecomp pure JAX * revert single halo extent change * update for latest jaxDecomp * remove fourrier_space in autoshmap * make normal_field work with single controller * format * make distributed pm work in single controller * merge bench_pm * update to leapfrog * add a strict dependency on jaxdecomp * global mesh no longer needed * kernels.py no longer uses global mesh * quick fix in distributed * pm.py no longer uses global mesh * painting.py no longer uses global mesh * update demo script * quick fix in kernels * quick fix in distributed * update demo * merge hugos LPT2 code * format * Small fix * format * remove duplicate get_ode_fn * update visualizer * update compensate CIC * By default check_rep is false for shard_map * remove experimental distributed code * update PGDCorrection and neural ode to use new fft3d * jaxDecomp pfft3d promotes to complex automatically * remove deprecated stuff * fix painting issue with read_cic * use jnp interp instead of jc interp * delete old slurms * add notebook examples * apply formatting * add distributed zeros * fix code in LPT2 * jit cic_paint * update notebooks * apply formating * get local shape and zeros can be used by users * add a user facing function to create uniform particle grid * use jax interp instead of jax_cosmo * use float64 for enmeshing * Allow applying weights with relative cic paint * Weights can be traced * remove script folder * update example notebooks * delete outdated design file * add readme for tutorials * update readme * fix small error * forgot particles in multi host * clarifying why cic_paint_dx is slower * clarifying the halo size dependence on the box size * ability to choose snapshots number with MultiHost script * Adding animation notebook * Put plotting in package * Add finite difference laplace kernel + powerspec functions from Hugo Co-authored-by: Hugo Simonfroy <hugo.simonfroy@gmail.com> * Put plotting utils in package * By default use absoulute painting with * update code * update notebooks * add tests * Upgrade setup.py to pyproject * Format * format tests * update test dependencies * add test workflow * fix deprecated FftType in jaxpm.kernels * Add aboucaud comments * JAX version is 0.4.35 until Diffrax new release * add numpy explicitly as dependency for tests * fix install order for tests * add numpy to be installed * enforce no build isolation for fastpm * pip install jaxpm test without build isolation * bump jaxdecomp version * revert test workflow * remove outdated tests --------- Co-authored-by: EiffL <fr.eiffel@gmail.com> Co-authored-by: Francois Lanusse <EiffL@users.noreply.github.com> Co-authored-by: Wassim KABALAN <wassim@apc.in2p3.fr> Co-authored-by: Hugo Simonfroy <hugo.simonfroy@gmail.com> Former-commit-id: 8c2e823d4669eac712089bf7f85ffb7912e8232d
155 lines
5 KiB
Python
155 lines
5 KiB
Python
import pytest
|
|
from diffrax import Dopri5, ODETerm, PIDController, SaveAt, diffeqsolve
|
|
from helpers import MSE, MSRE
|
|
from jax import numpy as jnp
|
|
|
|
from jaxpm.distributed import uniform_particles
|
|
from jaxpm.painting import cic_paint, cic_paint_dx
|
|
from jaxpm.pm import lpt, make_diffrax_ode
|
|
from jaxpm.utils import power_spectrum
|
|
|
|
_TOLERANCE = 1e-4
|
|
_PM_TOLERANCE = 1e-3
|
|
|
|
|
|
@pytest.mark.single_device
|
|
@pytest.mark.parametrize("order", [1, 2])
|
|
def test_lpt_absolute(simulation_config, initial_conditions, lpt_scale_factor,
|
|
fpm_lpt1_field, fpm_lpt2_field, cosmo, order):
|
|
|
|
mesh_shape, box_shape = simulation_config
|
|
cosmo._workspace = {}
|
|
particles = uniform_particles(mesh_shape)
|
|
|
|
# Initial displacement
|
|
dx, _, _ = lpt(cosmo,
|
|
initial_conditions,
|
|
particles,
|
|
a=lpt_scale_factor,
|
|
order=order)
|
|
|
|
fpm_ref_field = fpm_lpt1_field if order == 1 else fpm_lpt2_field
|
|
|
|
lpt_field = cic_paint(jnp.zeros(mesh_shape), particles + dx)
|
|
_, jpm_ps = power_spectrum(lpt_field, box_shape=box_shape)
|
|
_, fpm_ps = power_spectrum(fpm_ref_field, box_shape=box_shape)
|
|
|
|
assert MSE(lpt_field, fpm_ref_field) < _TOLERANCE
|
|
assert MSRE(jpm_ps, fpm_ps) < _TOLERANCE
|
|
|
|
|
|
@pytest.mark.single_device
|
|
@pytest.mark.parametrize("order", [1, 2])
|
|
def test_lpt_relative(simulation_config, initial_conditions, lpt_scale_factor,
|
|
fpm_lpt1_field, fpm_lpt2_field, cosmo, order):
|
|
|
|
mesh_shape, box_shape = simulation_config
|
|
cosmo._workspace = {}
|
|
# Initial displacement
|
|
dx, _, _ = lpt(cosmo, initial_conditions, a=lpt_scale_factor, order=order)
|
|
|
|
lpt_field = cic_paint_dx(dx)
|
|
|
|
fpm_ref_field = fpm_lpt1_field if order == 1 else fpm_lpt2_field
|
|
|
|
_, jpm_ps = power_spectrum(lpt_field, box_shape=box_shape)
|
|
_, fpm_ps = power_spectrum(fpm_ref_field, box_shape=box_shape)
|
|
|
|
assert MSE(lpt_field, fpm_ref_field) < _TOLERANCE
|
|
assert MSRE(jpm_ps, fpm_ps) < _TOLERANCE
|
|
|
|
|
|
@pytest.mark.single_device
|
|
@pytest.mark.parametrize("order", [1, 2])
|
|
def test_nbody_absolute(simulation_config, initial_conditions,
|
|
lpt_scale_factor, nbody_from_lpt1, nbody_from_lpt2,
|
|
cosmo, order):
|
|
|
|
mesh_shape, box_shape = simulation_config
|
|
cosmo._workspace = {}
|
|
particles = uniform_particles(mesh_shape)
|
|
|
|
# Initial displacement
|
|
dx, p, _ = lpt(cosmo,
|
|
initial_conditions,
|
|
particles,
|
|
a=lpt_scale_factor,
|
|
order=order)
|
|
|
|
ode_fn = ODETerm(make_diffrax_ode(cosmo, mesh_shape))
|
|
|
|
solver = Dopri5()
|
|
controller = PIDController(rtol=1e-8,
|
|
atol=1e-8,
|
|
pcoeff=0.4,
|
|
icoeff=1,
|
|
dcoeff=0)
|
|
|
|
saveat = SaveAt(t1=True)
|
|
|
|
y0 = jnp.stack([particles + dx, p])
|
|
|
|
solutions = diffeqsolve(ode_fn,
|
|
solver,
|
|
t0=lpt_scale_factor,
|
|
t1=1.0,
|
|
dt0=None,
|
|
y0=y0,
|
|
stepsize_controller=controller,
|
|
saveat=saveat)
|
|
|
|
final_field = cic_paint(jnp.zeros(mesh_shape), solutions.ys[-1, 0])
|
|
|
|
fpm_ref_field = nbody_from_lpt1 if order == 1 else nbody_from_lpt2
|
|
|
|
_, jpm_ps = power_spectrum(final_field, box_shape=box_shape)
|
|
_, fpm_ps = power_spectrum(fpm_ref_field, box_shape=box_shape)
|
|
|
|
assert MSE(final_field, fpm_ref_field) < _PM_TOLERANCE
|
|
assert MSRE(jpm_ps, fpm_ps) < _PM_TOLERANCE
|
|
|
|
|
|
@pytest.mark.single_device
|
|
@pytest.mark.parametrize("order", [1, 2])
|
|
def test_nbody_relative(simulation_config, initial_conditions,
|
|
lpt_scale_factor, nbody_from_lpt1, nbody_from_lpt2,
|
|
cosmo, order):
|
|
|
|
mesh_shape, box_shape = simulation_config
|
|
cosmo._workspace = {}
|
|
|
|
# Initial displacement
|
|
dx, p, _ = lpt(cosmo, initial_conditions, a=lpt_scale_factor, order=order)
|
|
|
|
ode_fn = ODETerm(
|
|
make_diffrax_ode(cosmo, mesh_shape, paint_absolute_pos=False))
|
|
|
|
solver = Dopri5()
|
|
controller = PIDController(rtol=1e-9,
|
|
atol=1e-9,
|
|
pcoeff=0.4,
|
|
icoeff=1,
|
|
dcoeff=0)
|
|
|
|
saveat = SaveAt(t1=True)
|
|
|
|
y0 = jnp.stack([dx, p])
|
|
|
|
solutions = diffeqsolve(ode_fn,
|
|
solver,
|
|
t0=lpt_scale_factor,
|
|
t1=1.0,
|
|
dt0=None,
|
|
y0=y0,
|
|
stepsize_controller=controller,
|
|
saveat=saveat)
|
|
|
|
final_field = cic_paint_dx(solutions.ys[-1, 0])
|
|
|
|
fpm_ref_field = nbody_from_lpt1 if order == 1 else nbody_from_lpt2
|
|
|
|
_, jpm_ps = power_spectrum(final_field, box_shape=box_shape)
|
|
_, fpm_ps = power_spectrum(fpm_ref_field, box_shape=box_shape)
|
|
|
|
assert MSE(final_field, fpm_ref_field) < _PM_TOLERANCE
|
|
assert MSRE(jpm_ps, fpm_ps) < _PM_TOLERANCE
|