fix: derive numba PSF kernel y/x shifts from the correct kernel axes

[Jammy2211]

Overview

Both numba PSF gathers in autoarray/inversion/inversion/imaging_numba/inversion_imaging_numba_util.py derive their kernel half-widths from the transposed kernel axes — the y shift is taken from kernel_native.shape[1] (the x axis) and the x shift from shape[0] (the y axis). For a square kernel the two are equal and nothing is observably wrong, which is why the existing tests pass. For a non-square odd PSF the gather is mis-centred along both axes, silently sampling the weight map / noise map off-centre.

This is reachable: kernels are validated as odd per axis, never as square (exc.KernelException("Convolver Convolver must be odd") in operators/convolver.py:268, and the same parity-only check in structures/grids/uniform_2d.py:1153). With the bounds guard from #456 now in place the mis-centred reads are clipped rather than reading uninitialized memory, so this is a silent wrong-answer bug, not a crash.

Split out from the OOB-gather fix (#456) under the one-prompt-one-task rule: separate defect, separate blast radius.

Confirmed reproducing on main (9e47505). With a 3x5 kernel on a fully-unmasked 5x5 image, the numba psf_weighted_data_from diverges from its zero-padded numpy twin by max|diff| = 1420.0; the numba output is the numpy output shifted, exactly as a transposed half-width predicts. The square-kernel control matches to 0.0.

Plan

• Swap the two kernel_shift_* right-hand sides in psf_weighted_data_from so the y shift comes from the kernel's y axis and the x shift from its x axis.
• Make the identical swap in psf_precision_value_from — together, so the psf_weighted_data and psf_precision_operator paths cannot disagree on kernel orientation.
• Correct the _reference_value helper in the existing precision-operator edge test, which currently replicates the same transposition and so is not the independent reference its comment claims.
• Extend the numba-vs-numpy equivalence test to a non-square odd kernel, and add non-square coverage to the precision-operator edge test.
• Add an explicit orientation guard test so a future transposition fails loudly rather than hiding behind a square kernel.

Detailed implementation plan

Work Classification

Library (PyAutoArray only). No workspace impact: this is an internal numba util and behaviour is bit-identical for square kernels, which is every PSF in the workspaces today.

Affected Repositories

• PyAutoArray (primary)

Branch Survey

Repository	Current Branch	Dirty?
./PyAutoArray	main	clean

Suggested branch: claude/numba-kernel-shift-axes-j5wo2p

Worktree root: not applicable — web-github session, the session clone is the working tree.

Orientation audit (repo-wide)

A scan for every half-width derivation in the library (shape[0] // 2 / shape[1] // 2) finds exactly four transposed lines, all in the numba util — 48/49 and 313/314. Everything else is correct and stays untouched:

• inversion_imaging_numba_util.py:788 convolve_with_kernel_native — ky, kx = psf_kernel.shape; cy, cx = ky // 2, kx // 2. Correct, and an in-file reference for the intended idiom.
• inversion_imaging_numba_util.py:212 psf_precision_operator_sparse_from — kernel_overlap_size is a symmetric product of both axes, orientation-agnostic.
• inversion_imaging_util.py:44-45 (numpy twin) — Ky, Kx = kernel_native.shape; ph, pw = Ky // 2, Kx // 2. Correct; the reference implementation.
• inversion_imaging_util.py:409/479/747 (JAX precision operator) — Ky, Kx = psf.shape, cy, cx = Ky // 2, Kx // 2. Correct.
• structures/arrays/array_2d_util.py:311-319 — resize centring, already y-from-shape[0]. Correct.

Implementation Steps

1. autoarray/inversion/inversion/imaging_numba/inversion_imaging_numba_util.py, psf_weighted_data_from lines 48-49 — swap to kernel_shift_y = -(kernel_native.shape[0] // 2) / kernel_shift_x = -(kernel_native.shape[1] // 2).
2. Same file, psf_precision_value_from lines 313-314 — the same swap. No other edit in that function: both the bounds guard (341-347) and the offset early-exit (319-324) consume these variables and become correct once they hold the right values.
3. test_autoarray/inversion/inversion/imaging/test_inversion_imaging_util.py, test__psf_precision_operator_sparse_from__edge_pixels — fix _reference_value's two shift lines so the reference is genuinely independent of the code under test, then cover a non-square kernel.
4. Same file, test__psf_weighted_data_from__unmasked_pixels_on_array_edge — add a non-square odd kernel case (asymmetric values, so a transposition cannot hide behind symmetry).
5. Add a dedicated orientation test asserting the tall-vs-wide kernel behaviour that only a correct axis mapping can produce.

Key Files

• autoarray/inversion/inversion/imaging_numba/inversion_imaging_numba_util.py — the two defect sites.
• autoarray/inversion/inversion/imaging/inversion_imaging_util.py — the correct zero-padded numpy reference the tests compare against.
• test_autoarray/inversion/inversion/imaging/test_inversion_imaging_util.py — the equivalence and edge tests.

Testing approach

New cases must fail on unpatched source (verified by reverting the source hunk) and pass with the fix. Then the full test_autoarray suite for regressions.

Original Prompt

Click to expand starting prompt

Numba PSF gathers derive the y/x kernel shifts from the wrong kernel axes

Type: bug
Target: autoarray
Repos:

• @PyAutoArray
  Difficulty: low
  Autonomy: supervised
  Priority: medium
  Status: formalised
  Filed: 2026-08-21 (backfilled from git)

Found 2026-08-21 while fixing
draft/bug/autoarray/numba_first_call_garbage_psf_weighted_data.md (the
out-of-bounds gather in psf_weighted_data_from). Split out under the
one-prompt-one-task rule: separate defect, separate blast radius.

Symptom

Both numba PSF gathers in
autoarray/inversion/inversion/imaging_numba/inversion_imaging_numba_util.py
compute their kernel half-widths from the transposed kernel axes:

kernel_shift_y = -(kernel_native.shape[1] // 2)   # shape[1] is x
kernel_shift_x = -(kernel_native.shape[0] // 2)   # shape[0] is y

at psf_weighted_data_from (line ~48) and psf_precision_value_from
(line ~294). The y shift must come from shape[0] and the x shift from
shape[1].

The zero-padded numpy twin
(imaging/inversion_imaging_util.py:psf_weighted_data_from) gets it right and
is the reference:

Ky, Kx = kernel_native.shape
ph, pw = Ky // 2, Kx // 2

Reachability

Harmless for square kernels (shape[0] == shape[1]), which is the common
case and why no test catches it. It is not unreachable: kernels are
validated as odd in each axis, not square — exc.KernelException("Convolver Convolver must be odd") in operators/convolver.py:268 and
structures/grids/uniform_2d.py:1153 check parity only. A non-square odd PSF
(e.g. 3x5) therefore mis-centres the gather, sampling the weight map / noise
map off-centre along both axes.

With the bounds guard now in place the mis-centred reads are clipped rather
than reading uninitialized memory, so this is a silent wrong-answer bug, not
a crash or a garbage-value bug.

Fix

Swap the two right-hand sides in both functions. Fix them together — they
must agree on kernel orientation, and correcting only one would make the
psf_weighted_data and psf_precision_operator paths disagree.

Acceptance

Extend the numba-vs-numpy equivalence test added by the OOB fix
(test_autoarray/inversion/inversion/imaging/test_inversion_imaging_util.py:: test__psf_weighted_data_from__unmasked_pixels_on_array_edge) to a non-square
odd kernel (e.g. 3x5). It passes today only because that test uses a square
kernel; with a non-square kernel the two implementations diverge.

[Jammy2211]

Shipped — PR #487, merged as 68a69ae.

Both gather sites fixed together (psf_weighted_data_from, psf_precision_value_from). Full test_autoarray: 1154 passed, 55 skipped; CI green on all three legs (3.12, 3.13, nojax).

Two findings beyond the issue as filed, both recorded in the PR:

1. In psf_precision_value_from the same variables drive the pair-overlap early-exit, so the transposition additionally dropped genuinely overlapping image-pixel pairs on one axis while admitting non-overlapping ones on the other. The swap corrects it; no separate change needed.
2. The existing _reference_value helper in test__psf_precision_operator_sparse_from__edge_pixels replicated the same transposition, so it was not the independent reference its comment claimed. Corrected — otherwise the new non-square case would have asserted the bug.

Characterisation, for the record. The defect is a pure rigid translation of the sampling window by (ky//2 − kx//2, kx//2 − ky//2), not a rotation or transpose — the kernel indexing kernel_native[k0_y, k0_x] was always correctly paired with the y/x steps; only the centring constant came from the wrong axis. Confirmed empirically: buggy(y,x) == fixed(y+Δy, x+Δx) to 0.000e+00, while transpose/rot90/rot180 hypotheses all leave residuals of order 10¹. Consequently square kernels are exactly inert at any size and either parity — the expression contains no parity term, so squareness is what matters, not oddness.

Follow-up check, no action needed. The "no kernel flip" (correlation) convention shared by the numba, numpy and JAX paths was verified correct against a dense ground-truth operator built from the definition of a PSF convolution, using a non-square asymmetric 3×5 PSF:

path	claims	vs dense truth
numpy psf_weighted_data_from	Hᵀ N⁻¹ d	1.78e-15
numba psf_weighted_data_from	Hᵀ N⁻¹ d	1.78e-15
numba psf_precision_operator_from	Hᵀ N⁻¹ H	3.55e-15
JAX ImagingSparseOperator.apply_operator	Hᵀ N⁻¹ H	7.11e-15

Correlation is right because the adjoint of a convolution is a correlation; the flipped alternative misses truth by 7.086e+00. JAX's Khat_r / Khat_flip_r pair is the correctly-assembled adjoint pair. Control: swapping JAX's crop offsets (cy,cx) → (cx,cy) shifts the result by 9.97e-01, so that path is genuinely orientation-sensitive and simply gets it right.

---

Generated by Claude Code