UPDATE: Latest JAX version is now 2025.11.5.1
This release marks the completion of two years work implementing JAX (https://docs.jax.dev/en/latest/notebooks/thinking_in_jax.html) in PyAutoGalaxy.
With JAX, any modeling analysis can be run on GPU, with speed up of ~x50 or more.
Core Release
The core PyAutoGalaxy API does not change significantly, however existing users redownload the new autogalaxy workspace, which has new configs and examples:
https://github.com/Jammy2211/autogalaxy_workspace
New user should checkout the start_here.ipynb notebook, which can be read via a Google Colab by clicking the hyperlink.
GPU Modeling Examples
The following Juypter Notebooks, which run via Google Colab, illustrate < 10 minute galaxy modeling for different science cases:
-
start_here_imaging.ipynb: Galaxy-scale strong galaxyes observed with CCD imaging (e.g. Hubble, James Webb).
-
start_here_interferometer.ipynb: Galaxy scale strong galaxyes observed with interferometer data (e.g. ALMA).
-
start_here_multi_wavelength.ipynb: Model multiple images (different wavelength imaging, imaging + interferometer) simultaneously.
Performance Of Features
-
Interferometer with many Visibilities: Above ~ GPU uv-plane analysis with hundreds of millions of visibilities and extremely high resolutions run in under and hour, a monumental speed up compared to CPU modeling.
-
Pixelized sources run ~x5 - x20 faster on modern HPC GPU clusters, with galaxy modeling times typically ~10 - 20 minutes. Pixelized source performance depends on the available GPU VRAM.