JaxPM

JAX-powered Cosmological Particle-Mesh N-body Solver

This project is currently in an early design phase. All inputs are welcome on the design document

Goals

Provide a modern infrastructure to support differentiable PM N-body simulations using JAX:

• Keep implementation simple and readable, in pure NumPy API
• Transparent distribution using builtin xmap
• Any order forward and backward automatic differentiation
• Support automated batching using vmap
• Compatibility with external optimizer libraries like optax

Open development and use

Current expectations are:

• This project is and will remain open source, and usable without any restrictions for any purposes
• Will be a simple publication on The Journal of Open Source Software
• Everyone is welcome to contribute, and can join the JOSS publication (until it is submitted to the journal).
• Anyone (including main contributors) can use this code as a framework to build and publish their own applications, with no expectation that they need to extend authorship to all jaxpm developers.

Contributors ✨

Thanks goes to these wonderful people (emoji key):

Francois Lanusse 🤔

Denise Lanzieri 💻

Wassim KABALAN 💻 🚇 👀

Hugo Simon-Onfroy 💻

This project follows the all-contributors specification. Contributions of any kind welcome!