guilhem_lavaux/JaxPM
0
0
Fork 1
mirror of https://github.com/DifferentiableUniverseInitiative/JaxPM.git synced 2025-02-21 09:17:12 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions
JAX-powered Cosmological Particle-Mesh N-body Solver
106 commits 7 branches 5 tags 132 MiB
Jupyter Notebook 99.6%
Python 0.4%
Find a file
Exact
Wassim KABALAN cb2a7ab17f
Fix pfft gradients (#34) 
* update jaxdecomp version and test gradients

* Prepare for DTO tests

* format
2024-12-22 12:47:42 -05:00
.github/workflows Pypi upload compatible version (#33) 2024-12-21 11:47:13 -05:00
jaxpm jaxdecomp proto (#21) 2024-12-20 05:44:02 -05:00
notebooks fixing formatting of notebook 2024-12-21 13:14:42 -05:00
tests Fix pfft gradients (#34) 2024-12-22 12:47:42 -05:00
.all-contributorsrc docs: update .all-contributorsrc [skip ci] 2024-12-07 19:54:33 +00:00
.gitignore Pypi upload compatible version (#33) 2024-12-21 11:47:13 -05:00
.pre-commit-config.yaml jaxdecomp proto (#21) 2024-12-20 05:44:02 -05:00
LICENSE Pypi upload compatible version (#33) 2024-12-21 11:47:13 -05:00
MANIFEST.in Pypi upload compatible version (#33) 2024-12-21 11:47:13 -05:00
pyproject.toml Fix pfft gradients (#34) 2024-12-22 12:47:42 -05:00
pytest.ini jaxdecomp proto (#21) 2024-12-20 05:44:02 -05:00
README.md minor typo fix 2024-12-21 15:28:20 -05:00
requirements-test.txt Pypi upload compatible version (#33) 2024-12-21 11:47:13 -05:00

README.md

JaxPM

Notebook PyPI version Tests All Contributors

JAX-powered Cosmological Particle-Mesh N-body Solver

Note

The new JaxPM v0.1.xx supports multi-GPU model distribution while remaining compatible with previous releases. These significant changes are still under development and testing, so please report any issues you encounter. For the older but more stable version, install:

pip install jaxpm==0.0.2

Install

Basic installation can be done using pip:

pip install jaxpm

For more advanced installation for optimized distribution on gpu clusters, please install jaxDecomp first. See instructions here.

Goals

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

  • Keep implementation simple and readable, in pure NumPy API
  • Any order forward and backward automatic differentiation
  • Support automated batching using vmap
  • Compatibility with external optimizer libraries like optax
  • Now fully distributable on multi-GPU and multi-node systems using jaxDecomp working withJAX v0.4.35

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.

Getting Started

To dive into JaxPM’s capabilities, please explore the notebook section for detailed tutorials and examples on various setups, from single-device simulations to multi-host configurations. You can find the notebooks' README here for a structured guide through each tutorial.

Contributors ✨

Thanks goes to these wonderful people (emoji key):

Francois Lanusse
Francois Lanusse
🤔		 Denise Lanzieri
Denise Lanzieri
💻		 Wassim KABALAN
Wassim KABALAN
💻 🚇 👀		 Hugo Simon-Onfroy
Hugo Simon-Onfroy
💻		 Alexandre Boucaud
Alexandre Boucaud
👀

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