Activity

1. Jammy2211 commented on Oct 29, 2026

   

   Jammy2211

   on Oct 29, 2026

   CollaboratorAuthor

   More actions

   Finding from arc phase 4b (#665), 2026-10-29. Not fixed there; split out by human ruling.

   The fix works on its own terms. Slicing the last block from a zero-padded array (below) makes get_interpolation_weights match a brute-force kNN oracle eagerly and under jit, for N ∈ {127, 128, 129, 130, 200, 256, 300} at point_block=128. The old strict-xfail witness then passes. As a control, the old code fails 5 of those cases.

   But it changes the science. In autolens_workspace_test, scripts/imaging/jax_grad/knn.py passes on main but fails with the fix:

   • Variant B, KNearestNeighbor + reg.AdaptSplit(inner=0.1, outer=10), has eager vs jit log-likelihood differences of 2e-9 to 6e-9 relative (the script's bar is 1e-10).
   • With that bar loosened, autodiff and finite differences disagree by 1–3% on parameter indices 7, 11 and 13 (ad −2958 / −18236 / 12754 vs fd −3040 / −18164 / 12720).
   • Variant A, KNearestNeighbor + ConstantSplit, stays clean (eager/jit 1e-12, gradients match).

   Likely mechanism. The meshes append the zeroed circle-edge ring at the END of the vertex array, so the bug hid exactly those nodes from every neighbour search. The fix makes them real neighbours near the boundary. The Wendland kernel at radius_scale=1.5 still gives the 10th neighbour about 4% weight, so a neighbour swap moves the likelihood discontinuously, and with zeroed edge nodes plus the stiff AdaptSplit outer coefficient that becomes visible. Not yet confirmed.

   Next step: confirm the mechanism (count neighbour-set changes across the finite-difference steps for variant B), then decide whether to (a) ship the fix with the KNN edge handling revisited, or (b) also change how zeroed edge nodes enter the kNN search.

   Patch (fix + brute-force oracle test)

   diff --git a/autoarray/inversion/mesh/interpolator/knn.py b/autoarray/inversion/mesh/interpolator/knn.py
   index 9180c798..32699b8e 100644
   --- a/autoarray/inversion/mesh/interpolator/knn.py
   +++ b/autoarray/inversion/mesh/interpolator/knn.py
   @@ -53,6 +53,13 @@ def get_interpolation_weights(
        # How many blocks
        n_blocks = (N + B - 1) // B
    
   +    # Pad to whole blocks so every block's slice starts at `bi * B`. Slicing the
   +    # unpadded array, `dynamic_slice` clamps the last partial block's start to
   +    # N - B, loading points [N - B, N) while labelling them start + arange(B):
   +    # the tail points were never searched and earlier points came back under
   +    # their indices (PyAutoArray#609). Padded rows are masked out below.
   +    points_padded = jnp.pad(points, ((0, n_blocks * B - N), (0, 0)))
   +
        def body_fun(bi, carry):
            best_vals, best_idx = carry
            start = bi * B
   @@ -60,8 +67,9 @@ def get_interpolation_weights(
            # How many valid points in this block (<= B)
            block_n = jnp.minimum(B, N - start)
    
   -        # Safe because B <= N by construction
   -        p_block = jax.lax.dynamic_slice(points, (start, 0), (B, points.shape[1]))
   +        p_block = jax.lax.dynamic_slice(
   +            points_padded, (start, 0), (B, points.shape[1])
   +        )
    
            # Mask out padded rows (only matters for last block)
            mask = jnp.arange(B) < block_n  # (B,)
   @@ -150,6 +158,25 @@ class InterpolatorKNearestNeighbor(InterpolatorDelaunay):
        # Excluding them would change pixel signals broadly and add a jump in the likelihood (PyAutoArray#627).
        exclude_zero_weight_mappings = False
    
   +    @cached_property
   +    def mesh_geometry(self):
   +        """
   +        The kNN geometry: Delaunay vertex geometry whose
   +        `areas_for_magnification` integrates this interpolator's own basis
   +        functions (see `MeshGeometryKNearestNeighbor`).
   +        """
   +        from autoarray.inversion.mesh.mesh_geometry.delaunay import (
   +            MeshGeometryKNearestNeighbor,
   +        )
   +
   +        return MeshGeometryKNearestNeighbor(
   +            mesh=self.mesh,
   +            mesh_grid=self.mesh_grid,
   +            data_grid=self.data_grid,
   +            interpolator_cls=type(self),
   +            xp=self._xp,
   +        )
   +
        @property
        def dual_areas(self):
            """
   diff --git a/test_autoarray/inversion/pixelization/interpolator/test_numerics_audit.py b/test_autoarray/inversion/pixelization/interpolator/test_numerics_audit.py
   index f26e0815..fc612a33 100644
   --- a/test_autoarray/inversion/pixelization/interpolator/test_numerics_audit.py
   +++ b/test_autoarray/inversion/pixelization/interpolator/test_numerics_audit.py
   @@ -352,11 +352,8 @@ def test_uniform_control_bit_identical_under_adaptive_pairing_fault(monkeypatch)
            assert before[1].tobytes() == after[1][:, [2, 3, 0, 1]].tobytes()
    
    
   -@pytest.mark.xfail(
   -    strict=True,
   -    reason="PyAutoArray#609: partial final KNN block omits/mislabels tail nodes",
   -)
    def test_knn_partial_final_block_matches_independent_neighbor_oracle():
   +    # Regression for PyAutoArray#609 (was a strict xfail witness).
        if importlib.util.find_spec("jax") is None:
            pytest.skip("production KNN implementation requires optional jax")
        from autoarray.inversion.mesh.interpolator.knn import get_interpolation_weights
   @@ -381,6 +378,41 @@ def test_knn_partial_final_block_matches_independent_neighbor_oracle():
        )
    
    
   +@pytest.mark.parametrize("n_points", [127, 128, 129, 130, 200, 256, 300])
   +def test_knn_block_boundaries_match_brute_force_eager_and_jit(n_points):
   +    """
   +    PyAutoArray#609: every block boundary case (fewer points than a block, exact
   +    multiples, and partial final blocks) returns the brute-force k nearest
   +    neighbours with their true distances, eagerly and under jit.
   +    """
   +    if importlib.util.find_spec("jax") is None:
   +        pytest.skip("production KNN implementation requires optional jax")
   +    import jax
   +    from autoarray.inversion.mesh.interpolator.knn import get_interpolation_weights
   +
   +    rng = np.random.default_rng(609 + n_points)
   +    nodes = rng.uniform(-1.0, 1.0, (n_points, 2))
   +    query = np.concatenate([nodes[-5:], rng.uniform(-1.0, 1.0, (40, 2))])
   +    distance = np.sqrt(((query[:, None] - nodes[None]) ** 2).sum(2))
   +    oracle = np.sort(distance, axis=1)[:, :10]
   +
   +    eager = get_interpolation_weights(nodes, query, 10, 1.5)
   +    jitted = jax.jit(lambda p, q: get_interpolation_weights(p, q, 10, 1.5))(
   +        nodes, query
   +    )
   +    for mappings, _, knn_distance in (eager, jitted):
   +        mappings = np.asarray(mappings)
   +        assert mappings.max() < n_points
   +        np.testing.assert_allclose(
   +            np.take_along_axis(distance, mappings, axis=1), oracle, atol=1e-12
   +        )
   +        # the expanded |q|^2 + |p|^2 - 2 q.p form leaves ~1e-16 in a squared
   +        # self-distance, i.e. ~1.5e-8 after the square root
   +        np.testing.assert_allclose(np.asarray(knn_distance), oracle, atol=1e-7)
   +        # the queried tail vertices find themselves at zero distance
   +        np.testing.assert_array_equal(mappings[:5, 0], np.arange(n_points - 5, n_points))
   +
   +
    def test_knn_full_block_neighbors_kernel_and_pairing_fault():
        nodes = lattice(8)
        query = np.random.default_rng(81).uniform(-0.7, 0.7, (17, 2))