Imaging.apply_mask: trim the native arrays to the mask plus a PSF half-width (trim_to_mask)

autoarray/config/general.yaml

@@ -1,3 +1,5 @@
+dataset:
+  trim_to_mask: true               # If True, Imaging.apply_mask cuts the native arrays down to the smallest centred frame holding the unmasked pixels plus a PSF half-width and a one-pixel buffer (Imaging.trimmed_to_mask_from), so the FFT convolution frame is set by the mask and PSF rather than by the size of the image the mask was cut from. The fit is unchanged (only masked pixels are removed, the central pixel stays put); the native shape of every array and model image shrinks. Set False to keep the native shape of the input image.
 psf:
   use_fft_default: true              # If True, PSFs are convolved using FFTs by default, which is faster and uses less memory in all cases except for very small PSFs, False uses direct convolution.
 inversion:

autoarray/dataset/imaging/dataset.py

@@ -12,9 +12,11 @@
 from autoarray.operators.convolver import ConvolverState
 from autoarray.operators.convolver import Convolver
 from autoarray.mask.mask_2d import Mask2D
+from autoarray.mask import mask_2d_util
 from autoarray import type as ty
 
 from autoarray import exc
+from autonerves import conf
 
 from autoarray.inversion.inversion.imaging import inversion_imaging_util
 
@@ -200,10 +202,7 @@ def __init__(
                 f"resolution."
             )
 
