GRaTeR-JAX

GRaTeR-JAX is a machine learning JAX-based implementation of the Generalized Radial Transporter (GRaTeR) framework (Augereau+ 1999), designed for modeling scattered light observations of debris disks. This repository provides tools for forward modeling, optimization, and parameter estimation of debris disk images using JAX's accelerated computations.

Features

• JAX-Based Optimization: Leverages JAX for fast, GPU/TPU-accelerated disk modeling.
• Scattered Light Debris Disk Modeling: Implements physical models of circumstellar debris disks.
• Differentiable Framework: Enables gradient-based optimization and probabilistic inference.
• Integration with Webbpsf: Supports JWST PSF convolutions for forward modeling. (Experimental JWST Notebook: in 'jwst' branch)

Installation

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

pip install grater-jax

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

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.

Check out GRaTeR Image Generator 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/             # Documentation and tutorial notebooks
|—— tests/            # Automated tests
│── pyproject.toml    # Installation file
│── README.md

Contributing

We welcome contributions! To contribute:

1. Fork the repository.
2. Create a feature branch:

   git checkout -b feature-branch

3. Commit your changes and push to your fork.
4. Open a pull request.

Citation

Please cite the DOI for this repository, as well as the accompanying paper (submitted to JOSS, currently available on arXiv):

@misc{kondapalli2026graterjaxacceleratedpackagedebris,
      title={GRaTer-JAX: An Accelerated Package for Debris Disk Modeling}, 
      author={Mihir Kondapalli and Briley L. Lewis and Jaren N. Ashcraft and Maxwell A. Millar-Blanchaer},
      year={2026},
      eprint={2607.25342},
      archivePrefix={arXiv},
      primaryClass={astro-ph.EP},
      url={https://arxiv.org/abs/2607.25342}, 
}

Acknowledgments

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

---