diff --git a/autoarray/inversion/inversion/interferometer/sparse.py b/autoarray/inversion/inversion/interferometer/sparse.py index 94be11c48..d30970853 100644 --- a/autoarray/inversion/inversion/interferometer/sparse.py +++ b/autoarray/inversion/inversion/interferometer/sparse.py @@ -14,6 +14,7 @@ from autoarray.inversion.mappers import mapper_util from autoarray.settings import Settings from autoarray.inversion.mappers.abstract import Mapper +from autoarray.structures.arrays.uniform_2d import Array2D from autoarray.structures.visibilities import Visibilities from autoarray.inversion.inversion import inversion_util @@ -513,6 +514,93 @@ def _curvature_matrix_func_list_and_mapper(self) -> np.ndarray: return curvature_matrix + @property + def mapped_reconstructed_data_dict( + self, + ) -> Dict[LinearObj, Array2D]: + """ + When constructing the simultaneous linear equations (via vectors and matrices) the quantities of each individual + linear object (e.g. their `mapping_matrix`) are combined into single ndarrays. This does not track which + quantities belong to which linear objects, therefore the linear equation's solutions (which are returned as + ndarrays) do not contain information on which linear object(s) they correspond to. + + This function converts an ndarray of a `reconstruction` to a dictionary of ndarrays containing each linear + object's reconstructed image (in real space, before the non-uniform fast Fourier transform), where the keys + are the instances of each linear object in the inversion. + + Each linear object's reconstructed image is `L s`, where `L` is its real-space mapping matrix and `s` its + block of the `reconstruction`. This is the mirror of the sparse `data_vector` (`Lᵀ d~`): + + - For a `Mapper` the image `A s` is computed from the mapper's sparse COO triplets + (`_sparse_triplets_slim_from`), as a scatter-add of `vals * s[cols]` over the slim image pixel + rows. The dense (masked pixels x source pixels) mapping matrix is therefore never built, which + on a 384 x 384 unmasked grid with a 48² source mesh is 2.7 GB. + - For an `AbstractLinearObjFuncList` the mapping matrix is small (one column per linear function), + so the dense `mapping_matrix @ reconstruction` product is used directly, as in the inherited + implementation. + + The result is used by `mapped_reconstructed_operated_data_dict` (model visibilities) and by fit + outputs such as galaxy image dictionaries. + """ + mapped_reconstructed_data_dict = {} + + reconstruction_dict = self.source_quantity_dict_from( + source_quantity=self.reconstruction + ) + + for linear_obj in self.linear_obj_list: + reconstruction = reconstruction_dict[linear_obj] + + if isinstance(linear_obj, Mapper): + mapped_reconstructed_data = self._mapped_reconstructed_data_mapper_from( + mapper=linear_obj, reconstruction=reconstruction + ) + else: + mapped_reconstructed_data = ( + inversion_util.mapped_reconstructed_data_via_mapping_matrix_from( + mapping_matrix=linear_obj.mapping_matrix, + reconstruction=reconstruction, + xp=self._xp, + ) + ) + + mapped_reconstructed_data = Array2D( + values=mapped_reconstructed_data, mask=self.mask + ) + + mapped_reconstructed_data_dict[linear_obj] = mapped_reconstructed_data + + return mapped_reconstructed_data_dict + + def _mapped_reconstructed_data_mapper_from( + self, mapper: Mapper, reconstruction + ) -> np.ndarray: + """ + Returns a mapper's reconstructed image, `A s`, computed from the sparse triplets of its mapping + operator `A` without building the dense mapping matrix. + + Entry `i` (a slim image pixel) is the sum over the non-zeros of `A`'s row `i` of `val * s[col]`, + i.e. a scatter-add of `vals * reconstruction[cols]` over the slim image pixel rows. This is the + transpose of the segment-sum performed by `_data_vector_mapper_from`. + + Parameters + ---------- + mapper + The mapper whose reconstructed image is computed. + reconstruction + The mapper's block of the `reconstruction` (its source pixel values `s`). + """ + rows, cols, vals = self._sparse_triplets_slim_from(mapper=mapper) + + contrib = vals * reconstruction[cols] + + n_slim = self.mask.pixels_in_mask + + if self._xp is np: + return np.bincount(rows, weights=contrib, minlength=n_slim) + + return self._xp.zeros(n_slim, dtype=contrib.dtype).at[rows].add(contrib) + @property def mapped_reconstructed_operated_data_dict( self, diff --git a/test_autoarray/inversion/inversion/interferometer/test_interferometer.py b/test_autoarray/inversion/inversion/interferometer/test_interferometer.py index 2361fb055..a683a8408 100644 --- a/test_autoarray/inversion/inversion/interferometer/test_interferometer.py +++ b/test_autoarray/inversion/inversion/interferometer/test_interferometer.py @@ -1898,3 +1898,221 @@ def test__interferometer_sparse_operator__data_vector_from_triplets_matches_dens rtol=1.0e-12, atol=1.0e-12 * np.abs(reference).max(), ) + + +def _rectangular_mapper_from(mask, shape, over_sample_size): + from autoarray.inversion.mesh.mesh.rectangular_rtu_adapt_density import ( + overlay_grid_from, + ) + + grid = aa.Grid2D.from_mask(mask=mask, over_sample_size=over_sample_size) + + interpolator = aa.mesh.RectangularUniform(shape=shape).interpolator_from( + source_plane_data_grid=grid, + source_plane_mesh_grid=aa.Grid2DIrregular( + overlay_grid_from(shape_native=shape, grid=grid) + ), + adapt_data=None, + ) + + return aa.Mapper( + interpolator=interpolator, regularization=aa.reg.Constant(coefficient=1.0) + ) + + +def _mapped_reconstructed_parity_setup(seed): + """ + A 12x12 circular-mask TransformerDFT dataset with its sparse operator, and the linear object lists + the sparse `mapped_reconstructed_data_dict` is checked on: a Rectangular mesh, a Delaunay mesh, + over-sampled (sub_size 2) Rectangular and Delaunay meshes, two mappers, and a mapper plus a linear + function list (in either order). + """ + mask = aa.Mask2D.circular(shape_native=(12, 12), pixel_scales=1.0, radius=4.0) + + rng = np.random.default_rng(seed=seed) + 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) + + rectangular = _rectangular_mapper_from(mask=mask, shape=(4, 4), over_sample_size=1) + delaunay = _over_sampled_mapper_from( + mask=mask, pixels=16, shape=(4, 4), over_sample_size=1 + ) + rectangular_sub_2 = _rectangular_mapper_from( + mask=mask, shape=(5, 5), over_sample_size=2 + ) + delaunay_sub_2 = _over_sampled_mapper_from( + mask=mask, pixels=9, shape=(3, 3), over_sample_size=2 + ) + linear_obj = aa.m.MockLinearObjFuncList( + parameters=2, + mapping_matrix=rng.normal(size=(mask.pixels_in_mask, 2)), + ) + + linear_obj_lists = [ + [rectangular], + [delaunay], + [rectangular_sub_2], + [delaunay_sub_2], + [delaunay_sub_2, rectangular], + [rectangular, linear_obj], + [linear_obj, delaunay_sub_2], + ] + + return mask, rng, dataset_sparse, linear_obj_lists + + +def _assert_mapped_reconstructed_data_dict_matches_dense( + image_dict, linear_obj_list, reconstruction, inversion, rtol +): + """ + Each entry of `image_dict` must equal the dense `mapping_matrix @ reconstruction` of its linear object, + where `reconstruction` is the (NumPy) reconstruction split per linear object. + """ + param_range_list = inversion.param_range_list_from(cls=aa.LinearObj) + + for linear_obj, (start, end) in zip(linear_obj_list, param_range_list): + reference = np.asarray(linear_obj.mapping_matrix) @ np.asarray( + reconstruction[start:end] + ) + + np.testing.assert_allclose( + np.asarray(image_dict[linear_obj]), + reference, + rtol=rtol, + atol=rtol * np.abs(reference).max(), + ) + + +def test__interferometer_sparse_operator__mapped_reconstructed_data_dict_matches_dense_mapping_matrix(): + """ + `InversionInterferometerSparse.mapped_reconstructed_data_dict` reconstructs each mapper's image by a + scatter-add over the mapper's sparse triplets instead of the dense `mapping_matrix @ reconstruction`. + It must reproduce the dense product for Rectangular and Delaunay meshes, over-sampled mappers, two + mappers and a mapper mixed with a linear function list, and its values must be wrapped as an + `Array2D` on the inversion's real-space mask (so the model visibilities are unchanged). + """ + mask, rng, dataset_sparse, linear_obj_lists = _mapped_reconstructed_parity_setup( + seed=7 + ) + + for linear_obj_list in linear_obj_lists: + inversion = aa.Inversion( + dataset=dataset_sparse, linear_obj_list=linear_obj_list + ) + + assert isinstance(inversion, aa.InversionInterferometerSparse) + + image_dict = inversion.mapped_reconstructed_data_dict + + for linear_obj in linear_obj_list: + assert isinstance(image_dict[linear_obj], aa.Array2D) + assert (image_dict[linear_obj].mask == mask).all() + + _assert_mapped_reconstructed_data_dict_matches_dense( + image_dict=image_dict, + linear_obj_list=linear_obj_list, + reconstruction=np.asarray(inversion.reconstruction), + inversion=inversion, + rtol=1.0e-10, + ) + + # The model visibilities consume the dict and must equal the transform of the dense image. + visibilities_dict = inversion.mapped_reconstructed_operated_data_dict + + for linear_obj in linear_obj_list: + reference = dataset_sparse.transformer.visibilities_from( + image=aa.Array2D( + values=np.asarray(linear_obj.mapping_matrix) + @ np.asarray(inversion.reconstruction_dict[linear_obj]), + mask=mask, + ) + ) + + np.testing.assert_allclose( + np.asarray(visibilities_dict[linear_obj]), + np.asarray(reference), + rtol=1.0e-10, + atol=1.0e-10 * np.abs(np.asarray(reference)).max(), + ) + + +def test__interferometer_sparse_operator__mapped_reconstructed_data_dict_matches_dense_mapping_matrix__jax_jit(): + """ + The JAX backend of the sparse mapped reconstructed image (a `.at[rows].add` scatter over image pixels) + must reproduce the dense `mapping_matrix @ reconstruction` under `jax.jit`, with the dirty image traced + so the reconstruction is computed inside the jitted function. + """ + jax = pytest.importorskip("jax") + + import jax.numpy as jnp + + mask, rng, dataset_sparse, linear_obj_lists = _mapped_reconstructed_parity_setup( + seed=9 + ) + + dirty_image = np.asarray(dataset_sparse.sparse_operator.dirty_image) + + def inversion_from(dirty_image, linear_obj_list, xp): + return aa.Inversion( + dataset=aa.DatasetInterface( + data=dataset_sparse.data, + noise_map=dataset_sparse.noise_map, + grids=dataset_sparse.grids, + transformer=dataset_sparse.transformer, + sparse_operator=dataset_sparse.sparse_operator, + sparse_dirty_image=dirty_image, + ), + linear_obj_list=linear_obj_list, + xp=xp, + ) + + for linear_obj_list in linear_obj_lists: + inversion_np = inversion_from(dirty_image, linear_obj_list, xp=np) + + assert isinstance(inversion_np, aa.InversionInterferometerSparse) + + def images_and_reconstruction_from(dirty_image): + inversion = inversion_from(dirty_image, linear_obj_list, xp=jnp) + + image_dict = inversion.mapped_reconstructed_data_dict + + return ( + [image_dict[linear_obj].array for linear_obj in linear_obj_list], + inversion.reconstruction, + ) + + image_list, reconstruction = jax.jit(images_and_reconstruction_from)( + jnp.asarray(dirty_image) + ) + + # A sanity check only: the NumPy and JAX linear solves differ by more than the matrices they are + # handed (see `test__interferometer_sparse_operator__numpy_inversion_matches_jax_inversion`), so the + # image is compared below against the dense product of the jitted reconstruction itself. + reference = np.asarray(inversion_np.reconstruction) + + np.testing.assert_allclose( + np.asarray(reconstruction), + reference, + rtol=0.0, + atol=1.0e-3 * np.abs(reference).max(), + ) + + _assert_mapped_reconstructed_data_dict_matches_dense( + image_dict=dict(zip(linear_obj_list, image_list)), + linear_obj_list=linear_obj_list, + reconstruction=np.asarray(reconstruction), + inversion=inversion_np, + rtol=1.0e-12, + )