pmwd is a differentiable cosmological particle-mesh forward model. The C2 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
pip install -e . # to install in editable/develop modeSee docs/examples.