particle mesh with derivatives

pmwd is a differentiable cosmological particle-mesh forward model. The C₂ symmetry of the name symbolizes the reversibility of the model, which helps to dramatically reduce the memory cost when used with the adjoint method. Based on JAX, pmwd is fully differentiable, and is highly performant on GPUs.

Particles align on the initial grid after evolving forward and then backward in time.

time_evolution.mp4

Optimizing the initial conditions by gradient descent to make some interesting projected patterns.

pmwd_optim.mp4

Installation

pip install -e .  # to install in editable/develop mode

Examples

See docs/examples.