SAJAX - Stellar Activity Grid for Exoplanets in JAX

SAJAX is a package that models contamination of exoplanet spectro-photometric time series by active regions (spots, faculae, and flares) on the stellar surface. Its core functionality builds on SAGE (Chakraborty et al. 2024), and it draws additional inspiration from ANTARESS (Bourrier et al. 2024).

The main innovation over SAGE is that SAJAX vectorises the spatial and spectral loops with jax.vmap, making it fast on both CPU and GPU - with no change to the calling code - and fully differentiable, enabling gradient-based inference with tools like NumPyro or Optax.

The full documentation can be found at sajax.readthedocs.io

Installation

pip install sajax

Or in development mode from a local clone. setup_env.sh builds a uv-managed .venv pinned by uv.lock, so every machine gets identical package versions:

git clone https://github.com/SamMerc/sajax.git
cd sajax
./setup_env.sh              # auto-detect: GPU if nvidia-smi sees a device

Useful flags: --cpu / --gpu to override the auto-detection (state --gpu explicitly when provisioning on a GPU-less HPC login node), --docs to add the sphinx dependencies, and --check to report what is already installed.

Run commands in the environment with uv run <cmd> (e.g. uv run pytest), or activate it with source .venv/bin/activate.

Repository layout

sajax/
├── sajax/
│   ├── __init__.py          # public API
│   ├── core.py              # JAX light-curve engine
│   ├── planet.py            # planet orbital dynamics
│   ├── geometry.py          # rotation matrices, coordinate transforms
├── docs/
│   ├── examples/
│   │   ├── introduction.ipynb
│   │   ├── inference.ipynb
├── tests/
│   ├── test_core.py
│   ├── test_planet.py
├── pyproject.toml
├── .gitignore
└── README.md