Fix pfft gradients (#34)

* update jaxdecomp version and test gradients

* Prepare for DTO tests

* format

README.md

@@ -1,9 +1,17 @@
 # JaxPM
+[![Notebook](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/DifferentiableUniverseInitiative/JaxPM/blob/main/notebooks/01-Introduction.ipynb)
 [![PyPI version](https://img.shields.io/pypi/v/jaxpm)](https://pypi.org/project/jaxpm/) [![Tests](https://github.com/DifferentiableUniverseInitiative/JaxPM/actions/workflows/tests.yml/badge.svg)](https://github.com/DifferentiableUniverseInitiative/JaxPM/actions/workflows/tests.yml) <!-- ALL-CONTRIBUTORS-BADGE:START - Do not remove or modify this section -->
 [![All Contributors](https://img.shields.io/badge/all_contributors-5-orange.svg?style=flat-square)](#contributors-)
 <!-- ALL-CONTRIBUTORS-BADGE:END -->
 JAX-powered Cosmological Particle-Mesh N-body Solver
 
+> ### Note
+> **The new JaxPM v0.1.xx** supports multi-GPU model distribution while remaining compatible with previous releases. These significant changes are still under development and testing, so please report any issues you encounter.
+> For the older but more stable version, install:
+> ```bash
+> pip install jaxpm==0.0.2
+> ```
+
 ## Install
 
 Basic installation can be done using pip:

pyproject.toml

@@ -11,7 +11,7 @@ readme = "README.md"
 requires-python = ">=3.9"
 license = { file = "LICENSE" }
 urls = { "Homepage" = "https://github.com/DifferentiableUniverseInitiative/JaxPM" }
-dependencies = ["jax_cosmo", "jax>=0.4.30", "jaxdecomp>=0.2.2"]
+dependencies = ["jax_cosmo", "jax>=0.4.35", "jaxdecomp>=0.2.3"]
 
 [tool.setuptools]
 packages = ["jaxpm"]

tests/test_gradients.py

@@ -0,0 +1,87 @@
+import jax
+import pytest
+from diffrax import (BacksolveAdjoint, Dopri5, ODETerm, PIDController,
+                     RecursiveCheckpointAdjoint, SaveAt, diffeqsolve)
+from helpers import MSE
+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
+
+
+@pytest.mark.single_device
+@pytest.mark.parametrize("order", [1, 2])
+@pytest.mark.parametrize("absolute_painting", [True, False])
+@pytest.mark.parametrize("adjoint", ['DTO', 'OTD'])
+def test_nbody_grad(simulation_config, initial_conditions, lpt_scale_factor,
+                    nbody_from_lpt1, nbody_from_lpt2, cosmo, order,
+                    absolute_painting, adjoint):
+
+    mesh_shape, _ = simulation_config
+    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.grad
+    def forward_model(initial_conditions, cosmo):
+
+        # Initial displacement
+        if absolute_painting:
+            particles = uniform_particles(mesh_shape)
+            dx, p, _ = lpt(cosmo,
+                           initial_conditions,
+                           particles,
+                           a=lpt_scale_factor,
+                           order=order)
+            ode_fn = ODETerm(make_diffrax_ode(cosmo, mesh_shape))
+            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,
+                                solver,
+                                t0=lpt_scale_factor,
+                                t1=1.0,
+                                dt0=None,
+                                y0=y0,
+                                adjoint=adjoint,
+                                stepsize_controller=controller,
+                                saveat=saveat)
+
+        if absolute_painting:
+            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)
+
+    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