Metadata-Version: 2.4
Name: grater-jax
Version: 0.1.0
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.10,>=3.9
Description-Content-Type: text/markdown
Requires-Dist: numpy<1.25,>=1.23
Requires-Dist: scipy<1.12,>=1.10
Requires-Dist: matplotlib<3.8,>=3.7
Requires-Dist: pandas<2.0,>=1.5
Requires-Dist: astropy<6.0,>=5.3
Requires-Dist: astroquery<0.5,>=0.4
Requires-Dist: poppy<1.1,>=1.0
Requires-Dist: stsynphot<2.0,>=1.3
Requires-Dist: synphot<1.4,>=1.3
Requires-Dist: pysiaf<0.23,>=0.21
Requires-Dist: emcee<4.0,>=2.2
Requires-Dist: arviz<0.16,>=0.15
Requires-Dist: corner<3.0,>=2.2
Requires-Dist: xarray<2024.0,>=2023.1
Requires-Dist: jax<0.5.0,>=0.4.23
Requires-Dist: jaxlib<0.5.0,>=0.4.23
Requires-Dist: jaxopt<0.9,>=0.8
Provides-Extra: dev
Requires-Dist: pytest<8.0,>=7.0; extra == "dev"
Requires-Dist: vip_hci<1.7,>=1.3; extra == "dev"
Provides-Extra: docs
Requires-Dist: sphinx==7.2.6; extra == "docs"
Requires-Dist: sphinx-rtd-theme<2.0,>=1.3; extra == "docs"
Requires-Dist: myst-parser<3.0,>=2.0; extra == "docs"
Requires-Dist: nbsphinx<1.0,>=0.9; extra == "docs"
Requires-Dist: photutils<2.0,>=1.8; extra == "docs"
Provides-Extra: webbpsf
Requires-Dist: webbpsf<1.1,>=1.0; extra == "webbpsf"
Requires-Dist: webbpsf_ext<1.2,>=1.0; extra == "webbpsf"
Requires-Dist: poppy<1.1,>=1.0; extra == "webbpsf"

# GRaTeR-JAX

**GRaTeR-JAX** is a machine learning JAX-based implementation of the **Generalized Radial Transporter (GRaTeR)** framework, designed for modeling scattered light disks in protoplanetary systems. 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 PSF convolution for telescope observations.

## Installation

To install GRaTeR-JAX and its dependencies, run:

```sh
git clone https://github.com/UCSB-Exoplanet-Polarimetry-Lab/GRaTeR-JAX.git
cd GRaTeR-JAX
pip install -e .

pip install -U <jax backend> ("jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html for cuda for example)
```

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

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

## 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/
│── disk_model/            # Code for disk modeling
│── optimization/          # Tools for statistical optimization and analysis
|── tutorials/             # Tutorial Jupyter notebooks
│── webbpsf-data           # PSF data for various instruments
│── PSFs/                  # PSF data for the disk model
│── environment.yml        # Dependencies
│── requirements.txt       # Pip dependencies
│── README.md              # This document
```

## 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. Additional thanks to Kellen Lawson for developing the Winnie package that this framework uses to model JWST PSFs.

---
