microjaxx 0.1.0

GPU-accelerated, differentiable microlensing modeling in JAX

microJAX is a GPU-accelerated, differentiable microlensing modeling library written in JAX.

microjax

microJAX is a fully‑differentiable, GPU‑accelerated software for modelling gravitational microlensing light curves produced by binary, and triple lens systems, using the image-centered ray shooting (ICRS) method (e.g., Bennett 2010). Written entirely in JAX, it delivers millisecond‑level evaluations of extended-source magnifications and exact gradients for every model parameter through the use of automatic differentiation, enabling gradient‑based Bayesian inference workflows such as Hamiltonian Monte Carlo (HMC) and variational inference.

This software is under active development and not yet feature complete.

---

✨ Key Features

Category	Description
Lens Systems	Supports point-source and finite-source magnification calculations for binary and triple lens systems
Extended Sources	Models uniform and limb-darkened source profiles
Computational Core	Implements the Image-Centered Ray Shooting (ICRS) algorithm in JAX, fully optimized for GPU acceleration
Root-Finding Engine	Uses a differentiable Ehrlich-Aberth method for complex polynomial roots with implicit gradients for stable optimization
Bayesian Inference	Provides a ready-to-use likelihood function compatible with NumPyro's HMC and variational inference frameworks

📦 Installation

From PyPI (recommended):

pip install microjaxx

Import name remains:

import microjax

Development install (from source):

# clone the repository
git clone https://github.com/ShotaMiyazaki94/microjax.git
cd microjax
pip install -e ".[dev]"

GPU support: JAX/JAXLIB with CUDA/ROCm depends on your environment. Please follow the official JAX installation guide to install the appropriate jaxlib for your accelerator.

---

Example output

Visualization of the ICRS method (binary-lens)	Triple-lens magnification and its gradients

Refer to the example directory for code that creates these plots.

📚 References

• Miyazaki & Kawahara (in prep.): microjax paper (expected within 2025!)
• Bennett (2010): Image-centred ray shooting (ICRS) method
• Cassan (2017): Hexadecapole approximations
• Sugiyama (2022): Fast FFT-based magnification evaluation with a single-lens extended source model

🤝 Contributing

Pull requests are welcome! Please see CONTRIBUTING.md for coding style, test suite, and CI guidelines. Bug reports can be filed via GitHub Issues.

Running Tests

CPU-only tests:

pytest -q

GPU-only (A100) tests are opt-in and skipped by default. To run them on an A100 machine:

export MICROJAX_GPU_TESTS=1
pytest -m gpu -q

These tests require JAX to detect an NVIDIA A100 (CUDA) device. If not available or the env var is not set, they are skipped.

📜 License

This project is licensed under the MIT License. If you use microJAX in academic work, please cite the upcoming Miyazaki et al. (2025) methods paper.

---