diff --git a/autoarray/inversion/mesh/interpolator/rectangular_kernel.py b/autoarray/inversion/mesh/interpolator/rectangular_kernel.py index ac897728..4b5bc2e4 100644 --- a/autoarray/inversion/mesh/interpolator/rectangular_kernel.py +++ b/autoarray/inversion/mesh/interpolator/rectangular_kernel.py @@ -57,6 +57,14 @@ KERNEL_CDF_DEFAULT_BANDWIDTH: float = 1.0 KERNEL_CDF_DEFAULT_KNOTS: int = 64 +# Queries per block of the forward-transform evaluation. The exact kernel sum +# broadcasts an (M, N, 2) array; blocking the query dimension caps the peak at +# (block, N, 2) — 512 queries × 15.4k points × 2 axes × 8 B ≈ 126 MB at +# production imaging scale (previously ~60 GB unblocked at M ≈ 246k). Values +# are identical to float precision: the sum over points is unchanged, blocks +# only tile the query axis. +KERNEL_FORWARD_BLOCK: int = 512 + _SQRT2 = np.sqrt(2.0) @@ -119,10 +127,34 @@ def create_transforms_kernel( span = hi - lo h = bandwidth * span / mesh_pixels - def F_raw(q): - t = (q[:, None, :] - points[None, :, :]) / h[None, None, :] + def F_raw_block(q_block): + t = (q_block[:, None, :] - points[None, :, :]) / h[None, None, :] return xp.sum(w[None, :, None] * _norm_cdf(t, xp), axis=1) + if xp.__name__.startswith("jax"): + + def F_raw(q): + import jax + + M = q.shape[0] + n_blocks = -(-M // KERNEL_FORWARD_BLOCK) + pad = n_blocks * KERNEL_FORWARD_BLOCK - M + q_padded = xp.pad(q, ((0, pad), (0, 0))) + blocks = q_padded.reshape(n_blocks, KERNEL_FORWARD_BLOCK, 2) + out = jax.lax.map(F_raw_block, blocks) + return out.reshape(n_blocks * KERNEL_FORWARD_BLOCK, 2)[:M] + + else: + + def F_raw(q): + return np.concatenate( + [ + F_raw_block(q[i : i + KERNEL_FORWARD_BLOCK]) + for i in range(0, q.shape[0], KERNEL_FORWARD_BLOCK) + ], + axis=0, + ) + # The unit square maps onto the data bounding box exactly, matching the # linear variant's convention (its empirical-CDF knots end at the extreme # points); the kernel tails outside [lo, hi] are absorbed by the rescale. diff --git a/test_autoarray/inversion/pixelization/mesh/test_rectangular_kernel.py b/test_autoarray/inversion/pixelization/mesh/test_rectangular_kernel.py index 30dbedd4..58963cc8 100644 --- a/test_autoarray/inversion/pixelization/mesh/test_rectangular_kernel.py +++ b/test_autoarray/inversion/pixelization/mesh/test_rectangular_kernel.py @@ -338,3 +338,25 @@ def __getattr__(self, item): assert areas.shape == (36,) assert np.all(np.isfinite(areas)) assert np.all(areas > 0.0) + + +def test__create_transforms_kernel__chunked_forward_is_block_size_invariant(): + """The forward transform blocks the query axis (KERNEL_FORWARD_BLOCK) to + cap peak memory; results must be identical to evaluating queries one at a + time — exercised across a query count spanning multiple blocks and not a + multiple of the block size.""" + from autoarray.inversion.mesh.interpolator.rectangular_kernel import ( + KERNEL_FORWARD_BLOCK, + ) + + rng = np.random.default_rng(7) + data_grid = rng.standard_normal((77, 2)) + + fwd, _ = create_transforms_kernel(data_grid, mesh_pixels=8, xp=np) + + q = rng.standard_normal((KERNEL_FORWARD_BLOCK + 137, 2)) + batched = fwd(q) + row_wise = np.vstack([fwd(q[i : i + 1]) for i in range(q.shape[0])]) + + assert batched.shape == (KERNEL_FORWARD_BLOCK + 137, 2) + assert batched == pytest.approx(row_wise, abs=0.0)