guilhem_lavaux/JaxPM
0
0
Fork 1
mirror of https://github.com/DifferentiableUniverseInitiative/JaxPM.git synced 2025-02-23 10:00:54 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Prepare for DTO tests

Browse source
This commit is contained in:
Wassim Kabalan 2024-12-21 23:24:19 +01:00
parent b132a0e2aa
commit a924458f0d
1 changed files with 56 additions and 86 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

122
tests/test_gradients.py
View file

@ -1,115 +1,85 @@
import jax
import pytest import pytest
from diffrax import Dopri5, ODETerm, PIDController, SaveAt, diffeqsolve from diffrax import Dopri5, ODETerm, PIDController, SaveAt, diffeqsolve , RecursiveCheckpointAdjoint, BacksolveAdjoint
from helpers import MSE from helpers import MSE
from jax import numpy as jnp from jax import numpy as jnp
from jaxpm.distributed import uniform_particles from jaxpm.distributed import uniform_particles
from jaxpm.painting import cic_paint, cic_paint_dx from jaxpm.painting import cic_paint, cic_paint_dx
from jaxpm.pm import lpt, make_diffrax_ode from jaxpm.pm import lpt, make_diffrax_ode
import jax
@pytest.mark.single_device @pytest.mark.single_device
@pytest.mark.parametrize("order", [1, 2]) @pytest.mark.parametrize("order", [1, 2])
def test_grad_relative(simulation_config, initial_conditions, @pytest.mark.parametrize("absolute_painting", [True, False])
lpt_scale_factor, nbody_from_lpt1, nbody_from_lpt2, @pytest.mark.parametrize("adjoint", ['DTO', 'OTD'])
cosmo, order): def test_nbody_grad(simulation_config, initial_conditions, lpt_scale_factor,
nbody_from_lpt1, nbody_from_lpt2, cosmo, order,
mesh_shape, _ = simulation_config absolute_painting , adjoint):
cosmo._workspace = {}
@jax.jit
@jax.grad
def forward_model(initial_conditions, cosmo):
# 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-7,
atol=1e-7,
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])
return MSE(final_field, nbody_from_lpt1 if order == 1 else nbody_from_lpt2)
bad_initial_conditions = initial_conditions + jax.random.normal(jax.random.PRNGKey(0), initial_conditions.shape) * 0.5
best_ic = forward_model(initial_conditions , cosmo)
bad_ic = forward_model(bad_initial_conditions, cosmo)
assert jnp.max(best_ic) < 1e-5
assert jnp.max(bad_ic) > 1e-5
@pytest.mark.single_device
@pytest.mark.parametrize("order", [1, 2])
def test_grad_absolute(simulation_config, initial_conditions,
lpt_scale_factor, nbody_from_lpt1, nbody_from_lpt2,
cosmo, order):
mesh_shape, _ = simulation_config mesh_shape, _ = simulation_config
cosmo._workspace = {} cosmo._workspace = {}
if adjoint == 'OTD':
pytest.skip("OTD adjoint not implemented yet (needs PFFT3D JVP)")
adjoint = RecursiveCheckpointAdjoint() if adjoint == 'DTO' else BacksolveAdjoint(solver=Dopri5())
@jax.jit @jax.jit
@jax.grad @jax.grad
def forward_model(initial_conditions, cosmo): def forward_model(initial_conditions, cosmo):
# Initial displacement # Initial displacement
if absolute_painting:
particles = uniform_particles(mesh_shape) particles = uniform_particles(mesh_shape)
dx, p, _ = lpt(cosmo, initial_conditions,particles, a=lpt_scale_factor, order=order) dx, p, _ = lpt(cosmo,
initial_conditions,
ode_fn = ODETerm( particles,
make_diffrax_ode(cosmo, mesh_shape, paint_absolute_pos=True)) a=lpt_scale_factor,
order=order)
solver = Dopri5() ode_fn = ODETerm(make_diffrax_ode(cosmo, mesh_shape))
controller = PIDController(rtol=1e-7,
atol=1e-7,
pcoeff=0.4,
icoeff=1,
dcoeff=0)
saveat = SaveAt(t1=True)
y0 = jnp.stack([particles + dx, p]) y0 = jnp.stack([particles + dx, p])
else:
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))
y0 = jnp.stack([dx, p])
solver = Dopri5()
controller = PIDController(rtol=1e-7,
atol=1e-7,
pcoeff=0.4,
icoeff=1,
dcoeff=0)
saveat = SaveAt(t1=True)
solutions = diffeqsolve(ode_fn, solutions = diffeqsolve(ode_fn,
solver, solver,
t0=lpt_scale_factor, t0=lpt_scale_factor,
t1=1.0, t1=1.0,
dt0=None, dt0=None,
y0=y0, y0=y0,
adjoint=adjoint,
stepsize_controller=controller, stepsize_controller=controller,
saveat=saveat) saveat=saveat)
if absolute_painting:
final_field = cic_paint(jnp.zeros(mesh_shape), solutions.ys[-1, 0]) final_field = cic_paint(jnp.zeros(mesh_shape), solutions.ys[-1, 0])
else:
final_field = cic_paint_dx(solutions.ys[-1, 0])
return MSE(final_field, nbody_from_lpt1 if order == 1 else nbody_from_lpt2) return MSE(final_field,
nbody_from_lpt1 if order == 1 else nbody_from_lpt2)
bad_initial_conditions = initial_conditions + jax.random.normal(
bad_initial_conditions = initial_conditions + jax.random.normal(jax.random.PRNGKey(0), initial_conditions.shape) * 0.5 jax.random.PRNGKey(0), initial_conditions.shape) * 0.5
best_ic = forward_model(initial_conditions , cosmo) best_ic = forward_model(initial_conditions, cosmo)
bad_ic = forward_model(bad_initial_conditions, cosmo) bad_ic = forward_model(bad_initial_conditions, cosmo)
assert jnp.max(best_ic) < 1e-5 assert jnp.max(best_ic) < 1e-5
assert jnp.max(bad_ic) > 1e-5 assert jnp.max(bad_ic) > 1e-5