Skip to content

perf: apply the interferometer sparse operator with rfft2/irfft2 #538

New issue
New issue
Closed
#540
Closed
perf: apply the interferometer sparse operator with rfft2/irfft2#538
#540

Description

@Jammy2211
Jammy2211
opened on Sep 8, 2026

Overview

InterferometerSparseOperator.apply_operator zero-pads a real batch to (2y, 2x) and applies the operator with the complex jnp.fft.fft2 / jnp.fft.ifft2 pair, taking the real part at the end; Khat = jnp.fft.fft2(nufft_precision_operator) is likewise the complex transform of a real array. Because both inputs are real, the real-transform pair rfft2 / irfft2 computes the same result exactly — not approximately — with roughly half the spectrum and half the butterflies. Measured 1.27–1.61× on the operator apply across every geometry and mesh tested, and the NumPy rfft2 route also beats the library's own JAX-CPU route by 1.10–1.30× at every cell.

Follow-up from autolens_profiling #226 (phase 2 of the retired epic numba-interferometer-revisit, ledger PyAutoMind/complete/archive/epics/numba_interferometer_likelihood_revisit.md). Numbers: autolens_profiling/results/notes/numba_interferometer_verdict.md and results/breakdown/interferometer/bakeoff_v2026.8.17.1.json.

Measured saving

Single thread, i9-10885H, autoarray 2026.8.17.1; same synthetic inputs, same batch size (128), scipy stack, so the comparison isolates the transform:

geometry (extent) mesh complex fft2 real rfft2 saving
sma (70²) Delaunay S=1500 1.7479 s 1.2660 s 1.38×
alma (140²) Delaunay S=1500 6.5003 s 4.3951 s 1.48×
alma_high (280²) Delaunay S=1500 17.7362 s 11.0346 s 1.61×
sma (70²) rect S=784 1.0124 s 0.7361 s 1.38×
alma (140²) rect S=784 3.3786 s 2.6557 s 1.27×
alma_high (280²) rect S=784 6.9360 s 4.8414 s 1.43×

Plan

  • Cache Khat as the real transform of the preload: jnp.fft.rfft2(nufft_precision_operator), shape (2y, x+1) instead of (2y, 2x) complex.
  • Apply the operator with rfft2 → multiply → irfft2(s=(2y, 2x)) → crop; the result is already real, so the trailing jnp.real goes away.
  • Change nothing else: the algebra is untouched, and the four consumers of apply_operator only ever call it.
  • Pin the change as exact — a new test compares apply_operator against an inline complex fft2/ifft2 reference at rtol=1e-10, on the small shared fixture and on a larger seeded case. A loosened tolerance would hide a bug.
  • Record in the docstrings why the real transform is exact here, with the measured saving and its source.
  • Confirm on the A100 before merging (see the /prm precondition below).

/prm precondition — A100 confirmation

The prompt requires: "GPU: the saving should hold or improve (half the FFT work and half the intermediate memory); confirm on the A100 before merging rather than assuming it."

There is no GPU in the development environment, so this PR ships on CPU parity plus the phase-2 measurements above. Before merge, run a short jax microbenchmark of apply_operator (280×280 preload, batch 128, old vs new) on RAL's A100 against a private merge-base checkout on PYTHONPATH (never the shared /mnt/ral/jnightin/PyAuto install), submitted through the project's hpc/sync push-submit gpu CLI, and paste the two timings into the PR before merging. The human decides at /prm whether to require it.

Detailed implementation plan

Work Classification

Library

Affected Repositories

  • PyAutoArray (primary)

Branch Survey

Repository Current Branch Dirty?
./PyAutoArray main (47a00e8c) clean

Suggested branch: feature/interferometer-apply-operator-rfft2

Worktree root: ~/Code/PyAutoLabs-wt/interferometer-apply-operator-rfft2/

