SuperResPM is a differentiable, FastPM-style cosmological solver implemented with Diffrax.
The goal is to make it a cosmology-ready code for field-level inferences.
- FastPM-like leapfrog solver using
diffrax - Differentiable with reversible adjoint for efficient gradient computation
- Compatible with
jaxpmforces andjax_cosmocosmology
git clone --recurse-submodules -j8 https://github.com/chto/SuperResPM.git
cd SuperResPM
sh install.sh
See notebook.