Metadata-Version: 2.4
Name: sajax
Version: 1.1
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: LICENSE
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)
[![License: MIT](https://img.shields.io/badge/License-MIT-purple.svg)](LICENCE)

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](https://github.com/chakrah/sage) ([Chakraborty et al. 2024](https://www.aanda.org/articles/aa/abs/2024/05/aa47727-23/aa47727-23.html)), and it draws additional inspiration from [ANTARESS](https://gitlab.unige.ch/spice_dune/antaress) ([Bourrier et al. 2024](https://www.aanda.org/articles/aa/full_html/2024/11/aa49203-24/aa49203-24.html)).

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