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

format tests

Browse source
This commit is contained in:
Wassim Kabalan 2024-12-08 22:46:20 +01:00
parent af29c4005d
commit 97f39bd051
3 changed files with 94 additions and 76 deletions
Download patch file Download diff file Expand all files Collapse all files

98
tests/test_against_fpm.py
View file

@ -1,21 +1,21 @@
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.distributed import uniform_particles
from helpers import MSE , MSRE
from diffrax import diffeqsolve, ODETerm, SaveAt, PIDController, Dopri5
from jaxpm.utils import power_spectrum
_TOLERANCE = 1e-4
_PM_TOLERANCE = 1e-3
@pytest.mark.parametrize("order", [1 , 2])
def test_lpt_absoulute(simulation_config, initial_conditions, lpt_scale_factor,
fpm_lpt1_field,fpm_lpt2_field, cosmo , order):
@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 = {}
@ -31,17 +31,18 @@ def test_lpt_absoulute(simulation_config, initial_conditions, lpt_scale_factor,
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)
_, 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):
fpm_lpt1_field, fpm_lpt2_field, cosmo, order):
mesh_shape, box_shape = simulation_config
cosmo._workspace = {}
# Initial displacement
@ -51,26 +52,31 @@ def test_lpt_relative(simulation_config, initial_conditions, lpt_scale_factor,
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)
_, 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):
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)
dx, p, _ = lpt(cosmo,
initial_conditions,
particles,
a=lpt_scale_factor,
order=order)
ode_fn = ODETerm(
make_diffrax_ode(cosmo, mesh_shape))
ode_fn = ODETerm(make_diffrax_ode(cosmo, mesh_shape))
solver = Dopri5()
controller = PIDController(rtol=1e-8,
@ -84,33 +90,36 @@ def test_nbody_absolute(simulation_config, initial_conditions, lpt_scale_factor,
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)
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])
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)
_, 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):
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)
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))
@ -127,23 +136,20 @@ def test_nbody_relative(simulation_config, initial_conditions, lpt_scale_factor,
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)
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])
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)
_, 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