forked from Aquila-Consortium/JaxPM_highres

JAX-powered Cosmological Particle-Mesh N-body Solver

Find a file

Francois Lanusse f245a1f685 Pypi upload compatible version (#33 ) * moving test dependencies separately * adding manifest to remove unecessary files * updating name of project * Fixing formatting * Adding badge for pypi version * Adding very simple install instructions		2024-12-21 11:47:13 -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	remove deprecated stuff	2024-12-21 09:58:42 -05:00
tests	jaxdecomp proto (#21 )	2024-12-20 05:44:02 -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	Pypi upload compatible version (#33 )	2024-12-21 11:47:13 -05:00
pytest.ini	jaxdecomp proto (#21 )	2024-12-20 05:44:02 -05:00
README.md	Pypi upload compatible version (#33 )	2024-12-21 11:47:13 -05:00
requirements-test.txt	Pypi upload compatible version (#33 )	2024-12-21 11:47:13 -05:00

README.md

JaxPM

JAX-powered Cosmological Particle-Mesh N-body Solver

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.