grater-jax 0.2.2

JAX-powered machine learning and modeling framework for GRaTeR disks.

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.

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:

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

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/             # 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:

   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.

---