diff --git a/autoarray/inversion/inversion/interferometer/abstract.py b/autoarray/inversion/inversion/interferometer/abstract.py index ad15fec9f..f42f65149 100644 --- a/autoarray/inversion/inversion/interferometer/abstract.py +++ b/autoarray/inversion/inversion/interferometer/abstract.py @@ -1,6 +1,7 @@ import numpy as np from typing import Dict, List, Optional, Union +from autoarray import exc from autoarray.dataset.interferometer.dataset import Interferometer from autoarray.inversion.inversion.dataset_interface import DatasetInterface from autoarray.inversion.inversion.abstract import AbstractInversion @@ -69,13 +70,41 @@ def operated_mapping_matrix_list(self) -> List[np.ndarray]: This is used to construct the simultaneous linear equations which reconstruct the data. This property returns the a list of each linear object's transformed mapping matrix. + + A linear object may have a `operated_mapping_matrix_override` property, which bypasses the `mapping_matrix` + computation and transformer operation and is directly placed in the `operated_mapping_matrix_list`. Because + the override bypasses the transformer it must already be in the data's visibility space, with (complex) + shape [total_visibilities, params] (e.g. computed via an analytic Fourier transform). """ - return [ - self.transformer.transform_mapping_matrix( - mapping_matrix=linear_obj.mapping_matrix, xp=self._xp - ) - for linear_obj in self.linear_obj_list - ] + operated_mapping_matrix_list = [] + + for linear_obj in self.linear_obj_list: + operated_mapping_matrix_override = linear_obj.operated_mapping_matrix_override + + if operated_mapping_matrix_override is not None: + expected_shape = (self.data.shape[0], linear_obj.params) + + if tuple(operated_mapping_matrix_override.shape) != expected_shape: + raise exc.InversionException( + f"The `operated_mapping_matrix_override` of a linear object input to an interferometer " + f"inversion has shape {tuple(operated_mapping_matrix_override.shape)} but shape " + f"{expected_shape} ([total_visibilities, params]) is required.\n\n" + f"For an interferometer dataset the override bypasses the transformer entirely and is " + f"placed directly in the `operated_mapping_matrix_list`, therefore it must be in the " + f"data's visibility space (unlike the real-space `mapping_matrix`, which the transformer " + f"maps to visibilities)." + ) + + operated_mapping_matrix_list.append(operated_mapping_matrix_override) + + else: + operated_mapping_matrix_list.append( + self.transformer.transform_mapping_matrix( + mapping_matrix=linear_obj.mapping_matrix, xp=self._xp + ) + ) + + return operated_mapping_matrix_list @property def mapped_reconstructed_data_dict( diff --git a/autoarray/inversion/inversion/interferometer/sparse.py b/autoarray/inversion/inversion/interferometer/sparse.py index 07a4a9ed5..6853ed902 100644 --- a/autoarray/inversion/inversion/interferometer/sparse.py +++ b/autoarray/inversion/inversion/interferometer/sparse.py @@ -1,6 +1,7 @@ import numpy as np from typing import Dict, List, Union +from autoarray import exc from autoarray.dataset.interferometer.dataset import Interferometer from autoarray.inversion.inversion.dataset_interface import DatasetInterface from autoarray.inversion.inversion.interferometer.abstract import ( @@ -45,6 +46,16 @@ def __init__( The linear objects used to reconstruct the data's observed values. If multiple linear objects are passed the simultaneous linear equations are combined and solved simultaneously. """ + for linear_obj in linear_obj_list: + if linear_obj.operated_mapping_matrix_override is not None: + raise exc.InversionException( + "A linear object with an `operated_mapping_matrix_override` was passed to the sparse " + "(w-tilde) interferometer inversion, which constructs its linear algebra without an " + "explicit operated mapping matrix and therefore cannot apply the override.\n\n" + "Use the mapping formalism instead (e.g. do not call `apply_sparse_operator` on the " + "interferometer dataset)." + ) + super().__init__( dataset=dataset, linear_obj_list=linear_obj_list, diff --git a/autoarray/inversion/linear_obj/linear_obj.py b/autoarray/inversion/linear_obj/linear_obj.py index 623812f84..52372adb2 100644 --- a/autoarray/inversion/linear_obj/linear_obj.py +++ b/autoarray/inversion/linear_obj/linear_obj.py @@ -122,7 +122,8 @@ def pixel_signals_from(self, signal_scale) -> np.ndarray: def operated_mapping_matrix_override(self) -> Optional[np.ndarray]: """ An `Inversion` takes the `mapping_matrix` of each linear object and combines it with the data's operators - (e.g. a PSF for `Imaging` data) to compute the `operated_mapping_matrix`. + (e.g. a PSF for `Imaging` data, the transformer for `Interferometer` data) to compute the + `operated_mapping_matrix`. If this property is overwritten this operation is not performed, with the `operated_mapping_matrix` output by this property automatically used instead. @@ -132,9 +133,19 @@ def operated_mapping_matrix_override(self) -> Optional[np.ndarray]: region which is blurred into the masked region which is linear solved for. This flux is outside the region that defines the `mapping_matrix` and thus this override is required to properly incorporate it. + Because the override bypasses the data's operators entirely, it must be in the data's space, which depends + on the dataset type being fitted: + + - `Imaging`: a real matrix of dimensions (total_mask_pixels, total_parameters), e.g. the PSF-convolved + image of each linear object's parameter. + + - `Interferometer`: a complex matrix of dimensions (total_visibilities, total_parameters), e.g. the + visibilities of each linear object's parameter computed via an analytic Fourier transform. The + transformer (NUFFT / DFT) is not applied to the override. + Returns ------- - An operated mapping matrix of dimensions (total_mask_pixels, total_parameters) which overrides the mapping + An operated mapping matrix in the data's space (see above) which overrides the mapping matrix calculations performed in the linear equation solvers. """ return None diff --git a/test_autoarray/inversion/inversion/interferometer/test_interferometer.py b/test_autoarray/inversion/inversion/interferometer/test_interferometer.py index ddba59baf..60b28d5b5 100644 --- a/test_autoarray/inversion/inversion/interferometer/test_interferometer.py +++ b/test_autoarray/inversion/inversion/interferometer/test_interferometer.py @@ -64,6 +64,179 @@ def test__fast_chi_squared( assert inversion.fast_chi_squared == pytest.approx(chi_squared, 1.0e-4) +def test__operated_mapping_matrix_list__override_is_honored(): + mask = aa.Mask2D( + mask=[ + [True, True, True, True, True, True, True], + [True, True, True, True, True, True, True], + [True, True, True, False, True, True, True], + [True, True, False, False, False, True, True], + [True, True, True, False, True, True, True], + [True, True, True, True, True, True, True], + [True, True, True, True, True, True, True], + ], + pixel_scales=2.0, + ) + + n_visibilities = 5 + rng = np.random.default_rng(seed=0) + data = aa.Visibilities( + visibilities=rng.normal(size=(n_visibilities, 2)).astype(np.float64) + ) + noise_map = aa.VisibilitiesNoiseMap( + visibilities=np.ones((n_visibilities, 2), dtype=np.float64) + ) + uv_wavelengths = rng.normal(size=(n_visibilities, 2)).astype(np.float64) + + dataset = aa.Interferometer( + data=data, + noise_map=noise_map, + uv_wavelengths=uv_wavelengths, + real_space_mask=mask, + transformer_class=aa.TransformerDFT, + ) + + mapping_matrix = np.ones((mask.pixels_in_mask, 1)) + override = (999.0 + 1.0j) * np.ones((n_visibilities, 1)) + + linear_obj_override = aa.m.MockLinearObjFuncList( + parameters=1, + mapping_matrix=mapping_matrix, + operated_mapping_matrix_override=override, + ) + linear_obj_no_override = aa.m.MockLinearObjFuncList( + parameters=1, + mapping_matrix=mapping_matrix, + ) + + inversion = aa.Inversion( + dataset=dataset, + linear_obj_list=[linear_obj_override, linear_obj_no_override], + ) + + operated_mapping_matrix_list = inversion.operated_mapping_matrix_list + + assert operated_mapping_matrix_list[0] == pytest.approx(override, 1.0e-8) + + transformed_mapping_matrix = dataset.transformer.transform_mapping_matrix( + mapping_matrix=mapping_matrix + ) + + assert operated_mapping_matrix_list[1] == pytest.approx( + transformed_mapping_matrix, 1.0e-8 + ) + + assert inversion.operated_mapping_matrix[:, 0] == pytest.approx( + override[:, 0], 1.0e-8 + ) + assert inversion.curvature_matrix.shape == (2, 2) + assert inversion.data_vector.shape == (2,) + + +def test__operated_mapping_matrix_override__wrong_shape_raises(): + mask = aa.Mask2D( + mask=[ + [True, True, True, True, True, True, True], + [True, True, True, True, True, True, True], + [True, True, True, False, True, True, True], + [True, True, False, False, False, True, True], + [True, True, True, False, True, True, True], + [True, True, True, True, True, True, True], + [True, True, True, True, True, True, True], + ], + pixel_scales=2.0, + ) + + n_visibilities = 7 + rng = np.random.default_rng(seed=0) + data = aa.Visibilities( + visibilities=rng.normal(size=(n_visibilities, 2)).astype(np.float64) + ) + noise_map = aa.VisibilitiesNoiseMap( + visibilities=np.ones((n_visibilities, 2), dtype=np.float64) + ) + uv_wavelengths = rng.normal(size=(n_visibilities, 2)).astype(np.float64) + + dataset = aa.Interferometer( + data=data, + noise_map=noise_map, + uv_wavelengths=uv_wavelengths, + real_space_mask=mask, + transformer_class=aa.TransformerDFT, + ) + + # A real-space shaped override (e.g. [total_mask_pixels, params]) is not valid for an + # interferometer inversion, whose override must be in visibility space. + linear_obj = aa.m.MockLinearObjFuncList( + parameters=1, + mapping_matrix=np.ones((mask.pixels_in_mask, 1)), + operated_mapping_matrix_override=np.ones((mask.pixels_in_mask, 1)), + ) + + inversion = aa.Inversion(dataset=dataset, linear_obj_list=[linear_obj]) + + with pytest.raises(aa.exc.InversionException): + inversion.operated_mapping_matrix_list + + +def test__operated_mapping_matrix_override__sparse_operator_raises(): + mask = aa.Mask2D( + mask=[ + [True, True, True, True, True, True, True], + [True, True, True, True, True, True, True], + [True, True, True, False, True, True, True], + [True, True, False, False, False, True, True], + [True, True, True, False, True, True, True], + [True, True, True, True, True, True, True], + [True, True, True, True, True, True, True], + ], + pixel_scales=2.0, + ) + + grid = aa.Grid2D.from_mask(mask=mask, over_sample_size=1) + + mesh = aa.mesh.Delaunay(pixels=9) + image_mesh = aa.image_mesh.Overlay(shape=(3, 3)) + image_mesh_grid = image_mesh.image_plane_mesh_grid_from(mask=mask, adapt_data=None) + + interpolator = mesh.interpolator_from( + source_plane_data_grid=grid, + source_plane_mesh_grid=image_mesh_grid, + ) + mapper = aa.Mapper(interpolator=interpolator) + + n_visibilities = 5 + rng = np.random.default_rng(seed=0) + data = aa.Visibilities( + visibilities=rng.normal(size=(n_visibilities, 2)).astype(np.float64) + ) + noise_map = aa.VisibilitiesNoiseMap( + visibilities=np.ones((n_visibilities, 2), dtype=np.float64) + ) + uv_wavelengths = rng.normal(size=(n_visibilities, 2)).astype(np.float64) + + dataset_sparse = aa.Interferometer( + data=data, + noise_map=noise_map, + uv_wavelengths=uv_wavelengths, + real_space_mask=mask, + transformer_class=aa.TransformerDFT, + ).apply_sparse_operator(use_jax=False) + + linear_obj = aa.m.MockLinearObjFuncList( + parameters=1, + mapping_matrix=np.ones((mask.pixels_in_mask, 1)), + operated_mapping_matrix_override=(999.0 + 1.0j) + * np.ones((n_visibilities, 1)), + ) + + with pytest.raises(aa.exc.InversionException): + aa.Inversion( + dataset=dataset_sparse, + linear_obj_list=[mapper, linear_obj], + ) + + def test__curvature_matrix__interferometer_sparse_operator__delaunay__identical_to_mapping(): mask = aa.Mask2D( mask=[