perf: interferometer sparse data vector without the dense mapping matrix

autoarray/inversion/inversion/interferometer/sparse.py

@@ -83,13 +83,17 @@ def data_vector(self) -> np.ndarray:
 
         The `data_vector` is computed as `Lᵀ d~`, where `L` is the real-space mapping matrix of every
         linear object stacked horizontally and `d~ = Re(Fᴴ W d)` is the dirty image cached on the
-        dataset's `sparse_operator`. This single expression covers every linear object type:
+        dataset's `sparse_operator`. Each linear object contributes its own block, concatenated in
+        `linear_obj_list` order (matching `param_range_list_from(cls=LinearObj)`):
 
-        - For a `Mapper` column `a`, `aᵀ d~` is the mapper's data vector entry.
+        - For a `Mapper` the block `Aᵀ d~` is computed from the mapper's sparse COO triplets
+          (`_data_vector_mapper_from`), as a segment-sum of `vals * d~[rows]` over source pixel
+          columns. The dense (masked pixels x source pixels) mapping matrix is therefore never
+          built in the likelihood, which at 91,507 masked pixels and 5,746 source pixels is 4.2 GB.
         - For an `AbstractLinearObjFuncList` column `b`, `bᵀ d~ = bᵀ Re(Fᴴ W d) = Re((F b)ᴴ W d)`,
           which is exactly the entry the mapping (dense) formalism computes via the linear
-          function's transformed mapping matrix. Linear function lists therefore require no
-          separate branch here.
+          function's transformed mapping matrix. Its mapping matrix is small (one column per
+          linear function), so it is used directly.
 
         The cached dirty image is that of the visibilities the `sparse_operator` was built from. When
         the inversion's data differs from them (e.g. a `DatasetInterface` whose data has the visibilities
@@ -101,7 +105,95 @@ def data_vector(self) -> np.ndarray:
         if dirty_image is None:
             dirty_image = self.dataset.sparse_operator.dirty_image
 
