store tuple in mask, simplify a lot of code

Author: Jammy2211

autoarray/mask/mask_2d.py

@@ -217,6 +217,9 @@ def __init__(
             xp=xp,
         )
 
+        slim_to_native = self.derive_indexes.native_for_slim.astype("int32")
+        self.slim_to_native_tuple = (slim_to_native[:, 0], slim_to_native[:, 1])
+
     @property
     def native_for_slim(self):
         return self.derive_indexes.native_for_slim

autoarray/structures/arrays/kernel_2d.py

@@ -122,25 +122,6 @@ 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],
-            )
-
         self.fft_shape = fft_shape
 
         self.mask_shape = None
@@ -585,18 +566,6 @@ def mapping_matrix_native_from(
             Contains contributions from both the main mapping matrix and, if provided,
             the blurring mapping matrix.
         """
-        slim_to_native_tuple = self.slim_to_native_tuple
-        if slim_to_native_tuple is None:
-            mask_flat = xp.logical_not(mask.array)
-
-            if xp.__name__.startswith("jax"):
-                slim_to_native_tuple = xp.nonzero(
-                    mask_flat, size=mapping_matrix.shape[0]
-                )
-            else:
-                slim_to_native = mask.derive_indexes.native_for_slim.astype("int32")
-                slim_to_native_tuple = (slim_to_native[:, 0], slim_to_native[:, 1])
-
         n_src = mapping_matrix.shape[1]
 
         # Allocate full native grid (ny, nx, n_src)
@@ -606,43 +575,22 @@ def mapping_matrix_native_from(
 
         # Scatter main mapping matrix into native cube
         if xp.__name__.startswith("jax"):
-            mapping_matrix_native = mapping_matrix_native.at[slim_to_native_tuple].set(
+            mapping_matrix_native = mapping_matrix_native.at[mask.slim_to_native_tuple].set(
                 mapping_matrix
             )
         else:
-            mapping_matrix_native[slim_to_native_tuple] = mapping_matrix
+            mapping_matrix_native[mask.slim_to_native_tuple] = mapping_matrix
+
         # Optionally scatter blurring mapping matrix
+
         if blurring_mapping_matrix is not None:
-            slim_to_native_blurring_tuple = self.slim_to_native_blurring_tuple
-
-            if slim_to_native_blurring_tuple is None:
-                if blurring_mask is None:
-                    raise ValueError(
-                        "blurring_mask must be provided if blurring_mapping_matrix is given "
-                        "and slim_to_native_blurring_tuple is None."
-                    )
-
-                if xp.__name__.startswith("jax"):
-                    mask_flat = xp.logical_not(blurring_mask.array)
-                    slim_to_native_blurring_tuple = xp.nonzero(
-                        mask_flat,
-                        size=blurring_mapping_matrix.shape[0],
-                    )
-                else:
-                    slim_to_native_blurring = (
-                        blurring_mask.derive_indexes.native_for_slim.astype("int32")
-                    )
-                    slim_to_native_blurring_tuple = (
-                        slim_to_native_blurring[:, 0],
-                        slim_to_native_blurring[:, 1],
-                    )
 
             if xp.__name__.startswith("jax"):
                 mapping_matrix_native = mapping_matrix_native.at[
-                    slim_to_native_blurring_tuple
+                    blurring_mask.slim_to_native_tuple
                 ].set(blurring_mapping_matrix)
             else:
-                mapping_matrix_native[slim_to_native_blurring_tuple] = (
+                mapping_matrix_native[blurring_mask.slim_to_native_tuple] = (
                     blurring_mapping_matrix
                 )
 
@@ -722,31 +670,17 @@ def convolved_image_from(self, image, blurring_image, jax_method="direct", xp=np
             mask_shape = self.mask_shape
             fft_psf = self.fft_psf
 
-        slim_to_native_tuple = self.slim_to_native_tuple
-        slim_to_native_blurring_tuple = self.slim_to_native_blurring_tuple
-
-        if slim_to_native_tuple is None:
-
-            mask_flat = xp.logical_not(image.mask.array)
-            slim_to_native_tuple = xp.nonzero(mask_flat, size=image.shape[0])
-
         # start with native image padded with zeros
         image_both_native = xp.zeros(image.mask.shape, dtype=image.dtype)
 
-        image_both_native = image_both_native.at[slim_to_native_tuple].set(
+        image_both_native = image_both_native.at[image.mask.slim_to_native_tuple].set(
             xp.asarray(image.array)
         )
 
         # add blurring contribution if provided
         if blurring_image is not None:
-            if slim_to_native_blurring_tuple is None:
-
-                mask_flat = xp.logical_not(blurring_image.mask.array)
-                slim_to_native_blurring_tuple = xp.nonzero(
-                    mask_flat, size=blurring_image.shape[0]
-                )
 
-            image_both_native = image_both_native.at[slim_to_native_blurring_tuple].set(
+            image_both_native = image_both_native.at[blurring_image.mask.slim_to_native_tuple].set(
                 xp.asarray(blurring_image.array)
             )
 
@@ -777,7 +711,7 @@ def convolved_image_from(self, image, blurring_image, jax_method="direct", xp=np
         )
 
         blurred_image = Array2D(
-            values=blurred_image_native[slim_to_native_tuple], mask=image.mask
+            values=blurred_image_native[image.mask.slim_to_native_tuple], mask=image.mask
         )
 
         if self.fft_shape is None:
@@ -880,13 +814,6 @@ def convolved_mapping_matrix_from(
             mask_shape = self.mask_shape
             fft_psf_mapping = self.fft_psf_mapping
 
-        slim_to_native_tuple = self.slim_to_native_tuple
-
-        if slim_to_native_tuple is None:
-
-            mask_flat = xp.logical_not(mask.array)
-            slim_to_native_tuple = xp.nonzero(mask_flat, size=mapping_matrix.shape[0])
-
         mapping_matrix_native = self.mapping_matrix_native_from(
             mapping_matrix=mapping_matrix,
             mask=mask,
@@ -1055,30 +982,18 @@ def convolved_image_via_real_space_from(
 
         import jax
 
-        slim_to_native_tuple = self.slim_to_native_tuple
-        slim_to_native_blurring_tuple = self.slim_to_native_blurring_tuple
-
-        if slim_to_native_tuple is None:
-            mask_flat = xp.logical_not(image.mask.array)
-            slim_to_native_tuple = xp.nonzero(mask_flat, size=image.shape[0])
-
         # start with native array padded with zeros
         image_native = xp.zeros(image.mask.shape, dtype=xp.asarray(image.array).dtype)
 
         # set image pixels
-        image_native = image_native.at[slim_to_native_tuple].set(
+        image_native = image_native.at[image.mask.slim_to_native_tuple].set(
             xp.asarray(image.array)
         )
 
         # add blurring contribution if provided
         if blurring_image is not None:
-            if slim_to_native_blurring_tuple is None:
 
-                slim_to_native_blurring_tuple = xp.nonzero(
-                    mask_flat,
-                    size=blurring_image.shape[0],
-                )
-            image_native = image_native.at[slim_to_native_blurring_tuple].set(
+            image_native = image_native.at[blurring_image.mask.slim_to_native_tuple].set(
                 xp.asarray(blurring_image.array)
             )
         else:
@@ -1094,7 +1009,7 @@ def convolved_image_via_real_space_from(
             image_native, kernel, mode="same", method=jax_method
         )
 
-        convolved_array_1d = convolve_native[slim_to_native_tuple]
+        convolved_array_1d = convolve_native[image.mask.slim_to_native_tuple]
 
         return Array2D(values=convolved_array_1d, mask=image.mask)
 
@@ -1146,16 +1061,6 @@ def convolved_mapping_matrix_via_real_space_from(
 
         import jax
 
-        slim_to_native_tuple = self.slim_to_native_tuple
-
-        if slim_to_native_tuple is None:
-
-            mask_flat = xp.logical_not(mask.array)
-            slim_to_native_tuple = xp.nonzero(
-                mask_flat,
-                size=mapping_matrix.shape[0],
-            )
-
         mapping_matrix_native = self.mapping_matrix_native_from(
             mapping_matrix=mapping_matrix,
             mask=mask,
@@ -1174,7 +1079,7 @@ def convolved_mapping_matrix_via_real_space_from(
         )
 
         # return slim form
-        return blurred_mapping_matrix_native[slim_to_native_tuple]
+        return blurred_mapping_matrix_native[mask.slim_to_native_tuple]
 
     def convolved_image_via_real_space_np_from(
         self, image: np.ndarray, blurring_image: Optional[np.ndarray] = None, xp=np
@@ -1207,32 +1112,16 @@ def convolved_image_via_real_space_np_from(
 
         from scipy.signal import convolve as scipy_convolve
 
-        slim_to_native_tuple = self.slim_to_native_tuple
-        slim_to_native_blurring_tuple = self.slim_to_native_blurring_tuple
-
-        if slim_to_native_tuple is None:
-            slim_to_native = image.mask.derive_indexes.native_for_slim.astype("int32")
-            slim_to_native_tuple = (slim_to_native[:, 0], slim_to_native[:, 1])
-
         # start with native array padded with zeros
         image_native = xp.zeros(image.mask.shape, dtype=xp.asarray(image.array).dtype)
 
         # set image pixels
-        image_native[slim_to_native_tuple] = xp.asarray(image.array)
+        image_native[image.mask.slim_to_native_tuple] = xp.asarray(image.array)
 
         # add blurring contribution if provided
         if blurring_image is not None:
-            if slim_to_native_blurring_tuple is None:
-
-                slim_to_native_blurring = (
-                    blurring_image.mask.derive_indexes.native_for_slim.astype("int32")
-                )
-                slim_to_native_blurring_tuple = (
-                    slim_to_native_blurring[:, 0],
-                    slim_to_native_blurring[:, 1],
-                )
 
-            image_native[slim_to_native_blurring_tuple] = xp.asarray(
+            image_native[blurring_image.mask.slim_to_native_tuple] = xp.asarray(
                 blurring_image.array
             )
         else:
@@ -1248,7 +1137,7 @@ def convolved_image_via_real_space_np_from(
             image_native, kernel, mode="same", method="auto"
         )
 
-        convolved_array_1d = convolve_native[slim_to_native_tuple]
+        convolved_array_1d = convolve_native[image.mask.slim_to_native_tuple]
 
         return Array2D(values=convolved_array_1d, mask=image.mask)
 
@@ -1290,13 +1179,6 @@ def convolved_mapping_matrix_via_real_space_np_from(
 
         from scipy.signal import convolve as scipy_convolve
 
-        slim_to_native_tuple = self.slim_to_native_tuple
-
-        if slim_to_native_tuple is None:
-
-            slim_to_native = mask.derive_indexes.native_for_slim.astype("int32")
-            slim_to_native_tuple = (slim_to_native[:, 0], slim_to_native[:, 1])
-
         mapping_matrix_native = self.mapping_matrix_native_from(
             mapping_matrix=mapping_matrix,
             mask=mask,
@@ -1314,4 +1196,4 @@ def convolved_mapping_matrix_via_real_space_np_from(
         )
 
         # return slim form
-        return blurred_mapping_matrix_native[slim_to_native_tuple]
+        return blurred_mapping_matrix_native[mask.slim_to_native_tuple]