diff --git a/autoarray/dataset/imaging/dataset.py b/autoarray/dataset/imaging/dataset.py index 13e4deffc..492effe5f 100644 --- a/autoarray/dataset/imaging/dataset.py +++ b/autoarray/dataset/imaging/dataset.py @@ -28,10 +28,11 @@ def _validate_convolve_over_sample_size( ) -> None: """ Validate a `convolve_over_sample_size` input against its matching - `over_sample_size`: it must be a plain int >= 1, and when above 1 the matching - over sample size must be uniform and equal to it, because the values convolved at - the fine resolution must first be evaluated at exactly that resolution (adaptive - over sampling makes 2D convolution ill-defined). + `over_sample_size`: it must be a plain int >= 1, and when above 1 every entry of + the matching over sample size (int or adaptive Array2D) must be divisible by it — + the k x s coupling, whereby values evaluated at per-pixel sizes k_i * s are + partially binned to the uniform s the convolution requires (see + `over_sample_util.binned_to_convolve_size_from`). """ if isinstance(convolve_over_sample_size, bool) or not isinstance( convolve_over_sample_size, (int, np.integer) @@ -51,24 +52,12 @@ def _validate_convolve_over_sample_size( if convolve_over_sample_size == 1: return - if isinstance(over_sample_size, (int, np.integer)) and not isinstance( - over_sample_size, bool - ): - if over_sample_size == convolve_over_sample_size: - return - raise exc.DatasetException( - f"convolve_over_sample_size_{name}={convolve_over_sample_size} requires " - f"over_sample_size_{name} to be equal to it, but " - f"{over_sample_size} was input. The values convolved at the fine " - f"resolution must be evaluated at that same resolution." - ) - - if not np.all(np.array(over_sample_size) == convolve_over_sample_size): + if not np.all(np.mod(np.array(over_sample_size), convolve_over_sample_size) == 0): raise exc.DatasetException( f"convolve_over_sample_size_{name}={convolve_over_sample_size} requires " - f"a uniform over_sample_size_{name} equal to it, but a non-uniform or " - f"unequal array was input. Adaptive over sampling makes 2D PSF " - f"convolution ill-defined." + f"every over_sample_size_{name} entry to be divisible by it (the k x s " + f"coupling: evaluate at k_i * s per pixel, partially bin to s, convolve), " + f"but {over_sample_size} was input." ) diff --git a/autoarray/inversion/inversion/imaging/abstract.py b/autoarray/inversion/inversion/imaging/abstract.py index 479fbd0b5..890279403 100644 --- a/autoarray/inversion/inversion/imaging/abstract.py +++ b/autoarray/inversion/inversion/imaging/abstract.py @@ -89,7 +89,9 @@ def _mapping_matrix_for_convolution_from(self, linear_obj) -> np.ndarray: return linear_obj.mapping_matrix if isinstance(linear_obj, Mapper): - return linear_obj.mapping_matrix_over_sampled + return linear_obj.mapping_matrix_over_sampled_for( + convolve_over_sample_size=self.psf.convolve_over_sample_size + ) raise exc.InversionException( "Oversampled PSF convolution (convolve_over_sample_size > 1) currently " diff --git a/autoarray/inversion/mappers/abstract.py b/autoarray/inversion/mappers/abstract.py index 31f7825df..1a79390c5 100644 --- a/autoarray/inversion/mappers/abstract.py +++ b/autoarray/inversion/mappers/abstract.py @@ -293,16 +293,47 @@ def mapping_matrix_over_sampled(self) -> np.ndarray: The matrix has shape [total_sub_pixels, source_pixels] and is the input format of `Convolver.convolved_mapping_matrix_from` when oversampled. """ - total_sub_pixels = int(self.over_sampler.sub_total) + return self.mapping_matrix_over_sampled_for( + convolve_over_sample_size=int(self.over_sampler.sub_size[0]) + ) + + def mapping_matrix_over_sampled_for( + self, convolve_over_sample_size: int + ) -> np.ndarray: + """ + The mapping matrix at the uniform convolution resolution ``s`` under the + k x s coupling: each data pixel's (possibly adaptive) k_i * s sub-samples + scatter directly into that pixel's s x s fine cells with fraction 1/k_i^2 — + the exact linear map from source pixels to the partially-binned fine image, + built without materialising the k_i * s resolution intermediate (linearity + makes scatter-then-nothing equal bin-after-build). + + With every evaluation size equal to s (k_i = 1) the parents are the identity + and the fractions unity, reproducing `mapping_matrix_over_sampled` exactly. + """ + from autoarray.operators.over_sampling.over_sample_util import ( + convolve_bin_segment_ids_from, + ) + + s = int(convolve_over_sample_size) + sub_size = np.array(self.over_sampler.sub_size).astype("int") + + parents = convolve_bin_segment_ids_from( + sub_size=sub_size, convolve_over_sample_size=s + ) + + n_fine = self.over_sampler.mask.pixels_in_mask * s**2 + k = sub_size // s + fraction = np.repeat(1.0 / k.astype("float") ** 2, s**2) return mapper_util.mapping_matrix_from( pix_indexes_for_sub_slim_index=self.pix_indexes_for_sub_slim_index, pix_size_for_sub_slim_index=self.pix_sizes_for_sub_slim_index, pix_weights_for_sub_slim_index=self.pix_weights_for_sub_slim_index, pixels=self.pixels, - total_mask_pixels=total_sub_pixels, - slim_index_for_sub_slim_index=np.arange(total_sub_pixels), - sub_fraction=np.ones(total_sub_pixels), + total_mask_pixels=n_fine, + slim_index_for_sub_slim_index=parents, + sub_fraction=fraction, use_mixed_precision=self.settings.use_mixed_precision, xp=self._xp, ) diff --git a/autoarray/operators/over_sampling/over_sample_util.py b/autoarray/operators/over_sampling/over_sample_util.py index dbe1d365b..87d79d842 100644 --- a/autoarray/operators/over_sampling/over_sample_util.py +++ b/autoarray/operators/over_sampling/over_sample_util.py @@ -125,6 +125,111 @@ def sub_slim_to_fine_slim_from(mask_2d: Mask2D, over_sample_size: int) -> np.nda return fine_slim_index_native[rows, cols].reshape(-1) +def convolve_bin_segment_ids_from( + sub_size: np.ndarray, convolve_over_sample_size: int +) -> np.ndarray: + """ + Returns the segment ids mapping every over-sampled evaluation sample to its cell on the uniform + fine grid used by oversampled PSF convolution. + + Evaluation over sampling and convolution over sampling compose via the k x s coupling: each pixel is + evaluated on its own (possibly adaptive) sub-grid of size ``sub_size_i = k_i * s`` and the evaluated + values are partially binned down to the uniform ``s x s`` block the oversampled Convolver requires + (see ``binned_to_convolve_size_from``). + + Samples are in per-pixel sub-block order (row-major within each pixel's block). Within pixel ``p``'s + block, the sample at (row, col) belongs to fine cell ``(row // k_p) * s + (col // k_p)``, and its + global segment id is ``p * s**2 +`` that cell. + + Parameters + ---------- + sub_size + The per-pixel evaluation sub-grid sizes (slim, one entry per unmasked pixel). Every entry must be + divisible by ``convolve_over_sample_size``. + convolve_over_sample_size + The uniform over sample size ``s`` of the PSF convolution. + + Returns + ------- + An integer array of shape [total evaluation samples] of segment ids into the uniform fine grid + (``total_unmasked_pixels * s**2`` segments). + """ + from autoarray import exc + + s = int(convolve_over_sample_size) + sub_size = np.asarray(sub_size).astype("int") + + if np.any(sub_size % s != 0): + raise exc.GridException( + f"Every over_sample_size entry must be divisible by " + f"convolve_over_sample_size={s} for oversampled PSF convolution, but " + f"sizes {np.unique(sub_size[sub_size % s != 0])} are not." + ) + + segment_ids = np.empty(int(np.sum(sub_size**2)), dtype="int64") + + offset = 0 + for p, n in enumerate(sub_size): + k = n // s + rows, cols = np.divmod(np.arange(n * n), n) + segment_ids[offset : offset + n * n] = p * s**2 + (rows // k) * s + (cols // k) + offset += n * n + + return segment_ids + + +def binned_to_convolve_size_from( + values, sub_size: np.ndarray, convolve_over_sample_size: int, xp=np +): + """ + Partially bin over-sampled evaluation values down to the uniform resolution required by oversampled + PSF convolution (the k x s coupling). + + Each pixel's evaluated block (size ``k_i * s`` per side, per-pixel sub-block order) is reduced to an + ``s x s`` block by the mean of each ``k_i x k_i`` group, producing values of length + ``total_unmasked_pixels * s**2`` in the per-pixel sub-block order the oversampled Convolver expects. + Trailing dimensions (e.g. the source axis of a mapping matrix) are supported. + + When every entry of ``sub_size`` equals ``convolve_over_sample_size`` the input is already at the + convolution resolution and is returned unchanged (the existing equal-sizes behaviour, byte-identical). + + Parameters + ---------- + values + The evaluated values, one per over-sampled sample in per-pixel sub-block order, with optional + trailing dimensions. + sub_size + The per-pixel evaluation sub-grid sizes (slim). Every entry must be divisible by + ``convolve_over_sample_size``. + convolve_over_sample_size + The uniform over sample size ``s`` of the PSF convolution. + """ + s = int(convolve_over_sample_size) + sub_size = np.asarray(sub_size).astype("int") + + if np.all(sub_size == s): + return values + + segment_ids = convolve_bin_segment_ids_from( + sub_size=sub_size, convolve_over_sample_size=s + ) + n_segments = sub_size.shape[0] * s**2 + + counts = np.bincount(segment_ids, minlength=n_segments).astype("float") + trailing_shape = values.shape[1:] + counts = counts.reshape((n_segments,) + (1,) * len(trailing_shape)) + + if xp.__name__.startswith("jax"): + import jax + + sums = jax.ops.segment_sum(values, segment_ids, n_segments) + return sums / counts + + sums = np.zeros((n_segments,) + trailing_shape) + np.add.at(sums, segment_ids, np.asarray(values)) + return sums / counts + + def total_sub_pixels_2d_from(sub_size: np.ndarray) -> int: """ Returns the total number of sub-pixels in unmasked pixels in a mask. diff --git a/test_autoarray/dataset/imaging/test_dataset.py b/test_autoarray/dataset/imaging/test_dataset.py index 87417f786..d880ab275 100644 --- a/test_autoarray/dataset/imaging/test_dataset.py +++ b/test_autoarray/dataset/imaging/test_dataset.py @@ -428,26 +428,39 @@ def test__convolve_over_sample_size__validation_and_plumbing(): convolve_over_sample_size_lp=2.0, ) - # over_sample_size must equal the convolve size when the latter is above 1. + # k x s coupling: every over_sample_size entry must be divisible by the + # convolve size — a non-divisible int raises, divisible ints and adaptive + # arrays are legal. with pytest.raises(aa.exc.DatasetException): aa.Imaging( data=data, noise_map=noise_map, psf=psf, - over_sample_size_lp=4, + over_sample_size_lp=3, convolve_over_sample_size_lp=2, ) - # Adaptive (non-uniform) over sampling is incompatible with oversampled convolution. sub_size_adaptive = np.full(fill_value=2, shape=data.shape_slim) sub_size_adaptive[0] = 4 + dataset_adaptive = aa.Imaging( + data=data, + noise_map=noise_map, + psf=psf, + over_sample_size_lp=aa.Array2D(values=sub_size_adaptive, mask=data.mask), + convolve_over_sample_size_lp=2, + ) + assert dataset_adaptive.convolve_over_sample_size_lp == 2 + + sub_size_bad = np.full(fill_value=2, shape=data.shape_slim) + sub_size_bad[0] = 3 + with pytest.raises(aa.exc.DatasetException): aa.Imaging( data=data, noise_map=noise_map, psf=psf, - over_sample_size_lp=aa.Array2D(values=sub_size_adaptive, mask=data.mask), + over_sample_size_lp=aa.Array2D(values=sub_size_bad, mask=data.mask), convolve_over_sample_size_lp=2, ) diff --git a/test_autoarray/inversion/inversion/imaging/test_imaging.py b/test_autoarray/inversion/inversion/imaging/test_imaging.py index 68bc33fbe..6d3a7b984 100644 --- a/test_autoarray/inversion/inversion/imaging/test_imaging.py +++ b/test_autoarray/inversion/inversion/imaging/test_imaging.py @@ -242,3 +242,28 @@ def test__oversampled_psf__linear_func_and_preload_guards( # The kernel-native preload fast path raises too. with pytest.raises(exc.InversionException): inversion.data_linear_func_matrix_dict + + +def test__mapping_matrix_over_sampled_for__kxs__full_bin_reproduces_mapping_matrix( + rectangular_mapper_7x7_3x3, +): + # By linearity, the k x s matrix mean-binned s^2 -> 1 must equal the regular + # (sub_fraction folded) mapping matrix exactly: the fold of the folds is the + # full fold. The fixture's over sampler is uniform size 2, so s=1 puts k=2 + # (adaptive-free k x s) and s=2 puts k=1 (the 2b identity). + mapper = rectangular_mapper_7x7_3x3 + mapping_matrix = np.array(mapper.mapping_matrix) + n_pix = mapper.mask.pixels_in_mask + + # k=1 (s equals the evaluation size): identical to the 2b property. + m_s2 = np.array(mapper.mapping_matrix_over_sampled_for(convolve_over_sample_size=2)) + assert m_s2 == pytest.approx(np.array(mapper.mapping_matrix_over_sampled), abs=0) + + # k=2 (s=1): one row per image pixel — must equal mapping_matrix exactly. + m_s1 = np.array(mapper.mapping_matrix_over_sampled_for(convolve_over_sample_size=1)) + assert m_s1.shape == mapping_matrix.shape + assert m_s1 == pytest.approx(mapping_matrix, abs=1.0e-14) + + # s=2 rows mean-binned to image resolution also reproduce mapping_matrix. + binned = m_s2.reshape(n_pix, 4, mapping_matrix.shape[1]).mean(axis=1) + assert binned == pytest.approx(mapping_matrix, abs=1.0e-14) diff --git a/test_autoarray/operators/over_sample/test_over_sample_util.py b/test_autoarray/operators/over_sample/test_over_sample_util.py index 7245acdd5..e794fa368 100644 --- a/test_autoarray/operators/over_sample/test_over_sample_util.py +++ b/test_autoarray/operators/over_sample/test_over_sample_util.py @@ -336,3 +336,60 @@ def test__sub_slim_to_fine_slim_from__is_bijection_and_round_trips(): fine_slim[perm] = values_sub assert (fine_slim[perm] == values_sub).all() + + +def test__binned_to_convolve_size_from__uniform_k__equals_manual_reshape_mean(): + # 3 pixels, evaluation size 4 (k=2, s=2): each pixel's 4x4 block bins to 2x2 + # by the mean of each 2x2 group. + s, n = 2, 4 + rng = np.random.default_rng(5) + values = rng.random(3 * n * n) + + binned = util.over_sample.binned_to_convolve_size_from( + values=values, sub_size=np.full(3, n), convolve_over_sample_size=s + ) + + manual = ( + values.reshape(3, n, n) + .reshape(3, s, n // s, s, n // s) + .mean(axis=(2, 4)) + .reshape(-1) + ) + + assert binned == pytest.approx(manual, abs=1.0e-14) + + +def test__binned_to_convolve_size_from__adaptive_k__and_identity_fast_path(): + s = 2 + sub_size = np.array([4, 2]) # k = [2, 1] + values = np.arange(16 + 4, dtype="float") + + binned = util.over_sample.binned_to_convolve_size_from( + values=values, sub_size=sub_size, convolve_over_sample_size=s + ) + + manual_p0 = values[:16].reshape(2, 2, 2, 2).mean(axis=(1, 3)).reshape(-1) + manual_p1 = values[16:] + + assert binned == pytest.approx(np.concatenate([manual_p0, manual_p1]), abs=1.0e-14) + + # equal sizes: returned unchanged (identity object, byte-identical behaviour) + same = util.over_sample.binned_to_convolve_size_from( + values=values[:8], sub_size=np.full(2, s), convolve_over_sample_size=s + ) + assert same is values[:8] or (same == values[:8]).all() + + # trailing dims (mapping-matrix rows) + values_2d = np.stack([values, 2 * values], axis=1) + binned_2d = util.over_sample.binned_to_convolve_size_from( + values=values_2d, sub_size=sub_size, convolve_over_sample_size=s + ) + assert binned_2d[:, 0] == pytest.approx(binned, abs=1.0e-14) + assert binned_2d[:, 1] == pytest.approx(2 * binned, abs=1.0e-14) + + +def test__convolve_bin_segment_ids_from__divisibility_guard(): + with pytest.raises(aa.exc.GridException): + util.over_sample.convolve_bin_segment_ids_from( + sub_size=np.array([4, 3]), convolve_over_sample_size=2 + )