PM++ documentation#
PM++ is a differentiable particle-mesh cosmology simulator built with JAX. It supports the complete path from Gaussian initial modes to Lagrangian perturbation theory (LPT), N-body evolution, density fields, summary statistics, and gradients. Its distributed runtime scales the same model across a one-dimensional mesh of accelerators.
For cosmologists, this means that larger volumes or higher particle counts can be simulated without giving up derivatives with respect to the initial modes and cosmological parameters. Those derivatives enable gradient-based reconstruction, inference, and sensitivity studies at scales that can be difficult to fit on one accelerator.
Note
PM++ is designed primarily for multi-accelerator simulations. For single-GPU work, PMWD is a closely related alternative. Some maintained tests directly validate matched PM++ and PMWD forward and gradient calculations. For the configurations covered by those tests, PM++ agrees with PMWD down to machine precision.
Build a multi-device mesh, evolve a small periodic box, and check mass conservation in about ten minutes.
Connect the equations to the arrays, JAX transforms, communication, and custom adjoints used by PM++.
Study reproducible examples pre-executed on multiple GPUs.
Choose a path#
New user: install PM++, then run the first simulation.
Scientific user: use the user guide to choose initial conditions, schedules, analysis products, and gradient targets.
Method developer: read How PM++ works, then consult the API reference.