Merge pull request #649 from PyAutoLabs/feature/imaging-sparse-triplets-masking

fix: mask invalid mapper triplets on the imaging sparse path; backend-dispatched scatter-add helpers

Author: Jammy2211

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

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)

autoarray/inversion/inversion/interferometer/sparse.py

@@ -180,31 +180,18 @@ 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,
             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,
             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

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,
         )

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