feat: support oversampled PSF convolution via convolve_over_sample_size

[Jammy2211]

Overview

Currently PSF convolution is restricted to the same pixel scale as the observed image. For accurate model-image generation, especially for steep light profiles (e.g. Sersic n>4) and small PSFs, evaluating the light profile and PSF on a finer grid before downsampling produces noticeably better-conditioned likelihoods. This task adds an integer convolve_over_sample_size to the Convolver and corresponding _lp / _pixelization flavours to the Imaging dataset so callers can opt into higher-resolution PSF convolution end-to-end.

Default-behaviour invariant: all new kwargs default to 1. With the defaults, the code path is byte-for-byte equivalent to the pre-change behaviour — no existing call site, test, or downstream script changes its result unless the caller explicitly sets convolve_over_sample_size_* > 1. This is the primary correctness gate for the library PR.

Plan

• Extend Convolver with convolve_over_sample_size (default 1) and treat the held PSF kernel as oversampled when > 1; bin the convolved output back to native resolution before returning.
• Add convolve_over_sample_size_lp (default 1) and convolve_over_sample_size_pixelization (default 1) to the Imaging dataset and propagate fine-resolution grids through GridsDataset.
• Reject undefined combinations (adaptive over_sample_size_* + oversampled PSF; sparse operator + oversampled PSF) with clear exceptions at construction time.
• Wire oversampled convolution through OperateImage.blurred_image_2d_from so model-image blurring uses the higher-resolution path when opted in.
• Wire the inversion mapping-matrix path (mapping.py) only; sparse.py is explicitly deferred.
• Add unit tests in PyAutoArray for the Convolver, dataset guards, and mapping path. Existing test suite must remain green at default kwargs.
• Extend autolens_workspace_test/scripts/imaging/convolution.py with a numerical cross-check and add a new convolution_oversampled.py against a closed-form Gaussian-convolution answer.
• Extend autolens_workspace/scripts/imaging/simulator.py to demonstrate end-to-end use of an oversampled PSF in a simulation.

Detailed implementation plan

Affected Repositories

• PyAutoArray (primary)
• PyAutoGalaxy
• autolens_workspace_test (after library PR)
• autolens_workspace (after library PR)

Work Classification

Both — library first, workspace follows.

Branch Survey

Repository	Current Branch	Dirty?
./PyAutoArray	main	clean
./PyAutoGalaxy	main	clean
./autolens_workspace_test	main	dirty (pre-existing dataset binary regen, unrelated)
./autolens_workspace	main	dirty (pre-existing dataset binary regen, unrelated)

Suggested branch: feature/psf-oversampling
Worktree root: ~/Code/PyAutoLabs-wt/psf-oversampling/ (created later by /start_library)

Default-behaviour invariant (correctness gate)

All new kwargs default to 1. The implementation must preserve byte-for-byte equivalence to existing behaviour at the defaults:

• Convolver(kernel=..., convolve_over_sample_size=1) returns identical results to Convolver(kernel=...) — no bin-down branch executes.
• Imaging(..., convolve_over_sample_size_lp=1, convolve_over_sample_size_pixelization=1) produces identical grids and identical inversion behaviour to today's Imaging(...) — no fine-resolution grid is constructed when both are 1.
• OperateImage.blurred_image_2d_from consumes the same grids it does today when the supplied psf.convolve_over_sample_size == 1.
• The new exception paths only trigger when convolve_over_sample_size_* > 1 is combined with adaptive over-sample or a sparse operator.

The CI signal for this gate is the pre-existing PyAutoArray and PyAutoGalaxy test suites continuing to pass without modification — any green-to-red transition in those suites means the default path has drifted.

Implementation Steps

PyAutoArray

