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(