refactor kernel mapping for speed

Author: Jammy2211

autoarray/dataset/imaging/dataset.py

@@ -159,9 +159,26 @@ def __init__(
                     """
                 )
 
-        if psf is not None and use_normalized_psf:
+        if psf is not None:
+
+            if not data.mask.is_all_false:
+
+                image_mask = data.mask
+                blurring_mask = data.mask.derive_mask.blurring_from(
+                    kernel_shape_native=psf.shape_native
+                )
+
+            else:
+
+                image_mask = None
+                blurring_mask = None
+
             psf = Kernel2D.no_mask(
-                values=psf.native._array, pixel_scales=psf.pixel_scales, normalize=True
+                values=psf.native._array,
+                pixel_scales=psf.pixel_scales,
+                normalize=use_normalized_psf,
+                image_mask=image_mask,
+                blurring_mask=blurring_mask,
             )
 
         self.psf = psf

autoarray/inversion/pixelization/mappers/rectangular.py

@@ -97,6 +97,24 @@ def pix_sub_weights(self) -> PixSubWeights:
         dimension of the array `pix_indexes_for_sub_slim_index` 1 and all entries in `pix_weights_for_sub_slim_index`
         are equal to 1.0.
         """
+        # from autoarray.geometry import geometry_util
+        #
+        # mappings = geometry_util.grid_pixel_indexes_2d_slim_from(
+        #     grid_scaled_2d_slim=np.array(self.source_plane_data_grid.over_sampled),
+        #     shape_native=self.source_plane_mesh_grid.shape_native,
+        #     pixel_scales=self.source_plane_mesh_grid.pixel_scales,
+        #     origin=self.source_plane_mesh_grid.origin,
+        # ).astype("int")
+        #
+        # mappings = mappings.reshape((len(mappings), 1))
+        #
+        # return PixSubWeights(
+        #     mappings=mappings,
+        #     sizes=np.ones(len(mappings), dtype="int"),
+        #     weights=np.ones(
+        #         (len(self.source_plane_data_grid.over_sampled), 1), dtype="int"
+        #     ),
+        # )
 
         mappings, weights = (
             mapper_util.rectangular_mappings_weights_via_interpolation_from(

autoarray/mask/mask_2d_util.py

@@ -50,9 +50,14 @@ def native_index_for_slim_index_2d_from(
 
     native_index_for_slim_index_2d = native_index_for_slim_index_2d_from(mask_2d=mask_2d)
     """
-    if isinstance(mask_2d, np.ndarray):
-        return np.stack(np.nonzero(~mask_2d.astype(bool))).T
-    return jnp.stack(jnp.nonzero(~mask_2d.astype(bool))).T
+
+    if isinstance(mask_2d, jnp.ndarray):
+        # JAX branch (assume jnp.ndarray)
+        rows, cols = jnp.where(~mask_2d.astype(bool))
+        return jnp.stack([rows, cols], axis=1)
+
+    rows, cols = np.where(~mask_2d.astype(bool))
+    return np.stack([rows, cols], axis=1)
 
 
 def mask_2d_centres_from(

autoarray/plot/wrap/base/output.py

@@ -108,7 +108,7 @@ def savefig(self, filename: str, output_path: str, format: str):
             plt.savefig(
                 path.join(output_path, f"{filename}.{format}"),
                 bbox_inches=self.bbox_inches,
-                pad_inches=0,
+                pad_inches=0.1,
             )
         except ValueError as e:
             logger.info(

autoarray/structures/arrays/kernel_2d.py

@@ -24,6 +24,8 @@ def __init__(
         header=None,
         normalize: bool = False,
         store_native: bool = False,
+        image_mask=None,
+        blurring_mask=None,
         *args,
         **kwargs,
     ):
@@ -56,6 +58,25 @@ def __init__(
 
         self.stored_native = self.native
 
+        self.slim_to_native_tuple = None
+
+        if image_mask is not None:
+
+            slim_to_native = image_mask.derive_indexes.native_for_slim.astype("int32")
+            self.slim_to_native_tuple = (slim_to_native[:, 0], slim_to_native[:, 1])
+
+        self.slim_to_native_blurring_tuple = None
+
+        if blurring_mask is not None:
+
+            slim_to_native_blurring = (
+                blurring_mask.derive_indexes.native_for_slim.astype("int32")
+            )
+            self.slim_to_native_blurring_tuple = (
+                slim_to_native_blurring[:, 0],
+                slim_to_native_blurring[:, 1],
+            )
+
     @classmethod
     def no_mask(
         cls,
@@ -64,6 +85,8 @@ def no_mask(
         shape_native: Tuple[int, int] = None,
         origin: Tuple[float, float] = (0.0, 0.0),
         normalize: bool = False,
+        image_mask=None,
+        blurring_mask=None,
     ):
         """
         Create a Kernel2D (see *Kernel2D.__new__*) by inputting the kernel values in 1D or 2D, automatically
@@ -91,7 +114,13 @@ def no_mask(
             pixel_scales=pixel_scales,
             origin=origin,
         )
-        return Kernel2D(values=values, mask=values.mask, normalize=normalize)
+        return Kernel2D(
+            values=values,
+            mask=values.mask,
+            normalize=normalize,
+            image_mask=image_mask,
+            blurring_mask=blurring_mask,
+        )
 
     @classmethod
     def full(
@@ -540,29 +569,41 @@ def convolve_image(self, image, blurring_image, jax_method="direct"):
             kernels that are more than about 5x5. Default is `fft`.
         """
 
-        slim_to_native = jnp.nonzero(
-            jnp.logical_not(image.mask.array), size=image.shape[0]
-        )
-        slim_to_native_blurring = jnp.nonzero(
-            jnp.logical_not(blurring_image.mask.array), size=blurring_image.shape[0]
-        )
+        slim_to_native_tuple = self.slim_to_native_tuple
+        slim_to_native_blurring_tuple = self.slim_to_native_blurring_tuple
 
-        expanded_array_native = jnp.zeros(image.mask.shape)
+        if slim_to_native_tuple is None:
 
-        expanded_array_native = expanded_array_native.at[slim_to_native].set(
-            image.array
+            slim_to_native_tuple = jnp.nonzero(
+                jnp.logical_not(image.mask.array), size=image.shape[0]
+            )
+
+        if slim_to_native_blurring_tuple is None:
+
+            slim_to_native_blurring_tuple = jnp.nonzero(
+                jnp.logical_not(blurring_image.mask.array), size=blurring_image.shape[0]
+            )
+
+        # make sure dtype matches what you want
+        expanded_array_native = jnp.zeros(
+            image.mask.shape, dtype=jnp.asarray(image.array).dtype
         )
-        expanded_array_native = expanded_array_native.at[slim_to_native_blurring].set(
-            blurring_image.array
+
+        # set using a tuple of index arrays
+        expanded_array_native = expanded_array_native.at[slim_to_native_tuple].set(
+            jnp.asarray(image.array)
         )
+        expanded_array_native = expanded_array_native.at[
+            slim_to_native_blurring_tuple
+        ].set(jnp.asarray(blurring_image.array))
 
         kernel = self.stored_native.array
 
         convolve_native = jax.scipy.signal.convolve(
             expanded_array_native, kernel, mode="same", method=jax_method
         )
 
-        convolved_array_1d = convolve_native[slim_to_native]
+        convolved_array_1d = convolve_native[slim_to_native_tuple]
 
         return Array2D(values=convolved_array_1d, mask=image.mask)
 
@@ -583,24 +624,77 @@ def convolve_image_no_blurring(self, image, mask, jax_method="direct"):
             kernels that are more than about 5x5. Default is `fft`.
         """
 
-        slim_to_native = jnp.nonzero(jnp.logical_not(mask.array), size=image.shape[0])
+        slim_to_native_tuple = self.slim_to_native_tuple
+
+        if slim_to_native_tuple is None:
+
+            slim_to_native_tuple = jnp.nonzero(
+                jnp.logical_not(mask.array), size=image.shape[0]
+            )
 
+        # make sure dtype matches what you want
         expanded_array_native = jnp.zeros(mask.shape)
 
+        # set using a tuple of index arrays
         if isinstance(image, np.ndarray) or isinstance(image, jnp.ndarray):
-            expanded_array_native = expanded_array_native.at[slim_to_native].set(image)
+            expanded_array_native = expanded_array_native.at[slim_to_native_tuple].set(
+                image
+            )
         else:
-            expanded_array_native = expanded_array_native.at[slim_to_native].set(
-                image.array
+            expanded_array_native = expanded_array_native.at[slim_to_native_tuple].set(
+                jnp.asarray(image.array)
+            )
+
+        kernel = self.stored_native.array
+
+        convolve_native = jax.scipy.signal.convolve(
+            expanded_array_native, kernel, mode="same", method=jax_method
+        )
+
+        convolved_array_1d = convolve_native[slim_to_native_tuple]
+
+        return Array2D(values=convolved_array_1d, mask=mask)
+
+    def convolve_image_no_blurring_for_mapping(self, image, mask, jax_method="direct"):
+        """
+        For a given 1D array and blurring array, convolve the two using this psf.
+
+        Parameters
+        ----------
+        image
+            1D array of the values which are to be blurred with the psf's PSF.
+        blurring_image
+            1D array of the blurring values which blur into the array after PSF convolution.
+        jax_method
+            If JAX is enabled this keyword will indicate what method is used for the PSF
+            convolution. Can be either `direct` to calculate it in real space or `fft`
+            to calculated it via a fast Fourier transform. `fft` is typically faster for
+            kernels that are more than about 5x5. Default is `fft`.
+        """
+
+        slim_to_native_tuple = self.slim_to_native_tuple
+
+        if slim_to_native_tuple is None:
+
+            slim_to_native_tuple = jnp.nonzero(
+                jnp.logical_not(mask.array), size=image.shape[0]
             )
 
+        # make sure dtype matches what you want
+        expanded_array_native = jnp.zeros(mask.shape)
+
+        # set using a tuple of index arrays
+        expanded_array_native = expanded_array_native.at[slim_to_native_tuple].set(
+            image
+        )
+
         kernel = self.stored_native.array
 
         convolve_native = jax.scipy.signal.convolve(
             expanded_array_native, kernel, mode="same", method=jax_method
         )
 
-        convolved_array_1d = convolve_native[slim_to_native]
+        convolved_array_1d = convolve_native[slim_to_native_tuple]
 
         return Array2D(values=convolved_array_1d, mask=mask)
 
@@ -612,6 +706,6 @@ def convolve_mapping_matrix(self, mapping_matrix, mask, jax_method="direct"):
         image
             1D array of the values which are to be blurred with the psf's PSF.
         """
-        return jax.vmap(self.convolve_image_no_blurring, in_axes=(1, None, None))(
-            mapping_matrix, mask, jax_method
-        ).T
+        return jax.vmap(
+            self.convolve_image_no_blurring_for_mapping, in_axes=(1, None, None)
+        )(mapping_matrix, mask, jax_method).T

autoarray/structures/grids/grid_2d_util.py

@@ -253,25 +253,29 @@ def grid_2d_slim_via_mask_from(
     centres_scaled = geometry_util.central_scaled_coordinate_2d_from(
         shape_native=mask_2d.shape, pixel_scales=pixel_scales, origin=origin
     )
-    if isinstance(mask_2d, jnp.ndarray):
 
-        centres_scaled = jnp.array(centres_scaled)
-        pixel_scales = jnp.array(pixel_scales)
+    # JAX branch
+    if isinstance(mask_2d, jnp.ndarray):
+        centres_scaled = jnp.asarray(centres_scaled)
+        pixel_scales = jnp.asarray(pixel_scales)
         sign = jnp.array([-1.0, 1.0])
-        return (
-            (jnp.stack(jnp.nonzero(~mask_2d.astype(bool))).T - centres_scaled)
-            * sign
-            * pixel_scales
-        )
 
-    centres_scaled = np.array(centres_scaled)
-    pixel_scales = np.array(pixel_scales)
+        # use jnp.where instead of jnp.nonzero
+        rows, cols = jnp.where(~mask_2d.astype(bool))
+        indices = jnp.stack([rows, cols], axis=1)  # shape (N_unmasked, 2)
+
+        # (indices - centre) -> pixel offsets; apply sign and scale to get physical coords
+        return (indices - centres_scaled) * sign * pixel_scales
+
+    # NumPy branch (kept consistent)
+    centres_scaled = np.asarray(centres_scaled)
+    pixel_scales = np.asarray(pixel_scales)
     sign = np.array([-1.0, 1.0])
-    return (
-        (np.stack(np.nonzero(~mask_2d.astype(bool))).T - centres_scaled)
-        * sign
-        * pixel_scales
-    )
+
+    rows, cols = np.where(~mask_2d.astype(bool))
+    indices = np.stack([rows, cols], axis=1)
+
+    return (indices - centres_scaled) * sign * pixel_scales
 
 
 def grid_2d_via_mask_from(