From 05f250645ff97c915d78e7fd7808fb7cfc107cfd Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 21 Aug 2026 01:19:51 +0000 Subject: [PATCH] fix: bounds-guard the kernel gather in psf_weighted_data_from MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `psf_weighted_data_from` indexed the weight map directly at `ip0_y + k0_y + kernel_shift_y`, with no check that the position lands on the array. numba `@jit()` does not bounds-check array reads, so for any unmasked pixel within `kernel_shape // 2` of the array edge the gather silently read uninitialized heap memory instead of raising IndexError — producing contributions of order 1e299 that overflowed downstream into a NaN figure of merit. A negative index is unsafe in the same way: it wraps to the opposite edge and convolves in unrelated pixels. Because the values read are whatever the allocator happened to leave next to the weight map, the corruption is heap-state dependent, which explains both reported symptoms: deterministic garbage on the first call after a cold-cache compile (numba's compilation churns the heap), and intermittent corruption in forked multiprocessing workers, where each worker has a different heap layout. The sibling `psf_precision_value_from` was already hardened against exactly this; `psf_weighted_data_from` was missed. This applies the same guard, so kernel positions off the array contribute zero — matching the zero-padded numpy reference in `inversion_imaging_util.psf_weighted_data_from`. The regression test compares the numba and numpy implementations on a mask whose unmasked pixels reach the array edge. Every existing test in that module masks a one-pixel border, so none of them exercised this path. Co-Authored-By: Claude Claude-Session: https://claude.ai/code/session_013unzp382r79c4BN8g4ckb2 --- .../inversion_imaging_numba_util.py | 25 +++++++++-- .../imaging/test_inversion_imaging_util.py | 45 +++++++++++++++++++ 2 files changed, 67 insertions(+), 3 deletions(-) diff --git a/autoarray/inversion/inversion/imaging_numba/inversion_imaging_numba_util.py b/autoarray/inversion/inversion/imaging_numba/inversion_imaging_numba_util.py index ccffcadef..3fa785048 100644 --- a/autoarray/inversion/inversion/imaging_numba/inversion_imaging_numba_util.py +++ b/autoarray/inversion/inversion/imaging_numba/inversion_imaging_numba_util.py @@ -61,9 +61,28 @@ def psf_weighted_data_from( for k0_y in range(kernel_native.shape[0]): for k0_x in range(kernel_native.shape[1]): - weight_value = weight_map_native[ - ip0_y + k0_y + kernel_shift_y, ip0_x + k0_x + kernel_shift_x - ] + iy = ip0_y + k0_y + kernel_shift_y + ix = ip0_x + k0_x + kernel_shift_x + + # numba @jit() does not bounds-check array reads. Without this + # guard, a kernel position that lands off the weight-map array + # (e.g. a mask pixel within `kernel_shift` of the array edge) + # silently reads uninitialized memory, producing astronomical + # or non-finite contributions that poison `psf_weighted_data` + # and the data vector computed from it. A negative index is + # just as unsafe: it wraps to the opposite edge and quietly + # convolves in unrelated pixels. Positions off the array + # contribute zero, matching the zero-padded reference in + # `inversion_imaging_util.psf_weighted_data_from`. + if ( + iy < 0 + or iy >= weight_map_native.shape[0] + or ix < 0 + or ix >= weight_map_native.shape[1] + ): + continue + + weight_value = weight_map_native[iy, ix] if not np.isnan(weight_value): value += kernel_native[k0_y, k0_x] * weight_value diff --git a/test_autoarray/inversion/inversion/imaging/test_inversion_imaging_util.py b/test_autoarray/inversion/inversion/imaging/test_inversion_imaging_util.py index d98d135ba..a5256a40d 100644 --- a/test_autoarray/inversion/inversion/imaging/test_inversion_imaging_util.py +++ b/test_autoarray/inversion/inversion/imaging/test_inversion_imaging_util.py @@ -36,6 +36,51 @@ def test__psf_weighted_noise_imaging_from(): ) +def test__psf_weighted_data_from__unmasked_pixels_on_array_edge(): + """ + Regression test: an unmasked pixel within `kernel_shape // 2` of the array + edge drives the kernel off the weight map. + + numba `@jit()` does not bounds-check array reads, so those positions + silently returned uninitialized memory (values of order 1e299) rather than + raising, poisoning `psf_weighted_data` and the data vector built from it. + Because the values read depend on whatever the allocator left next to the + weight map, the corruption was heap-state dependent: deterministic on the + first call after a cold-cache compile, and intermittent in forked + multiprocessing workers. + + The zero-padded numpy implementation is the reference — kernel positions + off the array contribute zero. Every other test in this module masks a + one-pixel border, so none of them exercise this path. + """ + + image = np.arange(1.0, 26.0).reshape(5, 5) + noise_map = np.ones((5, 5)) + + kernel = np.array([[0.0, 1.0, 2.0], [3.0, 4.0, 1.0], [2.0, 0.0, 1.0]]) + + # Every pixel unmasked, so the border pixels push the kernel off the array. + native_index_for_slim_index = np.array( + [[y, x] for y in range(5) for x in range(5)] + ) + + psf_weighted_data = aa.util.inversion_imaging_numba.psf_weighted_data_from( + image_native=image, + noise_map_native=noise_map, + kernel_native=kernel, + native_index_for_slim_index=native_index_for_slim_index, + ) + + psf_weighted_data_numpy = aa.util.inversion_imaging.psf_weighted_data_from( + weight_map_native=image / noise_map**2.0, + kernel_native=kernel, + native_index_for_slim_index=native_index_for_slim_index, + ) + + assert np.all(np.isfinite(psf_weighted_data)) + assert psf_weighted_data == pytest.approx(psf_weighted_data_numpy, 1.0e-8) + + def test__psf_weighted_data_from(): mask = aa.Mask2D(