.. image:: assets/logo.svg?raw=true :width: 42 em :align: center :alt: logo
pmwd
is a differentiable cosmological particle-mesh forward model.
The C\ :sub:2
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.
.. raw:: html
Optimizing the initial conditions by gradient descent to make some interesting projected patterns.
.. raw:: html
.. code:: sh
pip install -e . # to install in editable/develop mode
.. pip install pmwd pip install -e .[dev] # to install development dependencies
See docs/examples <docs/examples>
_.
.. code:: sh
XLA_PYTHON_CLIENT_MEM_FRACTION=.05 python -m pytest --cov --cov-report=term-missing:skip-covered --durations=5 -n 16
where XLA_PYTHON_CLIENT_MEM_FRACTION=.05
makes JAX preallocate 5% of
currently-available GPU memory, instead of the default 90%.
.. code:: sh
CUDA_VISIBLE_DEVICES= python -m pytest --cov --cov-report=term-missing:skip-covered --durations=5 -n 16
disables CUDA (to run tests on CPUs).
.. code:: sh
python -m pytest --durations=5 --benchmark-columns=mean,ops,rounds,iterations tests/benchmark.py
We refer the users to the following references for ... Please cite the following papers:
.. code:: bibtex
.. include:: CITATIONS.bib