From bdf2aec55651dcb404f7de6d8aa0e72ab4f41c70 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.tracer_to_inversion passes sparse_dirty_image via autogalaxy's sparse_dirty_image_from over the tracer galaxies. Test: lens Sersic + source MGE, sparse vs dense to 1e-8, with an all-linear control. Refs PyAutoLabs/PyAutoArray#575 Co-Authored-By: Claude Opus 5.5 --- autolens/interferometer/fit_interferometer.py | 12 ++++- .../interferometer/test_fit_interferometer.py | 49 +++++++++++++++++++ 2 files changed, 60 insertions(+), 1 deletion(-) diff --git a/autolens/interferometer/fit_interferometer.py b/autolens/interferometer/fit_interferometer.py index 99060015e..18f2864f1 100644 --- a/autolens/interferometer/fit_interferometer.py +++ b/autolens/interferometer/fit_interferometer.py @@ -24,6 +24,7 @@ import autogalaxy as ag from autogalaxy.abstract_fit import AbstractFitInversion +from autogalaxy.interferometer.fit_interferometer import sparse_dirty_image_from from autolens.lens.tracer import Tracer from autolens.lens.to_inversion import TracerToInversion @@ -132,12 +133,21 @@ def profile_subtracted_visibilities(self) -> aa.Visibilities: @property def tracer_to_inversion(self) -> TracerToInversion: + profile_subtracted_visibilities = self.profile_subtracted_visibilities + dataset = aa.DatasetInterface( - data=self.profile_subtracted_visibilities, + data=profile_subtracted_visibilities, noise_map=self.noise_map, 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.tracer.galaxies, + visibilities=profile_subtracted_visibilities, + noise_map=self.noise_map, + xp=self._xp, + ), ) return TracerToInversion( diff --git a/test_autolens/interferometer/test_fit_interferometer.py b/test_autolens/interferometer/test_fit_interferometer.py index 46cef27ce..0a4465a9a 100644 --- a/test_autolens/interferometer/test_fit_interferometer.py +++ b/test_autolens/interferometer/test_fit_interferometer.py @@ -1,6 +1,7 @@ import numpy as np import pytest +import autoarray as aa import autolens as al @@ -386,3 +387,51 @@ def test__model_visibilities_of_planes_list(interferometer_7): + fit.galaxy_model_visibilities_dict[galaxy_pix_1].array, 1.0e-4, ) + + +def test__fit_figure_of_merit__sparse_operator__lens_light_profile_and_source_mge__matches_dense( + interferometer_7, +): + """ + With the sparse operator applied, a lens ordinary light profile plus a source MGE (linear Gaussians) + must reproduce the dense fit: the lens light's visibilities are subtracted before the inversion, so the + sparse data vector must use the dirty image of the profile-subtracted visibilities. + """ + dataset_sparse = interferometer_7.apply_sparse_operator(use_jax=False) + + lens = al.Galaxy( + redshift=0.5, + bulge=al.lp.Sersic(intensity=0.1, centre=(0.05, 0.05)), + mass=al.mp.Isothermal(centre=(0.0, 0.0), einstein_radius=1.0), + ) + source = al.Galaxy( + redshift=1.0, + bulge=al.lp_basis.Basis( + profile_list=[ + al.lp_linear.Gaussian(sigma=sigma, centre=(0.1, 0.1)) + for sigma in (0.3, 1.0, 3.0) + ] + ), + ) + + # The second case, with no lens light, is the all-linear control which uses the cached dirty image. + for galaxies, has_lens_light in ( + ([lens, source], True), + ([al.Galaxy(redshift=0.5, mass=lens.mass), source], False), + ): + tracer = al.Tracer(galaxies=galaxies) + + fit = al.FitInterferometer(dataset=interferometer_7, tracer=tracer) + fit_sparse = al.FitInterferometer(dataset=dataset_sparse, tracer=tracer) + + assert ( + fit_sparse.inversion.dataset.sparse_dirty_image is not None + ) is has_lens_light + + assert isinstance(fit_sparse.inversion, aa.InversionInterferometerSparse) + assert isinstance(fit.inversion, aa.InversionInterferometerMapping) + + 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) From acb136cdd918c12d7f89b7fac4313de54fee0a7e 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) tracer_to_inversion passes the cached ordinary light-profile image to autogalaxy's sparse_dirty_image_from (d~ - W~ i_p); profile_visibilities is built from the same cached image. Test asserts profile_image transforms to the fit's profile_visibilities. Refs PyAutoLabs/PyAutoArray#575 Co-Authored-By: Claude Opus 5.5 --- autolens/interferometer/fit_interferometer.py | 24 +++++++++++----- .../interferometer/test_fit_interferometer.py | 28 +++++++++++++++++-- 2 files changed, 42 insertions(+), 10 deletions(-) diff --git a/autolens/interferometer/fit_interferometer.py b/autolens/interferometer/fit_interferometer.py index 18f2864f1..f3eee8b5a 100644 --- a/autolens/interferometer/fit_interferometer.py +++ b/autolens/interferometer/fit_interferometer.py @@ -113,14 +113,27 @@ def _xp(self): return jnp return np + @cached_property + def profile_image(self) -> aa.Array2D: + """ + Returns the summed image of every ordinary (non-linear) light profile in the tracer, which is Fourier + transformed to the `profile_visibilities`. + """ + return self.tracer.image_2d_from(grid=self.grids.lp, xp=self._xp) + @property def profile_visibilities(self) -> aa.Visibilities: """ Returns the visibilities of every light profile in the tracer, which are computed by performing a Fourier transform to the sum of light profile images. """ - return self.tracer.visibilities_from( - grid=self.grids.lp, transformer=self.dataset.transformer, xp=self._xp + if self.tracer.has(cls=ag.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],) ) @property @@ -133,10 +146,8 @@ def profile_subtracted_visibilities(self) -> aa.Visibilities: @property def tracer_to_inversion(self) -> TracerToInversion: - profile_subtracted_visibilities = self.profile_subtracted_visibilities - dataset = aa.DatasetInterface( - data=profile_subtracted_visibilities, + data=self.profile_subtracted_visibilities, noise_map=self.noise_map, grids=self.grids, transformer=self.dataset.transformer, @@ -144,8 +155,7 @@ def tracer_to_inversion(self) -> TracerToInversion: sparse_dirty_image=sparse_dirty_image_from( dataset=self.dataset, galaxies=self.tracer.galaxies, - visibilities=profile_subtracted_visibilities, - noise_map=self.noise_map, + image=self.profile_image, xp=self._xp, ), ) diff --git a/test_autolens/interferometer/test_fit_interferometer.py b/test_autolens/interferometer/test_fit_interferometer.py index 0a4465a9a..693f39a92 100644 --- a/test_autolens/interferometer/test_fit_interferometer.py +++ b/test_autolens/interferometer/test_fit_interferometer.py @@ -424,9 +424,6 @@ def test__fit_figure_of_merit__sparse_operator__lens_light_profile_and_source_mg fit = al.FitInterferometer(dataset=interferometer_7, tracer=tracer) fit_sparse = al.FitInterferometer(dataset=dataset_sparse, tracer=tracer) - assert ( - fit_sparse.inversion.dataset.sparse_dirty_image is not None - ) is has_lens_light assert isinstance(fit_sparse.inversion, aa.InversionInterferometerSparse) assert isinstance(fit.inversion, aa.InversionInterferometerMapping) @@ -435,3 +432,28 @@ def test__fit_figure_of_merit__sparse_operator__lens_light_profile_and_source_mg fit.log_likelihood, rel=1.0e-8 ) assert fit_sparse.log_evidence == pytest.approx(fit.log_evidence, rel=1.0e-8) + + assert ( + fit_sparse.inversion.dataset.sparse_dirty_image is not None + ) is has_lens_light + + # 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 = tracer.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, + ) From ac333b175c660bd64e1d5920ddf2b489e532229d 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 the tracer has no ordinary light (PyAutoArray#575) profile_visibilities uses autogalaxy's Basis-aware _has_light_profile_non_linear over tracer.galaxies and returns zeros for an all-linear (e.g. MGE source, mass-only lens) tracer instead of forward-NUFFTing an all-zero image. Test: profile_visibilities are zeros and the transformer is not called. Refs PyAutoLabs/PyAutoArray#575 Co-Authored-By: Claude Opus 5.5 --- autolens/interferometer/fit_interferometer.py | 11 ++++- .../interferometer/test_fit_interferometer.py | 47 +++++++++++++++++++ 2 files changed, 56 insertions(+), 2 deletions(-) diff --git a/autolens/interferometer/fit_interferometer.py b/autolens/interferometer/fit_interferometer.py index f3eee8b5a..65e10adde 100644 --- a/autolens/interferometer/fit_interferometer.py +++ b/autolens/interferometer/fit_interferometer.py @@ -24,7 +24,10 @@ import autogalaxy as ag from autogalaxy.abstract_fit import AbstractFitInversion -from autogalaxy.interferometer.fit_interferometer import sparse_dirty_image_from +from autogalaxy.interferometer.fit_interferometer import ( + _has_light_profile_non_linear, + sparse_dirty_image_from, +) from autolens.lens.tracer import Tracer from autolens.lens.to_inversion import TracerToInversion @@ -126,8 +129,12 @@ def profile_visibilities(self) -> aa.Visibilities: """ Returns the visibilities of every light profile in the tracer, which are computed by performing a Fourier transform to the sum of light profile images. + + If the tracer has no ordinary (non-linear) light profile (e.g. its 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.tracer.has(cls=ag.LightProfile): + if _has_light_profile_non_linear(galaxies=self.tracer.galaxies): return self.dataset.transformer.visibilities_from( image=self.profile_image, xp=self._xp ) diff --git a/test_autolens/interferometer/test_fit_interferometer.py b/test_autolens/interferometer/test_fit_interferometer.py index 693f39a92..bd784777e 100644 --- a/test_autolens/interferometer/test_fit_interferometer.py +++ b/test_autolens/interferometer/test_fit_interferometer.py @@ -457,3 +457,50 @@ def test__fit_figure_of_merit__sparse_operator__lens_light_profile_and_source_mg rtol=1.0e-12, atol=1.0e-12, ) + + +def test__profile_visibilities__linear_light_only__zeros_without_fourier_transform( + interferometer_7, monkeypatch +): + """ + A tracer whose light is entirely linear (a lens with only mass and a source 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) + + source = al.Galaxy( + redshift=1.0, + bulge=al.lp_basis.Basis( + profile_list=[ + al.lp_linear.Gaussian(sigma=sigma, centre=(0.1, 0.1)) + for sigma in (0.3, 1.0, 3.0) + ] + ), + ) + + tracer = al.Tracer( + galaxies=[ + al.Galaxy( + redshift=0.5, + mass=al.mp.Isothermal(centre=(0.0, 0.0), einstein_radius=1.0), + ), + source, + ] + ) + + fit = al.FitInterferometer(dataset=interferometer_7, tracer=tracer) + + profile_visibilities = fit.profile_visibilities + + assert calls == [] + assert profile_visibilities.shape == interferometer_7.data.shape + assert np.all(profile_visibilities.array == 0.0)