fix: bounds-guard the kernel gather in psf_weighted_data_from

[Jammy2211]

Fixes the numba sparse-operator likelihood returning garbage / NaN (PyAutoMind draft/bug/autoarray/numba_first_call_garbage_psf_weighted_data.md).

Root cause — not a numba caching bug

The prompt's leading suspect was numba codegen/caching (first-call-after-compile executing wrong code). It isn't.

psf_weighted_data_from gathers the weight map at [ip0_y + k0_y + kernel_shift_y, ...] with no bounds check. numba @jit() does not bounds-check array reads, so for any unmasked pixel within kernel_shape // 2 of the array edge the gather reads uninitialized heap memory instead of raising IndexError.

Proof: compiling the shipped source unchanged under boundscheck=True raises IndexError: index is out of bounds for a mask reaching the array edge, and is clean for an interior-only mask.

This is also why the inputs were verified identical between call 1 and call 2 — they were. The function reads memory outside its inputs. Both reported symptoms follow:

• Cold-cache first call runs right after numba's compilation churned the heap, so the memory next to the freshly allocated weight map holds compiler garbage (~1e299). Warm cache: no compile, benign neighbour.
• Forked workers each have a different heap layout, so whether the neighbouring memory is poisonous varies per worker and per run (2/8 corrupted in one map, 0/24 in the next).

The sibling psf_precision_value_from was already hardened against exactly this, with a comment describing the failure mode almost verbatim; psf_weighted_data_from was missed. It was the only unguarded native-array read left in the module.

The fix

The same guard as the sibling: kernel positions off the array contribute zero, matching the zero-padded numpy reference in inversion_imaging_util.psf_weighted_data_from.

Numerics impact

The guard only ever changes what an off-array position contributes, so behaviour splits into four cases:

Case	Old	New	Changed?
Read in bounds	correct	correct	no — bit-identical
Negative index wrapping onto a masked pixel	NaN, skipped → 0	0	no
Negative index wrapping onto an unmasked pixel	adds an unrelated pixel	0	yes — old was wrong
Index past the array end	reads uninitialized memory (UB)	0	yes — old was undefined

Only the two broken cases change. Verified old-vs-new by extracting both versions of the function and compiling each:

• Interior masks (the production case): bit-identical. 8 array/kernel geometries (7x7 up to 31x31, kernels 3x3–9x9), 3 random trials each, compared with np.array_equal and raw-byte equality — not approx. Zero difference.
• Edge-touching masks: old disagrees with the zero-padded numpy reference on 19–342 pixels per case; new matches it exactly (rtol=1e-9).
• Change is confined to border pixels: with an edge-touching mask, interior pixels are still bit-identical; only pixels within kernel//2 of the edge move.

End-to-end at the inversion level (Inversion over the standard 7x7 fixture, pre-fix vs post-fix), every quantity is bit-identical by SHA over raw bytes — psf_weighted_data, data_vector, curvature_matrix, reconstruction, mapped_reconstructed_operated_data, log_det_curvature_reg_matrix_term, regularization_term. The sparse-operator inversion also still agrees with the independent mapping-matrix formalism to 6.7e-16 (that formalism never touches this code path).

Edge-touching masks are constructible — nothing upstream forbids them — and with the fix such an inversion agrees with the mapping-matrix formalism to 2.2e-16.

Tests

Regression test compares the numba and numpy implementations on a mask whose unmasked pixels reach the array edge. It fails without the fix (13/25 border pixels wrong) and passes with it. Every existing test in that module masks a one-pixel border, which is why this survived.

Full test_autoarray: 1034 passed. The 3 test_transformer.py pynufft failures are pre-existing — identical on a clean tree, from the optional pynufft dep not being installed.

One caveat on test design: the out-of-bounds read is undefined behaviour, so the garbage cannot be reproduced deterministically at unit scale — on small in-process arrays the neighbouring heap is usually benign. The test therefore asserts the correctness property (agreement with the zero-padded reference), which is deterministic, rather than trying to catch a specific garbage value.

Not verified here

The parallel_scaling/pixelization_numba.py corruption counters named as the acceptance probe need the profiling workspace and its datasets, which aren't available in this environment. The case that this fix covers the forked-worker symptom is mechanistic, not measured — worth re-running that probe against this branch before closing the prompt.

Found alongside, filed separately

Both numba gathers derive their kernel shifts from the transposed kernel axes (kernel_shift_y from shape[1]). Harmless for square kernels; kernels are validated as odd, not square, so a 3x5 PSF mis-centres the gather. Filed as PyAutoMind draft/bug/autoarray/numba_kernel_shift_axes_swapped.md rather than fixed here — psf_precision_value_from has the identical swap, and correcting one without the other would make the two paths disagree about kernel orientation.

🤖 Generated with Claude Code

https://claude.ai/code/session_013unzp382r79c4BN8g4ckb2

---

Generated by Claude Code

[Jammy2211]

Post-merge follow-up: I ran the acceptance probe this PR said it couldn't run. The fix is confirmed on the real dataset, but two claims in the description above need correcting.

Symptom 1 reproduced exactly

Calling psf_weighted_data_from on the real profiling dataset (autolens_profiling, dataset/imaging/euclid, mask radius 3.5", PSF 21x21):

	max abs(psf_weighted_data)	sum
pre-fix (1c33850)	4.901e300	1.333e301
post-fix	298.312	559433.417

The original bug report's numbers were "max abs = 4.8e299 on fit #1, 2.98e02 on fit #2". The post-fix value 298.31 = 2.98e02 matches the report's correct value exactly, and the pre-fix value reproduces the uninitialized-memory scale. 1244 of the 3841 unmasked pixels gather off the array.

Correction 1 — the affected-pixel figure

The description says 35.1%. The measured figure on the real dataset is 1244/3841 = 32.4%. I derived 35.1% from a synthetic 70x70 mask rather than the pipeline's actual 71x71 geometry.

Worth stating explicitly, since I briefly believed otherwise while checking this: the mask padding does not protect this path. apply_mask emits no padding warning and leaves data.native and data.mask at (71, 71). Only derive_mask.blurring_from(allow_padding=True) pads, to (89, 89), and that padded blurring mask feeds the dense convolver — not the sparse numba path. psf_weighted_data_from reads the unpadded (71, 71) array via data.mask.derive_indexes.native_for_slim, so the mask sits flush against the array edge and the 21x21 kernel reads past it.

Correction 2 — the named acceptance probe is a weak detector

parallel_scaling/pixelization_numba.py was run at P=2, 24 evals, 2 map repeats, cold NUMBA_CACHE_DIR, both pre-fix and post-fix. Both runs reported corrupt_evals_first_map = 0 and corrupt_evals_steady_maps = [0, 0], and both had a finite warm-up likelihood.

That is not evidence of no bug. The values an out-of-bounds read returns are whatever the allocator left next to the weight map, so in that process they happened to be benign. The pre-fix warm-up likelihood still drifted from the post-fix one in the 6th decimal (5860.175003698866 vs 5860.175922117387) — the same out-of-bounds reads, landing on small values instead of huge ones. This is exactly the intermittency the original report described (2/8 corrupted in one map, 0/24 in the next).

So those counters can read zero on a run where the bug is fully present, and should not be treated as the acceptance gate. The deterministic max abs(psf_weighted_data) comparison above is the reliable probe.

Everything else in the description stands: the bit-identical interior-mask result, the four-case numerics table, and the boundscheck=True proof are unaffected. Full detail recorded in PyAutoMind#257.

---

Generated by Claude Code