Metadata-Version: 2.4
Name: sajax
Version: 0.5.0
Summary: Stellar Activity Grid for Exoplanets — JAX implementation
Author-email: Samson Mercier <merci228@mit.edu>, Lucas Arthur <lmarthur@mit.edu>, Hritam Chakraborty <Hritam.Chakraborty@unige.ch>
License-Expression: MIT
Project-URL: Repository, https://github.com/SamMerc/sajax
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENCE
Requires-Dist: numpy>=1.24
Requires-Dist: jax>=0.4
Requires-Dist: jaxlib>=0.4
Requires-Dist: matplotlib>=3.7
Requires-Dist: scipy>=1.10
Requires-Dist: interpax>=0.3
Provides-Extra: dev
Requires-Dist: pytest>=7; extra == "dev"
Requires-Dist: pytest-cov; extra == "dev"
Requires-Dist: jupyter; extra == "dev"
Requires-Dist: ipykernel; extra == "dev"
Provides-Extra: docs
Requires-Dist: sphinx>=7.0; extra == "docs"
Requires-Dist: sphinx-book-theme; extra == "docs"
Requires-Dist: myst-nb; extra == "docs"
Requires-Dist: sphinx-copybutton; extra == "docs"
Requires-Dist: sphinx-autoapi; extra == "docs"
Requires-Dist: ipykernel; extra == "docs"
Provides-Extra: cuda
Requires-Dist: jax[cuda12]>=0.4; (sys_platform == "linux" and platform_machine == "x86_64") and extra == "cuda"
Dynamic: license-file

<p align="center">
  <img src="docs/_static/logo.png" width="300">
</p>

# SAJAX — Stellar Activity Grid for Exoplanets in JAX

![Tests](https://github.com/SamMerc/sajax/actions/workflows/tests.yml/badge.svg)
[![codecov](https://codecov.io/gh/SamMerc/sajax/branch/main/graph/badge.svg)](https://codecov.io/gh/SamMerc/sajax)

SAJAX is a JAX-accelerated reimplementation of
[SAGE](https://github.com/chakrah/sage) ([Chakraborty et al. 2024](https://www.aanda.org/articles/aa/abs/2024/05/aa47727-23/aa47727-23.html)), a code
that models stellar contamination of exoplanet transmission spectra from
active regions (ars, faculae) on the stellar surface.

The key innovation over plain SAGE is that SAJAX vectorises the spectral
loop with `jax.vmap`, making it fast on both CPU and GPU without any
change to the calling code, and fully differentiable — enabling
gradient-based inference with tools like NumPyro or Optax.

Documentation can be found at [sajax.readthedocs.io](http://sajax.readthedocs.io/)

## Installation

```bash
pip install sajax
```

Or in development mode from a local clone. `setup_env.sh` builds a
[uv](https://docs.astral.sh/uv/)-managed `.venv` pinned by `uv.lock`, so every
machine gets identical package versions:

```bash
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
│   │   ├── comparison.ipynb
│   │   ├── inference.ipynb
├── tests/
│   ├── test_core.py
│   ├── test_planet.py
├── pyproject.toml
├── .gitignore
└── README.md
```
