diff --git a/autoarray/inversion/inversion/interferometer/sparse.py b/autoarray/inversion/inversion/interferometer/sparse.py index 49fb23ba0..94be11c48 100644 --- a/autoarray/inversion/inversion/interferometer/sparse.py +++ b/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): """ diff --git a/test_autoarray/inversion/inversion/interferometer/test_interferometer.py b/test_autoarray/inversion/inversion/interferometer/test_interferometer.py index df33d8ae3..2361fb055 100644 --- a/test_autoarray/inversion/inversion/interferometer/test_interferometer.py +++ b/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(), + )