Imaging.apply_mask: trim the native arrays to the mask plus a PSF half-width, so the FFT frame follows the mask rather than the stamp

[samlange04]

While fitting JWST lenses I noticed that the size of the image a mask is cut from, not the mask, sets the cost of every PSF convolution in a fit. I would like to propose a small change to Imaging.apply_mask that removes that dependence, which I have implemented and run on real data, and get your view (in particular on the packaged default) before you look at the PR.

Symptom. ConvolverState sizes its FFT frame from the native array the mask lives on (full_shape = mask_shape + kernel - 1, with mask_shape the native shape; required_shape_for_kernel only ever pads). So the convolution of every model image, and of every column of a pixelization's mapping matrix, runs on the frame of the stamp the mask was cut from, however small the mask. A 1.5" mask on a 12" JWST stamp at 0.03"/px fits in about 136 px but convolves on 400 x 400: 160k pixels, 16x the area the fit needs. GPU memory for a pixelized search scales as batch x source pixels x frame, which is what set the batch size on a 48 GB A6000 for one of my lenses. Cropping the PSF does nothing here (37, 31 and 21 px kernels all give a 400 px frame); only the stamp size matters, and asking users to cut their stamps by hand before from_fits is the wrong fix, since the stamp is also what the lens-light and contaminant masks are drawn on.

Proposal. Imaging.apply_mask(mask, trim_to_mask=None) cuts the native arrays, after masking, to the smallest centred frame that holds every unmasked pixel a PSF half-width plus a one-pixel buffer from the edge:

• mask_2d_util.trimmed_shape_for_kernel(mask_2d, kernel_shape_native, buffer=1) returns that shape. It is centred on the mask's central pixel and keeps the parity of each dimension, so the central pixel (the origin of the scaled coordinate system, the resized_array_2d_from convention) stays on the same pixel: scaled coordinates, positions and priors are unchanged. A mask whose unmasked pixels sit closer to the edge than the half-width makes the frame grow, as the convolver would pad it.
• Imaging.trimmed_to_mask_from(buffer=1) resizes data, noise_map, both over-sampling size arrays and the mask (Array2D.resized_from, mask padded True), rebuilds the ConvolverState, carries the noise_covariance_matrix over unchanged (only masked pixels are removed, so the slim ordering is the same), and returns self when the dataset is already that size. A sparse_operator is discarded with the existing warning, as apply_mask does.
• apply_mask's trim_to_mask=None reads a new config key, general.yaml [dataset] trim_to_mask.

Nothing a fit sees changes: the slim arrays, the blurring region and the convolution are identical (to FFT rounding), and the mask and data are unchanged so search identifiers are unchanged: a trimmed chain resumes an untrimmed one and vice versa. On a synthetic 200 px stamp with a 1.6" mask and an 11 px PSF a parametric FitImaging log-likelihood agrees to 1e-11 between the two.

What a user does see change: dataset.shape_native, data.native, and every native model image or residual map a fit produces are the trimmed shape, and Mask2D.circular(shape_native=dataset.shape_native, ...) built after apply_mask no longer matches a mask built before it. Project code that builds Array2Ds or image meshes on the mask variable it passed to apply_mask has to use dataset.mask instead. In my projects that was two call sites (an adapt-image over-sampling map and a Hilbert image mesh), plus two places that combine a fit's arrays with the full stamp (a source-subtracted image written for a lens-light fit, and a sensitivity-mapping simulator), which pad back with resized_from, exactly, since the central pixel is shared.

Numbers (ConvolverState frame with the real masks):

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 (160k px) -> 180 x 180 (32k px)
JWST F444W	200 px, 0.06"	2.2"	41 x 41	200 -> 116	200 x 200 (40k px) -> 160 x 160 (26k px)

Per-call GPU timings still to come; the memory saving is what let me raise the batch size.

The question for you: should the packaged default be on or off? The branch ships it on (trim_to_mask: true), on the argument that nothing a fit sees changes and every masked fit gets a smaller frame for free, with false for anyone who wants the input image's native shape kept. The case for off is that native shapes change silently for every existing user whose code builds things on the pre-mask mask (the pattern above), which is a visible behaviour change in a point release. I am happy either way; the PR has a note for the reviewer to pick, and it is a one-line config change. The unit-test config keeps it off either way so fixture shapes are untouched.

PR: PyAutoLabs/PyAutoArray#675