Metadata-Version: 2.4
Name: grater-jax
Version: 0.2.3
Summary: JAX-powered machine learning and modeling framework for GRaTeR disks.
Author-email: Mihir Kondapalli <mihir.kondapalli@gmail.com>, Briley Lewis <brileylewis@g.ucla.edu>
License: MIT
Requires-Python: <3.13,>=3.9
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: numpy<2.0,>=1.23
Requires-Dist: scipy<=1.14,>=1.10
Requires-Dist: matplotlib<3.9,>=3.7
Requires-Dist: pandas<2.3,>=1.5
Requires-Dist: astropy<7.0,>=5.3
Requires-Dist: astroquery<0.5,>=0.4
Requires-Dist: poppy<1.2,>=1.0
Requires-Dist: stsynphot<1.5,>=1.3
Requires-Dist: synphot<1.5,>=1.3
Requires-Dist: pysiaf<0.24,>=0.21
Requires-Dist: emcee<3.2,>=3.1
Requires-Dist: arviz<0.18,>=0.15
Requires-Dist: corner<2.3,>=2.2
Requires-Dist: xarray<2024.10,>=2023.1
Requires-Dist: jax<=0.4.34,>=0.4.20
Requires-Dist: jaxlib<=0.4.34,>=0.4.20
Requires-Dist: jaxopt<0.9,>=0.8
Requires-Dist: tqdm<5.0,>=4.65
Provides-Extra: dev
Requires-Dist: pytest<9.0,>=7.0; extra == "dev"
Requires-Dist: vip_hci>=1.3; extra == "dev"
Provides-Extra: docs
Requires-Dist: sphinx<8.0,>=7.2; extra == "docs"
Requires-Dist: sphinx-rtd-theme<2.1,>=1.3; extra == "docs"
Requires-Dist: myst-parser<3.1,>=2.0; extra == "docs"
Requires-Dist: nbsphinx<0.10,>=0.9; extra == "docs"
Requires-Dist: photutils<1.14,>=1.8; extra == "docs"
Provides-Extra: webbpsf
Requires-Dist: webbpsf<1.4,>=1.0; extra == "webbpsf"
Requires-Dist: webbpsf_ext<1.4,>=1.0; extra == "webbpsf"
Requires-Dist: poppy<1.2,>=1.0; extra == "webbpsf"
Dynamic: license-file

# GRaTeR-JAX

[![Run Tests](https://github.com/UCSB-Exoplanet-Polarimetry-Lab/GRaTeR-JAX/actions/workflows/tests.yml/badge.svg)](https://github.com/UCSB-Exoplanet-Polarimetry-Lab/GRaTeR-JAX/actions/workflows/tests.yml)

**GRaTeR-JAX** is a machine learning JAX-based implementation of the **Generalized Radial Transporter (GRaTeR)** framework [(Augereau+ 1999)](https://arxiv.org/abs/astro-ph/9906429), designed for modeling scattered-light images of debris disks. This repository provides tools for forward modeling, optimization, and parameter estimation of scattered-light disk images using JAX's accelerated computations.

<img src="https://github.com/user-attachments/assets/c10f45e8-5449-4891-b6a7-33954cf6d954" width="300">

## Features

- **JAX-Based Optimization**: Leverages JAX for fast, GPU/TPU-accelerated disk modeling.
- **Scattered Light Disk Modeling**: Implements physical models of exoplanetary debris disks.
- **Differentiable Framework**: Enables gradient-based optimization and probabilistic inference.
- **Integration with Webbpsf**: Supports JWST PSF convolutions for forward modeling.

## Installation

To install GRaTeR-JAX and its dependencies, create a new Conda environment with Python and run:

```sh
pip install grater-jax
```

Make sure you have JAX installed with the correct backend for your hardware:

```sh
pip install --upgrade "jax[cpu]"  # or "jax[cuda]" for GPU
```

Highly recommended to install this on a fresh environment, just to be safe.

## Usage

Refer to the documentation at [grater-jax.readthedocs.io](https://grater-jax.readthedocs.io/en/latest/).

Check out [GRaTeR Image Generator](https://scattered-light-disks.vercel.app) to visualize how each of the parameters affect the disk model!

## Repository Structure

```
GRaTeR-JAX/
│── grater-jax/       # Package root for grater-jax
   │── disk_model/    # Code for disk modeling
   │── optimization/  # Tools for statistical optimization and analysis
|── docs/             # Documenation and tutorial notebooks
│── pyproject.toml    # Installation file
│── README.md
```

## Contributing

We welcome contributions! To contribute:

1. Fork the repository.
2. Create a feature branch:
   ```sh
   git checkout -b feature-branch
   ```
3. Commit your changes and push to your fork.
4. Open a pull request.

## Acknowledgments

Developed by the **UCSB Exoplanet Polarimetry Lab**. This work is inspired by previous implementations of GRaTeR and advances in JAX-based differentiable modeling.

---
