From 9b52fac56050b8d3b3d242a17abe6e5e60c8e024 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Wed, 21 Oct 2026 12:56:11 +0100 Subject: [PATCH 1/3] fix: mask invalid stencil slots in shared sparse_triplets_from (#643) sparse_triplets_from gains pix_sizes_for_sub: entries mapping_matrix_from drops (slot >= pix_sizes, or a negative source index) get value 0 and column 0 on both numpy and JAX branches. Mapper.sparse_triplets_data / sparse_triplets_curvature pass it, so every imaging sparse consumer is covered; the interferometer _sparse_triplets_slim_from uses the shared masking and drops its local copy (outputs bit-identical). Co-Authored-By: Claude Fable 5.1 --- .../inversion/interferometer/sparse.py | 19 +++------------- autoarray/inversion/mappers/abstract.py | 2 ++ autoarray/inversion/mappers/mapper_util.py | 22 +++++++++++++++++++ 3 files changed, 27 insertions(+), 16 deletions(-) diff --git a/autoarray/inversion/inversion/interferometer/sparse.py b/autoarray/inversion/inversion/interferometer/sparse.py index d49c81bf1..c897bf628 100644 --- a/autoarray/inversion/inversion/interferometer/sparse.py +++ b/autoarray/inversion/inversion/interferometer/sparse.py @@ -180,9 +180,7 @@ def _sparse_triplets_slim_from(self, mapper: Mapper, apply_primary_beam: bool = beamed operator `P A` (used by the data vector); `False` gives those of `A` (used for the true-sky reconstructed image). """ - xp = self._xp - - rows, cols, vals = mapper_util.sparse_triplets_from( + return 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, @@ -190,21 +188,10 @@ def _sparse_triplets_slim_from(self, mapper: Mapper, apply_primary_beam: bool = sub_fraction_slim=mapper.over_sampler.sub_fraction.array, return_rows_slim=True, slim_weights=self._primary_beam_slim if apply_primary_beam else None, - xp=xp, + pix_sizes_for_sub=mapper.pix_sizes_for_sub_slim_index, + xp=self._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 diff --git a/autoarray/inversion/mappers/abstract.py b/autoarray/inversion/mappers/abstract.py index e2cee3e57..b3c3e45fc 100644 --- a/autoarray/inversion/mappers/abstract.py +++ b/autoarray/inversion/mappers/abstract.py @@ -407,6 +407,7 @@ def sparse_triplets_data(self): slim_index_for_sub=self.slim_index_for_sub_slim_index, fft_index_for_masked_pixel=self.mask.fft_index_for_masked_pixel, sub_fraction_slim=self.over_sampler.sub_fraction.array, + pix_sizes_for_sub=self.pix_sizes_for_sub_slim_index, xp=self._xp, ) @@ -463,6 +464,7 @@ def sparse_triplets_curvature(self): slim_index_for_sub=self.slim_index_for_sub_slim_index, fft_index_for_masked_pixel=self.mask.fft_index_for_masked_pixel, sub_fraction_slim=self.over_sampler.sub_fraction.array, + pix_sizes_for_sub=self.pix_sizes_for_sub_slim_index, xp=self._xp, return_rows_slim=False, ) diff --git a/autoarray/inversion/mappers/mapper_util.py b/autoarray/inversion/mappers/mapper_util.py index abaef95df..9bfa6978b 100644 --- a/autoarray/inversion/mappers/mapper_util.py +++ b/autoarray/inversion/mappers/mapper_util.py @@ -276,6 +276,7 @@ def sparse_triplets_from( *, return_rows_slim: bool = True, slim_weights=None, + pix_sizes_for_sub=None, xp=np, ): """ @@ -303,6 +304,13 @@ def sparse_triplets_from( of `diag(w) A` rather than `A`. The interferometer sparse inversion passes the primary beam here so the curvature and data vector use the beamed response `P A`. `None` (default) applies no weighting. + pix_sizes_for_sub + Optional number of used stencil slots per sub-pixel (M_sub,), i.e. the mapper's + `pix_sizes_for_sub_slim_index`. When given, every entry that `mapping_matrix_from` drops + (slot index `>= pix_sizes_for_sub`, or a negative source pixel index such as the `-1` + padding Delaunay meshes write into unused slots) is given value 0 and column 0, so it + contributes nothing to any scatter-add and never indexes out of range. `None` (default) + returns the raw stencil entries unmasked. xp Backend module (np or jnp) @@ -340,6 +348,13 @@ def sparse_triplets_from( if slim_weights is not None: vals = vals * np.asarray(slim_weights, dtype=np.float64)[slim_rows] + if pix_sizes_for_sub is not None: + slot = np.tile(np.arange(P, dtype=np.int32), M_sub) + pix_sizes = np.asarray(pix_sizes_for_sub, dtype=np.int32) + valid = (slot < np.repeat(pix_sizes, P)) & (cols >= 0) + vals = np.where(valid, vals, 0.0) + cols = np.where(valid, cols, 0).astype(np.int32) + if return_rows_slim: return slim_rows, cols, vals @@ -370,6 +385,13 @@ def sparse_triplets_from( if slim_weights is not None: vals = vals * xp.asarray(slim_weights, dtype=xp.float64)[slim_rows] + if pix_sizes_for_sub is not None: + slot = xp.tile(xp.arange(P, dtype=xp.int32), M_sub) + pix_sizes = xp.asarray(pix_sizes_for_sub, dtype=xp.int32) + valid = (slot < xp.repeat(pix_sizes, P)) & (cols >= 0) + vals = xp.where(valid, vals, 0.0) + cols = xp.where(valid, cols, 0) + if return_rows_slim: return slim_rows, cols, vals From 3e53e0aa829d2b6c13a4293f7465d5840fb4d089 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Wed, 21 Oct 2026 12:56:11 +0100 Subject: [PATCH 2/3] fix: backend-dispatched scatter-add in imaging sparse helpers (#643) data_vector_via_psf_weighted_data_from and mapped_reconstructed_operated_data_via_sparse_operator_from take xp: numpy uses np.bincount with no JAX import, JAX keeps segment_sum with the import inside its branch. InversionImagingSparse passes xp=self._xp. Co-Authored-By: Claude Fable 5.1 --- .../imaging/inversion_imaging_util.py | 39 +++++++++++++++++-- .../inversion/inversion/imaging/sparse.py | 4 ++ 2 files changed, 40 insertions(+), 3 deletions(-) diff --git a/autoarray/inversion/inversion/imaging/inversion_imaging_util.py b/autoarray/inversion/inversion/imaging/inversion_imaging_util.py index 9b7763262..b4ce10d7c 100644 --- a/autoarray/inversion/inversion/imaging/inversion_imaging_util.py +++ b/autoarray/inversion/inversion/imaging/inversion_imaging_util.py @@ -73,6 +73,7 @@ def data_vector_via_psf_weighted_data_from( cols: np.ndarray, # (nnz,) int32 source pixel index vals: np.ndarray, # (nnz,) float64 mapping weights incl sub_fraction S: int, # number of source pixels + xp=np, ) -> np.ndarray: r""" Returns the data vector `D` from the `psf_weighted_data` matrix (see `psf_weighted_data_from`), which encodes the @@ -96,7 +97,16 @@ def data_vector_via_psf_weighted_data_from( The values of the sparse mapping matrix triplet representation, which map the image pixels to source pixels. S The number of source pixels. + xp + The array module: `np` scatter-adds with `np.bincount` (no JAX import, so this runs without + JAX installed); `jax.numpy` uses `jax.ops.segment_sum`. The triplets must be masked (see + `mapper_util.sparse_triplets_from(pix_sizes_for_sub=...)`) so every column is non-negative. """ + if xp is np: + w = np.asarray(psf_weighted_data)[np.asarray(rows)] # (nnz,) + contrib = np.asarray(vals) * w # (nnz,) + return np.bincount(np.asarray(cols), weights=contrib, minlength=S)[:S] # (S,) + from jax.ops import segment_sum w = psf_weighted_data[rows] # (nnz,) @@ -284,7 +294,32 @@ def mapped_reconstructed_operated_data_via_sparse_operator_from( vals, # (nnz,) fft_index_for_masked_pixel, data_shape: int, # y_shape * x_shape + xp=np, ): + """ + Returns the slim image `A s` of a reconstruction `s` mapped through a mapper's sparse mapping + operator `A`, whose COO triplets `(rows, cols, vals)` index rows on the rectangular FFT grid + (`Mapper.sparse_triplets_curvature`). + + The scatter-add is performed on the rectangular grid of `data_shape` and the masked pixels are + then gathered via `fft_index_for_masked_pixel`. With `xp=np` it uses `np.bincount` and imports + no JAX, so it runs without JAX installed; with `xp=jax.numpy` it uses `jax.ops.segment_sum`. + """ + num_segments = data_shape[0] * data_shape[1] + + if xp is np: + reconstruction = np.asarray(reconstruction, dtype=np.float64) + rows = np.asarray(rows, dtype=np.int64) + cols = np.asarray(cols, dtype=np.int64) + vals = np.asarray(vals, dtype=np.float64) + + contrib = vals * reconstruction[cols] # (nnz,) + image_rect = np.bincount(rows, weights=contrib, minlength=num_segments)[ + :num_segments + ] # (M_rect,) + + return image_rect[np.asarray(fft_index_for_masked_pixel)] # (M_pix,) + import jax.numpy as jnp from jax.ops import segment_sum @@ -294,9 +329,7 @@ def mapped_reconstructed_operated_data_via_sparse_operator_from( vals = jnp.asarray(vals, dtype=jnp.float64) contrib = vals * reconstruction[cols] # (nnz,) - image_rect = segment_sum( - contrib, rows, num_segments=data_shape[0] * data_shape[1] - ) # (M_rect,) + image_rect = segment_sum(contrib, rows, num_segments=num_segments) # (M_rect,) image_slim = image_rect[fft_index_for_masked_pixel] # (M_pix,) return image_slim diff --git a/autoarray/inversion/inversion/imaging/sparse.py b/autoarray/inversion/inversion/imaging/sparse.py index 78d57d433..b628d18f1 100644 --- a/autoarray/inversion/inversion/imaging/sparse.py +++ b/autoarray/inversion/inversion/imaging/sparse.py @@ -96,6 +96,7 @@ def _data_vector_mapper(self) -> np.ndarray: cols=cols, vals=vals, S=mapper.params, + xp=self._xp, ) ) param_range = mapper_param_range[mapper_index] @@ -149,6 +150,7 @@ def _data_vector_x1_mapper(self) -> np.ndarray: cols=cols, vals=vals, S=linear_obj.params, + xp=self._xp, ) @property @@ -174,6 +176,7 @@ def _data_vector_multi_mapper(self) -> np.ndarray: cols=cols, vals=vals, S=mapper.params, + xp=self._xp, ) ) @@ -532,6 +535,7 @@ def _mapped_reconstructed_data_dict_from( vals=vals, fft_index_for_masked_pixel=self.mask.fft_index_for_masked_pixel, data_shape=self.mask.shape_native, + xp=self._xp, ) mapped = Array2D(values=mapped, mask=self.mask) From dd26822498dc402ca2426d5536f16fd6e12ffe35 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Wed, 21 Oct 2026 13:04:21 +0100 Subject: [PATCH 3/3] test: imaging sparse invalid-slot parity and nojax helpers (#643) - sparse_triplets_from zeroes a -1 index and unused slots with non-zero weights (numpy and JAX), and its triplets scatter to mapping_matrix_from; - InversionImagingSparse on a mapper with -1 indices, a negative index in a used slot and weighted unused slots matches the mapping inversion (data_vector, curvature_matrix, mapped reconstructed data) on numpy and JAX backends; - the mapper triplets and both imaging scatter-add helpers run with jax, jaxlib and nufftax blocked (subprocess import hook). Co-Authored-By: Claude Fable 5.1 --- .../inversion/imaging/test_imaging.py | 264 ++++++++++++++++++ .../pixelization/mappers/test_mapper_util.py | 142 ++++++++++ 2 files changed, 406 insertions(+) diff --git a/test_autoarray/inversion/inversion/imaging/test_imaging.py b/test_autoarray/inversion/inversion/imaging/test_imaging.py index 761e854a4..f956dbdb9 100644 --- a/test_autoarray/inversion/inversion/imaging/test_imaging.py +++ b/test_autoarray/inversion/inversion/imaging/test_imaging.py @@ -1,3 +1,9 @@ +import importlib.util +import os +import subprocess +import sys +import textwrap + import autoarray as aa from autoarray.inversion.inversion.imaging.inversion_imaging_util import ( @@ -299,3 +305,261 @@ def counted(self, *args, **kwargs): inversion.mapped_reconstructed_data assert len(calls) == 1 + + +requires_jax = pytest.mark.skipif( + importlib.util.find_spec("jax") is None, + reason="the imaging sparse operator is built with jax", +) + + +def _mapper_with_invalid_slots_from(mapper): + """ + Returns a copy of `mapper` whose stencil table carries every kind of entry + `mapping_matrix_from` drops, all with non-zero weights: + + - an unused slot (slot >= pix_sizes) padded with index -1 (odd sub-pixels of a Delaunay row); + - a negative index inside a used slot (even sub-pixels, whose pix_size is bumped over a -1 slot); + - an extra unused slot holding a real source pixel index (every sub-pixel). + + The dense mapping matrix is unchanged by these entries, so a correct sparse path must be too. + """ + pix_indexes = np.asarray(mapper.pix_indexes_for_sub_slim_index) + pix_weights = np.asarray(mapper.pix_weights_for_sub_slim_index, dtype=float) + pix_sizes = np.asarray(mapper.pix_sizes_for_sub_slim_index).copy() + + M_sub, P = pix_indexes.shape + + unused = np.arange(P)[None, :] >= pix_sizes[:, None] + + assert unused.any() + assert (pix_indexes[unused] == -1).all() + + pix_weights = np.where(unused, 0.7, pix_weights) + + bump = (np.arange(M_sub) % 2 == 0) & (pix_sizes < P) + pix_sizes[bump] += 1 + + pix_indexes = np.concatenate( + [pix_indexes, (np.arange(M_sub) % mapper.params)[:, None]], axis=1 + ) + pix_weights = np.concatenate([pix_weights, np.full((M_sub, 1), 0.9)], axis=1) + + class MapperWithInvalidSlots(type(mapper)): + @property + def pix_indexes_for_sub_slim_index(self): + return pix_indexes + + @property + def pix_weights_for_sub_slim_index(self): + return pix_weights + + @property + def pix_sizes_for_sub_slim_index(self): + return pix_sizes + + corrupted = MapperWithInvalidSlots( + interpolator=mapper.interpolator, + regularization=mapper.regularization, + image_plane_mesh_grid=mapper.image_plane_mesh_grid, + ) + + assert (corrupted.pix_indexes_for_sub_slim_index < 0).any() + assert np.allclose(corrupted.mapping_matrix, mapper.mapping_matrix) + + return corrupted + + +@requires_jax +def test__sparse__mapper_with_invalid_stencil_slots__matches_mapping_inversion( + masked_imaging_7x7, delaunay_mapper_9_3x3 +): + mapper = _mapper_with_invalid_slots_from(delaunay_mapper_9_3x3) + + settings = aa.Settings(use_positive_only_solver=False) + + inversion_sparse = aa.Inversion( + dataset=masked_imaging_7x7.apply_sparse_operator(), + linear_obj_list=[mapper], + settings=settings, + ) + inversion_mapping = aa.Inversion( + dataset=masked_imaging_7x7, + linear_obj_list=[mapper], + settings=settings, + ) + + assert isinstance(inversion_sparse, InversionImagingSparse) + + assert np.asarray(inversion_sparse.data_vector) == pytest.approx( + np.asarray(inversion_mapping.data_vector), 1.0e-6 + ) + assert np.asarray(inversion_sparse.curvature_matrix) == pytest.approx( + np.asarray(inversion_mapping.curvature_matrix), 1.0e-6 + ) + + reconstruction = np.asarray(inversion_sparse.reconstruction) + + assert reconstruction == pytest.approx( + np.asarray(inversion_mapping.reconstruction), 1.0e-4 + ) + + mapped = np.asarray(inversion_sparse.mapped_reconstructed_data_dict[mapper].array) + + assert mapped == pytest.approx( + np.asarray(mapper.mapping_matrix) @ reconstruction, 1.0e-8 + ) + assert mapped == pytest.approx( + np.asarray(inversion_mapping.mapped_reconstructed_data_dict[mapper].array), + 1.0e-4, + ) + assert np.asarray( + inversion_sparse.mapped_reconstructed_operated_data_dict[mapper].array + ) == pytest.approx( + np.asarray( + inversion_mapping.mapped_reconstructed_operated_data_dict[mapper].array + ), + 1.0e-4, + ) + + +@requires_jax +def test__sparse__mapper_with_invalid_stencil_slots__matches_mapping_inversion__jax( + masked_imaging_7x7, delaunay_mapper_9_3x3 +): + import jax + + jax.config.update("jax_enable_x64", True) + + import jax.numpy as jnp + + mapper = _mapper_with_invalid_slots_from(delaunay_mapper_9_3x3) + mapper_jax = type(mapper)( + interpolator=mapper.interpolator, + regularization=mapper.regularization, + image_plane_mesh_grid=mapper.image_plane_mesh_grid, + xp=jnp, + ) + + settings = aa.Settings(use_positive_only_solver=False) + + inversion_sparse = aa.Inversion( + dataset=masked_imaging_7x7.apply_sparse_operator(), + linear_obj_list=[mapper_jax], + settings=settings, + xp=jnp, + ) + inversion_mapping = aa.Inversion( + dataset=masked_imaging_7x7, + linear_obj_list=[mapper], + settings=settings, + ) + + assert isinstance(inversion_sparse, InversionImagingSparse) + + assert np.asarray(inversion_sparse.data_vector) == pytest.approx( + np.asarray(inversion_mapping.data_vector), 1.0e-6 + ) + assert np.asarray(inversion_sparse.curvature_matrix) == pytest.approx( + np.asarray(inversion_mapping.curvature_matrix), 1.0e-6 + ) + + reconstruction = np.asarray(inversion_sparse.reconstruction) + + assert np.asarray( + inversion_sparse.mapped_reconstructed_data_dict[mapper_jax].array + ) == pytest.approx(np.asarray(mapper.mapping_matrix) @ reconstruction, 1.0e-8) + + +def test__sparse_helpers__import_and_run_with_jax_blocked(): + """ + The mapper triplets and the two imaging scatter-add helpers (`data_vector_via_psf_weighted_data_from`, + `mapped_reconstructed_operated_data_via_sparse_operator_from`) must work in the `unittest-nojax` + leg. JAX is in `sys.modules` for the rest of a normal test session, so a subprocess hides `jax`, + `jaxlib` and `nufftax` from the import system and checks them against the dense mapping matrix + on a Delaunay mapper, whose unused stencil slots carry `-1` indices. + + The `ImagingSparseOperator` itself (FFT precompute and curvature blocks) is JAX-only, so the full + `InversionImagingSparse` is not run here. + """ + script = textwrap.dedent( + """ + import importlib.machinery + import sys + + _find_spec = importlib.machinery.PathFinder.find_spec.__func__ + + + def _find_spec_blocking_jax(cls, name, path=None, target=None): + if name.split(".")[0] in ("jax", "jaxlib", "nufftax"): + return None + return _find_spec(cls, name, path, target) + + + importlib.machinery.PathFinder.find_spec = classmethod(_find_spec_blocking_jax) + + # Guard the guard: if the block stopped working this test would pass vacuously. + try: + import jax + except ModuleNotFoundError: + pass + else: + raise AssertionError("jax was not blocked") + + import numpy as np + + import autoarray as aa + from autoarray import fixtures + from autoarray.inversion.inversion.imaging import inversion_imaging_util + + mapper = fixtures.make_delaunay_mapper_9_3x3() + mask = fixtures.make_mask_2d_7x7() + + assert (np.asarray(mapper.pix_indexes_for_sub_slim_index) < 0).any() + + mapping_matrix = np.asarray(mapper.mapping_matrix) + + rows, cols, vals = mapper.sparse_triplets_data + + rng = np.random.default_rng(seed=1) + psf_weighted_data = rng.normal(size=mapping_matrix.shape[0]) + + data_vector = inversion_imaging_util.data_vector_via_psf_weighted_data_from( + psf_weighted_data=psf_weighted_data, + rows=rows, + cols=cols, + vals=vals, + S=mapper.params, + xp=np, + ) + + assert np.allclose(data_vector, mapping_matrix.T @ psf_weighted_data) + + rows, cols, vals = mapper.sparse_triplets_curvature + + reconstruction = rng.normal(size=mapper.params) + + mapped = inversion_imaging_util.mapped_reconstructed_operated_data_via_sparse_operator_from( + reconstruction=reconstruction, + rows=rows, + cols=cols, + vals=vals, + fft_index_for_masked_pixel=mask.fft_index_for_masked_pixel, + data_shape=mask.shape_native, + xp=np, + ) + + assert np.allclose(mapped, mapping_matrix @ reconstruction) + + assert not any( + name.split(".")[0] in ("jax", "jaxlib", "nufftax") for name in sys.modules + ) + """ + ) + + environment = dict(os.environ) + environment["PYAUTO_DISABLE_JAX"] = "1" + + subprocess.run( + [sys.executable, "-c", script], check=True, env=environment, timeout=600 + ) diff --git a/test_autoarray/inversion/pixelization/mappers/test_mapper_util.py b/test_autoarray/inversion/pixelization/mappers/test_mapper_util.py index 9eae96393..db9449afe 100644 --- a/test_autoarray/inversion/pixelization/mappers/test_mapper_util.py +++ b/test_autoarray/inversion/pixelization/mappers/test_mapper_util.py @@ -914,3 +914,145 @@ def test__mapper_pixel_signals_from__delaunay_excludes_zero_weights(): mapper, adapt_data, exclude_zero_weight_mappings=True ), ) + + +def _invalid_slot_table(): + """ + A hand-built stencil table with every kind of entry `mapping_matrix_from` drops: + + - row 0, slot 2: an unused slot (slot >= pix_sizes) padded with index -1 and a non-zero weight; + - row 1, slot 1: a negative index inside a used slot, with a non-zero weight; + - row 1, slot 2: an unused slot holding a real-looking index (1) and a non-zero weight. + """ + pix_indexes = np.array([[0, 1, -1], [2, -1, 1], [1, 2, 0]]) + pix_weights = np.array([[0.5, 0.5, 0.3], [0.6, 0.4, 0.9], [0.2, 0.3, 0.5]]) + pix_sizes = np.array([2, 2, 3]) + slim_index_for_sub = np.array([0, 1, 1]) + fft_index_for_masked_pixel = np.array([3, 5]) + sub_fraction_slim = np.array([1.0, 0.5]) + + return ( + pix_indexes, + pix_weights, + pix_sizes, + slim_index_for_sub, + fft_index_for_masked_pixel, + sub_fraction_slim, + ) + + +def _assert_invalid_slots_zeroed(rows, cols, vals): + ( + pix_indexes, + pix_weights, + pix_sizes, + slim_index_for_sub, + fft_index_for_masked_pixel, + sub_fraction_slim, + ) = _invalid_slot_table() + + rows, cols, vals = np.asarray(rows), np.asarray(cols), np.asarray(vals) + + assert (cols >= 0).all() + assert cols.tolist() == [0, 1, 0, 2, 0, 0, 1, 2, 0] + assert vals == pytest.approx( + np.array([0.5, 0.5, 0.0, 0.3, 0.0, 0.0, 0.1, 0.15, 0.25]), 1.0e-12 + ) + + # The masked triplets scatter to exactly the dense mapping matrix. + mapping_matrix = aa.util.mapper.mapping_matrix_from( + pix_indexes_for_sub_slim_index=pix_indexes, + pix_size_for_sub_slim_index=pix_sizes, + pix_weights_for_sub_slim_index=pix_weights, + pixels=3, + total_mask_pixels=2, + slim_index_for_sub_slim_index=slim_index_for_sub, + sub_fraction=sub_fraction_slim, + ) + + dense = np.zeros((2, 3)) + np.add.at(dense, (rows, cols), vals) + + assert dense == pytest.approx(mapping_matrix, 1.0e-12) + + +def test__sparse_triplets_from__pix_sizes_for_sub_zeroes_invalid_entries(): + ( + pix_indexes, + pix_weights, + pix_sizes, + slim_index_for_sub, + fft_index_for_masked_pixel, + sub_fraction_slim, + ) = _invalid_slot_table() + + rows, cols, vals = aa.util.mapper.sparse_triplets_from( + pix_indexes_for_sub=pix_indexes, + pix_weights_for_sub=pix_weights, + slim_index_for_sub=slim_index_for_sub, + fft_index_for_masked_pixel=fft_index_for_masked_pixel, + sub_fraction_slim=sub_fraction_slim, + pix_sizes_for_sub=pix_sizes, + ) + + _assert_invalid_slots_zeroed(rows, cols, vals) + + # Rectangular FFT-grid rows are unaffected by the masking. + rows_rect, cols_rect, vals_rect = aa.util.mapper.sparse_triplets_from( + pix_indexes_for_sub=pix_indexes, + pix_weights_for_sub=pix_weights, + slim_index_for_sub=slim_index_for_sub, + fft_index_for_masked_pixel=fft_index_for_masked_pixel, + sub_fraction_slim=sub_fraction_slim, + return_rows_slim=False, + pix_sizes_for_sub=pix_sizes, + ) + + assert rows_rect.tolist() == fft_index_for_masked_pixel[rows].tolist() + assert (cols_rect == cols).all() + assert (vals_rect == vals).all() + + # Without `pix_sizes_for_sub` the raw stencil entries are returned unmasked. + _, cols_raw, vals_raw = aa.util.mapper.sparse_triplets_from( + pix_indexes_for_sub=pix_indexes, + pix_weights_for_sub=pix_weights, + slim_index_for_sub=slim_index_for_sub, + fft_index_for_masked_pixel=fft_index_for_masked_pixel, + sub_fraction_slim=sub_fraction_slim, + ) + + assert (cols_raw < 0).sum() == 2 + assert vals_raw[5] == pytest.approx(0.45, 1.0e-12) + + +@pytest.mark.skipif( + importlib.util.find_spec("jax") is None, + reason="the JAX branch of sparse_triplets_from needs jax", +) +def test__sparse_triplets_from__pix_sizes_for_sub_zeroes_invalid_entries__jax(): + import jax + + jax.config.update("jax_enable_x64", True) + + import jax.numpy as jnp + + ( + pix_indexes, + pix_weights, + pix_sizes, + slim_index_for_sub, + fft_index_for_masked_pixel, + sub_fraction_slim, + ) = _invalid_slot_table() + + rows, cols, vals = aa.util.mapper.sparse_triplets_from( + pix_indexes_for_sub=pix_indexes, + pix_weights_for_sub=pix_weights, + slim_index_for_sub=slim_index_for_sub, + fft_index_for_masked_pixel=fft_index_for_masked_pixel, + sub_fraction_slim=sub_fraction_slim, + pix_sizes_for_sub=pix_sizes, + xp=jnp, + ) + + _assert_invalid_slots_zeroed(rows, cols, vals)