-        if (
-            sparse_operator is not None
-            and convolve_over_sample_size_pixelization > 1
-        ):
+        if sparse_operator is not None and convolve_over_sample_size_pixelization > 1:
             raise exc.DatasetException(
                 "convolve_over_sample_size_pixelization > 1 is incompatible with the "
                 "sparse linear algebra formalism (sparse_operator), whose PSF "
@@ -227,15 +226,13 @@ def __init__(
             if ((self.noise_map.native <= 0.0) * np.invert(self.noise_map.mask)).any():
                 zero_entries = np.argwhere(self.noise_map.native <= 0.0)
 
-                raise exc.DatasetException(
-                    f"""
+                raise exc.DatasetException(f"""
                     A value in the noise-map of the dataset is {np.min(self.noise_map)}. 
 
                     This is less than or equal to zero, and therefore an ill-defined value which must be corrected.
 
                     The 2D indexes of the arrays in the native noise map array are {zero_entries}.
-                    """
-                )
+                    """)
 
         convolve_over_sample_size = max(
             self.convolve_over_sample_size_lp,
@@ -371,7 +368,9 @@ def from_fits(
         noise_map = Array2D.from_fits(
             file_path=noise_map_path, hdu=noise_map_hdu, pixel_scales=pixel_scales
         )
-        noise_map, pixel_scales = cap_array_2d_for_small_datasets(noise_map, pixel_scales)
+        noise_map, pixel_scales = cap_array_2d_for_small_datasets(
+            noise_map, pixel_scales
+        )
 
         if psf_path is not None:
 
@@ -413,7 +412,9 @@ def from_fits(
             convolve_over_sample_size_pixelization=convolve_over_sample_size_pixelization,
         )
 
-    def apply_mask(self, mask: Mask2D) -> "Imaging":
+    def apply_mask(
+        self, mask: Mask2D, trim_to_mask: Optional[bool] = None
+    ) -> "Imaging":
         """
         Apply a mask to the imaging dataset, whereby the mask is applied to the image data, noise-map and other
         quantities one-by-one.
@@ -423,6 +424,12 @@ def apply_mask(self, mask: Mask2D) -> "Imaging":
         and columns corresponding to masked pixels are removed so it stays consistent with the
         remaining unmasked pixels. The PSF `ConvolverState` is recomputed for the new mask.
 
+        With `trim_to_mask` the native arrays are then cut down to the smallest centred frame that
+        holds the unmasked pixels plus a PSF half-width and a one-pixel buffer (see
+        `trimmed_to_mask_from`), so a small mask on a large image does not pay for the empty frame
+        in every FFT convolution. The default comes from the config (`general.yaml`
+        `[dataset] trim_to_mask`, packaged True).
+
         The `apply_mask` function cannot be called multiple times — a new mask cannot expand the
         unmasked region beyond what was already unmasked, as the underlying data has already been
         trimmed. An exception is raised if this is attempted. If you wish to apply a different mask,
@@ -432,6 +439,9 @@ def apply_mask(self, mask: Mask2D) -> "Imaging":
         ----------
         mask
             The 2D mask that is applied to the image.
+        trim_to_mask
+            If True, the masked dataset is trimmed with `trimmed_to_mask_from`; if None, the
+            config setting `[dataset] trim_to_mask` decides.
 
         Returns
         -------
@@ -489,6 +499,86 @@ def apply_mask(self, mask: Mask2D) -> "Imaging":
             f"IMAGING - Data masked, contains a total of {mask.pixels_in_mask} image-pixels"
         )
 
+        if trim_to_mask is None:
+            trim_to_mask = conf.instance["general"]["dataset"]["trim_to_mask"]
+
+        if trim_to_mask:
+            dataset = dataset.trimmed_to_mask_from()
+
+        return dataset
+
+    def trimmed_to_mask_from(self, buffer: int = 1) -> "Imaging":
+        """
+        Return the dataset with its native arrays cut down to the smallest frame that holds every
+        unmasked pixel at least a PSF half-width plus `buffer` pixels from the edge.
+
+        A lens or galaxy is usually fitted inside a mask that covers a small part of the image it
+        was cut from, and the FFT convolution of every model image (and of every column of a
+        pixelization's mapping matrix) runs on a frame that is set by the native array, not by the
+        mask. For a 1.5" mask on a 12" JWST stamp that frame is 16 times the area the fit needs.
+        This method removes the surplus: the frame it keeps is the one `ConvolverState` already
+        pads to for the blurring region (`mask_2d_util.trimmed_shape_for_kernel`) plus a margin of
+        `buffer` masked pixels on every side, so the FFT shape falls to close to the minimum the
+        mask and kernel allow without being cut exactly at the blurring region's edge.
+
+        Nothing a fit sees changes. The trim removes masked pixels only, so the slim `data`,
+        `noise_map`, over-sampling sizes and `noise_covariance_matrix` are identical; the central
+        pixel stays on the same pixel (the parity of each dimension is preserved), so scaled
+        coordinates, positions and priors are unchanged; the blurring region is kept whole, so the
+        convolution is the same to floating-point rounding. What changes is the native shape of
+        every array (and of the images a fit produces), the PSF `ConvolverState`, and the memory
+        and time of each convolution. A mask whose unmasked pixels sit closer to the edge than the
+        PSF half-width is padded instead, as the convolver would pad it internally.
+
+        Applied automatically by `apply_mask` when the config setting `[dataset] trim_to_mask` is
+        True. A `sparse_operator` is discarded, as by `apply_mask`.
+
+        Parameters
+        ----------
+        buffer
+            Extra masked pixels kept on every side beyond the PSF half-width (default 1).
+
+        Returns
+        -------
+        Imaging
+            A new `Imaging` dataset on the trimmed frame, or this dataset if it is already that size.
+        """
+        kernel_shape = self.psf.kernel.shape_native if self.psf is not None else (1, 1)
+
+        new_shape = mask_2d_util.trimmed_shape_for_kernel(
+            mask_2d=np.array(self.data.mask),
+            kernel_shape_native=kernel_shape,
+            buffer=buffer,
+        )
+
+        if tuple(new_shape) == tuple(self.data.shape_native):
+            return self
+
+        _warn_sparse_operator_discarded(
+            sparse_operator=self.sparse_operator, method_name="trimmed_to_mask_from"
+        )
+
+        def trimmed(array):
+            return array.resized_from(new_shape=new_shape, mask_pad_value=True)
+
+        dataset = Imaging(
+            data=trimmed(self.data),
+            noise_map=trimmed(self.noise_map),
+            psf=self.psf,
+            psf_setup_state=True,
+            noise_covariance_matrix=self.noise_covariance_matrix,
+            over_sample_size_lp=trimmed(self.over_sample_size_lp),
+            over_sample_size_pixelization=trimmed(self.over_sample_size_pixelization),
+            convolve_over_sample_size_lp=self.convolve_over_sample_size_lp,
+            convolve_over_sample_size_pixelization=self.convolve_over_sample_size_pixelization,
+            check_noise_map=False,
+        )
+
+        logger.info(
+            f"IMAGING - Data trimmed to the mask, native shape {tuple(self.data.shape_native)} -> "
+            f"{tuple(new_shape)} (PSF kernel {tuple(kernel_shape)})"
+        )
+
         return dataset
 
     def apply_noise_scaling(
@@ -665,9 +755,7 @@ def apply_sparse_operator(
         """
 
         if self.psf is not None and self.psf.convolve_over_sample_size > 1:
-            raise exc.DatasetException(
-                _SPARSE_OPERATOR_CONVOLVE_OVER_SAMPLE_SIZE_ERROR
-            )
+            raise exc.DatasetException(_SPARSE_OPERATOR_CONVOLVE_OVER_SAMPLE_SIZE_ERROR)
 
         logger.info(
             "IMAGING - Setting Up Sparse Operator For low Memory Pixelizations."
@@ -715,9 +803,7 @@ def apply_sparse_operator_cpu(
             enabling efficient pixelized source reconstruction on CPU hardware.
         """
         if self.psf is not None and self.psf.convolve_over_sample_size > 1:
-            raise exc.DatasetException(
-                _SPARSE_OPERATOR_CONVOLVE_OVER_SAMPLE_SIZE_ERROR
-            )
+            raise exc.DatasetException(_SPARSE_OPERATOR_CONVOLVE_OVER_SAMPLE_SIZE_ERROR)
 
         try:
             import numba
@@ -767,4 +853,3 @@ def apply_sparse_operator_cpu(
             check_noise_map=False,
             sparse_operator=sparse_operator,
         )
-

autoarray/mask/mask_2d_util.py

@@ -515,6 +515,90 @@ def required_shape_for_kernel(
     return new_y, new_x
 
 
+def trimmed_shape_for_kernel(
+    mask_2d: np.ndarray,
+    kernel_shape_native: Tuple[int, int],
+    buffer: int = 1,
+) -> Tuple[int, int]:
+    """
+    Return the smallest shape, centred on the mask's central pixel and of the same parity as the
+    mask in each dimension, that holds every unmasked pixel at least a kernel half-width plus
+    `buffer` pixels from the array edge.
+
+    This is the frame a masked dataset can be cut down to without changing anything a fit sees:
+    every unmasked pixel is kept, the blurring region a kernel of this footprint reaches is kept,
+    and because the parity of each dimension is preserved the central pixel (the origin of the
+    scaled coordinate system) stays on the same pixel, so scaled coordinates, positions and
+    priors are unchanged. Cutting an array to this shape with ``resized_array_2d_from`` therefore
+    only removes masked pixels.
+
+    The `buffer` keeps a margin of masked pixels around the blurring region, so the frame is not
+    cut exactly at the mask plus the half-width; the default of one pixel costs nothing a fit can
+    measure. Pass ``buffer=0`` for the exact minimum.
+
+    The returned shape can be larger than the mask's own in a dimension where unmasked pixels sit
+    closer to the edge than the kernel half-width plus the buffer; the frame is then padded, as
+    ``required_shape_for_kernel`` would.
+
+    For an even dimension the central pixel is the lower of the central pair (the convention of
+    ``resized_array_2d_from``), so one more pixel lies after it than before it.
+
+    Parameters
+    ----------
+    mask_2d
+        2D boolean array where False is unmasked and True is masked.
+    kernel_shape_native
+        (ky, kx) footprint of the convolution kernel.
+    buffer
+        Extra masked pixels kept on every side beyond the kernel half-width.
+
+    Returns
+    -------
+    trimmed_shape
+        The (ny, nx) centred shape.
+
+    Examples
+    --------
+    >>> mask = np.full((9, 9), True)
+    >>> mask[3:6, 3:6] = False
+    >>> trimmed_shape_for_kernel(mask, kernel_shape_native=(3, 3))
+    (7, 7)
+    >>> trimmed_shape_for_kernel(mask, kernel_shape_native=(3, 3), buffer=0)
+    (5, 5)
+    """
+    mask_2d = np.asarray(mask_2d, dtype=bool)
+
+    ky, kx = kernel_shape_native
+    if ky <= 0 or kx <= 0:
+        raise ValueError(
+            f"kernel_shape_native must be positive, got {kernel_shape_native}."
+        )
+    if buffer < 0:
+        raise ValueError(f"buffer must be non-negative, got {buffer}.")
+
+    unmasked = np.where(~mask_2d)
+    if unmasked[0].size == 0:
+        raise ValueError(
+            "trimmed_shape_for_kernel requires at least one unmasked pixel."
+        )
+
+    shape = []
+    for length, indexes, pad in zip(
+        mask_2d.shape, unmasked, (ky // 2 + buffer, kx // 2 + buffer)
+    ):
+        centre = (length - 1) // 2
+        before = max(centre - int(indexes.min()), 0) + pad
+        after = max(int(indexes.max()) - centre, 0) + pad
+        if length % 2:
+            half = max(before, after)
+            shape.append(2 * half + 1)
+        else:
+            half = max(before, after - 1)
+            shape.append(2 * half + 2)
+
+    return tuple(shape)
+
+
 def blurring_mask_2d_from(
     mask_2d: np.ndarray,
     kernel_shape_native: Tuple[int, int],

test_autoarray/config/general.yaml

@@ -1,5 +1,7 @@
 analysis:
   n_cores: 1
+dataset:
+  trim_to_mask: false              # If True, Imaging.apply_mask trims the native arrays to the mask plus a PSF half-width and a one-pixel buffer. Off for unit tests so fixture shapes are kept.
 psf:
   use_fft_default: false              # If True, PSFs are convolved using FFTs by default, which is faster and uses less memory in all cases except for very small PSFs, False uses direct convolution. Real space used for unit tests.
 grid: