astrojax 0.4.0

Astrodynamics written in JAX for massively parallel computation

Astrojax - Blazing fast 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.