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.

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

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

See docs/examples.