Jammy2211/PyAutoLens: JAX Interferometry More Speed Up + Shapelets

This code speeds up lens modeling with interferometers even more, with uv-plane lens modeling on a HPC GPU taking mere hours for datasets with 100,000,000+ visibilities and extremely high resolution, and run times under 30 minutes for more modest datasets.

This release also includes mature JAX support for shapelets thanks to @Chocologism.

Interferometer Speed Up

https://github.com/Jammy2211/PyAutoArray/pull/202

https://github.com/Jammy2211/autolens_workspace/tree/release/notebooks/interferometer

The fast interferometry JAX GPU implementation (https://github.com/Jammy2211/PyAutoArray/pull/201) uses a preload of the NUFFT in order to compute the curvature_matrix.

This curvature_preload matrix is computed once, before lens modeling, and reused throughout lens modeling. For high number of visibilities and resolution real space mask, this calculation can take minutes or hours on a CPU.

The previous PR did not convert this calculation to JAX or run on GPU.

This pull request makes the W-Tilde curvature preload computation support JAX and GPU, with profiling suggeting at least x100 speed up for high resolution datasets, with the calculation taking under a minute for the highest resolution / visibilities ALMA datasets tested.

It also includes utilities for safely saving and loading precomputed data, which check metadata to ensure the loaded data matches the data being analysed.

Shapelets

https://github.com/Jammy2211/PyAutoGalaxy/pull/259

https://github.com/Jammy2211/autolens_workspace/tree/release/notebooks/imaging/features/advanced/shapelets

Full JAX support for elliptical polar and Cartesian shapelets, with the elliptical power shapelet the recommended default!