data vector stuff all jaxd

Author: Jammy2211

autoarray/dataset/abstract/w_tilde.py

@@ -52,17 +52,4 @@ def fft_index_for_masked_pixel(self) -> np.ndarray:
         - This method is intentionally backend-agnostic and can be used by both
           imaging and interferometer curvature pipelines.
         """
-
-        # Boolean mask defined on the rectangular FFT grid
-        # True  = masked pixel
-        # False = unmasked pixel
-        mask_fft = self.fft_mask
-
-        # Coordinates of unmasked pixels in the FFT grid
-        ys, xs = np.where(~mask_fft)
-
-        # Width of the FFT grid (number of columns)
-        width = mask_fft.shape[1]
-
-        # Convert (y, x) coordinates to flat row-major indices
-        return (ys * width + xs).astype(np.int32)
+        self.fft_mask.fft_index_for_masked_pixel

autoarray/dataset/imaging/dataset.py

@@ -505,6 +505,7 @@ def apply_w_tilde(
         """
 
         w_tilde = WTildeImaging(
+            data=self.data,
             noise_map=self.noise_map,
             psf=self.psf,
             fft_mask=self.mask,

autoarray/dataset/imaging/w_tilde.py

@@ -12,6 +12,7 @@
 class WTildeImaging(AbstractWTilde):
     def __init__(
         self,
+        data: np.ndarray,
         noise_map: np.ndarray,
         psf: np.ndarray,
         fft_mask: np.ndarray,
@@ -41,9 +42,13 @@ def __init__(
             fft_mask=fft_mask
         )
 
+        self.data = data
         self.noise_map = noise_map
         self.psf = psf
 
+        self.data_native = data.native
+        self.noise_map_native = noise_map.native
+
     @property
     def psf_operator_matrix_dense(self):
 

autoarray/dataset/interferometer/w_tilde.py

@@ -250,8 +250,6 @@ def __init__(
             curvature_preload=self.curvature_preload, batch_size=batch_size
         )
 
-
-
     def save_curvature_preload(
         self,
         file: Union[str, Path],

autoarray/inversion/inversion/imaging/inversion_imaging_util.py

@@ -1,38 +1,5 @@
 import numpy as np
 
-def pixel_triplets_from_subpixel_arrays_from(
-    pix_indexes_for_sub,          # (M_sub, P)
-    pix_weights_for_sub,          # (M_sub, P)
-    slim_index_for_sub,           # (M_sub,)
-    fft_index_for_masked_pixel,   # (N_unmasked,)
-    sub_fraction_slim,            # (N_unmasked,)
-):
-    """
-    Build sparse source→image mapping triplets (rows, cols, vals)
-    for a fixed-size interpolation stencil.
-
-    Assumptions:
-    - Every subpixel maps to exactly P source pixels
-    - All entries in pix_indexes_for_sub are valid
-    - No padding / ragged rows needed
-    """
-    import jax.numpy as jnp
-
-
-    M_sub, P = pix_indexes_for_sub.shape
-
-    sub_ids = jnp.repeat(jnp.arange(M_sub, dtype=jnp.int32), P)
-
-    cols = pix_indexes_for_sub.reshape(-1).astype(jnp.int32)
-    vals = pix_weights_for_sub.reshape(-1).astype(jnp.float64)
-
-    slim_rows = slim_index_for_sub[sub_ids].astype(jnp.int32)
-    rows = fft_index_for_masked_pixel[slim_rows].astype(jnp.int32)
-
-    vals = vals * sub_fraction_slim[slim_rows].astype(jnp.float64)
-    return rows, cols, vals
-
-
 def psf_operator_matrix_dense_from(
     kernel_native: np.ndarray,
     native_index_for_slim_index: np.ndarray,  # shape (N_pix, 2), native (y,x) coords of masked pixels
@@ -159,6 +126,29 @@ def w_tilde_data_imaging_from(
     return xp.sum(patches * kernel_native[None, :, :], axis=(1, 2))  # (N,)
 
 
+def data_vector_via_w_tilde_from(
+    w_tilde_data: np.ndarray,     # (M_pix,) float64
+    rows: np.ndarray,             # (nnz,) int32  each triplet's data pixel (slim index)
+    cols: np.ndarray,                  # (nnz,) int32  source pixel index
+    vals: np.ndarray,                  # (nnz,) float64 mapping weights incl sub_fraction
+    S: int,                             # number of source pixels
+) -> np.ndarray:
+    """
+    Replacement for numba data_vector_via_w_tilde_data_imaging_from using triplets.
+
+    Computes:
+        D[p] = sum_{triplets t with col_t=p} vals[t] * w_tilde_data_slim[slim_rows[t]]
+
+    Returns:
+        (S,) float64
+    """
+    from jax.ops import segment_sum
+
+    w = w_tilde_data[rows]          # (nnz,)
+    contrib = vals * w                         # (nnz,)
+    return segment_sum(contrib, cols, num_segments=S)  # (S,)
+
+
 def data_vector_via_blurred_mapping_matrix_from(
     blurred_mapping_matrix: np.ndarray, image: np.ndarray, noise_map: np.ndarray
 ) -> np.ndarray:

autoarray/inversion/inversion/imaging/w_tilde.py

@@ -58,10 +58,11 @@ def __init__(
     @cached_property
     def w_tilde_data(self):
         return inversion_imaging_util.w_tilde_data_imaging_from(
-            image_native=self.data.native.array,
-            noise_map_native=self.noise_map.native.array,
-            kernel_native=self.psf.native.array,
+            image_native=self.w_tilde.data_native.array,
+            noise_map_native=self.w_tilde.noise_map_native.array,
+            kernel_native=self.psf.stored_native,
             native_index_for_slim_index=self.data.mask.derive_indexes.native_for_slim,
+            xp=self._xp
         )
 
     @property
@@ -83,15 +84,16 @@ def _data_vector_mapper(self) -> np.ndarray:
         mapper_param_range = self.param_range_list_from(cls=AbstractMapper)
 
         for mapper_index, mapper in enumerate(mapper_list):
+
+            rows, cols, vals = mapper.pixel_triplets
+
             data_vector_mapper = (
-                inversion_imaging_numba_util.data_vector_via_w_tilde_data_imaging_from(
+                inversion_imaging_util.data_vector_via_w_tilde_data_imaging_from(
                     w_tilde_data=self.w_tilde_data,
-                    data_to_pix_unique=np.array(
-                        mapper.unique_mappings.data_to_pix_unique
-                    ),
-                    data_weights=np.array(mapper.unique_mappings.data_weights),
-                    pix_lengths=np.array(mapper.unique_mappings.pix_lengths),
-                    pix_pixels=mapper.params,
+                    rows=rows,
+                    cols=cols,
+                    vals=vals,
+                    S=mapper.total_params,
                 )
             )
             param_range = mapper_param_range[mapper_index]
@@ -131,12 +133,14 @@ def _data_vector_x1_mapper(self) -> np.ndarray:
         """
         linear_obj = self.linear_obj_list[0]
 
-        return inversion_imaging_numba_util.data_vector_via_w_tilde_data_imaging_from(
+        rows, cols, vals = linear_obj.pixel_triplets
+
+        return inversion_imaging_util.data_vector_via_w_tilde_from(
             w_tilde_data=self.w_tilde_data,
-            data_to_pix_unique=linear_obj.unique_mappings.data_to_pix_unique,
-            data_weights=linear_obj.unique_mappings.data_weights,
-            pix_lengths=linear_obj.unique_mappings.pix_lengths,
-            pix_pixels=linear_obj.params,
+            rows=rows,
+            cols=cols,
+            vals=vals,
+            S=linear_obj.params,
         )
 
     @property

autoarray/inversion/pixelization/mappers/abstract.py

@@ -268,6 +268,20 @@ def mapping_matrix(self) -> np.ndarray:
             xp=self._xp,
         )
 
+    @cached_property
+    def pixel_triplets(self):
+
+        rows, cols, vals = mapper_util.pixel_triplets_from_subpixel_arrays_from(
+            pix_indexes_for_sub=self.pix_indexes_for_sub_slim_index,
+            pix_weights_for_sub=self.pix_weights_for_sub_slim_index,
+            slim_index_for_sub=self.slim_index_for_sub_slim_index,
+            fft_index_for_masked_pixel=self.mapper_grids.mask.fft_index_for_masked_pixel,
+            sub_fraction_slim=self.over_sampler.sub_fraction.array,
+            xp=self._xp
+        )
+
+        return rows, cols, vals
+
     def pixel_signals_from(self, signal_scale: float, xp=np) -> np.ndarray:
         """
         Returns the signal in each pixelization pixel, where this signal is an estimate of the expected signal
@@ -410,7 +424,6 @@ def extent_from(
             extent=self.source_plane_mesh_grid.geometry.extent
         )
 
-
 class PixSubWeights:
     def __init__(self, mappings: np.ndarray, sizes: np.ndarray, weights: np.ndarray):
         """

autoarray/inversion/pixelization/mappers/mapper_util.py

@@ -442,6 +442,102 @@ def adaptive_pixel_signals_from(
     return pixel_signals**signal_scale
 
 
+import numpy as np
+
+def pixel_triplets_from_subpixel_arrays_from(
+    pix_indexes_for_sub,          # (M_sub, P)
+    pix_weights_for_sub,          # (M_sub, P)
+    slim_index_for_sub,           # (M_sub,)
+    fft_index_for_masked_pixel,   # (N_unmasked,)
+    sub_fraction_slim,            # (N_unmasked,)
+    *,
+    xp=np,
+):
+    """
+    Build sparse source→image mapping triplets (rows, cols, vals)
+    for a fixed-size interpolation stencil.
+
+    This supports both:
+      - NumPy (xp=np)
+      - JAX  (xp=jax.numpy)
+
+    Parameters
+    ----------
+    pix_indexes_for_sub
+        Source pixel indices for each subpixel (M_sub, P)
+    pix_weights_for_sub
+        Interpolation weights for each subpixel (M_sub, P)
+    slim_index_for_sub
+        Mapping subpixel -> slim image pixel index (M_sub,)
+    fft_index_for_masked_pixel
+        Mapping slim pixel -> rectangular FFT-grid pixel index (N_unmasked,)
+    sub_fraction_slim
+        Oversampling normalization per slim pixel (N_unmasked,)
+    xp
+        Backend module (np or jnp)
+
+    Returns
+    -------
+    rows : (nnz,) int32
+        Rectangular FFT grid row index per mapping entry
+    cols : (nnz,) int32
+        Source pixel index per mapping entry
+    vals : (nnz,) float64
+        Mapping weight per entry including sub_fraction normalization
+    """
+
+    # ------------------------------------------------------------
+    # Put everything on the right backend
+    # ------------------------------------------------------------
+    pix_indexes_for_sub = xp.asarray(pix_indexes_for_sub)
+    pix_weights_for_sub = xp.asarray(pix_weights_for_sub)
+    slim_index_for_sub = xp.asarray(slim_index_for_sub)
+    fft_index_for_masked_pixel = xp.asarray(fft_index_for_masked_pixel)
+    sub_fraction_slim = xp.asarray(sub_fraction_slim)
+
+    # dtypes (important for JAX scatter / indexing performance)
+    pix_indexes_for_sub = pix_indexes_for_sub.astype(xp.int32)
+    pix_weights_for_sub = pix_weights_for_sub.astype(xp.float64)
+    slim_index_for_sub = slim_index_for_sub.astype(xp.int32)
+    fft_index_for_masked_pixel = fft_index_for_masked_pixel.astype(xp.int32)
+    sub_fraction_slim = sub_fraction_slim.astype(xp.float64)
+
+    # ------------------------------------------------------------
+    # Dimensions
+    # ------------------------------------------------------------
+    M_sub, P = pix_indexes_for_sub.shape
+
+    # ------------------------------------------------------------
+    # Subpixel IDs repeated P times (fixed stencil)
+    # ------------------------------------------------------------
+    sub_ids = xp.repeat(xp.arange(M_sub, dtype=xp.int32), P)  # (M_sub*P,)
+
+    # ------------------------------------------------------------
+    # Flatten interpolation stencil
+    # ------------------------------------------------------------
+    cols = pix_indexes_for_sub.reshape(-1).astype(xp.int32)      # (nnz,)
+    vals = pix_weights_for_sub.reshape(-1).astype(xp.float64)    # (nnz,)
+
+    # ------------------------------------------------------------
+    # subpixel -> slim image pixel
+    # ------------------------------------------------------------
+    slim_rows = slim_index_for_sub[sub_ids].astype(xp.int32)     # (nnz,)
+
+    # ------------------------------------------------------------
+    # slim pixel -> FFT rectangular pixel
+    # ------------------------------------------------------------
+    rows = fft_index_for_masked_pixel[slim_rows].astype(xp.int32)
+
+    # ------------------------------------------------------------
+    # Oversampling normalization
+    # ------------------------------------------------------------
+    vals = vals * sub_fraction_slim[slim_rows].astype(xp.float64)
+
+    return rows, cols, vals
+
+
+
+
 def mapping_matrix_from(
     pix_indexes_for_sub_slim_index: np.ndarray,
     pix_size_for_sub_slim_index: np.ndarray,

autoarray/mask/mask_2d.py

@@ -619,6 +619,48 @@ def from_fits(
     def shape_native(self) -> Tuple[int, ...]:
         return self.shape
 
+    @property
+    def fft_index_for_masked_pixel(self) -> np.ndarray:
+        """
+        Return a mapping from masked-pixel (slim) indices to flat indices
+        on the rectangular FFT grid.
+
+        This array is used to translate between:
+
+            - "masked pixel space" (a compact 1D indexing over unmasked pixels)
+            - the 2D rectangular grid on which FFT-based convolutions are performed
+
+        The FFT grid is assumed to be rectangular and already suitable for FFTs
+        (e.g. padded and centered appropriately). Masked pixels are present on
+        this grid but are ignored in computations via zero-weighting.
+
+        Returns
+        -------
+        np.ndarray
+            A 1D array of shape (N_unmasked,), where element `i` gives the flat
+            (row-major) index into the FFT grid corresponding to the `i`-th
+            unmasked pixel in slim ordering.
+
+        Notes
+        -----
+        - The slim ordering is defined as the order returned by `np.where(~mask)`.
+        - The flat FFT index is computed assuming row-major (C-style) ordering:
+              flat_index = y * width + x
+        - This method is intentionally backend-agnostic and can be used by both
+          imaging and interferometer curvature pipelines.
+        """
+        # Boolean mask defined on the rectangular FFT grid
+        mask_fft = self
+
+        # Coordinates of unmasked pixels in the FFT grid
+        ys, xs = np.where(~mask_fft)
+
+        # Width of the FFT grid (number of columns)
+        width = mask_fft.shape[1]
+
+        # Convert (y, x) coordinates to flat row-major indices
+        return (ys * width + xs).astype(np.int32)
+
     def trimmed_array_from(self, padded_array, image_shape) -> Array2D:
         """
         Map a padded 1D array of values to its original 2D array, trimming all edge values.