fix: Preserve KNN neighbors in a partial final point block

[Jammy2211]

Overview

The numerical audit (#603) found a block-tail indexing defect in the shared KNN neighbor search. With 130 nodes and the default 128-node block, queries at the last two nodes select earlier nodes instead of themselves. This affects KNearestNeighbor and KNNBarycentric.

Plan

• Preserve the audit's regression and reproduce the tail-node failure.
• Align final-block coordinates, validity masks and emitted point indices.
• Compare neighbor selection with a brute-force oracle across block boundaries and JAX modes.
• Remove the strict expected-failure marker only once the repair passes.

Detailed implementation plan

Affected repositories: PyAutoArray (primary), autolens_workspace_test. Suggested branch: feature/knn-partial-point-block.

1. Fix autoarray/inversion/mesh/interpolator/knn.py:get_interpolation_weights. A final lax.dynamic_slice start is clamped backward but the mask/indices assume the requested start. Prefer static padding to a whole number of blocks plus a true global-index validity mask, or a consistent clamped-start scheme with no duplicate candidate corruption. Preserve bounded memory and existing full-block outputs.
2. Use deterministic unique-distance datasets at N below/equal/above B, including129/130/257 at default128. Compare eager/JIT selected indices and distances against brute force. Cover both public interpolation variants and verify mappings refer to the coordinates used to calculate weights.
3. Promote the strict expected-failure coverage in audit#603 after the fix; run focused NumPy tests and JAX workspace checks, full library suite and applicable smoke. No production repair belongs in the audit PR.

Classification: Both. Filing only: no worktree or implementation authorized by this bug registration. Queue behind the currently active audit; reuse the normal start-dev/start-library process when scheduled. User explicitly requested that audit-discovered defects be filed separately.

Original Prompt

Audit finding and authorized follow-up

Fix KNN neighbor search for a partial final point block

Type: bug
Target: PyAutoArray
Repos:

• PyAutoArray
• autolens_workspace_test
  Difficulty: small
  Autonomy: supervised
  Priority: high
  Status: formalised
  Consequence: glance
  Witness: At N=130 with point_block=128, exact tail-node queries select indices128/129 with zero distance; eager/JIT results match a brute-force oracle at block boundaries for both KNN variants.
  Review-minutes: 5
  Unattended: ready

Fix KNN neighbor search for a partial final point block

Type: bug
Target: autoarray
Repos:

• PyAutoArray
• autolens_workspace_test
  Difficulty: small
  Autonomy: supervised
  Priority: high
  Witness: At N=130 with point_block=128, exact tail-node queries select indices128/129 with zero distance; eager/JIT results match a brute-force oracle at block boundaries for both KNN variants.
  Review-minutes: 5

Found while executing PyAutoArray#603, the approved final numerics audit. User authorization from the original audit prompt: "Any bug found is filed separately; this task is the audit, not the fixes."

In autoarray/inversion/mesh/interpolator/knn.py:get_interpolation_weights, lax.dynamic_slice clamps a requested final-block start when N is not divisible by point_block. For N=130 and default point_block=128, the last requested start128 actually slices from2. The code still labels candidate rows with start128 and masks only2 rows, excluding or mislabelling the true final nodes.

Minimal witness on untouched main: points (0,0) through (129,0), query exact points128 and129, k_neighbors3. Expected first neighbors128 and129 with zero distance. Actual selected indices are [127,126,125] for both, with distances[1,2,3] and[2,3,4]. Both InterpolatorKNearestNeighbor and InterpolatorKNNBarycentric share the selection path.

Fix in a separate task by padding point blocks and masking padded rows, or otherwise keeping slice start, point coordinates and emitted indices consistent. Test N below/equal/above block size and nonmultiples (including129/130/257), compare eager/JIT neighbor indices/distances with a brute-force oracle on unique-distance data, cover both interpolation variants, and promote the audit's strict expected-failure regression to ordinary passing coverage. Preserve memory-bounded block behavior and full-block results. No production repair is part of audit#603.