Implementation Steps

  1. autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py, InterferometerSparseOperator.from_nufft_precision_operator (~line 621): Khat = jnp.fft.rfft2(nufft_precision_operator) → shape (2y, x+1). Update the field comment to Khat: "jax.Array" # (2y, x+1), rfft2 of the real preload.
  2. Same file, apply_operator (~line 688): Fhat = jnp.fft.rfft2(F_pad); Ghat = Fhat * Khat[None]; G_pad = jnp.fft.irfft2(Ghat, s=(2 * y_shape, 2 * x_shape)) (already real — drop jnp.real); crop [:, :y_shape, :x_shape].
  3. Leave the four consumers alone — curvature_matrix_diag_from, curvature_matrix_off_diag_from, curvature_matrix_func_list_from, operated_matrix_slim_from only call apply_operator.
  4. Docstrings: one paragraph on why rfft2/irfft2 is exact for a real preload and a real padded batch, plus the measured saving, citing autolens_profiling docs: refactor docstrings for autoarray/mask package #226 / the verdict note.
  5. Tests, test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py: keep the existing dense-oracle tests (_operator_dense, pytest.approx(..., 1e-8)). Add test__apply_operator__rfft2_matches_complex_fft2_reference building the operator from the shared 7×7 / K=5 fixture and a larger seeded case (e.g. 12×12 mask, K=64), comparing apply_operator(np.eye(M)) against an inline complex fft2/ifft2 reference at rtol=1e-10, atol=1e-10·max|G|.
  6. Existing test_interferometer.py comparisons against the mapping path (1e-4) remain the end-to-end gate.

Compatibility notes

  • The Khat shape change is invisible to serialisation: the class has no __getstate__ and the operator rides as pytree aux data via simulator.py:56 no_flatten.
  • Workspace .npy caches store the raw preload, not Khat, so they are unaffected.

Verification

  • New rfft2-vs-fft2 pin passes at rtol=1e-10.
  • pytest test_autoarray/inversion/inversion/interferometer test_autoarray/dataset/interferometer test_autoarray/operators -q green.
  • The full test_autoarray/inversion tree green once.

Key Files

  • autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py — InterferometerSparseOperator (~line 538): from_nufft_precision_operator (~621) and apply_operator (~688).
  • test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py — parity pins.
  • test_autoarray/inversion/inversion/interferometer/test_interferometer.py — end-to-end sparse-vs-mapping gate.

Related

Original Prompt

Click to expand starting prompt

InterferometerSparseOperator.apply_operator pads real input to complex — use rfft2/irfft2

Type: feature
Target: autoarray
Repos:

  • PyAutoArray
    Themes:
  • interferometer
  • jax-performance
  • likelihood-profiling
    Difficulty: small
    Autonomy: supervised
    Priority: high
    Epic: numba-interferometer-revisit
    Filed: 2026-09-07

Follow-up from autolens_profiling#226 (phase 2 of numba-interferometer-revisit); the
numbers are in autolens_profiling/results/notes/numba_interferometer_verdict.md and
results/breakdown/interferometer/bakeoff_v2026.8.17.1.json.

What

autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py
(InterferometerSparseOperator.apply_operator, ~689-728) zero-pads a real batch to
(2y, 2x) and applies the operator with complex jnp.fft.fft2 / jnp.fft.ifft2, taking
the real part at the end. Khat = jnp.fft.fft2(nufft_precision_operator) is likewise the
complex transform of a real array. Both halve-able: the real-transform pair
rfft2 / irfft2 computes the same result exactly (not approximately) with roughly half
the spectrum and half the butterflies.

Measured saving (single thread, i9-10885H, autoarray 2026.8.17.1)

Same synthetic inputs, same batch size (128), scipy stack, so the comparison isolates the
transform:

