Skip to content

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

Merged
Jammy2211 merged 1 commit into
mainfrom
feature/fix-interferometer-sparse-curvature
May 16, 2026
Merged

Jammy2211 merged 1 commit into
mainfrom
feature/fix-interferometer-sparse-curvature

Conversation

@Jammy2211

Copy link
Copy Markdown
Collaborator

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

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

Replaces the InterferometerSparseOperator curvature path with one that
mirrors ImagingSparseOperator.curvature_matrix_diag_from(rows, cols,
vals, *, S) and uses extent-flat row indices for the W~ operator's
(2*y_ext, 2*x_ext) grid. The previous path used native-flat indices
from fft_index_for_masked_pixel, so most JAX scatter writes silently
fell out-of-bounds for any real_space_mask with extent < native shape.
This fixed both the 34% Frobenius gap on Delaunay (Pmax=3 barycentric,
issue #314) and the previously-documented ~0.4% Pmax=1 "numerical
reformulation" gap — both reduce to the same indexing bug.

Removes the defensive NotImplementedError guard added in #315, and
converts the raise-test into a sparse-vs-mapping parity assertion at
rtol=1e-4.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
@Jammy2211

Copy link
Copy Markdown
Collaborator Author

Workspace PR: PyAutoLabs/autolens_workspace_test#98

@Jammy2211
Jammy2211 merged commit aa418a5 into main May 16, 2026
6 checks passed
@Jammy2211
Jammy2211 deleted the feature/fix-interferometer-sparse-curvature branch May 16, 2026 08:31
@Jammy2211 Jammy2211 removed the pending-release PR queued for the next release build label Sep 4, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Interferometer + Delaunay sparse-operator curvature math is wrong (silent ~34% error)

1 participant