diff --git a/autogalaxy/interferometer/fit_interferometer.py b/autogalaxy/interferometer/fit_interferometer.py index bb6ef154..5e5fdcb4 100644 --- a/autogalaxy/interferometer/fit_interferometer.py +++ b/autogalaxy/interferometer/fit_interferometer.py @@ -24,6 +24,89 @@ 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], + 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 `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 + light is entirely linear return `None` and pay no extra cost. + + Parameters + ---------- + dataset + The interferometer dataset being fitted, whose `sparse_operator` is used. + galaxies + The galaxies of the fit, checked for ordinary light profiles. + 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`). + """ + if dataset.sparse_operator is None: + return None + + if not _has_light_profile_non_linear(galaxies=galaxies): + return None + + image = getattr(image, "array", image) + + 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, + )[:, 0] + + return xp.asarray(dataset.sparse_operator.dirty_image) - operated_image class FitInterferometer(aa.FitInterferometer, AbstractFitInversion): @@ -91,14 +174,31 @@ 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. + + 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`. """ - return self.galaxies.visibilities_from( - grid=self.grids.lp, transformer=self.dataset.transformer, xp=self._xp + if _has_light_profile_non_linear(galaxies=self.galaxies): + 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 @@ -116,6 +216,12 @@ 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, + image=self.profile_image, + xp=self._xp, + ), ) return GalaxiesToInversion( diff --git a/test_autogalaxy/interferometer/test_fit_interferometer.py b/test_autogalaxy/interferometer/test_fit_interferometer.py index 68439a85..ff967bec 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,161 @@ 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) + + # 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, +): + """ + 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): + _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 + + +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