microjaxx 0.1.1

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.

---

Documentation

For detailed guides, tutorials, and the API reference, visit the hosted microJAX documentation. The rendered HTML bundle also lives in docs/ if you prefer browsing locally.

---

📦 Installation

From PyPI:

pip install microjaxx

Notes:

• PyPI package name: microjaxx
• Python import name: microjax

import microjax

Development install (from source, recommended):

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

• JAX installation (CPU/GPU): https://jax.readthedocs.io/en/latest/installation.html

---

🚀 Quickstart

Compute an extended-source binary-lens magnification light curve using the image-centered ray shooting (ICRS) method:

Note: mag_binary also works on CPU but is very slow.

import jax
import jax.numpy as jnp
from microjax.inverse_ray.lightcurve import mag_binary
jax.config.update("jax_enable_x64", True)

# Binary-lens parameters
s, q = 1.0, 0.01            # separation and mass ratio (m2/m1)
rho = 0.02                  # source radius (Einstein units)
tE, u0 = 30.0, 0.0          # Einstein time [days], impact parameter
alpha = jnp.deg2rad(10.0)  # trajectory angle
t0 = 0.0

# Source trajectory
N_points = 1000
t = t0 + jnp.linspace(-tE, tE, N_points)
tau = (t - t0)/tE
y1 = -u0*jnp.sin(alpha) + tau*jnp.cos(alpha)
y2 =  u0*jnp.cos(alpha) + tau*jnp.sin(alpha)
w_points = jnp.array(y1 + y2 * 1j, dtype=complex)

# Extended-source magnification (binary lens)
mu = mag_binary(w_points, rho, s=s, q=q)

For point-source magnification, use:

Note: mag_point_source runs on CPU (and GPU), so it works without a GPU.

from microjax.point_source import mag_point_source
mu_point = mag_point_source(w_points, nlenses=2, s=s, q=q)

---

Example output

Visualization of the ICRS method (binary-lens)	Triple-lens magnification and its gradients	Compare with VBBL (uniform source, binary-lens)

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

Note: Finite-source calculation with microJAX is extremely slow without a GPU. So, the latter two examples are significantly slower on a CPU.

---

⚠️ Known Limitations

• Finite-source magnification trades memory/runtime for accuracy through resolution parameters; tune these settings to match your GPU's available memory and throughput.
• For numerical stability and agreement across libraries, enable 64-bit precision in JAX (jax_enable_x64=True).
• Triple-lens hexadecapole/ghost-image test is not yet implemented: triple-lens calculations fall back to full contour integration everywhere, which can be substantially slower.
• GPU tests are opt-in; run them explicitly with pytest -m gpu. If JAX cannot see a CUDA GPU, those tests are skipped.

📚 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 tests are opt-in and skipped by default. To run them on a CUDA-capable machine:

# optionally: export JAX_PLATFORMS=cuda
pytest -m gpu -q

These tests require JAX to detect a CUDA device. If not available, they are skipped. No additional environment flag is required beyond the -m gpu marker.

📜 License

This project is licensed under the MIT License. Third-party components bundled in the tree and their respective licenses are listed in third_party/README.md. If you use microJAX in academic work, please cite the upcoming Miyazaki et al. (2025) methods paper.

---