Merge pull request #378 from PyAutoLabs/feature/kernel-forward-chunking

refactor: chunk the kernel-CDF forward transform (512-query blocks)

Author: Jammy2211

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.

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)