1. autoarray/operators/convolver.py — Convolver.__init__: add convolve_over_sample_size: int = 1, validate positive int, store on self. The held kernel is interpreted as oversampled when > 1. The supplied mask in ConvolverState must be the oversampled mask (caller's responsibility). After FFT or real-space convolution in convolved_image_from / convolved_mapping_matrix_from (and their real-space and _np siblings), bin the output by N×N (sum) back to native resolution and return an Array2D paired with the native mask. The bin-down branch is gated behind if self.convolve_over_sample_size > 1 so the default path is unchanged. Add a shape sanity assertion in state_from that fires when the input mask shape is inconsistent with convolve_over_sample_size × native_mask_shape.
2. autoarray/dataset/imaging/dataset.py — Imaging.__init__ and from_fits: add convolve_over_sample_size_lp: int = 1, convolve_over_sample_size_pixelization: int = 1. Raise DatasetException for: adaptive over_sample_size_* combined with convolve_over_sample_size_* > 1; sparse operator combined with convolve_over_sample_size_* > 1. Build the fine mask from data.mask for ConvolverState setup when psf_setup_state and convolve_over_sample_size > 1.
3. autoarray/dataset/grids.py — GridsDataset: accept the two new ints; expose lp_oversampled, lp_blurring_oversampled, pixelization_oversampled grids at fine resolution only when > 1. Native grids unchanged. When both are 1, no oversampled-grid attribute is constructed.
4. autoarray/operators/over_sampling/over_sample_util.py — small helpers for Mask2D upscale-by-N and fine-to-native sum-reduce of a flat (slim) array. Reuse existing utilities where possible.
5. autoarray/inversion/inversion/imaging/abstract.py and mapping.py — audit operated_mapping_matrix_list to ensure the supplied linear_obj.mapping_matrix is on the fine grid when convolve_over_sample_size_pixelization > 1. Add an explicit shape-vs-mask assertion at the call site so a misconfiguration fails loudly. sparse.py left untouched (gated upstream by dataset construction).

PyAutoGalaxy

6. autogalaxy/operate/image.py — OperateImage.blurred_image_2d_from: when psf.convolve_over_sample_size > 1, expect grid and blurring_grid to already be the fine-resolution grids (caller responsibility). The Convolver bins down internally, so the body change is small. LightProfileOperated branch: fine-grid eval + manual bin-down (no convolution), then add to the blurred image. When psf.convolve_over_sample_size == 1, the existing code path is used unchanged.
7. autogalaxy/imaging/fit_imaging.py and PyAutoLens FitImaging — when dataset.convolve_over_sample_size_lp > 1, pass dataset.grids.lp_oversampled to the blurring path instead of dataset.grids.lp. Default-1 path uses today's dataset.grids.lp unchanged.

Library tests

8. PyAutoArray/test_autoarray/operators/test_convolver.py — regression test with convolve_over_sample_size=1 matches existing behaviour byte-for-byte; new test with =2 agrees with a brute-force scipy.signal.convolve2d + reshape sum reference for delta-function inputs.
9. PyAutoArray/test_autoarray/dataset/imaging/test_dataset.py — kwarg acceptance, default-1 invariance (existing tests untouched), all four exception paths.
10. PyAutoArray/test_autoarray/inversion/inversion/imaging/test_mapping.py — small inversion with fixed mapper at convolve_over_sample_size_pixelization=2, asserting operated_mapping_matrix shape and values for one column against a brute-force reference.

Workspace (after library PR ships)

11. autolens_workspace_test/scripts/imaging/convolution.py — extend with a section comparing native vs. oversampled PSF convolution + bin-down for the same physical scenario, asserting agreement to ~1e-3 absolute on the data-frame image.
12. autolens_workspace_test/scripts/imaging/convolution_oversampled.py — new. Build an analytic Gaussian-PSF + Gaussian-light scenario whose downsampled blurred image has a closed-form expected value; assert numerical equality.
13. autolens_workspace/scripts/imaging/simulator.py — extend with a section demonstrating oversampled PSF simulation end-to-end.

Key Files

• PyAutoArray/autoarray/operators/convolver.py — Convolver oversample kwarg + bin-down
• PyAutoArray/autoarray/dataset/imaging/dataset.py — Imaging kwargs + guards
• PyAutoArray/autoarray/dataset/grids.py — fine-resolution grid attrs
• PyAutoArray/autoarray/inversion/inversion/imaging/mapping.py + abstract.py — assertion-only audit
• PyAutoGalaxy/autogalaxy/operate/image.py — fine-grid blurring path
• autolens_workspace_test/scripts/imaging/convolution.py — extend
• autolens_workspace_test/scripts/imaging/convolution_oversampled.py — new
• autolens_workspace/scripts/imaging/simulator.py — extend

Testing approach

• Library pytest must remain green at default convolve_over_sample_size=1 (regression — primary correctness gate).
• New unit tests cover Convolver bin-down, dataset construction guards, and inversion mapping path.
• Workspace numerical test (convolution_oversampled.py) provides the analytic ground-truth check.

Original Prompt

Click to expand starting prompt

A point spreadh function is used to blur images via 2d convolution.

This blurring occurs predominantly in the package @PyAutoArray/autoarray/operators/convolver.py.

The source code currently requires PSF blurring to occur at the same resolution (pixel scale) as the
image, meaning the PSF is always the same resolution as the image.

However, for modeling, convolution can be performed at a higher resolution than the image, which allows for more accurate
blurring and modeling of the image. This requires us to have an oversampled PSF, which is a PSF that has a higher
resolution than the image.

For modeling, where images are generated PSF blurring happens in @PyAutoGalaxy/autogalaxy/operate/image.py.

Modeling can always evaluate images using a hgiher resolition grid, blurring them with the PSF at high
resolution and then downsample to the observed image resolution. Oversampling is implemented in
@PyAutoArray/autoarray/operators/over_sampling.

Note that over sampling often uses an adaptive sub-szie, which means that 2D covnolution with a PSF is not
well defined. for now, we will assume adaptive over sampling is not used.

I want us to be able to append the Convolver class with a convolve_over_sample_size integer, which specifies the over sample size of the PSF.
This will allow us to perform convolution at a higher resolution than the image, which will improve the accuracy of the blurring and modeling of the image.
For example, if convolve_over_sample_size is 2, then the PSF will be oversampled by a factor of 2, meaning it will have a resolution that is 2 times higher than the image.

This, in turn, means out imaging object @PyAutoArray/autoarray/dataset/imaging/dataset.py will need to be
extended to include the convolve_over_sample_size_lp and convolve_over_sample_size_pixelization attributes, which will
specify the over sample size of the PSF for the lensing and pixelization operations, respectively.

class Imaging(AbstractDataset):
def init(
self,
data: Array2D,
noise_map: Optional[Array2D] = None,
psf: Optional[Convolver] = None,
psf_setup_state: bool = False,
noise_covariance_matrix: Optional[np.ndarray] = None,
over_sample_size_lp: Union[int, Array2D] = 4,
over_sample_size_pixelization: Union[int, Array2D] = 4,
use_normalized_psf: Optional[bool] = True,
check_noise_map: bool = True,
sparse_operator: Optional[ImagingSparseOperator] = None,
):

Also,read through the @PyAutoArray/autoarray/inversion/inversion/imaging package, and parents, to see
how PSF convolution enters this. I think we can get it to work in PyAutoArray/autoarray/inversion/inversion/imaging/mapping.py,
and will leave work in PyAutoArray/autoarray/inversion/inversion/imaging/sparse.py to future work.

This is a complex task, therefore I think we should extend @autolens_workspace_test/scripts/imaging/convolution.py
with a numerical test.

We should then build on this test in a separate test file using a simple over sampled PSF, to get a numerical
result we can test the source code against.

@autolens_workspace/scripts/imaging/simulator.py is a good example we can build on to show how to use
over sampled PSFs in a real simulation. We can extend this script to show how to use over sampled PSFs in a real simulation.

Come up with a plan to implement over sampled PSFs.