Skip to content

feat: route interferometer numba vs FFT curvature by a cost model - #632

Merged
Jammy2211 merged 2 commits into
mainfrom
feature/interferometer-numba-route-gate
Oct 15, 2026
Merged

Jammy2211 merged 2 commits into
mainfrom
feature/interferometer-numba-route-gate

Conversation

@Jammy2211

Copy link
Copy Markdown
Collaborator

Summary

The NumPy single-mapper interferometer inversion chooses between the numba direct_conv curvature route (InversionInterferometerSparseNumba) and the FFT route (InversionInterferometerSparse). Until now the factory used a fixed gate of 60 mean non-zeros per source column, calibrated on the autolens_profiling#226 geometries. On the Aris 512x512 mask (91,507 px, extent 341x341, FFT grid 682x682) the measured crossover is ~130-150, so at 2,577 vertices (106 nnz/col) the gate picked the FFT route: 44.1 s/call against 34.2 s for direct_conv.

The default routing is now a two-constant cost model over quantities known at routing time (source pixels S, total mapping non-zeros, the extent rectangle from dataset.sparse_operator):

  • t_fft = c_fft · S · [rfft2 + irfft2 work of the (2ny, 2nx) padded extent]. The transform work is scipy/pocketfft's own planner estimate (cost_guess: n × sum of prime factors, with the Bluestein switch), so awkward lengths cost what they really cost: 682 = 2·11·31 and 166 = 2·83 are 2-5x dearer per element than 560 = 2^4·5·7.
  • t_dc = c_dc · nnz · ny · nx (the kernel scatters the W~ row over the whole extent once per non-zero; flat in S at fixed nnz, as measured: 34→41 s over 2,577→5,746 vertices).
  • c_fft = 7.62e-10, c_dc = 1.77e-09: least squares on log seconds over 9 paired single-thread timings (6 docs: refactor docstrings for autoarray/mask package #226 bake-off cells + the 3 Aris 512 points). The predicted faster route matches the measured one at all nine. Measured/predicted t_fft/t_dc ratio residuals are 0.68-1.47 (log-rms 0.26). Implied crossovers: ~51/58/58 nnz/col at the 70²/140²/280² bake-off extents (measured ~60 Delaunay / ~77 rect), ~166 at the Aris 512 extent (measured between 106 and 188).
  • The FFT work form was checked, not fitted, with an rfft2+irfft2 micro-benchmark at padded lengths 140-682. Seconds per model unit varied 1.8x across lengths, against 5.8x for seconds per N log2 N.

API Changes

  • Settings.interferometer_numba_nnz_per_source_max: the packaged default changes from 60.0 to null, which means the cost model. A positive value is still a hard gate, so explicit A/B settings keep their meaning, and 0 is still the kill switch.
  • New Settings(interferometer_route=None | "numba" | "fft") (config key interferometer_route: null) forces a route for benchmarking. "numba" overrides the gate and the kill switch but still requires an inversion the kernel can represent.
  • factory._use_interferometer_numba (private) gains an extent_shape kwarg. New helpers curvature_route_seconds_from, fft_route_work_from, direct_conv_work_from and _fft_length_cost are in inversion_interferometer_numba_util.
  • Behaviour: a NumPy single-mapper interferometer fit may now take the numba route at a higher nnz/col on large extents, or the FFT route at a lower one on small smooth extents. The likelihood is unchanged (both routes are pinned equal by existing parity tests); only speed changes.
    See full details below.

Test Plan

  • pytest test_autoarray -x in the worktree: 2055 passed, 4 xfailed
  • Downstream suites with the branch on PYTHONPATH: PyAutoGalaxy 1357 passed; PyAutoLens 832 passed, 1 xfailed
  • New routing tests (test_factory.py): every measured point (Aris 512 ×3, Aris 125 ×3, docs: refactor docstrings for autoarray/mask package #226 bake-off ×6) routes to the measured faster route under the default; crossover moves with extent; override semantics (hard gate, kill switch, forced route, missing extent); FFT-length cost penalises awkward lengths.
  • The new tests need no JAX. Verified passing with a fake jax package that raises ModuleNotFoundError on PYTHONPATH (60/60 in the three touched test files).
  • Smoke: pyauto-heart smoke --root <task worktree>: 165 curated scripts PASS / 0 FAIL across autofit, autogalaxy, autolens, autolens_test, euclid and howtolens (incl. autolens interferometer/modeling.py, interferometer/features/pixelization/delaunay.py, datacube/modeling_array_free.py, autolens_test interferometer/jax_likelihood/*, datacube/shared_preloads.py, misc/jax_assertions/fit_interferometer_sparse_operator.py) + 6 notebooks, no failures; autocti / autocti_test environments not built locally (GSL headers missing for arcticpy, environment-only, no interferometer code path)
  • Live route check on the original Aris datasets (table below)

Witness

On both the Aris 125x125 (5,496 px) and 512x512 (91,507 px) masks, the default routing picks the faster of numba direct_conv vs NumPy FFT curvature at every tested vertex count (1.4k-5.7k), asserted by a routing test.

test__interferometer_route__default_cost_model_picks_the_measured_faster_route asserts this. As a live check, the branch's factory was run on the original Aris datasets, eager, routing only (fit.inversion class; nothing was re-timed because the machine is shared). Timings are from the #628 prompt table and PROGRESS3/4.

mask vertices nnz/col measured FFT s measured direct_conv s faster predicted t_fft / t_dc (s) chosen (this branch) old gate 60
512x512 (341² extent) 1,459 187.6 28.3 38.4 FFT 49.9 / 56.3 InversionInterferometerSparse (FFT) ✓ FFT
512x512 2,577 106.2 44.1 34.2 numba 88.2 / 56.3 InversionInterferometerSparseNumba ✓ FFT ✗ (re-run with --gate 60 reproduces it)
512x512 5,746 47.7 97.3 41.2 numba 196.6 / 56.4 InversionInterferometerSparseNumba ✓ numba
125x125 (83² extent) 1,405 11.7 JAX FFT 1.96 (2.8 cores) 0.40 numba 5.30 / 0.20 InversionInterferometerSparseNumba ✓ numba
125x125 2,947 5.6 JAX FFT 6.84 (2.9 cores) 2.08 numba 11.1 / 0.20 InversionInterferometerSparseNumba ✓ numba
125x125 5,526 3.0 JAX FFT 12.0 (3.6 cores) 6.12 numba 20.8 / 0.20 InversionInterferometerSparseNumba ✓ numba

The predictions cover the curvature build only. The 125 measured times are whole likelihood calls, and the dense solve common to both routes dominates them at large S.

Decision taken

The FFT work term is pocketfft's cost_guess (factorisation-aware), not N log2 N. The issue's plan sketched extent·log2(extent). With that form no single constant ratio routes both the bake-off cells and the Aris 512 points correctly: fitted on all nine points it mis-routes Aris 512 at 2,577 vertices (the bug being fixed), and fitted on the 512 points alone it mis-routes two bake-off cells. The factorisation term explains the gap, because the Aris FFT lengths (682, 166) are awkward. The micro-benchmark above confirms it.

  • Rejected alternatives: the N log2 N form above, and a fixed per-mask-size table of gates. A table does not generalise to an unmeasured extent and needs re-measuring per mask.
  • Deviation from the issue's test sketch: "125x125 Delaunay 50 → numba, 80 → FFT" is tested at the 140² extent (docs: refactor docstrings for autoarray/mask package #226 alma), where the ~60 crossover was actually measured. The 60/77 crossovers come from the docs: refactor docstrings for autoarray/mask package #226 70²-280² geometries, not the Aris 125 mask. On that mask (extent 83², FFT length 166 = 2·83, pocketfft's Bluestein case) the model predicts a crossover of ~309. Only low-nnz points (≤ 11.7) were measured there, and all of them are numba.
  • Revert (restores the pre-feat: scale the interferometer numba route gate with mask size #628 behaviour exactly): set interferometer_numba_nnz_per_source_max: 60.0 in autoarray/config/general.yaml (or per run, Settings(interferometer_numba_nnz_per_source_max=60.0)), or git revert <merge sha>.

Heart RED development override

Heart is RED. The current verdict (2026-10-15 ~21:20) is RELEASE READINESS ✗ RED score 55 with reasons "release validation FAILED (stage integrate)" (red) and "PyAutoFit: open PR 8d old" (warning). Both are organism-scope and unrelated to this branch. The live human authorized the development-only override in session 3e3dea38 (AskUserQuestion answer "Yes, override for development (Recommended)") for the three PyAutoArray Aris prompts, including this one. The full reason list at launch is recorded on #628 and in PyAutoMind active.md. Permitted: commit, push, PR-open. Merge is a separate human act with all required checks green. Heart remains RED for release purposes. This branch does not claim to fix Heart.

Validation checklist (--auto run — plan was not pre-approved)

  • Effective level: supervised (header: supervised, cap: feature → supervised; Difficulty large); decide-and-flag at ship (one decision, above)
  • Plan: on the issue (feat: scale the interferometer numba route gate with mask size #628), written at filing, unmodified since
  • Gate: tests PyAutoArray 2055 passed / 4 xfailed; downstream PyAutoGalaxy 1357 passed, PyAutoLens 832 passed / 1 xfailed · smoke 165/165 scripts PASS (autocti envs not built: local GSL) · review FINDINGS (1 low, doc-only: wrong crossover figures in two docstrings) → fixed in 46befdb → re-review CLEAN (independent Opus subagent). Claim dispositions: hard-gate/kill-switch semantics → basis-cited test__interferometer_route__settings_overrides + test__factory__gate_of_zero_disables_the_numba_path; interferometer_route forcing → basis-cited test__factory__interferometer_route_forces_the_route; Witness → basis-cited 12-point routing test + independent recompute of the constants (exact match). Review note: the Aris 125 'numba faster' labels compare against JAX FFT on 2.8-3.6 cores; the model predicts single-thread NumPy FFT slower still · Heart RED under the recorded development override
  • Human: plan sound in hindsight?
  • Human: diff matches plan (no scope creep)?
  • Human: merge, amend, or reject — then log the outcome

Campaign: PyAutoPulse#62, task interferometer_large_mask_aris. Phase 2 (an autolens_profiling re-measurement record on both masks with the new default) follows after this PR merges.

Closes #628

Full API Changes (for automation & release notes)

Added

  • Settings(interferometer_route: Optional[str] = None) and property Settings.interferometer_route: None | "numba" | "fft"; invalid values raise ValueError at construction.
  • Config key general.inversion.interferometer_route: null.
  • inversion_interferometer_numba_util.curvature_route_seconds_from(source_pixels, nnz, extent_shape) -> (t_fft, t_direct_conv)
  • inversion_interferometer_numba_util.fft_route_work_from(source_pixels, extent_shape), direct_conv_work_from(nnz, extent_shape), _fft_length_cost(n)
  • Constants FFT_ROUTE_SECONDS_PER_UNIT = 7.62e-10, DIRECT_CONV_SECONDS_PER_UNIT = 1.77e-09

Changed Signature

  • factory._use_interferometer_numba(linear_obj_list, settings=None, xp=np, extent_shape=None) (private; new optional kwarg)

Changed Behaviour

  • Packaged interferometer_numba_nnz_per_source_max default 60.0 → null. Settings().interferometer_numba_nnz_per_source_max now returns None (the cost model) rather than 60.0. The property's return type is Optional[float].
  • The default interferometer route is chosen by the cost model; explicit positive gate / 0 unchanged.

Migration

  • To keep the old fixed gate: Settings(interferometer_numba_nnz_per_source_max=60.0) or the same key in a workspace general.yaml.

🤖 Generated with Claude Code

Jammy2211 and others added 2 commits October 15, 2026 21:14
The numba direct_conv route gate was a fixed 60 non-zeros per source column,
calibrated at the autolens_profiling#226 extents (70^2-280^2). On the Aris
512x512 mask (extent 341^2) it sent 2,577 vertices (106 nnz/col) to the FFT
route at 44.1 s/call against 34.2 s for direct_conv.

The default routing is now a two-constant cost model over quantities known at
routing time: t_fft = c_fft * S * [rfft2+irfft2 work of the (2ny,2nx) padded
extent, pocketfft cost_guess incl. factorisation/Bluestein], t_dc = c_dc * nnz
* ny * nx. Constants fitted by least squares on log seconds to the 9 paired
timings (6 bake-off cells of autolens_profiling#226 + 3 Aris 512 points);
predicted faster route matches the measured one at all nine, ratio residuals
0.68-1.47.

Settings: interferometer_numba_nnz_per_source_max packaged default -> null
(cost model); positive value = hard gate (unchanged); 0 = kill switch
(unchanged). New interferometer_route ("numba" | "fft" | None) forces a route
for benchmarking. _use_interferometer_numba takes extent_shape, passed from
dataset.sparse_operator (y_shape, x_shape).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@Jammy2211 Jammy2211 added pending-release PR queued for the next release build decision-taken autonomous run took one flagged decision (AUTONOMY.md decide-and-flag) labels Oct 15, 2026
@Jammy2211
Jammy2211 merged commit 17caf22 into main Oct 15, 2026
3 checks passed
@Jammy2211
Jammy2211 deleted the feature/interferometer-numba-route-gate branch October 15, 2026 20:53
@Jammy2211

Copy link
Copy Markdown
Collaborator Author

Workspace PR: PyAutoLabs/autolens_profiling#420 (phase-2 Aris routing record)

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

Labels

decision-taken autonomous run took one flagged decision (AUTONOMY.md decide-and-flag)

Projects

None yet

Development

Successfully merging this pull request may close these issues.

feat: scale the interferometer numba route gate with mask size

1 participant