From a7f9fbc7b34e8d8dafedc7cf45d1768ea914440f Mon Sep 17 00:00:00 2001 From: samlange04 Date: Mon, 2 Nov 2026 22:31:57 +0800 Subject: [PATCH 1/2] Imaging.apply_mask: trim the native arrays to the mask plus a PSF half-width A fit's FFT convolution frame is set by the native array, not by the mask, so a small mask on a large image pays for the empty frame in every convolution (a 1.5" mask on a 12" JWST stamp: 16x the area the fit needs). apply_mask now takes trim_to_mask (default from the new general.yaml [dataset] trim_to_mask, packaged false) and cuts data, noise map, over-sampling sizes and mask to the smallest centred frame that keeps every unmasked pixel a PSF half-width from the edge (mask_2d_util.trimmed_shape_for_kernel, parity preserved so the central pixel and the scaled coordinates stay put). Only masked pixels are removed: the slim arrays, the noise covariance matrix and the convolution are unchanged. Imaging.trimmed_to_mask_from does the same on an already masked dataset. Co-Authored-By: Claude Fable 5.1 --- autoarray/config/general.yaml | 2 + autoarray/dataset/imaging/dataset.py | 88 +++++++++++++++++- autoarray/mask/mask_2d_util.py | 69 ++++++++++++++ test_autoarray/config/general.yaml | 2 + .../dataset/imaging/test_dataset.py | 89 +++++++++++++++++++ test_autoarray/mask/test_mask_2d_util.py | 58 ++++++++++++ 6 files changed, 307 insertions(+), 1 deletion(-) diff --git a/autoarray/config/general.yaml b/autoarray/config/general.yaml index 4f6b67df7..e63066838 100644 --- a/autoarray/config/general.yaml +++ b/autoarray/config/general.yaml @@ -1,3 +1,5 @@ +dataset: + trim_to_mask: false # If True, Imaging.apply_mask cuts the native arrays down to the smallest centred frame holding the unmasked pixels plus a PSF half-width (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. 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: diff --git a/autoarray/dataset/imaging/dataset.py b/autoarray/dataset/imaging/dataset.py index 3ac1c7d8d..cc1e3e616 100644 --- a/autoarray/dataset/imaging/dataset.py +++ b/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 @@ -413,7 +415,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 +427,11 @@ 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 (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`). + 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 +441,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 +501,80 @@ 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) -> "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 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`), so the FFT + shape falls to the minimum the mask and kernel allow. + + 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`. + + 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 + ) + + 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( diff --git a/autoarray/mask/mask_2d_util.py b/autoarray/mask/mask_2d_util.py index 00b444fbd..bfe8fe4d6 100644 --- a/autoarray/mask/mask_2d_util.py +++ b/autoarray/mask/mask_2d_util.py @@ -515,6 +515,75 @@ 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], +) -> 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 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 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; 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. + + Returns + ------- + trimmed_shape + The (ny, nx) centred shape. + + Examples + -------- + >>> mask = np.full((7, 7), True) + >>> mask[2:5, 2:5] = False + >>> trimmed_shape_for_kernel(mask, kernel_shape_native=(3, 3)) + (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}." + ) + + 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, kx // 2)): + 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], diff --git a/test_autoarray/config/general.yaml b/test_autoarray/config/general.yaml index b07c331e1..aeed5854b 100644 --- a/test_autoarray/config/general.yaml +++ b/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. 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: diff --git a/test_autoarray/dataset/imaging/test_dataset.py b/test_autoarray/dataset/imaging/test_dataset.py index 4ecc9fe81..321499382 100644 --- a/test_autoarray/dataset/imaging/test_dataset.py +++ b/test_autoarray/dataset/imaging/test_dataset.py @@ -5,6 +5,7 @@ import pytest import autoarray as aa +from autonerves import conf from autoarray import exc from autoarray.inversion.inversion.imaging_numba.sparse import ( @@ -268,6 +269,94 @@ def test__apply_mask__data_noise_map_psf_correctly_masked( assert type(masked_imaging_7x7.psf) == aa.Convolver +def test__apply_mask__trim_to_mask__native_cut_to_mask_plus_psf_half_width__slim_unchanged( + imaging_7x7, mask_2d_7x7 +): + untrimmed = imaging_7x7.apply_mask(mask=mask_2d_7x7, trim_to_mask=False) + trimmed = imaging_7x7.apply_mask(mask=mask_2d_7x7, trim_to_mask=True) + + # The 3x3 unmasked block plus the 3x3 PSF's half-width of 1 on each side. + assert trimmed.data.shape_native == (5, 5) + assert trimmed.noise_map.shape_native == (5, 5) + assert trimmed.over_sample_size_lp.shape_native == (5, 5) + assert trimmed.over_sample_size_pixelization.shape_native == (5, 5) + assert trimmed.mask.shape_native == (5, 5) + assert trimmed.psf._state.fft_shape[0] <= untrimmed.psf._state.fft_shape[0] + + # Only masked pixels were removed. + assert (trimmed.data.slim == untrimmed.data.slim).all() + assert (trimmed.noise_map.slim == untrimmed.noise_map.slim).all() + assert (trimmed.over_sample_size_lp.slim == untrimmed.over_sample_size_lp.slim).all() + assert (trimmed.mask == untrimmed.mask.resized_from(new_shape=(5, 5), pad_value=True)).all() + assert (trimmed.data.native[1:4, 1:4] == untrimmed.data.native[2:5, 2:5]).all() + + # The central pixel stays on the same pixel: scaled coordinates are unchanged. + assert (trimmed.grids.lp.slim == untrimmed.grids.lp.slim).all() + assert trimmed.mask.origin == untrimmed.mask.origin + + # The convolution gives the same blurred image. + image = aa.Array2D(values=np.arange(9, dtype=float), mask=untrimmed.mask) + blurring_untrimmed = aa.Array2D( + values=np.arange(untrimmed.grids.blurring.shape[0], dtype=float), + mask=untrimmed.grids.blurring.mask, + ) + blurring_trimmed = aa.Array2D( + values=blurring_untrimmed.slim.array, mask=trimmed.grids.blurring.mask + ) + convolved_untrimmed = untrimmed.psf.convolved_image_from( + image=image, blurring_image=blurring_untrimmed + ) + convolved_trimmed = trimmed.psf.convolved_image_from( + image=aa.Array2D(values=np.arange(9, dtype=float), mask=trimmed.mask), + blurring_image=blurring_trimmed, + ) + assert np.asarray(convolved_trimmed) == pytest.approx(np.asarray(convolved_untrimmed), 1e-10) + + +def test__apply_mask__trim_to_mask__default_from_config( + imaging_7x7, mask_2d_7x7, monkeypatch +): + # The test config sets [dataset] trim_to_mask: false. + masked = imaging_7x7.apply_mask(mask=mask_2d_7x7) + + assert masked.data.shape_native == (7, 7) + + monkeypatch.setitem(conf.instance["general"]["dataset"], "trim_to_mask", True) + + masked = imaging_7x7.apply_mask(mask=mask_2d_7x7) + + assert masked.data.shape_native == (5, 5) + + +def test__trimmed_to_mask_from__already_minimal__returns_self__covariance_kept(): + image = aa.Array2D.ones(shape_native=(2, 2), pixel_scales=(1.0, 1.0)) + noise_map = aa.Array2D.ones(shape_native=(2, 2), pixel_scales=(1.0, 1.0)) + noise_covariance_matrix = np.array( + [[1.0, 0.1, 0.0, 0.0], [0.1, 2.0, 0.0, 0.0], [0.0, 0.0, 3.0, 0.2], [0.0, 0.0, 0.2, 4.0]] + ) + psf = aa.Convolver( + kernel=aa.Array2D.no_mask(values=[[1.0]], pixel_scales=(1.0, 1.0)) + ) + imaging = aa.Imaging( + data=image, + noise_map=noise_map, + psf=psf, + noise_covariance_matrix=noise_covariance_matrix, + over_sample_size_lp=1, + ) + + mask = aa.Mask2D(mask=[[True, False], [False, True]], pixel_scales=(1.0, 1.0)) + + masked = imaging.apply_mask(mask=mask, trim_to_mask=False) + + assert masked.trimmed_to_mask_from() is masked + + trimmed = imaging.apply_mask(mask=mask, trim_to_mask=True) + + assert trimmed.data.shape_native == (2, 2) + assert (trimmed.noise_covariance_matrix == masked.noise_covariance_matrix).all() + + def test__apply_noise_scaling__masked_pixel_data_zeroed_and_noise_set_to_noise_value( imaging_7x7, mask_2d_7x7 ): diff --git a/test_autoarray/mask/test_mask_2d_util.py b/test_autoarray/mask/test_mask_2d_util.py index 89d6d6bc1..c5e6cc639 100644 --- a/test_autoarray/mask/test_mask_2d_util.py +++ b/test_autoarray/mask/test_mask_2d_util.py @@ -1104,3 +1104,61 @@ def test__buffed_mask_2d_from(): ] ) ).all() + + +def test__trimmed_shape_for_kernel__centred_square_region__shape_is_region_plus_half_kernel(): + mask = np.full((7, 7), True) + mask[2:5, 2:5] = False + + assert util.mask_2d.trimmed_shape_for_kernel(mask, kernel_shape_native=(3, 3)) == (5, 5) + assert util.mask_2d.trimmed_shape_for_kernel(mask, kernel_shape_native=(1, 1)) == (3, 3) + assert util.mask_2d.trimmed_shape_for_kernel(mask, kernel_shape_native=(5, 3)) == (7, 5) + + +def test__trimmed_shape_for_kernel__parity_preserved_and_centre_kept(): + # Even frame: the central pixel is the lower of the central pair (3 of 0..7), and the + # unmasked 2x2 block sits on pixels 3 and 4, so it is symmetric about the pixel boundary. + mask = np.full((8, 8), True) + mask[3:5, 3:5] = False + + shape = util.mask_2d.trimmed_shape_for_kernel(mask, kernel_shape_native=(3, 3)) + + assert shape == (4, 4) + + # Cutting to that shape about the centre keeps every unmasked pixel one pixel from the edge. + trimmed = util.array_2d.resized_array_2d_from( + array_2d=mask.astype(float), resized_shape=shape, pad_value=1.0 + ).astype(bool) + assert trimmed.sum() == mask.sum() - (64 - 16) + assert util.mask_2d.min_false_distance_to_edge(trimmed) == (1, 1) + + # Odd frame, region offset from the centre: the frame grows symmetrically so the + # central pixel stays put, rather than re-centring on the region. + mask = np.full((9, 9), True) + mask[1:3, 5:8] = False + + shape = util.mask_2d.trimmed_shape_for_kernel(mask, kernel_shape_native=(3, 3)) + + assert shape == (9, 9) + assert shape[0] % 2 == 1 and shape[1] % 2 == 1 + + +def test__trimmed_shape_for_kernel__region_at_edge__frame_pads_rather_than_cuts(): + mask = np.full((5, 5), True) + mask[0, 2] = False + + shape = util.mask_2d.trimmed_shape_for_kernel(mask, kernel_shape_native=(5, 5)) + + # The pixel at row 0 needs two rows above it: 2 before the centre + 2 pad, so 9 rows. + assert shape == (9, 5) + + +def test__trimmed_shape_for_kernel__invalid_inputs_raise(): + mask = np.full((5, 5), True) + mask[2, 2] = False + + with pytest.raises(ValueError): + util.mask_2d.trimmed_shape_for_kernel(mask, kernel_shape_native=(0, 3)) + + with pytest.raises(ValueError): + util.mask_2d.trimmed_shape_for_kernel(np.full((5, 5), True), kernel_shape_native=(3, 3)) From aaa5b58286880527b3cad1aab37a0d9d3c08f541 Mon Sep 17 00:00:00 2001 From: samlange04 Date: Tue, 3 Nov 2026 19:43:19 +0800 Subject: [PATCH 2/2] Trim to mask: packaged default True; keep a one-pixel buffer beyond the PSF half-width trim_to_mask is now on by default: nothing a fit sees changes, every masked fit gets a smaller FFT frame, and a user who wants the input image's native shape kept sets general.yaml [dataset] trim_to_mask false. The unit-test config keeps it false so fixture shapes are unchanged. trimmed_shape_for_kernel and Imaging.trimmed_to_mask_from take buffer (default 1): the frame keeps that many extra masked pixels on every side beyond the kernel half-width, so it is not cut exactly at the blurring region's edge. buffer=0 gives the exact minimum. Co-Authored-By: Claude Fable 5.1 --- autoarray/config/general.yaml | 2 +- autoarray/dataset/imaging/dataset.py | 53 +++++----- autoarray/mask/mask_2d_util.py | 29 ++++-- test_autoarray/config/general.yaml | 2 +- .../dataset/imaging/test_dataset.py | 99 ++++++++++++++----- test_autoarray/mask/test_mask_2d_util.py | 69 ++++++++++--- 6 files changed, 181 insertions(+), 73 deletions(-) diff --git a/autoarray/config/general.yaml b/autoarray/config/general.yaml index e63066838..53aa6595a 100644 --- a/autoarray/config/general.yaml +++ b/autoarray/config/general.yaml @@ -1,5 +1,5 @@ dataset: - trim_to_mask: false # If True, Imaging.apply_mask cuts the native arrays down to the smallest centred frame holding the unmasked pixels plus a PSF half-width (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. + 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: diff --git a/autoarray/dataset/imaging/dataset.py b/autoarray/dataset/imaging/dataset.py index cc1e3e616..3784cf8b7 100644 --- a/autoarray/dataset/imaging/dataset.py +++ b/autoarray/dataset/imaging/dataset.py @@ -202,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 " @@ -229,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, @@ -373,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: @@ -428,9 +425,10 @@ def apply_mask( 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 (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`). + 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 @@ -509,18 +507,19 @@ def apply_mask( return dataset - def trimmed_to_mask_from(self) -> "Imaging": + 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 from the edge. + 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`), so the FFT - shape falls to the minimum the mask and kernel allow. + 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 @@ -534,17 +533,22 @@ def trimmed_to_mask_from(self) -> "Imaging": 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) - ) + 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 + mask_2d=np.array(self.data.mask), + kernel_shape_native=kernel_shape, + buffer=buffer, ) if tuple(new_shape) == tuple(self.data.shape_native): @@ -751,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." @@ -801,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 @@ -853,4 +853,3 @@ def apply_sparse_operator_cpu( check_noise_map=False, sparse_operator=sparse_operator, ) - diff --git a/autoarray/mask/mask_2d_util.py b/autoarray/mask/mask_2d_util.py index bfe8fe4d6..1d9fd1caf 100644 --- a/autoarray/mask/mask_2d_util.py +++ b/autoarray/mask/mask_2d_util.py @@ -518,11 +518,12 @@ def required_shape_for_kernel( 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 from - the array edge. + 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, @@ -531,8 +532,12 @@ def trimmed_shape_for_kernel( 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; the frame is then padded, as + 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 @@ -544,6 +549,8 @@ def trimmed_shape_for_kernel( 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 ------- @@ -552,9 +559,11 @@ def trimmed_shape_for_kernel( Examples -------- - >>> mask = np.full((7, 7), True) - >>> mask[2:5, 2:5] = False + >>> 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) @@ -564,13 +573,19 @@ def trimmed_shape_for_kernel( 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.") + 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, kx // 2)): + 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 diff --git a/test_autoarray/config/general.yaml b/test_autoarray/config/general.yaml index aeed5854b..fb830367c 100644 --- a/test_autoarray/config/general.yaml +++ b/test_autoarray/config/general.yaml @@ -1,7 +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. + 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: diff --git a/test_autoarray/dataset/imaging/test_dataset.py b/test_autoarray/dataset/imaging/test_dataset.py index 321499382..7e0d8b93a 100644 --- a/test_autoarray/dataset/imaging/test_dataset.py +++ b/test_autoarray/dataset/imaging/test_dataset.py @@ -220,7 +220,9 @@ def test__from_fits__small_datasets_env_unset__shape_unchanged(tmp_path, monkeyp assert dataset.pixel_scales == (0.08, 0.08) -def test__output_to_fits__round_trips_data_psf_noise_map_correctly(imaging_7x7, tmp_path): +def test__output_to_fits__round_trips_data_psf_noise_map_correctly( + imaging_7x7, tmp_path +): from autoarray.dataset.plot.imaging_plots import fits_imaging @@ -269,26 +271,45 @@ def test__apply_mask__data_noise_map_psf_correctly_masked( assert type(masked_imaging_7x7.psf) == aa.Convolver -def test__apply_mask__trim_to_mask__native_cut_to_mask_plus_psf_half_width__slim_unchanged( +def _imaging_11x11_from(imaging_7x7): + """The 7x7 fixture padded to 11x11 (two masked pixels on every side), so the 3x3 mask plus the + 3x3 PSF's half-width and the one-pixel buffer (7x7) is a genuine trim.""" + return aa.Imaging( + data=imaging_7x7.data.resized_from(new_shape=(11, 11), mask_pad_value=True), + noise_map=imaging_7x7.noise_map.resized_from( + new_shape=(11, 11), mask_pad_value=True + ), + psf=imaging_7x7.psf, + over_sample_size_lp=1, + ) + + +def test__apply_mask__trim_to_mask__native_cut_to_mask_plus_psf_half_width_plus_buffer__slim_unchanged( imaging_7x7, mask_2d_7x7 ): - untrimmed = imaging_7x7.apply_mask(mask=mask_2d_7x7, trim_to_mask=False) - trimmed = imaging_7x7.apply_mask(mask=mask_2d_7x7, trim_to_mask=True) - - # The 3x3 unmasked block plus the 3x3 PSF's half-width of 1 on each side. - assert trimmed.data.shape_native == (5, 5) - assert trimmed.noise_map.shape_native == (5, 5) - assert trimmed.over_sample_size_lp.shape_native == (5, 5) - assert trimmed.over_sample_size_pixelization.shape_native == (5, 5) - assert trimmed.mask.shape_native == (5, 5) + imaging_11x11 = _imaging_11x11_from(imaging_7x7) + mask_2d_11x11 = mask_2d_7x7.resized_from(new_shape=(11, 11), pad_value=True) + + untrimmed = imaging_11x11.apply_mask(mask=mask_2d_11x11, trim_to_mask=False) + trimmed = imaging_11x11.apply_mask(mask=mask_2d_11x11, trim_to_mask=True) + + # The 3x3 unmasked block plus the 3x3 PSF's half-width of 1 and the buffer of 1 on each side. + assert untrimmed.data.shape_native == (11, 11) + assert trimmed.data.shape_native == (7, 7) + assert trimmed.noise_map.shape_native == (7, 7) + assert trimmed.over_sample_size_lp.shape_native == (7, 7) + assert trimmed.over_sample_size_pixelization.shape_native == (7, 7) + assert trimmed.mask.shape_native == (7, 7) assert trimmed.psf._state.fft_shape[0] <= untrimmed.psf._state.fft_shape[0] # Only masked pixels were removed. assert (trimmed.data.slim == untrimmed.data.slim).all() assert (trimmed.noise_map.slim == untrimmed.noise_map.slim).all() - assert (trimmed.over_sample_size_lp.slim == untrimmed.over_sample_size_lp.slim).all() - assert (trimmed.mask == untrimmed.mask.resized_from(new_shape=(5, 5), pad_value=True)).all() - assert (trimmed.data.native[1:4, 1:4] == untrimmed.data.native[2:5, 2:5]).all() + assert ( + trimmed.over_sample_size_lp.slim == untrimmed.over_sample_size_lp.slim + ).all() + assert (trimmed.mask == mask_2d_7x7).all() + assert (trimmed.data.native[2:5, 2:5] == untrimmed.data.native[4:7, 4:7]).all() # The central pixel stays on the same pixel: scaled coordinates are unchanged. assert (trimmed.grids.lp.slim == untrimmed.grids.lp.slim).all() @@ -310,29 +331,56 @@ def test__apply_mask__trim_to_mask__native_cut_to_mask_plus_psf_half_width__slim image=aa.Array2D(values=np.arange(9, dtype=float), mask=trimmed.mask), blurring_image=blurring_trimmed, ) - assert np.asarray(convolved_trimmed) == pytest.approx(np.asarray(convolved_untrimmed), 1e-10) + assert np.asarray(convolved_trimmed) == pytest.approx( + np.asarray(convolved_untrimmed), 1e-10 + ) + + # The buffer is a margin of masked pixels around the blurring region, not part of it. + exact = untrimmed.trimmed_to_mask_from(buffer=0) + assert exact.data.shape_native == (5, 5) + assert ( + trimmed.grids.blurring.mask.resized_from(new_shape=(5, 5), pad_value=True) + == exact.grids.blurring.mask + ).all() def test__apply_mask__trim_to_mask__default_from_config( imaging_7x7, mask_2d_7x7, monkeypatch ): - # The test config sets [dataset] trim_to_mask: false. - masked = imaging_7x7.apply_mask(mask=mask_2d_7x7) + imaging_11x11 = _imaging_11x11_from(imaging_7x7) + mask_2d_11x11 = mask_2d_7x7.resized_from(new_shape=(11, 11), pad_value=True) - assert masked.data.shape_native == (7, 7) + # The test config sets [dataset] trim_to_mask: false (the packaged default is true). + masked = imaging_11x11.apply_mask(mask=mask_2d_11x11) + + assert masked.data.shape_native == (11, 11) monkeypatch.setitem(conf.instance["general"]["dataset"], "trim_to_mask", True) - masked = imaging_7x7.apply_mask(mask=mask_2d_7x7) + masked = imaging_11x11.apply_mask(mask=mask_2d_11x11) + + assert masked.data.shape_native == (7, 7) + - assert masked.data.shape_native == (5, 5) +def test__apply_mask__trim_to_mask__fixture_already_at_buffered_minimum__unchanged( + imaging_7x7, mask_2d_7x7 +): + # 3x3 mask + half-width 1 + buffer 1 on each side is exactly the 7x7 fixture. + trimmed = imaging_7x7.apply_mask(mask=mask_2d_7x7, trim_to_mask=True) + + assert trimmed.data.shape_native == (7, 7) def test__trimmed_to_mask_from__already_minimal__returns_self__covariance_kept(): image = aa.Array2D.ones(shape_native=(2, 2), pixel_scales=(1.0, 1.0)) noise_map = aa.Array2D.ones(shape_native=(2, 2), pixel_scales=(1.0, 1.0)) noise_covariance_matrix = np.array( - [[1.0, 0.1, 0.0, 0.0], [0.1, 2.0, 0.0, 0.0], [0.0, 0.0, 3.0, 0.2], [0.0, 0.0, 0.2, 4.0]] + [ + [1.0, 0.1, 0.0, 0.0], + [0.1, 2.0, 0.0, 0.0], + [0.0, 0.0, 3.0, 0.2], + [0.0, 0.0, 0.2, 4.0], + ] ) psf = aa.Convolver( kernel=aa.Array2D.no_mask(values=[[1.0]], pixel_scales=(1.0, 1.0)) @@ -349,12 +397,17 @@ def test__trimmed_to_mask_from__already_minimal__returns_self__covariance_kept() masked = imaging.apply_mask(mask=mask, trim_to_mask=False) - assert masked.trimmed_to_mask_from() is masked + # Without the buffer the 2x2 frame is already the minimum for a 1x1 kernel. + assert masked.trimmed_to_mask_from(buffer=0) is masked + # With the default buffer the unmasked pixels sit on the edge, so the frame is padded, + # as the convolver would pad it; the slim data and the covariance are unchanged. trimmed = imaging.apply_mask(mask=mask, trim_to_mask=True) - assert trimmed.data.shape_native == (2, 2) + assert trimmed.data.shape_native == (4, 4) + assert (trimmed.data.slim == masked.data.slim).all() assert (trimmed.noise_covariance_matrix == masked.noise_covariance_matrix).all() + assert trimmed.trimmed_to_mask_from() is trimmed def test__apply_noise_scaling__masked_pixel_data_zeroed_and_noise_set_to_noise_value( diff --git a/test_autoarray/mask/test_mask_2d_util.py b/test_autoarray/mask/test_mask_2d_util.py index c5e6cc639..0ebfa6270 100644 --- a/test_autoarray/mask/test_mask_2d_util.py +++ b/test_autoarray/mask/test_mask_2d_util.py @@ -1106,13 +1106,34 @@ def test__buffed_mask_2d_from(): ).all() -def test__trimmed_shape_for_kernel__centred_square_region__shape_is_region_plus_half_kernel(): - mask = np.full((7, 7), True) - mask[2:5, 2:5] = False +def test__trimmed_shape_for_kernel__centred_square_region__shape_is_region_plus_half_kernel_plus_buffer(): + mask = np.full((9, 9), True) + mask[3:6, 3:6] = False + + # 3x3 region + half-kernel on each side + the default one-pixel buffer on each side. + assert util.mask_2d.trimmed_shape_for_kernel(mask, kernel_shape_native=(3, 3)) == ( + 7, + 7, + ) + assert util.mask_2d.trimmed_shape_for_kernel(mask, kernel_shape_native=(1, 1)) == ( + 5, + 5, + ) + assert util.mask_2d.trimmed_shape_for_kernel(mask, kernel_shape_native=(5, 3)) == ( + 9, + 7, + ) - assert util.mask_2d.trimmed_shape_for_kernel(mask, kernel_shape_native=(3, 3)) == (5, 5) - assert util.mask_2d.trimmed_shape_for_kernel(mask, kernel_shape_native=(1, 1)) == (3, 3) - assert util.mask_2d.trimmed_shape_for_kernel(mask, kernel_shape_native=(5, 3)) == (7, 5) + # buffer=0 is the exact minimum: region + half-kernel. + assert util.mask_2d.trimmed_shape_for_kernel( + mask, kernel_shape_native=(3, 3), buffer=0 + ) == (5, 5) + assert util.mask_2d.trimmed_shape_for_kernel( + mask, kernel_shape_native=(1, 1), buffer=0 + ) == (3, 3) + assert util.mask_2d.trimmed_shape_for_kernel( + mask, kernel_shape_native=(3, 3), buffer=2 + ) == (9, 9) def test__trimmed_shape_for_kernel__parity_preserved_and_centre_kept(): @@ -1123,14 +1144,20 @@ def test__trimmed_shape_for_kernel__parity_preserved_and_centre_kept(): shape = util.mask_2d.trimmed_shape_for_kernel(mask, kernel_shape_native=(3, 3)) - assert shape == (4, 4) + # 2 unmasked + 1 half-kernel + 1 buffer on each side. + assert shape == (6, 6) - # Cutting to that shape about the centre keeps every unmasked pixel one pixel from the edge. + # Cutting to that shape about the centre keeps every unmasked pixel a half-width plus the + # buffer from the edge. trimmed = util.array_2d.resized_array_2d_from( array_2d=mask.astype(float), resized_shape=shape, pad_value=1.0 ).astype(bool) - assert trimmed.sum() == mask.sum() - (64 - 16) - assert util.mask_2d.min_false_distance_to_edge(trimmed) == (1, 1) + assert trimmed.sum() == mask.sum() - (64 - 36) + assert util.mask_2d.min_false_distance_to_edge(trimmed) == (2, 2) + + assert util.mask_2d.trimmed_shape_for_kernel( + mask, kernel_shape_native=(3, 3), buffer=0 + ) == (4, 4) # Odd frame, region offset from the centre: the frame grows symmetrically so the # central pixel stays put, rather than re-centring on the region. @@ -1139,8 +1166,11 @@ def test__trimmed_shape_for_kernel__parity_preserved_and_centre_kept(): shape = util.mask_2d.trimmed_shape_for_kernel(mask, kernel_shape_native=(3, 3)) - assert shape == (9, 9) + assert shape == (11, 11) assert shape[0] % 2 == 1 and shape[1] % 2 == 1 + assert util.mask_2d.trimmed_shape_for_kernel( + mask, kernel_shape_native=(3, 3), buffer=0 + ) == (9, 9) def test__trimmed_shape_for_kernel__region_at_edge__frame_pads_rather_than_cuts(): @@ -1149,8 +1179,12 @@ def test__trimmed_shape_for_kernel__region_at_edge__frame_pads_rather_than_cuts( shape = util.mask_2d.trimmed_shape_for_kernel(mask, kernel_shape_native=(5, 5)) - # The pixel at row 0 needs two rows above it: 2 before the centre + 2 pad, so 9 rows. - assert shape == (9, 5) + # The pixel at row 0 needs two rows above it plus the buffer: 2 before the centre + 2 pad + # + 1 buffer, so 11 rows; the columns need the half-width plus buffer, 7. + assert shape == (11, 7) + assert util.mask_2d.trimmed_shape_for_kernel( + mask, kernel_shape_native=(5, 5), buffer=0 + ) == (9, 5) def test__trimmed_shape_for_kernel__invalid_inputs_raise(): @@ -1161,4 +1195,11 @@ def test__trimmed_shape_for_kernel__invalid_inputs_raise(): util.mask_2d.trimmed_shape_for_kernel(mask, kernel_shape_native=(0, 3)) with pytest.raises(ValueError): - util.mask_2d.trimmed_shape_for_kernel(np.full((5, 5), True), kernel_shape_native=(3, 3)) + util.mask_2d.trimmed_shape_for_kernel( + mask, kernel_shape_native=(3, 3), buffer=-1 + ) + + with pytest.raises(ValueError): + util.mask_2d.trimmed_shape_for_kernel( + np.full((5, 5), True), kernel_shape_native=(3, 3) + )