fix(interferometer): correct sparse curvature for Pmax > 1 (Delaunay)

[Jammy2211]

Summary

Fixes the silent ~34% Frobenius mismatch between the interferometer sparse-operator curvature matrix and the mapping path on Delaunay (Pmax = 3 barycentric) meshes (issue #314). Replaces the defensive NotImplementedError guard from #315 with a real math fix. As a bonus, the same bug was producing a quieter ~0.4% gap on Pmax = 1 (rectangular) — both reduce to the same indexing mistake and are fixed together.

API Changes

The interferometer sparse operator now exposes the same API surface as the imaging operator: apply_operator(F) + curvature_matrix_diag_from(rows, cols, vals, *, S). The old curvature_matrix_via_sparse_operator_from(pix_indexes, pix_weights, pix_pixels, fft_index) shape-array signature is gone. The NotImplementedError guard on InversionInterferometerSparse.curvature_matrix_diag (added in #315) is also gone. New Mask2D.extent_index_for_masked_pixel cached_property exposes the slim → unmasked-extent-flat index — the interferometer counterpart of fft_index_for_masked_pixel.

See full details below.

Test Plan

• pytest test_autoarray/inversion/inversion/interferometer/test_interferometer.py — the new test__curvature_matrix__interferometer_sparse_operator__delaunay__identical_to_mapping asserts sparse-vs-mapping parity at rtol=1e-4 on a Delaunay mapper.
• pytest test_autoarray/inversion/ (full inversion suite) — 166 passed locally.
• pytest test_autoarray/mask/ (mask suite) — 105 passed locally, covers the new extent_index_for_masked_pixel via the integration parity test.
• Follow-up workspace PR will pick up the one Pmax=1 call-site in autolens_workspace_test/scripts/jax_assertions/sparse_operators.py and the now-closed-gap literal in rectangular_sparse.py.

Root cause

The interferometer's W~ operator lives on the unmasked-extent rectangular grid: M = shape_native_masked_pixels[0] * shape_native_masked_pixels[1]. The previous implementation used mask.fft_index_for_masked_pixel as row indices for the (M, batch_size) scatter buffer — but fft_index_for_masked_pixel returns native-flat indices in [0, native_y * native_x). For any real-space mask where the unmasked extent is smaller than the native shape (i.e. any non-trivial circular mask), most rows fell out of bounds; JAX's default .at[i, j].add(v) silently drops out-of-bounds writes, so most contributions vanished. The result was very-nearly-zero curvature for Delaunay (34% Frobenius from the correct value) and a quieter mostly-correct-but-not-quite curvature for Pmax = 1 (the documented 0.4% gap).

The fix adds Mask2D.extent_index_for_masked_pixel (slim → extent-flat), and InversionInterferometerSparse.curvature_matrix_diag plumbs it through mapper_util.sparse_triplets_from(...) to produce extent-flat row indices that land in the operator's actual (M, B) scatter buffer. The sparse-operator class itself is then a near-clone of ImagingSparseOperator (same scatter / apply_operator / segment_sum / width-mask / S_pad shape), differing only in the actual W~ apply — FFT-conv with Khat for interferometer, vs H^T N^{-1} H for imaging.

Scripts Changed

None in this library PR. Workspace follow-up coming in a separate /ship_workspace PR for autolens_workspace_test:

• scripts/jax_assertions/sparse_operators.py — switch the Pmax=1 call to the new curvature_matrix_diag_from(rows, cols, vals, S) API.
• scripts/jax_likelihood_functions/interferometer/rectangular_sparse.py — Path A literal -3152.03184792 → -3164.286252 (now matches Path B and Path C exactly); the Path B docstring's "~0.4% numerical reformulation" comment is replaced with the actual behaviour (mathematically exact, agrees to ~1e-13).
• notebooks/jax_likelihood_functions/interferometer/rectangular_sparse.ipynb — corresponding literal update.

Full API Changes (for automation & release notes)

Added

• Mask2D.extent_index_for_masked_pixel: np.ndarray — cached_property. 1D int32 array of shape (N_unmasked,). Maps slim masked-pixel index to the flat row-major index on the unmasked-extent rectangular FFT grid (extent_y * extent_x). Use this for the interferometer sparse path. The existing fft_index_for_masked_pixel (native-flat) remains correct for the imaging path.
• InterferometerSparseOperator.apply_operator(Fbatch_flat) -> jax.Array — FFT-conv method (Re(IFFT(FFT(F_pad) * Khat))[:y, :x]). Previously inlined inside curvature_matrix_via_sparse_operator_from.
• InterferometerSparseOperator.col_offsets: jax.Array — (batch_size,) int32. Width-mask helper, populated by from_nufft_precision_operator.
• InterferometerSparseOperator.curvature_matrix_diag_from(rows, cols, vals, *, S) -> jax.Array — replaces the old method; see Renamed / Signature Change.

Removed

• InterferometerSparseOperator.curvature_matrix_via_sparse_operator_from(pix_indexes_for_sub_slim_index, pix_weights_for_sub_slim_index, pix_pixels, fft_index_for_masked_pixel) — see Renamed / Signature Change.
• The NotImplementedError guard on InversionInterferometerSparse.curvature_matrix_diag (and the from autoarray.inversion.mesh.mesh.delaunay import Delaunay import that supported it) — guard is no longer needed; the math is now correct.

Renamed / Signature Change

• InterferometerSparseOperator.curvature_matrix_via_sparse_operator_from(pix_indexes_for_sub_slim_index, pix_weights_for_sub_slim_index, pix_pixels, fft_index_for_masked_pixel) →
  InterferometerSparseOperator.curvature_matrix_diag_from(rows, cols, vals, *, S). Same name and call shape as ImagingSparseOperator.curvature_matrix_diag_from. Triplets must use extent-flat row indexing; callers should produce them via mapper_util.sparse_triplets_from(..., fft_index_for_masked_pixel=mask.extent_index_for_masked_pixel, return_rows_slim=False).

Changed Behaviour

• InversionInterferometerSparse.curvature_matrix_diag now returns the correct curvature for arbitrary Pmax (including Delaunay barycentric). Previously raised NotImplementedError on Delaunay (PR fix(interferometer-sparse): guard against Delaunay mappers (issue #314) #315) and silently produced an indexed-out-of-bounds result with a quieter ~0.4% gap on Pmax = 1.

Migration

If you held a direct reference to InterferometerSparseOperator.curvature_matrix_via_sparse_operator_from, replace the call with:

```python

Before:

op.curvature_matrix_via_sparse_operator_from(
pix_indexes_for_sub_slim_index=mapper.pix_indexes_for_sub_slim_index,
pix_weights_for_sub_slim_index=mapper.pix_weights_for_sub_slim_index,
pix_pixels=mapper.params,
fft_index_for_masked_pixel=mask.fft_index_for_masked_pixel,
)

After:

from autoarray.inversion.mappers import mapper_util
rows, cols, vals = mapper_util.sparse_triplets_from(
pix_indexes_for_sub=mapper.pix_indexes_for_sub_slim_index,
pix_weights_for_sub=mapper.pix_weights_for_sub_slim_index,
slim_index_for_sub=mapper.slim_index_for_sub_slim_index,
fft_index_for_masked_pixel=mask.extent_index_for_masked_pixel, # extent-flat, not native-flat
sub_fraction_slim=mapper.over_sampler.sub_fraction.array,
return_rows_slim=False,
)
op.curvature_matrix_diag_from(rows, cols, vals, S=mapper.params)
```

This pattern mirrors ImagingSparseOperator.curvature_matrix_diag_from's call site — see inversion/inversion/imaging/sparse.py:288 for the imaging counterpart.

Closes #314. Supersedes the defensive guard from #315.

🤖 Generated with Claude Code

[Jammy2211]

Workspace PR: PyAutoLabs/autolens_workspace_test#98