geometry (extent) mesh complex fft2 real rfft2 saving
sma (70²) Delaunay S=1500 1.7479 s 1.2660 s 1.38×
alma (140²) Delaunay S=1500 6.5003 s 4.3951 s 1.48×
alma_high (280²) Delaunay S=1500 17.7362 s 11.0346 s 1.61×
sma (70²) rect S=784 1.0124 s 0.7361 s 1.38×
alma (140²) rect S=784 3.3786 s 2.6557 s 1.27×
alma_high (280²) rect S=784 6.9360 s 4.8414 s 1.43×

rfft2 on the NumPy stack also beats the library's own JAX-CPU route by 1.10-1.30× at every
one of those cells, i.e. the transform change is worth more than the backend.

Scope

  • Change from_nufft_precision_operator to cache Khat = jnp.fft.rfft2(...) and
    apply_operator to rfft2 → multiply → irfft2(s=(2y, 2x)) → crop.
  • The parity gate is the existing interferometer inversion tests plus an explicit pin of
    curvature_matrix before/after at rtol=1e-10; the change is exact, so a loosened
    tolerance would be hiding a bug.
  • GPU: the saving should hold or improve (half the FFT work and half the intermediate
    memory); confirm on the A100 before merging rather than assuming it.
  • Nothing else in the interferometer path changes — the algebra is untouched.

Activity

  1. Jammy2211 commented on Sep 8, 2026

    @Jammy2211
    Jammy2211
    on Sep 8, 2026
    CollaboratorAuthor

    Library PR Created

    API changes

    None — internal only. Khat changes shape/meaning ((2y, 2x) complex fft2 →
    (2y, x + 1) rfft2) but it is an internal cache field; no public signature moves and no
    numerical result changes beyond round-off (new exactness pin: max abs error 2.8e-14 at
    rtol=1e-10).

    Downstream workspace impact — none (option iii)

    grep across autofit_workspace/, autogalaxy_workspace/, autolens_workspace/,
    autolens_workspace_test/, euclid_strong_lens_modeling_pipeline/ and HowToLens/ finds a
    single consumer, autolens_workspace_test/scripts/misc/jax_assertions/sparse_operators.py,
    which calls from_nufft_precision_operator(...) with an unchanged signature and pins at
    rtol=1e-4. Khat is not serialised (pytree aux data, no __getstate__) and the workspace
    .npy caches hold the raw preload, so nothing is invalidated. No migration, no new demo.

    Open precondition

    The prompt asks for the saving to be confirmed on the A100 before merging. No GPU was
    available in this session; the PR body carries the exact RAL recipe as a /prm precondition
    for the human to require or waive.

    Next: /prm (merge gate). Task 2 (interferometer-preload-nufft-type1, #539) is stacked on
    this branch.

  2. Jammy2211
    closed this as completedin #540on Sep 8, 2026

  3. Jammy2211 commented on Sep 8, 2026

    @Jammy2211
    Jammy2211
    on Sep 8, 2026
    CollaboratorAuthor

    Shipped

    Merged in #540 (merge commit 7a4cb700, head a67828e3).

    InterferometerSparseOperator.apply_operator now uses the exact real-transform pair
    jnp.fft.rfft2 / jnp.fft.irfft2 instead of padding the real batch to complex and
    running fft2 / ifft2 with a .real at the end. Khat is likewise built with
    rfft2, giving it shape (2y, x + 1) rather than (2y, 2x).

    • Exact, not approximate: the new path is pinned against the previous complex fft2
      route at rtol=1e-10; measured agreement is 2.8e-14.
    • Measured speedup 1.27-1.61x over the complex path (autolens_profiling#226).
    • Tests: test_inversion_interferometer_util.py extended (+109/-29 lines overall diff),
      full PyAutoArray suite green on unittest 3.12, unittest 3.13 and unittest-nojax.

    Open verification item: the A100 confirmation was not run before merge — the human
    waived it at merge time, so it stays recorded as outstanding in the completion record.

    Pending release: PyAutoArray@https://github.com/PyAutoLabs/PyAutoArray/pull/540 — merged
    is not released; the obligation is carried into the PyAutoMind completion record.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions