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

don't send static cosmo into ODETerm

Browse source
This commit is contained in:
Wassim Kabalan 2025-01-21 11:18:00 +01:00
parent 6ab26ea1ec
commit 6aa2ccf7b5
5 changed files with 9 additions and 9 deletions
Download patch file Download diff file Expand all files Collapse all files

4
tests/test_gradients.py
View file

@ -39,7 +39,7 @@ def test_nbody_grad(simulation_config, initial_conditions, lpt_scale_factor,
particles,
a=lpt_scale_factor,
order=order)
ode_fn = ODETerm(make_diffrax_ode(cosmo, mesh_shape))
ode_fn = ODETerm(make_diffrax_ode(mesh_shape))
y0 = jnp.stack([particles + dx, p])
else:
@ -48,7 +48,7 @@ def test_nbody_grad(simulation_config, initial_conditions, lpt_scale_factor,
a=lpt_scale_factor,
order=order)
ode_fn = ODETerm(
make_diffrax_ode(cosmo, mesh_shape, paint_absolute_pos=False))
make_diffrax_ode(mesh_shape, paint_absolute_pos=False))
y0 = jnp.stack([dx, p])
solver = Dopri5()