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).
pip install astrojax
# or
uv add astrojaxIf 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.
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 installEvery 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 |
The code in this repository is licensed under the MIT License. See LICENSE for details.