Merge pull request #460 from PyAutoLabs/feature/interferometer-operated-override

Honor operated_mapping_matrix_override in interferometer inversions

Author: Jammy2211

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(

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,

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

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=[