PyAutoGalaxy JAX

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.

May 2025

• Results workflow API, which generates .csv, .png and .fits files of large libraries of results for quick and efficient inspection:

https://github.com/Jammy2211/autolens_workspace/tree/main/notebooks/results/workflow

• Visualization now outputs .fits files corresponding to each subplot, which more concisely contain all information of a fit and are used by the above workflow API.

• Visualization Simplified, removing customization of individual image outputs.

• Remove Analysis summing API, replacing all dataset combinations with AnalysisFactor and FactorGraphModel API used for graphical modeling:

https://github.com/Jammy2211/autolens_workspace/blob/main/notebooks/advanced/multi/modeling/start_here.ipynb

• Pixelized source reconstruction output as a .csv file which can be loaded and interpolated for better source science analysis.

• Latent variable API bug fixes and now used in some test example scripts.

January 2025

The main updates are visualization of Delaunay mesh's using Delaunah triangles and a significant refactoring of over sampling, with the primary motivation to make the code much less complex for the ongoing JAX implementation.

What's Changed

• Feature/disable noise by @Jammy2211 in #211
• feature/delaunay_visual by @Jammy2211 in #210
• feature/inversion noise map by @Jammy2211 in #212
• feature/operate deflections api by @rhayes777 in #195
• Revert "feature/operate deflections api" by @rhayes777 in #213
• Feature/over sampling refactor by @Jammy2211 in #214

Full Changelog: 2024.11.13.2...2025.1.18.7

November 2024 update

Small bug fixes and optimizations for Euclid lens modeling pipeline.

November 2024

Minor release with stability updates and one main feature.

• Extra Galaxies API for modeling multiple galaxies at once: https://github.com/Jammy2211/autogalaxy_workspace/blob/release/notebooks/features/extra_galaxies.ipynb

September 2024

This release updates all projects to support Python 3.12, with support tested for Python 3.9 - 3.12 and 3.11 regarded as most stable.

This includes many project dependency updates:

https://github.com/rhayes777/PyAutoFit/blob/main/requirements.txt
https://github.com/rhayes777/PyAutoFit/blob/main/optional_requirements.txt
https://github.com/Jammy2211/PyAutoGalaxy/blob/main/requirements.txt
https://github.com/Jammy2211/PyAutoGalaxy/blob/main/optional_requirements.txt

Workspace Restructure:

This release has a workspace restructure, which is now grouped at a high level by tasks (e.g. modeling, simulators) rather than datasets:

https://github.com/Jammy2211/autogalaxy_workspace

The readthedocs have been greatly simplified and include a new user guide to help navitgate the new workspace:

https://pyautogalaxy.readthedocs.io/en/latest/overview/overview_2_new_user_guide.html

PyAutoGalaxy:

• Improved Cosmology wrapper to support new astropy and easier to use in models: #193
• Ellipse Fitting: https://github.com/Jammy2211/autogalaxy_workspace/tree/release/notebooks/advanced/misc/ellipse

PyAutoFit:

https://github.com/rhayes777/PyAutoFit/pulls?q=is%3Apr+is%3Aclosed

• Improvements to HowToFit lectures: rhayes777/PyAutoFit#1022
• Support for NumPy arrays in model composition and prior creation, for example creating an ndarray of input shape where each value is a free parameter in the seach: rhayes777/PyAutoFit#1021
• Name of optimize searches renamed to mle, for maximum likelihood estimator, with improvements to visualization: rhayes777/PyAutoFit#1029
• Improvement to sensitivity mapping functionality and results: https://github.com/rhayes777/PyAutoFit/pulls?q=is%3Apr+is%3Aclosed
• More improvements to JAX Pytree interface, documentation still to come.

May 2024

PyAutoFit:

• Nautilus now outputs results on the fly: rhayes777/PyAutoFit#961
• Output latent samples of a model-fit, which are parameters derived from a model which may be marginalized over:

PR: rhayes777/PyAutoFit#994
Example: https://github.com/Jammy2211/autofit_workspace/blob/release/notebooks/cookbooks/analysis.ipynb

• model.info file displays complex models in a more concise and readable way: rhayes777/PyAutoFit#1012
• All samples with a weight below an input value are now removed from samples.csv to save hard disk space: rhayes777/PyAutoFit#979
• Documentation describing autofit scientific workflow: rhayes777/PyAutoFit#1011
• Refactor visualization into stand alone module: rhayes777/PyAutoFit#995
• Refactor how results are returned after a search: rhayes777/PyAutoFit#989
• Improved parallelism logging: rhayes777/PyAutoFit#1009
• Likelihood consistency check now performed internally: rhayes777/PyAutoFit#987
• Generation of initial search samples is now performed in parallel: rhayes777/PyAutoFit#997
• No longer store search_internal on hard-disk. simplifying source code internals: rhayes777/PyAutoFit#938
• Multiple small bug fixes and improvements to interface.