-        return self._xp.dot(self.mapping_matrix.T, dirty_image)
+        dirty_image = self._xp.asarray(
+            getattr(dirty_image, "array", dirty_image), dtype=self._xp.float64
+        )
+
+        data_vector_list = []
+
+        for linear_obj in self.linear_obj_list:
+            if isinstance(linear_obj, Mapper):
+                data_vector_list.append(
+                    self._data_vector_mapper_from(
+                        mapper=linear_obj, dirty_image=dirty_image
+                    )
+                )
+            else:
+                data_vector_list.append(
+                    self._xp.dot(linear_obj.mapping_matrix.T, dirty_image)
+                )
+
+        if len(data_vector_list) == 1:
+            return data_vector_list[0]
+
+        return self._xp.concatenate(data_vector_list)
+
+    def _sparse_triplets_slim_from(self, mapper: Mapper):
+        """
+        Returns the sparse COO triplets `(rows, cols, vals)` of a mapper's real-space mapping
+        operator `A`, with `rows` indexed on the slim (masked-pixel) grid the dirty image lives on.
+
+        This is the sibling of `_sparse_triplets_curvature_from`, whose rows are extent-flat
+        indices on the `W~` operator's rectangular grid and therefore cannot index the slim dirty
+        image.
+
+        The triplets reproduce `mapper_util.mapping_matrix_from` exactly: `vals` carry the
+        `sub_fraction` normalisation of each sub-pixel's slim pixel, and stencil slots that
+        `mapping_matrix_from` drops (slot index `>= pix_sizes_for_sub_slim_index`, or a negative
+        source pixel index) are given zero weight and column 0, so they contribute nothing.
+
+        Parameters
+        ----------
+        mapper
+            The mapper whose mapping operator is expressed in sparse triplet form.
+        """
+        xp = self._xp
+
+        rows, cols, vals = mapper_util.sparse_triplets_from(
+            pix_indexes_for_sub=mapper.pix_indexes_for_sub_slim_index,
+            pix_weights_for_sub=mapper.pix_weights_for_sub_slim_index,
+            slim_index_for_sub=mapper.slim_index_for_sub_slim_index,
+            fft_index_for_masked_pixel=self.mask.extent_index_for_masked_pixel,
+            sub_fraction_slim=mapper.over_sampler.sub_fraction.array,
+            return_rows_slim=True,
+            xp=xp,
+        )
+
+        pix_indexes = xp.asarray(mapper.pix_indexes_for_sub_slim_index)
+        M_sub, P = pix_indexes.shape
+
+        pix_sizes = xp.asarray(mapper.pix_sizes_for_sub_slim_index, dtype=xp.int32)
+        slot = xp.tile(xp.arange(P, dtype=xp.int32), M_sub)
+        valid = (slot < xp.repeat(pix_sizes, P)) & (cols >= 0)
+
+        vals = xp.where(valid, vals, 0.0)
+        cols = xp.where(valid, cols, 0)
+
+        return rows, cols, vals
+
+    def _data_vector_mapper_from(self, mapper: Mapper, dirty_image) -> np.ndarray:
+        """
+        Returns a mapper's block of the `data_vector`, `Aᵀ d~`, computed from the sparse triplets
+        of its mapping operator `A` without building the dense mapping matrix.
+
+        Entry `j` is the sum over the non-zeros of `A`'s column `j` of `val * d~[row]`, i.e. a
+        segment-sum of `vals * dirty_image[rows]` over the source pixel columns.
+
+        Parameters
+        ----------
+        mapper
+            The mapper whose data vector block is computed.
+        dirty_image
+            The slim (masked-pixel indexed) dirty image `d~`.
+        """
+        rows, cols, vals = self._sparse_triplets_slim_from(mapper=mapper)
+
+        weights = vals * dirty_image[rows]
+
+        if self._xp is np:
+            return np.bincount(cols, weights=weights, minlength=mapper.params)
+
+        return self._xp.zeros(mapper.params, dtype=weights.dtype).at[cols].add(weights)
 
     def _sparse_triplets_curvature_from(self, mapper: Mapper):
         """

test_autoarray/inversion/inversion/interferometer/test_interferometer.py

@@ -1739,3 +1739,162 @@ def test__inversion_interferometer_mask__is_the_real_space_mask_for_every_datase
         assert isinstance(inversion, aa.InversionInterferometerSparse)
         assert (inversion.mask == mask).all()
         assert inversion.mask.pixel_scales == mask.pixel_scales
+
+
+def _over_sampled_mapper_from(mask, pixels, shape, over_sample_size):
+    grid = aa.Grid2D.from_mask(mask=mask, over_sample_size=over_sample_size)
+
+    mesh = aa.mesh.Delaunay(pixels=pixels)
+    image_mesh_grid = aa.image_mesh.Overlay(shape=shape).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,
+    )
+
+    return aa.Mapper(
+        interpolator=interpolator, regularization=aa.reg.Constant(coefficient=1.0)
+    )
+
+
+def _dense_data_vector_from(inversion, dirty_image):
+    return np.asarray(inversion.mapping_matrix).T @ np.asarray(dirty_image)
+
+
+def test__interferometer_sparse_operator__data_vector_from_triplets_matches_dense_mapping_matrix():
+    """
+    `InversionInterferometerSparse.data_vector` computes each mapper block from the mapper's sparse
+    triplets (a segment-sum over source columns) instead of the dense `mapping_matrix.T @ d~`. It must
+    reproduce the dense product exactly for an over-sampled mapper (where the triplets carry the
+    `sub_fraction` normalisation), for a mapper mixed with a linear function list (in either order), and
+    for a `DatasetInterface` whose `sparse_dirty_image` differs from the operator's cached dirty image.
+    """
+    mask = aa.Mask2D.circular(shape_native=(12, 12), pixel_scales=1.0, radius=4.0)
+
+    rng = np.random.default_rng(seed=3)
+    n_visibilities = 32
+
+    dataset = aa.Interferometer(
+        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),
+        real_space_mask=mask,
+        transformer_class=aa.TransformerDFT,
+    )
+    dataset_sparse = dataset.apply_sparse_operator(use_jax=False)
+    dirty_image = np.asarray(dataset_sparse.sparse_operator.dirty_image)
+
+    mapper = _over_sampled_mapper_from(
+        mask=mask, pixels=16, shape=(4, 4), over_sample_size=2
+    )
+    linear_obj = aa.m.MockLinearObjFuncList(
+        parameters=2,
+        mapping_matrix=rng.normal(size=(mask.pixels_in_mask, 2)),
+    )
+
+    for linear_obj_list in ([mapper], [mapper, linear_obj], [linear_obj, mapper]):
+        inversion = aa.Inversion(
+            dataset=dataset_sparse, linear_obj_list=linear_obj_list
+        )
+
+        assert type(inversion) is aa.InversionInterferometerSparse
+
+        reference = _dense_data_vector_from(inversion, dirty_image)
+
+        np.testing.assert_allclose(
+            np.asarray(inversion.data_vector),
+            reference,
+            rtol=1.0e-12,
+            atol=1.0e-12 * np.abs(reference).max(),
+        )
+
+    # A `DatasetInterface` with its own `sparse_dirty_image` must use it in place of the operator's.
+    sparse_dirty_image = rng.normal(size=mask.pixels_in_mask)
+
+    inversion = aa.Inversion(
+        dataset=aa.DatasetInterface(
+            data=dataset.data,
+            noise_map=dataset.noise_map,
+            grids=dataset.grids,
+            transformer=dataset.transformer,
+            sparse_operator=dataset_sparse.sparse_operator,
+            sparse_dirty_image=sparse_dirty_image,
+        ),
+        linear_obj_list=[mapper, linear_obj],
+    )
+
+    reference = _dense_data_vector_from(inversion, sparse_dirty_image)
+
+    np.testing.assert_allclose(
+        np.asarray(inversion.data_vector),
+        reference,
+        rtol=1.0e-12,
+        atol=1.0e-12 * np.abs(reference).max(),
+    )
+
+    assert (
+        np.abs(_dense_data_vector_from(inversion, dirty_image) - reference).max()
+        > 1.0e-4 * np.abs(reference).max()
+    )
+
+
+def test__interferometer_sparse_operator__data_vector_from_triplets_matches_dense_mapping_matrix__jax():
+    """
+    The JAX backend of the triplet data vector (a scatter-add over source columns) must reproduce the
+    dense `mapping_matrix.T @ d~` for an over-sampled mapper and a mapper mixed with a linear function list.
+    """
+    pytest.importorskip("jax")
+
+    import jax.numpy as jnp
+
+    mask = aa.Mask2D.circular(shape_native=(12, 12), pixel_scales=1.0, radius=4.0)
+
+    rng = np.random.default_rng(seed=5)
+    n_visibilities = 32
+
+    dataset = aa.Interferometer(
+        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),
+        real_space_mask=mask,
+        transformer_class=aa.TransformerDFT,
+    )
+    dataset_sparse = dataset.apply_sparse_operator(use_jax=False)
+    dirty_image = np.asarray(dataset_sparse.sparse_operator.dirty_image)
+
+    mapper = _over_sampled_mapper_from(
+        mask=mask, pixels=16, shape=(4, 4), over_sample_size=2
+    )
+    linear_obj = aa.m.MockLinearObjFuncList(
+        parameters=2,
+        mapping_matrix=rng.normal(size=(mask.pixels_in_mask, 2)),
+    )
+
+    for linear_obj_list in ([mapper], [linear_obj, mapper]):
+        inversion_np = aa.Inversion(
+            dataset=dataset_sparse, linear_obj_list=linear_obj_list, xp=np
+        )
+        inversion_jax = aa.Inversion(
+            dataset=dataset_sparse, linear_obj_list=linear_obj_list, xp=jnp
+        )
+
+        assert isinstance(inversion_jax, aa.InversionInterferometerSparse)
+
+        reference = _dense_data_vector_from(inversion_np, dirty_image)
+
+        np.testing.assert_allclose(
+            np.asarray(inversion_jax.data_vector),
+            reference,
+            rtol=1.0e-12,
+            atol=1.0e-12 * np.abs(reference).max(),
+        )