Skip to content

perf(interferometer): cache curvature_matrix / data_vector on sparse + mapping inversions #581

Closed
#582
Closed
perf(interferometer): cache curvature_matrix / data_vector on sparse + mapping inversions#581
#582

Description

@Jammy2211
Jammy2211
opened on Sep 27, 2026

Overview

On the NumPy path, one figure_of_merit of InversionInterferometerSparse / InversionInterferometerSparseNumba builds the curvature matrix F twice and the data vector D four times, because data_vector, curvature_matrix and curvature_matrix_diag are plain @property. The repeat is 43-50 % of every alma / alma_high RAL CPU call (autolens_profiling #326 / PR #328). This caches them per inversion instance with autonerves.cached_property, bringing the interferometer inversions in line with the imaging ones (which already cache these under jit / vmap / grad).

Plan

  • Make data_vector, curvature_matrix, curvature_matrix_diag cached per instance on the sparse interferometer inversion, its numba subclass override, and (for parity with imaging) the dense mapping interferometer inversion.
  • No change to consumers (curvature_reg_matrix, fast_chi_squared, reconstruction); an inversion is built per fit, so no invalidation is needed.
  • Add an evaluation-count test (F == 1, D == 1 per likelihood) for Sparse, SparseNumba and Mapping, run red on unfixed source first, and a JAX jit + vmap + grad sparse-path case if none exists.
  • Verify: full test_autoarray, PyAutoGalaxy / PyAutoLens interferometer suites, the phase-1 harness on sma (counter {1,1}, log evidence bit-identical), JAX harness unchanged.
Detailed implementation plan

Affected Repositories

  • PyAutoArray (primary)
  • autolens_profiling (after-measurement, separate small PR after the library merges)

Branch Survey

Repository Current Branch Dirty?
./PyAutoArray main clean

PyAutoArray is also claimed by pointsolver-step0-gather (PR #580, PointSolver step-0 files); file sets disjoint, parallel claim human-approved 2026-09-27.

Suggested branch: feature/interferometer-sparse-cache

Implementation Steps

  1. autoarray/inversion/inversion/interferometer/sparse.py: data_vector, curvature_matrix, curvature_matrix_diag -> @cached_property (from autonerves import cached_property, as in imaging/sparse.py).
  2. autoarray/inversion/inversion/interferometer_numba/sparse.py: the curvature_matrix_diag override -> @cached_property; check base-class resolution so the cache holds on the subclass.
  3. autoarray/inversion/inversion/interferometer/mapping.py: data_vector, curvature_matrix -> @cached_property (parity with imaging/mapping.py).
  4. Tests in test_autoarray/inversion/inversion/interferometer*/: counting wrapper on the underlying util / kernel functions, assert one evaluation each per likelihood; JAX jit + vmap + grad sparse-path case.
  5. After-measurement (autolens_profiling, post-merge): re-run the RAL CPU numba rows (alma + alma_high both meshes, sma) with evaluations_per_figure_of_merit = {1, 1}.

Key Files

  • autoarray/inversion/inversion/interferometer/sparse.py — sparse interferometer inversion (F, D, diag)
  • autoarray/inversion/inversion/interferometer_numba/sparse.py — numba curvature diag override
  • autoarray/inversion/inversion/interferometer/mapping.py — dense mapping inversion
  • autoarray/inversion/inversion/abstract.py — consumers curvature_reg_matrix, reconstruction (unchanged)
  • autoarray/inversion/inversion/interferometer/abstract.py — consumer fast_chi_squared (unchanged)

Original Prompt

Click to expand starting prompt

Interferometer likelihood campaign: cache curvature_matrix / data_vector on the NumPy sparse interferometer inversions

Type: feature
Target: PyAutoArray
Repos:

  • PyAutoArray
  • autolens_profiling
    Themes:
  • interferometer
  • pixelization
  • numba
    Difficulty: small
    Autonomy: supervised
    Priority: medium
    Status: draft
    Consequence: glance
    Witness: On the NumPy path (FitInterferometer(..., xp=np)), one figure_of_merit of InversionInterferometerSparse and InversionInterferometerSparseNumba evaluates curvature_matrix_diag exactly once and data_vector exactly once (the evaluations_per_figure_of_merit counter in scripts/misc/likelihood_breakdown/interferometer_pixelized_numpy.py reads {curvature_matrix_diag: 1, data_vector: 1}, down from {2, 4}), the log evidence is bit-identical to before on the sma / alma rows, the JAX path (jit + vmap + grad) is unaffected, and the RAL CPU alma Delaunay numba full call drops from 2297 ms towards ~1.3 s.
    Review-minutes: 10
    Epic: interferometer-likelihood-campaign
    Filed: 2026-09-27

Source: autolens_profiling#326 phase 1 (PR autolens_profiling#328), RAL CPU rows
results/breakdown/interferometer/{,sma/,alma_high/}*_numba_hpc_ral_cpu_fp64.json
(arms.*.evaluations_per_figure_of_merit and arms.*.sub_rows).

Why

The library-dispatch breakdown counts property evaluations on one prior-median
figure_of_merit: on both NumPy sparse classes F is built twice and D four
times. They are plain @property on the sparse classes (PyAutoArray main 14d6336):

  • autoarray/inversion/inversion/interferometer/sparse.py:70-71 data_vector — @property
  • autoarray/inversion/inversion/interferometer/sparse.py:132-133 curvature_matrix — @property
  • autoarray/inversion/inversion/interferometer/sparse.py:190-191 curvature_matrix_diag — @property
  • autoarray/inversion/inversion/interferometer_numba/sparse.py:184-185 numba curvature_matrix_diag override — @property

Consumers: curvature_reg_matrix (inversion/abstract.py:368-389, itself a
cached_property) takes F once; fast_chi_squared
(inversion/interferometer/abstract.py:162, reads self.curvature_matrix at :181
and self.data_vector at :189) takes F a second time; reconstruction
(inversion/abstract.py:660, self.data_vector at :699/:724/:739/:755 depending on
branch) plus fast_chi_squared account for the four D evaluations.

Measured redundant cost (one extra F + three extra D) from the RAL CPU rows,
1 thread, median full call:

row arm F alone D alone full call redundant share
sma Delaunay numba 76 ms 8 ms 357 ms 101 ms 28%
sma Delaunay NumPy FFT 625 ms 9 ms 1440 ms 652 ms 45%
alma Delaunay numba 913 ms 36 ms 2297 ms 1021 ms 44%
alma rect numba 1281 ms 32 ms 3164 ms 1376 ms 43%
alma rect NumPy FFT 2251 ms 31 ms 5140 ms 2343 ms 46%
alma_high Delaunay numba 13139 ms 122 ms 27284 ms 13505 ms 49%
alma_high rect NumPy FFT 8579 ms 118 ms 18244 ms 8933 ms 49%

So roughly half of every alma / alma_high NumPy figure_of_merit is recomputation.
This changes the numba-vs-JAX-CPU comparison in the campaign decision matrix
(at alma_high numba 27.3 s vs JAX-CPU 8.3 s would become ~13.8 s).

What

Make curvature_matrix (or curvature_matrix_diag) and data_vector on the
sparse interferometer inversions cache per instance on the NumPy path (the
autonerves.cached_property used elsewhere in these classes), checking the
e819fa1 (2025-11-05) property sweep's reason for removing such caches does not
apply (a new inversion is built per fit, so no invalidation is needed), and that
JAX tracing is unaffected. Re-run the RAL CPU numba rows as the after-measurement.

Activity

  1. Jammy2211 commented on Sep 27, 2026

    @Jammy2211
    Jammy2211
    on Sep 27, 2026
    CollaboratorAuthor

    Library PR Created

    Summary: data_vector, curvature_matrix and curvature_matrix_diag on the sparse, numba and mapping interferometer inversions are now autonerves.cached_property. On one NumPy figure of merit, F and D evaluations drop from {2, 4} to {1, 1}, and the log evidence is bit-identical. Laptop sma: Delaunay numba full call 492 → 343 ms; NumPy FFT 1783 → 923 ms.

    API changes: None — internal changes only.

    Workspace impact: option (iii), none. Nothing in PyAutoGalaxy, PyAutoLens, the workspaces or the euclid pipeline assigns or patches these attributes. Evidence against this branch:

    • PyAutoGalaxy interferometer tests: 40 passed
    • PyAutoLens interferometer tests: 35 passed
    • autolens_workspace_test JAX sparse scripts pass: fit_interferometer_sparse_operator.py, rectangular_sparse.py, jax_grad/gradient.py

    The one exception is the autolens_profiling harness counter, which wraps these attributes via property(... .fget). It needs the patch below in the after-measurement PR.

    Next: human /prm 582. Then the autolens_profiling after-measurement PR: apply the harness patch, and re-run the RAL CPU numba rows (sma / alma / alma_high, both meshes) on the merged main.

    Harness patch (scripts/misc/likelihood_breakdown/interferometer_pixelized_numpy.py) 
    diff --git a/scripts/misc/likelihood_breakdown/interferometer_pixelized_numpy.py b/scripts/misc/likelihood_breakdown/interferometer_pixelized_numpy.py
index de7d80f..516ec3a 100644
--- a/scripts/misc/likelihood_breakdown/interferometer_pixelized_numpy.py
+++ b/scripts/misc/likelihood_breakdown/interferometer_pixelized_numpy.py
@@ -89,6 +89,7 @@ is the headline (``steps`` / ``total_step_by_step``), the FFT arm is under
 
 from __future__ import annotations
 
+import functools
 import gc
 import json
 import os
@@ -409,13 +410,15 @@ def run(
             (InversionInterferometerSparse, "data_vector"),
         ):
             original = cls.__dict__[name]
+            _fn = original.fget if isinstance(original, property) else original.func
 
-            def getter(self, _original=original, _name=name):
+            @functools.wraps(_fn)
+            def getter(self, _fn=_fn, _name=name):
                 counts[_name] += 1
-                return _original.fget(self)
+                return _fn(self)
 
             patched.append((cls, name, original))
-            setattr(cls, name, property(getter))
+            setattr(cls, name, type(original)(getter))
         try:
             fom = float(fit_from(median_instance, arm).figure_of_merit)
         finally:
@@ -530,8 +533,8 @@ def run(
         inv = fit_from(median_instance, arm).inversion
         mapper = inv.cls_list_from(cls=Mapper)[0]
         rows_arm = {
-            "F alone (curvature_matrix_diag)": repeat_s(lambda inv=inv: inv.curvature_matrix_diag),
-            "D alone (data_vector = Lᵀ d~)": repeat_s(lambda inv=inv: inv.data_vector),
+            "F alone (curvature_matrix_diag)": repeat_s(lambda inv=inv: (inv.__dict__.pop("curvature_matrix_diag", None), inv.curvature_matrix_diag)),
+            "D alone (data_vector = Lᵀ d~)": repeat_s(lambda inv=inv: (inv.__dict__.pop("data_vector", None), inv.data_vector)),
         }
         if arm == "numpy_fft":
             rows_arm["Sparse triplets alone (extent grid)"] = repeat_s(
  2. Jammy2211
    added a commit that references this issue on Sep 27, 2026

    Merge pull request #582 from PyAutoLabs/feature/interferometer-sparse…

    e281abf
  3. Jammy2211
    closed this as completedin #582on Sep 27, 2026

  4. Jammy2211 commented on Sep 27, 2026

    @Jammy2211
    Jammy2211
    on Sep 27, 2026
    CollaboratorAuthor

    Shipped

    Merged in #582 (merge e281abf3, 2026-09-27; 3/3 unittest jobs green). Pending release.

    What changed

    • autoarray/inversion/inversion/interferometer/sparse.py: data_vector, curvature_matrix and curvature_matrix_diag are now @cached_property, so each is evaluated once per fit instead of twice; the numba path gains a curvature_matrix_diag override.
    • autoarray/inversion/inversion/interferometer/mapping.py: same caching.
    • Evaluation-count tests went red (F=2, D=2) on unfixed source and green (1, 1) after the fix.

    Verification

    • test_autoarray 1736 passed; PyAutoGalaxy interferometer 40, PyAutoLens 35; 3 workspace_test JAX sparse scripts pass.
    • Laptop sma: Delaunay numba 492 → 343 ms, NumPy FFT 1783 → 923 ms; figure of merit unchanged; JAX jaxpr 305 → 272 eqns.

    Not in this PR (re-filed)

    • The workspace after-measurement (RAL CPU numba rows at sma/alma/alma_high, both meshes, plus the harness counter fix from the comment above) is re-filed as PyAutoMind draft/research/autolens_profiling/interferometer_sparse_cache_after_measurement.md.
    • JAX grad vs finite-difference follow-up: draft/research/autoarray/interferometer_sparse_jax_grad_vs_finite_difference_1pct.md.
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