A simple Particle-Mesh implementation in JAX

Navigation

Verified details

These details have been verified by PyPI
Project links
GitHub Statistics
Maintainers
Avatar for ASKabalan from gravatar.com ASKabalan Avatar for eiffL from gravatar.com eiffL

Unverified details

These details have not been verified by PyPI
Meta
  • License: MIT License
  • Author: JaxPM developers
  • Requires: Python >=3.9
Report project as malware

Project description

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!

Project details

Verified details

These details have been verified by PyPI
Project links
GitHub Statistics
Maintainers
Avatar for ASKabalan from gravatar.com ASKabalan Avatar for eiffL from gravatar.com eiffL

Unverified details

These details have not been verified by PyPI
Meta
  • License: MIT License
  • Author: JaxPM developers
  • Requires: Python >=3.9

Release history Release notifications | RSS feed

This version

0.1.6

0.1.5

0.1.4

0.1.3 yanked

0.0.2

0.0.1

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

jaxpm-0.1.6.tar.gz (27.2 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

jaxpm-0.1.6-py3-none-any.whl (24.8 kB view details)

Uploaded Python 3

File details

Details for the file jaxpm-0.1.6.tar.gz.

File metadata

  • Download URL: jaxpm-0.1.6.tar.gz
  • Upload date:
  • Size: 27.2 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/6.1.0 CPython/3.12.9

File hashes

Hashes for jaxpm-0.1.6.tar.gz
Algorithm Hash digest
SHA256 1957ab86ffe4929968488a83ce7ad172d30e12baa47c1c611bc31ae09cbfff61
MD5 761c3bc63ece7c832057446a1d6a42a6
BLAKE2b-256 d1ec236c9b772f4d950b647ef91755953c21f336dcf90fb687a88ebfe03ef7a3

See more details on using hashes here.

Provenance

The following attestation bundles were made for jaxpm-0.1.6.tar.gz:

Publisher: python-publish.yml on DifferentiableUniverseInitiative/JaxPM

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file jaxpm-0.1.6-py3-none-any.whl.

File metadata

  • Download URL: jaxpm-0.1.6-py3-none-any.whl
  • Upload date:
  • Size: 24.8 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/6.1.0 CPython/3.12.9

File hashes

Hashes for jaxpm-0.1.6-py3-none-any.whl
Algorithm Hash digest
SHA256 dc1b9aa917757989f3c14e55908f71d93b296a8f3bc9f38701acd98f322d7144
MD5 451b9aad332848d2617c93c20c3e35d8
BLAKE2b-256 36e4692fae60c9009fa97160e20ef1b9f2c88157d26ca60d1d690a8c98f5a63e

See more details on using hashes here.

Provenance

The following attestation bundles were made for jaxpm-0.1.6-py3-none-any.whl:

Publisher: python-publish.yml on DifferentiableUniverseInitiative/JaxPM

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page