PyAutoGalaxy:

• Remove Plane object and replace with Galaxies object
• Shapelets improvements: #173
• Adaptive over sampling of grids for a pixelization: #168
• BasisPlotter which plots each basis (e.g. each Gaussian of an MGE): #173
• Plot mappings between source and image plane of a pixelization as lines: #172
• For multi-wavelength datasets model offsets between each dataset: #171
• Modeling of background sky: #170
• Improvements to use of adapt images for adaptive pixelizations: #160
• Improved angle conversions for computing errors on mass profile and shear angles from ell_comps: #169
• Remove sub_size from all classes (e.g. Array2D, Mask2D) to simplify API.
• MaternKernel added: #148

January 2024 (2024.1.27.4)

• Log10 plots implemented in 1D and 2D, which show certain quantities (e.g. galaxy convergence) more clear and include contours showing log10 values:

• Improved subplots including addition of log10 panels:

• Pixelization API now has separate entry for an image_mesh, defining how the source pixel centres are computed (E.g. using a KMeans clustering) and the mesh is now just the method (e.g. Delaunay):

pixelization = al.Pixelization(
    image_mesh=al.image_mesh.Overlay(shape=(25, 25)),
    mesh=al.mesh.Delaunay(),
    regularization=al.reg.Constant(coefficient=1.0),
)

• Implemented Hilbert image-mesh which is a significant improvement on KMeans clustering for creating the distribution of source pixels for a pixelization and inversion.

• Simplified adapt_dataset API to now only pass via an AdaptImage class, which is not passed as Galaxy attributes anymore but instead through the FitImaging object.

• Removed SetupAdapt object and pass image_mesh_pixels as an integer through pipelines.

• Added Exponential / Gaussian smoothing kernels for regularization documented in Vernardos 2022 (https://arxiv.org/abs/2202.09378)

October 2023 (2023.10.23.3)

• Support for Python 3.11 by updating requirement on core libraries (e.g. numpy, scipy, scikit-learn).
• Fix issues with sqlite database following switch from .pickle outputs to .json / .fits / .csv.
• Database use of Samples object much more efficient.
• Methods to output classes to hard-disk (e.g. output_to_json, from_json, to_dict) are now all handled and called from autoconf.
• Fix bug where nautilus parallel fits sometimes crashed.
• Fix bug where nautilus single CPU fits did not work.

September (v2023.9.18.4)

This release implements two major changes to PyAutoGalaxy:

Nautilus:

For the past ~3 years, model fitting has used the nested sampling algorithm Dynesty.

Recently, a new nested sampler, Nautilus (https://nautilus-sampler.readthedocs.io/en/stable/), was released, which uses machine-learning based techniques to improve sampling.

Extensive testing of modeling with Nautilus has revealed that it:

• Speeds up the fitting of simple models by ~x2 - x3.
• Speeds up the fitting of complex models by ~x3 - x5+.
• Is more robust and reliable (e.g less likely to infer a local maxima, can fit more complex lens models).
• Controlled predominantly by just one parameter n_live, so is simpler to use than dynesty.
• Parallelization using Python multiprocessing is more efficient than dynesty and now supports proper error handling.

Nautilus is therefore now the default modeler, with all workspace examples updated accordingly.

NOTE: Nautilus does not currently support on-the-fly output and to get the results of a lens model mid-fit a user can instead cancel the run (e.g. via Ctrl + C) and restart it, where the maximum likelihood model will be output.

Results Output

Result metadata was previously output as .pickle files, which were not human readable and depended on project imports, hurting backwards compatibility.

All metadata is now output as human readable .json files and dataset as .fits files, making it a lot more straight forward for a user to interpret how data is stored internally within PyAutoGalaxy:

Here is an example of the search.json file:

All internal functionality (e.g. the sqlite database) has been updated to use these files.

All workspace documentation has been updated accordingly.

Other:

• imaging/modeling/features split to make linear light profiles and multi gaussian expansion more visible.
• Improved HowToGalaxy tutorial 5 on linear light profiles.
• Power law with multipole parameterization updated, now supports multipoles of any order (#115).
• Update certain requirements (e.g. PyYAML) to mitigate installation issues (rhayes777/PyAutoConf#41).
• Lots of quality-of-life improvements thoughout the code bases.