Skip to content

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

Open
samlange04 wants to merge 2 commits into
PyAutoLabs:mainfrom
samlange04:feature/trim-to-mask
Open

samlange04 wants to merge 2 commits into
PyAutoLabs:mainfrom
samlange04:feature/trim-to-mask

Conversation

@samlange04

Copy link
Copy Markdown
Collaborator

Discussion: https://github.com/orgs/PyAutoLabs/discussions/43

Summary

ConvolverState sizes the FFT frame from the native array the mask lives on, so a small mask on a large image pays for the empty frame in every convolution of every model image and every column of a pixelization's mapping matrix (a 1.5" mask on a 12" JWST stamp at 0.03"/px: 400 x 400 = 160k px for a fit that needs ~136 px, 16x the area; GPU memory for a pixelized search goes as batch x source pixels x frame). This PR lets Imaging.apply_mask cut the native arrays down to the mask plus a PSF half-width, after masking.

  • mask_2d_util.trimmed_shape_for_kernel(mask_2d, kernel_shape_native, buffer=1): the smallest shape, centred on the mask's central pixel and of the same parity per axis, holding every unmasked pixel at least a kernel half-width plus buffer pixels from the edge. Parity preservation keeps the central pixel on the same pixel (the resized_array_2d_from convention), so scaled coordinates, positions and priors are unchanged. A mask whose unmasked pixels sit closer to the edge than that grows the frame, as the convolver would pad.
  • Imaging.trimmed_to_mask_from(buffer=1): resizes data, noise_map, over_sample_size_lp, over_sample_size_pixelization and the mask to that shape (Array2D.resized_from, mask padded True), rebuilds the ConvolverState, carries noise_covariance_matrix over unchanged (only masked pixels are removed, so the slim ordering is the same), discards a sparse_operator with the existing warning, returns self when already that size.
  • Imaging.apply_mask(mask, trim_to_mask=None): None reads the new config key general.yaml [dataset] trim_to_mask; True/False override it.

Nothing a fit sees changes: slim arrays, blurring region and convolution are identical (to FFT rounding; a parametric FitImaging log-likelihood agrees to 1e-11 on a synthetic test), and the mask and data are unchanged so search identifiers are unchanged (a trimmed chain resumes an untrimmed one and vice versa). What changes is the native shape of every array and of the images a fit produces, so code that builds arrays or image meshes on the mask it passed to apply_mask must use dataset.mask instead.

Note for the reviewer: the packaged default

Please confirm whether you want trim_to_mask on or off by default. As submitted, autoarray/config/general.yaml ships it on (true): nothing a fit sees changes and every masked fit gets the smaller frame; a user who wants the input image's native shape kept sets false. The case for off is that native shapes change silently for every existing user whose code builds things on the pre-mask mask, a visible behaviour change in a point release. It is a one-line change to the yaml (and its comment) either way; the unit-test config keeps it false regardless so the fixture shapes are untouched. Say which and I will update the PR.

The one-pixel buffer is a margin so the frame is not cut exactly at the blurring region's edge; buffer=0 is the exact minimum and is tested.

Tests

  • test_mask_2d_util.py: shape = region + half-kernel + buffer (and buffer=0, buffer=2); even-frame parity and centre kept, min_false_distance_to_edge equal to half-width + buffer; an off-centre region grows the frame symmetrically rather than re-centring; a region at the edge pads; invalid kernel, negative buffer and an empty mask raise.
  • test_dataset.py: the 7x7 fixture padded to 11x11 trims to 7x7 with slim data, noise map, over-sampling, grids, mask origin and the convolved image unchanged; the buffer is a margin outside the blurring region; the config default is honoured (monkeypatched); a dataset already at the minimum returns self with its covariance kept, and one whose unmasked pixels touch the edge is padded.
  • Full test_autoarray/: 2092 passed, the same 4 pre-existing failures as main (matern kernel, kernel JAX gradients, two numba NaN-weight tests).

Numbers

lens stamp mask PSF native before -> after FFT frame before -> after
JWST F150W 400 px, 0.03" 1.5" 37 x 37 400 -> 138 400 x 400 -> 180 x 180
JWST F444W 200 px, 0.06" 2.2" 41 x 41 200 -> 116 200 x 200 -> 160 x 160

🤖 Generated with Claude Code

samlange04 and others added 2 commits November 2, 2026 22:31
…f-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 <noreply@anthropic.com>
…he 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 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant