guilhem_lavaux/JaxPM
0
0
Fork 1
mirror of https://github.com/DifferentiableUniverseInitiative/JaxPM.git synced 2025-04-17 00:20:55 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

update readme

Browse source
This commit is contained in:
Wassim KABALAN 2024-10-30 01:59:33 +01:00
parent f0c43f8d76
commit a067954d0a
1 changed files with 6 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
README.md
View file

@ -4,16 +4,15 @@
<!-- ALL-CONTRIBUTORS-BADGE:END --> <!-- ALL-CONTRIBUTORS-BADGE:END -->
JAX-powered Cosmological Particle-Mesh N-body Solver 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](https://github.com/DifferentiableUniverseInitiative/JaxPM/blob/main/design.md)**
## Goals ## Goals
Provide a modern infrastructure to support differentiable PM N-body simulations using JAX: Provide a modern infrastructure to support differentiable PM N-body simulations using JAX:
- Keep implementation simple and readable, in pure NumPy API - Keep implementation simple and readable, in pure NumPy API
- Transparent distribution using builtin `xmap`
- Any order forward and backward automatic differentiation - Any order forward and backward automatic differentiation
- Support automated batching using `vmap` - Support automated batching using `vmap`
- Compatibility with external optimizer libraries like `optax` - Compatibility with external optimizer libraries like `optax`
- Now fully distributable on **multi-GPU and multi-node** systems using [jaxDecomp](https://github.com/DifferentiableUniverseInitiative/jaxDecomp) working with the latex `JAX v0.4.35`
## Open development and use ## Open development and use
@ -23,6 +22,10 @@ Current expectations are:
- Everyone is welcome to contribute, and can join the JOSS publication (until it is submitted to the journal). - 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. - 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](notebooks/README.md) for a structured guide through each tutorial.
## Contributors ✨ ## Contributors ✨