From 5444382816da6edcf15e4978a66433a76422c084 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Sat, 26 Sep 2026 11:53:14 +0100 Subject: [PATCH 1/3] fix: sparse interferometer fits use the dirty image of profile-subtracted visibilities (PyAutoArray#575) FitInterferometer passes sparse_dirty_image (the noise-weighted adjoint transform of profile_subtracted_visibilities, as apply_sparse_operator builds it) whenever the dataset has a sparse operator and a galaxy has an ordinary light profile (Basis contents checked, so an all-linear MGE keeps the cached image at no cost). Shared helper sparse_dirty_image_from is reused by PyAutoLens. Tests: sparse vs dense log_likelihood / log_evidence to 1e-8 with and without an ordinary light profile. Refs PyAutoLabs/PyAutoArray#575 Co-Authored-By: Claude Opus 5.5 --- .../interferometer/fit_interferometer.py | 86 +++++++++++++++++ .../interferometer/test_fit_interferometer.py | 92 +++++++++++++++++++ 2 files changed, 178 insertions(+) diff --git a/autogalaxy/interferometer/fit_interferometer.py b/autogalaxy/interferometer/fit_interferometer.py index bb6ef1549..d1a53238a 100644 --- a/autogalaxy/interferometer/fit_interferometer.py +++ b/autogalaxy/interferometer/fit_interferometer.py @@ -24,6 +24,85 @@ from autogalaxy.galaxy.galaxy import Galaxy from autogalaxy.galaxy.galaxies import Galaxies from autogalaxy.galaxy.to_inversion import GalaxiesToInversion +from autogalaxy.profiles.basis import Basis +from autogalaxy.profiles.light.abstract import LightProfile +from autogalaxy.profiles.light.linear import LightProfileLinear + + +def _has_light_profile_non_linear(galaxies: List[Galaxy]) -> bool: + """ + Returns whether any galaxy has an ordinary (non-linear) light profile, whose visibilities are subtracted + from the data before an inversion. + + A `Basis` is itself a (non-linear) `LightProfile` class, so its contents are checked: a basis made only of + linear light profiles (e.g. an MGE) contributes nothing to the subtracted visibilities. + """ + for galaxy in galaxies: + for light_profile in galaxy.cls_list_from( + cls=LightProfile, cls_filtered=LightProfileLinear + ): + if not isinstance(light_profile, Basis): + return True + + if any( + not isinstance(profile, LightProfileLinear) + for profile in light_profile.light_profile_list + ): + return True + + return False + + +def sparse_dirty_image_from( + dataset: aa.Interferometer, + galaxies: List[Galaxy], + visibilities: aa.Visibilities, + noise_map: aa.VisibilitiesNoiseMap, + xp=np, +) -> Optional[np.ndarray]: + """ + Returns the dirty image the sparse (w-tilde) inversion of an interferometer fit must use to form its data + vector, or `None` if the dirty image cached on the dataset's `sparse_operator` is already the correct one. + + The `sparse_operator` caches the noise-weighted dirty image `Re(Fᴴ W d)` of the dataset's visibilities, + as computed by `Interferometer.apply_sparse_operator`. A fit with ordinary (non-linear) light profiles + inverts the `profile_subtracted_visibilities` instead, so the cached image would give a data vector + inconsistent with the data the chi-squared is computed from. This function computes the dirty image of + the input (profile-subtracted) visibilities with the same weighting (one adjoint NUFFT / DFT). + + Whether this is needed is decided structurally (is there a sparse operator, do any galaxies have an + ordinary light profile), so the branch is fixed at trace time and is safe under `jax.jit`. Fits whose + light is entirely linear return `None` and pay no extra cost. + + Parameters + ---------- + dataset + The interferometer dataset being fitted, whose `sparse_operator` and `transformer` are used. + galaxies + The galaxies of the fit, checked for ordinary light profiles. + visibilities + The visibilities fitted by the inversion (the profile-subtracted visibilities). + noise_map + The noise-map of the visibilities, whose real and imaginary parts weight the dirty image. + xp + The array module (`numpy` or `jax.numpy`). + """ + if dataset.sparse_operator is None: + return None + + if not _has_light_profile_non_linear(galaxies=galaxies): + return None + + visibilities = getattr(visibilities, "array", visibilities) + noise_map = getattr(noise_map, "array", noise_map) + + return dataset.transformer.image_from( + visibilities=aa.Visibilities( + visibilities=visibilities.real * noise_map.real**-2.0 + + 1j * visibilities.imag * noise_map.imag**-2.0 + ), + xp=xp, + ).array class FitInterferometer(aa.FitInterferometer, AbstractFitInversion): @@ -116,6 +195,13 @@ def galaxies_to_inversion(self) -> GalaxiesToInversion: grids=self.grids, transformer=self.dataset.transformer, sparse_operator=self.dataset.sparse_operator, + sparse_dirty_image=sparse_dirty_image_from( + dataset=self.dataset, + galaxies=self.galaxies, + visibilities=self.profile_subtracted_visibilities, + noise_map=self.noise_map, + xp=self._xp, + ), ) return GalaxiesToInversion( diff --git a/test_autogalaxy/interferometer/test_fit_interferometer.py b/test_autogalaxy/interferometer/test_fit_interferometer.py index 68439a85a..d61bd6a1e 100644 --- a/test_autogalaxy/interferometer/test_fit_interferometer.py +++ b/test_autogalaxy/interferometer/test_fit_interferometer.py @@ -1,6 +1,7 @@ import numpy as np import pytest +import autoarray as aa import autogalaxy as ag @@ -509,3 +510,94 @@ def test__model_visibilities_of_galaxies_list__matches_galaxy_model_visibilities assert fit.model_visibilities_of_galaxies_list[2].array == pytest.approx( fit.galaxy_model_visibilities_dict[galaxy_pix].array, 1.0e-4 ) + + +def _assert_sparse_fit_matches_dense(dataset, dataset_sparse, galaxies): + fit = ag.FitInterferometer(dataset=dataset, galaxies=galaxies) + fit_sparse = ag.FitInterferometer(dataset=dataset_sparse, galaxies=galaxies) + + assert isinstance(fit.inversion, aa.InversionInterferometerMapping) + assert isinstance(fit_sparse.inversion, aa.InversionInterferometerSparse) + + assert fit_sparse.log_likelihood == pytest.approx(fit.log_likelihood, rel=1.0e-8) + assert fit_sparse.log_evidence == pytest.approx(fit.log_evidence, rel=1.0e-8) + + +def test__fit_figure_of_merit__sparse_operator__linear_light_only__matches_dense( + interferometer_7, +): + """ + A fit whose light is entirely linear (e.g. an MGE) inverts the unsubtracted visibilities, so the dirty + image cached on the sparse operator is used as-is and the sparse fit reproduces the dense one. + """ + dataset_sparse = interferometer_7.apply_sparse_operator(use_jax=False) + + galaxies = [ + ag.Galaxy( + redshift=0.5, + bulge=ag.lp_linear.Gaussian(sigma=0.5, centre=(0.05, 0.05)), + disk=ag.lp_linear.Gaussian(sigma=1.5, centre=(0.05, 0.05)), + ) + ] + + galaxies_basis = [ + ag.Galaxy( + redshift=0.5, + bulge=ag.lp_basis.Basis( + profile_list=[ + ag.lp_linear.Gaussian(sigma=0.5, centre=(0.05, 0.05)), + ag.lp_linear.Gaussian(sigma=1.5, centre=(0.05, 0.05)), + ] + ), + ) + ] + + for galaxies in (galaxies, galaxies_basis): + fit_sparse = ag.FitInterferometer(dataset=dataset_sparse, galaxies=galaxies) + + assert fit_sparse.inversion.dataset.sparse_dirty_image is None + + _assert_sparse_fit_matches_dense( + dataset=interferometer_7, dataset_sparse=dataset_sparse, galaxies=galaxies + ) + + +def test__fit_figure_of_merit__sparse_operator__light_profile_and_linear_light__matches_dense( + interferometer_7, +): + """ + An ordinary light profile's visibilities are subtracted before the inversion, so the sparse data vector + must be formed from the dirty image of the profile-subtracted visibilities rather than the one cached on + the sparse operator (which is of the unsubtracted visibilities). + """ + dataset_sparse = interferometer_7.apply_sparse_operator(use_jax=False) + + galaxies = [ + ag.Galaxy( + redshift=0.5, + bulge=ag.lp.Sersic(intensity=0.1, centre=(0.05, 0.05)), + disk=ag.lp_linear.Gaussian(sigma=1.0, centre=(0.05, 0.05)), + ) + ] + + # The same ordinary light profile inside a `Basis` alongside linear ones. + galaxies_basis = [ + ag.Galaxy( + redshift=0.5, + bulge=ag.lp_basis.Basis( + profile_list=[ + ag.lp.Sersic(intensity=0.1, centre=(0.05, 0.05)), + ag.lp_linear.Gaussian(sigma=1.0, centre=(0.05, 0.05)), + ] + ), + ) + ] + + for galaxies in (galaxies, galaxies_basis): + fit_sparse = ag.FitInterferometer(dataset=dataset_sparse, galaxies=galaxies) + + assert fit_sparse.inversion.dataset.sparse_dirty_image is not None + + _assert_sparse_fit_matches_dense( + dataset=interferometer_7, dataset_sparse=dataset_sparse, galaxies=galaxies + ) From 068daa9db0933c32b4e48f919046771f69c44e00 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Sat, 26 Sep 2026 12:17:28 +0100 Subject: [PATCH 2/3] perf: sparse dirty-image correction via W~ instead of an adjoint NUFFT (PyAutoArray#575) sparse_dirty_image_from now returns d~ - W~ i_p, where i_p is the ordinary light-profile image the fit's profile_visibilities transform. W~ is the operator the sparse curvature matrix already uses, so this adds no new assumption and replaces a type-1 NUFFT over every visibility with one FFT convolution (alma CPU jit, MGE + lens Sersic: 4.54 s -> 1.30 s; dense 13.8 s). FitInterferometer caches profile_image and builds profile_visibilities from it. Tests assert profile_image transforms to the fit's profile_visibilities. (The TransformerDFT.image_from xp change in PyAutoArray, an incidental fix for a latent jit crash, is kept.) Refs PyAutoLabs/PyAutoArray#575 Co-Authored-By: Claude Opus 5.5 --- .../interferometer/fit_interferometer.py | 64 ++++++++++++------- .../interferometer/test_fit_interferometer.py | 29 +++++++-- 2 files changed, 65 insertions(+), 28 deletions(-) diff --git a/autogalaxy/interferometer/fit_interferometer.py b/autogalaxy/interferometer/fit_interferometer.py index d1a53238a..45eac417d 100644 --- a/autogalaxy/interferometer/fit_interferometer.py +++ b/autogalaxy/interferometer/fit_interferometer.py @@ -56,19 +56,25 @@ def _has_light_profile_non_linear(galaxies: List[Galaxy]) -> bool: def sparse_dirty_image_from( dataset: aa.Interferometer, galaxies: List[Galaxy], - visibilities: aa.Visibilities, - noise_map: aa.VisibilitiesNoiseMap, + image: aa.Array2D, xp=np, ) -> Optional[np.ndarray]: """ Returns the dirty image the sparse (w-tilde) inversion of an interferometer fit must use to form its data vector, or `None` if the dirty image cached on the dataset's `sparse_operator` is already the correct one. - The `sparse_operator` caches the noise-weighted dirty image `Re(Fᴴ W d)` of the dataset's visibilities, - as computed by `Interferometer.apply_sparse_operator`. A fit with ordinary (non-linear) light profiles - inverts the `profile_subtracted_visibilities` instead, so the cached image would give a data vector - inconsistent with the data the chi-squared is computed from. This function computes the dirty image of - the input (profile-subtracted) visibilities with the same weighting (one adjoint NUFFT / DFT). + The `sparse_operator` caches the noise-weighted dirty image `d~ = Re(Fᴴ W d)` of the dataset's visibilities + `d`, as computed by `Interferometer.apply_sparse_operator`. A fit with ordinary (non-linear) light profiles + inverts the profile-subtracted visibilities `d - F i_p` instead, where `i_p` is the image of those light + profiles, so the cached image would give a data vector inconsistent with the data the chi-squared is + computed from. By linearity, the dirty image of the profile-subtracted visibilities is: + + Re(Fᴴ W (d - F i_p)) = d~ - W~ i_p + + where `W~ = Re(Fᴴ W F)` is the operator the sparse inversion already uses for its curvature matrix. Applying + `W~` is one FFT convolution on the real-space grid, far cheaper than an adjoint NUFFT / DFT over every + visibility, and it relies on no assumption the sparse curvature matrix does not already make (equal real and + imaginary noise, enforced by `Interferometer.apply_sparse_operator`). Whether this is needed is decided structurally (is there a sparse operator, do any galaxies have an ordinary light profile), so the branch is fixed at trace time and is safe under `jax.jit`. Fits whose @@ -77,13 +83,12 @@ def sparse_dirty_image_from( Parameters ---------- dataset - The interferometer dataset being fitted, whose `sparse_operator` and `transformer` are used. + The interferometer dataset being fitted, whose `sparse_operator` is used. galaxies The galaxies of the fit, checked for ordinary light profiles. - visibilities - The visibilities fitted by the inversion (the profile-subtracted visibilities). - noise_map - The noise-map of the visibilities, whose real and imaginary parts weight the dirty image. + image + The image `i_p` of the ordinary light profiles on the fit's `grids.lp`, which is exactly the image that + is Fourier transformed to the fit's `profile_visibilities`. xp The array module (`numpy` or `jax.numpy`). """ @@ -93,16 +98,15 @@ def sparse_dirty_image_from( if not _has_light_profile_non_linear(galaxies=galaxies): return None - visibilities = getattr(visibilities, "array", visibilities) - noise_map = getattr(noise_map, "array", noise_map) + image = getattr(image, "array", image) - return dataset.transformer.image_from( - visibilities=aa.Visibilities( - visibilities=visibilities.real * noise_map.real**-2.0 - + 1j * visibilities.imag * noise_map.imag**-2.0 - ), + operated_image = dataset.sparse_operator.operated_matrix_slim_from( + matrix_slim=image[:, None], + extent_index_for_masked_pixel=dataset.real_space_mask.extent_index_for_masked_pixel, xp=xp, - ).array + )[:, 0] + + return xp.asarray(dataset.sparse_operator.dirty_image) - operated_image class FitInterferometer(aa.FitInterferometer, AbstractFitInversion): @@ -170,14 +174,27 @@ def __init__( self.adapt_images = adapt_images self.settings = settings or aa.Settings() + @functools.cached_property + def profile_image(self) -> aa.Array2D: + """ + Returns the summed image of every ordinary (non-linear) light profile of every galaxy, which is Fourier + transformed to the `profile_visibilities`. + """ + return self.galaxies.image_2d_from(grid=self.grids.lp, xp=self._xp) + @functools.cached_property def profile_visibilities(self) -> aa.Visibilities: """ Returns the visibilities of every light profile of every galaxy, which are computed by performing a Fourier transform to the sum of light profile images. """ - return self.galaxies.visibilities_from( - grid=self.grids.lp, transformer=self.dataset.transformer, xp=self._xp + if self.galaxies.has(cls=LightProfile): + return self.dataset.transformer.visibilities_from( + image=self.profile_image, xp=self._xp + ) + + return aa.Visibilities.zeros( + shape_slim=(self.dataset.transformer.uv_wavelengths.shape[0],) ) @functools.cached_property @@ -198,8 +215,7 @@ def galaxies_to_inversion(self) -> GalaxiesToInversion: sparse_dirty_image=sparse_dirty_image_from( dataset=self.dataset, galaxies=self.galaxies, - visibilities=self.profile_subtracted_visibilities, - noise_map=self.noise_map, + image=self.profile_image, xp=self._xp, ), ) diff --git a/test_autogalaxy/interferometer/test_fit_interferometer.py b/test_autogalaxy/interferometer/test_fit_interferometer.py index d61bd6a1e..44f9804a9 100644 --- a/test_autogalaxy/interferometer/test_fit_interferometer.py +++ b/test_autogalaxy/interferometer/test_fit_interferometer.py @@ -522,6 +522,27 @@ def _assert_sparse_fit_matches_dense(dataset, dataset_sparse, galaxies): assert fit_sparse.log_likelihood == pytest.approx(fit.log_likelihood, rel=1.0e-8) assert fit_sparse.log_evidence == pytest.approx(fit.log_evidence, rel=1.0e-8) + # The image `i_p` the sparse dirty image is corrected with (`d~ - W~ i_p`) must be exactly the image the + # fit's `profile_visibilities` are the Fourier transform of. + profile_visibilities = ag.Galaxies(galaxies=galaxies).visibilities_from( + grid=dataset_sparse.grids.lp, transformer=dataset_sparse.transformer + ) + + np.testing.assert_allclose( + dataset_sparse.transformer.visibilities_from( + image=fit_sparse.profile_image + ).array, + profile_visibilities.array, + rtol=1.0e-12, + atol=1.0e-12, + ) + np.testing.assert_allclose( + fit_sparse.profile_visibilities.array, + profile_visibilities.array, + rtol=1.0e-12, + atol=1.0e-12, + ) + def test__fit_figure_of_merit__sparse_operator__linear_light_only__matches_dense( interferometer_7, @@ -594,10 +615,10 @@ def test__fit_figure_of_merit__sparse_operator__light_profile_and_linear_light__ ] for galaxies in (galaxies, galaxies_basis): - fit_sparse = ag.FitInterferometer(dataset=dataset_sparse, galaxies=galaxies) - - assert fit_sparse.inversion.dataset.sparse_dirty_image is not None - _assert_sparse_fit_matches_dense( dataset=interferometer_7, dataset_sparse=dataset_sparse, galaxies=galaxies ) + + fit_sparse = ag.FitInterferometer(dataset=dataset_sparse, galaxies=galaxies) + + assert fit_sparse.inversion.dataset.sparse_dirty_image is not None From 36d1b43605444ea2a213204d71360e53936cf76b Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Sat, 26 Sep 2026 15:45:58 +0100 Subject: [PATCH 3/3] perf: skip the zero-image Fourier transform when galaxies have no ordinary light (PyAutoArray#575) An MGE Basis counts as a LightProfile, so profile_visibilities forward-NUFFTed an all-zero image on every call of an MGE-only fit. It now returns zeros unless _has_light_profile_non_linear (Basis-aware, structural, jit-safe) and only then evaluates profile_image. alma CPU jit, MGE-only sparse: 0.995 s -> 0.579 s. Test: MGE-only profile_visibilities are zeros and the transformer is not called. Refs PyAutoLabs/PyAutoArray#575 Co-Authored-By: Claude Opus 5.5 --- .../interferometer/fit_interferometer.py | 6 ++- .../interferometer/test_fit_interferometer.py | 46 +++++++++++++++++++ 2 files changed, 51 insertions(+), 1 deletion(-) diff --git a/autogalaxy/interferometer/fit_interferometer.py b/autogalaxy/interferometer/fit_interferometer.py index 45eac417d..5e5fdcb46 100644 --- a/autogalaxy/interferometer/fit_interferometer.py +++ b/autogalaxy/interferometer/fit_interferometer.py @@ -187,8 +187,12 @@ def profile_visibilities(self) -> aa.Visibilities: """ Returns the visibilities of every light profile of every galaxy, which are computed by performing a Fourier transform to the sum of light profile images. + + If the galaxies have no ordinary (non-linear) light profile (e.g. their light is entirely an MGE of + linear Gaussians), the image is all zeros and the Fourier transform is skipped. This is decided + structurally, so it is safe under `jax.jit`. """ - if self.galaxies.has(cls=LightProfile): + if _has_light_profile_non_linear(galaxies=self.galaxies): return self.dataset.transformer.visibilities_from( image=self.profile_image, xp=self._xp ) diff --git a/test_autogalaxy/interferometer/test_fit_interferometer.py b/test_autogalaxy/interferometer/test_fit_interferometer.py index 44f9804a9..ff967bec8 100644 --- a/test_autogalaxy/interferometer/test_fit_interferometer.py +++ b/test_autogalaxy/interferometer/test_fit_interferometer.py @@ -622,3 +622,49 @@ def test__fit_figure_of_merit__sparse_operator__light_profile_and_linear_light__ fit_sparse = ag.FitInterferometer(dataset=dataset_sparse, galaxies=galaxies) assert fit_sparse.inversion.dataset.sparse_dirty_image is not None + + +def test__profile_visibilities__linear_light_only__zeros_without_fourier_transform( + interferometer_7, monkeypatch +): + """ + A fit whose light is entirely linear (e.g. an MGE `Basis` of linear Gaussians) has an all-zero ordinary + light image, so `profile_visibilities` must be zeros without performing a Fourier transform. + """ + calls = [] + + visibilities_from = interferometer_7.transformer.visibilities_from + + def spy(*args, **kwargs): + calls.append(1) + return visibilities_from(*args, **kwargs) + + monkeypatch.setattr(interferometer_7.transformer, "visibilities_from", spy) + + galaxies = [ + ag.Galaxy( + redshift=0.5, + bulge=ag.lp_basis.Basis( + profile_list=[ + ag.lp_linear.Gaussian(sigma=0.5, centre=(0.05, 0.05)), + ag.lp_linear.Gaussian(sigma=1.5, centre=(0.05, 0.05)), + ] + ), + ) + ] + + fit = ag.FitInterferometer(dataset=interferometer_7, galaxies=galaxies) + + profile_visibilities = fit.profile_visibilities + + assert calls == [] + assert profile_visibilities.shape == interferometer_7.data.shape + assert np.all(profile_visibilities.array == 0.0) + + # With an ordinary light profile added, the transform is performed. + galaxies[0].disk = ag.lp.Sersic(intensity=0.1, centre=(0.05, 0.05)) + + fit = ag.FitInterferometer(dataset=interferometer_7, galaxies=galaxies) + + assert np.any(fit.profile_visibilities.array != 0.0) + assert len(calls) == 1