Astrojax - Accelerated astrodynamics in JAX
Astrodynamics written in JAX for massively parallel simulation.
The goal of this project is to provide a high-performance astrodynamics library that can be used for research and education. The library is built on top of JAX, which allows for automatic differentiation and GPU/TPU acceleration. The goal is to provide a proof-of-concept implementation of common astrodynamics algorithms that can be used as a starting point for further development. It is not intended to be a full-featured, high-accuracy astrodynamics library (at least not yet).
Install
pip install astrojax
# or
uv add astrojax
If you have a GPU and want to take advantage of JAX's acceleration can then install the package with the appropriate JAX extras:
pip install astrojax[cuda13]
# or
uv add astrojax[cuda13]
Astrojax supports CUDA 12 and CUDA 13.
Quickstart
This project uses just as a command runner and uv for Python package management.
# Install just (macOS)
brew install just
# Install uv
curl -LsSf https://astral.sh/uv/install.sh | sh
# Clone and install everything (dev deps + extras)
just install
Development
Every recipe has an equivalent raw command you can run directly.
| Task | just recipe |
Raw command |
|---|---|---|
| Install dev deps + extras | just install |
uv sync --dev --extra extras |
| Run tests | just test |
uv run pytest tests/ -v |
| Test with coverage | just test-cov |
uv run pytest --cov=astrojax --cov-report=term-missing |
| Format code | just fmt |
uv run ruff format |
| Lint (auto-fix) | just lint |
uv run ruff check --fix |
| Type check | just typecheck |
uvx ty check |
| All quality checks | just check |
Runs fmt + lint + typecheck |
| Build docs | just docs-build |
uv run zensical build --clean |
| Serve docs | just docs-serve |
uv run zensical serve --clean |
License
The code in this repository is licensed under the MIT License. See LICENSE for details.
1 maintainer
Metadata
1 maintainer
Release files for astrojax 0.8.0
For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.
Source distribution (sdist)
| File | Size | Uploaded | |
|---|---|---|---|
| astrojax-0.8.0.tar.gz | 4.8 MB | Details |
Built distribution (wheel)
| File | Interpreter | ABI | Platform | Reset |
|---|---|---|---|---|
| astrojax-0.8.0-py3-none-any.whl | Python 3 | none | any | Details |
Total release size: 9.7 MB
Release files / astrojax-0.8.0.tar.gz
| Download URL | astrojax-0.8.0.tar.gz |
|---|---|
| Size | 4.8 MB |
| Tags | Source |
|
SHA-256 checksum How to use checksums |
4ac05283df00e519be16e67f10079abd83db5e6491ac767055b932a15b4577e1
|
|
BLAKE2b-256 checksum How to use checksums |
d3960863a1ed27e2269b87dd3df813a2193069d7c2340f7492260dc717a3ecc1
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
Yes |
| Uploaded via |
twine/7.0.0 CPython/3.13.14
|
Provenance
Provenance describes where a file came from. On PyPI, provenance is shared via attestations, which provide a verifiable record of the build or publishing details. View details, limitations and caveats.
PyPI Publish Attestation
PyPI verified that this artifact, at this checksum, originated from the publisher listed below.
Signed by GitHub Actions, verified by PyPI on Jul 31, 2026.
Transparency logRelease files / astrojax-0.8.0-py3-none-any.whl
| Download URL | astrojax-0.8.0-py3-none-any.whl |
|---|---|
| Size | 4.9 MB |
| Tags | Python 3 |
|
SHA-256 checksum How to use checksums |
e984625ed32e70b3d364b5b2b8d23ea7f3faef7584278d5f507299ed890c1f03
|
|
BLAKE2b-256 checksum How to use checksums |
706885d4b2f1bbad1021d2012b670b5ffe8d87c981c6ce322b2c807cc8884496
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
Yes |
| Uploaded via |
twine/7.0.0 CPython/3.13.14
|
Provenance
Provenance describes where a file came from. On PyPI, provenance is shared via attestations, which provide a verifiable record of the build or publishing details. View details, limitations and caveats.
PyPI Publish Attestation
PyPI verified that this artifact, at this checksum, originated from the publisher listed below.
Signed by GitHub Actions, verified by PyPI on Jul 31, 2026.
Transparency log
