How PM++ Works#
This section describes the algorithm implemented in src/pmpp, from the
cosmological parameters and random modes to a differentiable, distributed
particle-mesh simulation. It develops the equations together with the array
layouts and discrete operators that evaluate them
[Hockney and Eastwood, 1988, Li et al., 2024].
The chapters first define the code boundaries, then follow the forward simulation and its reverse pass:
Python package layout maps implementation modules to their scientific and runtime responsibilities.
System architecture defines the state, units, JAX program structure, and end-to-end dataflow.
Initial modes and LPT constructs the transfer and growth tables, Gaussian density modes, and first- or second-order particle initial conditions.
Particle-mesh force derives CIC assignment, the periodic Poisson solve, spectral forces, the Pallas kernels, and their transposes.
Distributed runtime explains slab ownership, mesh halos, particle migration, distributed real FFTs, and static-capacity invariants.
Integration and discrete adjoint derives the growth-matched drift and kick factors and the reverse-time custom VJP.
Optional CUDA routing documents the typed FFI boundary, route records, stable merge, and route transpose.
The underlying PM and discrete-adjoint formulation follows Li et al. [2024]. PM++ keeps those mathematical operators while adding distributed particle ownership, mesh-halo communication, distributed FFT layouts, and optional local CUDA routing.
Conventions#
The simulation is periodic. Bold lower-case symbols denote vectors, hats
denote Fourier coefficients, and a prime denotes differentiation with respect
to \(\ln a\). The discrete Fourier transforms use the default normalization of
jax.numpy.fft.rfftn and
jax.numpy.fft.irfftn. The forward transform is unnormalized and
the inverse transform carries \(1/N_\mathrm{mesh}\).
PM++ separates three kinds of quantities:
physical state, such as cosmological parameters, particle displacement, velocity, and acceleration
discrete representation, such as mesh indices, Fourier layouts, masks, and padded particle slots
static program structure, such as shapes, the device mesh, capacities, and selected kernels.
Keeping these distinctions explicit is important. A mathematical field may be unchanged while its sharding changes, and a particle may change owner while its physical trajectory remains continuous.