Skip to content

Commit ec06f69

Browse files
Jammy2211claude
andcommitted
feat: route interferometer numba vs FFT curvature by a cost model (#628)
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>
1 parent 0377eea commit ec06f69

6 files changed

Lines changed: 579 additions & 63 deletions

File tree

‎autoarray/config/general.yaml‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,8 @@ inversion:
1818
reconstruction_vmax_factor: 0.5 # Plots of an Inversion's reconstruction use the reconstructed data's bright value multiplied by this factor.
1919
log_det_method: cholesky # How the Bayesian-evidence log-determinant terms are computed. "cholesky" (default) is the historical 2*sum(log(diag(cholesky(M)))); "slogdet" uses logabsdet of slogdet(M), which is identical where M is positive-definite but finite (not NaN) where the Cholesky fails, for gradient-based searches (opt-in, non-default; does not change the default evidence). Under "slogdet" the kernel regularization schemes (Matern/Gaussian/Exponential) also compute the regularization log-det analytically from a Cholesky of their covariance instead of factorizing the formed inverse. See PyAutoArray#391.
2020
regularization_term_method: matmul # How the Bayesian-evidence regularization term s^T H s is computed. "matmul" (default) is the historical s @ (H @ s) against the explicitly formed regularization matrix; "cho_solve" evaluates coefficient * s^T C^-1 s for the kernel schemes (Matern/Gaussian/Exponential/MaternAdapt) via one Cholesky solve of their covariance C, avoiding the explicit inverse whose round-off is amplified by cond(C) (~1e9 on clustered traced mesh vertices). Opt-in, non-default; does not change the default evidence. Schemes with no such factorization fall back to the formed matrix.
21-
interferometer_numba_nnz_per_source_max: 60.0 # Geometry gate for the numba `direct_conv` interferometer curvature path, in mean non-zeros per source column (mapper.pix_sizes_for_sub_slim_index.sum() / mapper.params). At or below this the factory routes a NumPy (xp=np) single-mapper interferometer inversion to InversionInterferometerSparseNumba, which is 2-7x faster than the JAX/FFT route while the mapping operator stays sparse; above it the FFT route wins and is used. Measured crossovers are ~60 (Delaunay) and ~77 (rectangular) on the reference CPU (autolens_profiling#226 verdict section 2) -- this is a machine-dependent constant, so re-measure before tuning. Set 0 to disable the numba path entirely.
21+
interferometer_numba_nnz_per_source_max: null # How a NumPy (xp=np) single-mapper interferometer inversion chooses between the numba `direct_conv` curvature path (InversionInterferometerSparseNumba) and the FFT one (InversionInterferometerSparse). null (default) = cost model: predicts both routes' single-thread time from source pixels S, total mapping non-zeros and the extent rectangle (FFT ~ S x rfft2+irfft2 cost of the padded (2ny,2nx) extent, direct_conv ~ nnz x ny x nx) and picks the faster; calibrated on autolens_profiling#226 and the Aris 512x512 mask (PyAutoArray#628). A positive value is a hard gate in mean non-zeros per source column (mapper.pix_sizes_for_sub_slim_index.sum() / mapper.params): numba at or below it, FFT above (the pre-#628 behaviour, default was 60.0). Set 0 to disable the numba path entirely.
22+
interferometer_route: null # Force the interferometer curvature route for benchmarking: null (default) routes by interferometer_numba_nnz_per_source_max; "fft" always uses the FFT route; "numba" uses the numba direct_conv route whenever it can represent the inversion (NumPy, one mapper, no linear functions, no over-sampling), overriding the gate and the 0 kill switch.
2223
numba:
2324
use_numba: true
2425
cache: true

‎autoarray/inversion/inversion/factory.py‎

Lines changed: 59 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import numpy as np
2-
from typing import List, Union
2+
from typing import List, Optional, Tuple, Union
33

44
from autoarray.dataset.imaging.dataset import Imaging
55
from autoarray.dataset.interferometer.dataset import Interferometer
@@ -214,6 +214,10 @@ def inversion_interferometer_from(
214214
linear_obj_list=linear_obj_list,
215215
settings=settings,
216216
xp=xp,
217+
extent_shape=(
218+
dataset.sparse_operator.y_shape,
219+
dataset.sparse_operator.x_shape,
220+
),
217221
):
218222
return InversionInterferometerSparseNumba(
219223
dataset=dataset,
@@ -243,6 +247,7 @@ def _use_interferometer_numba(
243247
linear_obj_list: List[LinearObj],
244248
settings: Settings = None,
245249
xp=np,
250+
extent_shape: Optional[Tuple[int, int]] = None,
246251
) -> bool:
247252
"""
248253
Whether an interferometer inversion is routed to the numba `direct_conv` curvature
@@ -258,12 +263,25 @@ def _use_interferometer_numba(
258263
The conditions are, in order of cost to evaluate:
259264
260265
- `xp is np` -- the kernel is numba, with no JAX path.
261-
- `settings.interferometer_numba_nnz_per_source_max > 0` -- `0` is the kill switch.
266+
- `settings.interferometer_route` -- `"fft"` never routes to numba; `"numba"` forces it
267+
whenever the kernel can represent the inversion (the conditions below that are not
268+
about speed), bypassing the geometry gate and the kill switch. For benchmarking.
269+
- `settings.interferometer_numba_nnz_per_source_max` is not `0` -- `0` is the kill
270+
switch.
262271
- exactly one `Mapper` and no `AbstractLinearObjFuncList` -- the kernel builds a single
263272
mapper-mapper block and has no off-diagonal or function blocks.
264273
- no over-sampling (`sub_fraction == 1`) -- the kernel uses the mapper's weights as-is.
265-
- the mapper's mean non-zeros per source column is at or below the gate -- above it the
266-
FFT route is faster (see `Settings.interferometer_numba_nnz_per_source_max`).
274+
- the geometry gate, which picks the faster route:
275+
276+
- `interferometer_numba_nnz_per_source_max` positive -- a hard gate: numba at or
277+
below that many mean non-zeros per source column (the pre-#628 behaviour, kept
278+
so explicit A/B settings are honoured).
279+
- `interferometer_numba_nnz_per_source_max` `None` (the packaged default) -- the
280+
cost model `inversion_interferometer_numba_util.curvature_route_seconds_from`,
281+
which predicts both routes' time from the source pixels, the mapping operator's
282+
non-zeros and the extent rectangle `extent_shape`, and routes to numba when
283+
`direct_conv` is predicted no slower. Without an `extent_shape` the model cannot
284+
be evaluated and the FFT route is used.
267285
- `import numba` succeeds.
268286
269287
NaN mapping weights (an overflowing or degenerate `DelaunayNN` Sibson row) are
@@ -277,19 +295,29 @@ def _use_interferometer_numba(
277295
linear_obj_list
278296
The linear objects reconstructing the data.
279297
settings
280-
The inversion settings, whose `interferometer_numba_nnz_per_source_max` is the
281-
geometry gate. `None` uses the packaged defaults.
298+
The inversion settings, whose `interferometer_route` and
299+
`interferometer_numba_nnz_per_source_max` select the route. `None` uses the
300+
packaged defaults.
282301
xp
283302
The array module the inversion runs on.
303+
extent_shape
304+
The `(ny, nx)` extent rectangle the `W~` operator lives on, i.e. the dataset's
305+
`sparse_operator.y_shape, sparse_operator.x_shape`
306+
(`= mask.shape_native_masked_pixels`). Needed only by the default cost model.
284307
"""
285308
if xp is not np:
286309
return False
287310

288311
settings = settings if settings is not None else Settings()
289312

313+
route = settings.interferometer_route
314+
315+
if route == "fft":
316+
return False
317+
290318
nnz_max = settings.interferometer_numba_nnz_per_source_max
291319

292-
if nnz_max is None or nnz_max <= 0:
320+
if route is None and nnz_max is not None and nnz_max <= 0:
293321
return False
294322

295323
if any(
@@ -312,12 +340,31 @@ def _use_interferometer_numba(
312340
if not np.all(sub_fraction == 1.0):
313341
return False
314342

315-
nnz_per_source_column = (
316-
inversion_interferometer_numba_util.nnz_per_source_column_from(mapper=mapper)
317-
)
343+
if route is None:
344+
if nnz_max is not None:
345+
nnz_per_source_column = (
346+
inversion_interferometer_numba_util.nnz_per_source_column_from(
347+
mapper=mapper
348+
)
349+
)
318350

319-
if nnz_per_source_column > nnz_max:
320-
return False
351+
if nnz_per_source_column > nnz_max:
352+
return False
353+
354+
else:
355+
if extent_shape is None:
356+
return False
357+
358+
t_fft, t_direct_conv = (
359+
inversion_interferometer_numba_util.curvature_route_seconds_from(
360+
source_pixels=int(mapper.params),
361+
nnz=int(np.asarray(mapper.pix_sizes_for_sub_slim_index).sum()),
362+
extent_shape=extent_shape,
363+
)
364+
)
365+
366+
if t_direct_conv > t_fft:
367+
return False
321368

322369
try:
323370
import numba # noqa: F401

‎autoarray/inversion/inversion/interferometer_numba/inversion_interferometer_numba_util.py‎

Lines changed: 213 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -33,9 +33,9 @@ def curvature_direct_conv(
3333
Cost `O(nnz·M + S·nnz)` with `M = ny·nx`, versus the historic pair loop's `O(N² P²)`:
3434
the convolution replaces the `N²` pixel-pair space with the `M`-cell extent rectangle,
3535
which is the whole point of the extent-grid form. It beats the FFT route (the JAX and
36-
NumPy `InterferometerSparseOperator` paths) while the source columns stay sparse --
37-
measured crossover ~60 non-zeros per source column on Delaunay meshes, ~77 on
38-
rectangular ones (autolens_profiling issue #226).
36+
NumPy `InterferometerSparseOperator` paths) while the source columns stay sparse; where
37+
the two cross depends on the extent as well as the column density, and the factory
38+
decides with :func:`curvature_route_seconds_from` (PyAutoArray#628).
3939
4040
The result is a complete, symmetric `[pix_pixels, pix_pixels]` matrix: the kernel
4141
loops the full column x row space rather than halving on symmetry, so no mirroring
@@ -284,11 +284,13 @@ def nnz_per_source_column_from(mapper) -> float:
284284
"""
285285
The mean number of non-zero mapping weights per source column of `A`.
286286
287-
This is the geometry the `direct_conv` kernel's cost scales with, and therefore the
288-
quantity the factory gates on: the kernel's convolution step costs `O(nnz · M)`,
289-
while the FFT route's cost is set by the number of source *columns* rather than their
290-
density, so the two cross at a roughly fixed non-zeros-per-column value (~60 on
291-
Delaunay meshes, ~77 on rectangular ones).
287+
The kernel's convolution step costs `O(nnz · M)` while the FFT route's cost is set by
288+
the number of source *columns* rather than their density, so at a fixed extent the two
289+
cross at a fixed non-zeros-per-column value (~60 on Delaunay meshes, ~77 on
290+
rectangular ones at the autolens_profiling#226 geometries). That value moves with the
291+
extent, so the default routing uses :func:`curvature_route_seconds_from`; this ratio is
292+
what an explicit `Settings.interferometer_numba_nnz_per_source_max` hard gate compares
293+
against.
292294
293295
`nnz = pix_sizes_for_sub_slim_index.sum()` is the total number of valid triplets and
294296
`mapper.params` the number of source pixels, so this is simply their ratio.
@@ -300,3 +302,206 @@ def nnz_per_source_column_from(mapper) -> float:
300302
return 0.0
301303

302304
return nnz / params
305+
306+
307+
def _fft_length_cost(n: int) -> float:
308+
"""
309+
The work of one length-`n` complex FFT, as `scipy.fft`'s pocketfft backend estimates it.
310+
311+
This mirrors pocketfft's own planner heuristic (`util::cost_guess`): `n` times the sum
312+
of its prime factors, with a 1.1 penalty on factors above 5, so a smooth length such as
313+
`560 = 2^4 * 5 * 7` costs `560 * 20` while an awkward one such as
314+
`682 = 2 * 11 * 31` costs `682 * 48.2` -- 2.4x more per element although
315+
`log2(682)` and `log2(560)` differ by 3 %. That factorisation penalty is what makes the
316+
FFT route relatively slower on the Aris masks (extents 83 and 341, so FFT lengths
317+
`166 = 2 * 83` and `682`) than `n log n` predicts, and is the reason the route gate is
318+
a cost model rather than a fixed non-zeros-per-column constant.
319+
320+
For a length with a large prime factor (`p^2 > n`, `n >= 50`) pocketfft switches to
321+
Bluestein's algorithm when that is cheaper -- two FFTs of the smallest 11-smooth length
322+
`>= 2n - 1`, with pocketfft's 1.5 fudge factor -- and so does this estimate.
323+
"""
324+
n = int(n)
325+
326+
def cost_guess(m: int) -> float:
327+
m_in = m
328+
result = 0.0
329+
330+
while m % 4 == 0:
331+
result += 2.0
332+
m //= 4
333+
334+
while m % 2 == 0:
335+
result += 2.0
336+
m //= 2
337+
338+
x = 3
339+
while x * x <= m:
340+
while m % x == 0:
341+
result += float(x) if x <= 5 else 1.1 * float(x)
342+
m //= x
343+
x += 2
344+
345+
if m > 1:
346+
result += float(m) if m <= 5 else 1.1 * float(m)
347+
348+
return result * float(m_in)
349+
350+
def largest_prime_factor(m: int) -> int:
351+
largest = 1
352+
x = 2
353+
while x * x <= m:
354+
while m % x == 0:
355+
largest = x
356+
m //= x
357+
x += 1
358+
return m if m > 1 else largest
359+
360+
def good_size(m: int) -> int:
361+
"""The smallest `2^a 3^b 5^c 7^d 11^e >= m` (pocketfft's `good_size_cmplx`)."""
362+
best = 2 * m
363+
f11 = 1
364+
while f11 < best:
365+
f7 = f11
366+
while f7 < best:
367+
f5 = f7
368+
while f5 < best:
369+
f3 = f5
370+
while f3 < best:
371+
f2 = f3
372+
while f2 < m:
373+
f2 *= 2
374+
best = min(best, f2)
375+
f3 *= 3
376+
f5 *= 5
377+
f7 *= 7
378+
f11 *= 11
379+
return best
380+
381+
direct = cost_guess(n)
382+
383+
if n < 50:
384+
return direct
385+
386+
p = largest_prime_factor(n)
387+
388+
if p * p <= n:
389+
return direct
390+
391+
bluestein = 1.5 * 2.0 * cost_guess(good_size(2 * n - 1))
392+
393+
return min(direct, bluestein)
394+
395+
396+
def fft_route_work_from(source_pixels: int, extent_shape) -> float:
397+
"""
398+
The FFT route's curvature work, in pocketfft cost units.
399+
400+
`InterferometerSparseOperator.curvature_matrix_diag_from` (NumPy branch) applies `W~`
401+
to every source column once: an `rfft2` and an `irfft2` of the `(2ny, 2nx)` padded
402+
extent, in batches of `batch_size` columns. Batching changes the memory traffic, not
403+
the number of transforms, so the work is `S` times one forward plus one inverse 2D
404+
real transform. A 2D real transform of shape `(a, b)` is `a` real transforms of length
405+
`b` (half the cost of a complex one) followed by `b // 2 + 1` complex transforms of
406+
length `a`.
407+
408+
The sparse `Aᵀ G` projection (`S · nnz` multiply-adds) is left out: the `direct_conv`
409+
kernel does the identical gather, so it cancels in the comparison.
410+
"""
411+
ny, nx = int(extent_shape[0]), int(extent_shape[1])
412+
a, b = 2 * ny, 2 * nx
413+
414+
per_transform = 0.5 * a * _fft_length_cost(b) + (b // 2 + 1) * _fft_length_cost(a)
415+
416+
return float(source_pixels) * 2.0 * per_transform
417+
418+
419+
def direct_conv_work_from(nnz: int, extent_shape) -> float:
420+
"""
421+
The `direct_conv` kernel's curvature work, in multiply-adds.
422+
423+
:func:`curvature_direct_conv` scatters a shifted copy of the `W~` kernel over the whole
424+
`(ny, nx)` extent rectangle for every non-zero of the mapping operator, so the
425+
convolution step is `nnz · ny · nx` multiply-adds whatever the number of source pixels
426+
(its `S · nnz` gather is shared with the FFT route and left out of both).
427+
"""
428+
ny, nx = int(extent_shape[0]), int(extent_shape[1])
429+
430+
return float(nnz) * float(ny) * float(nx)
431+
432+
433+
# Seconds per unit of `fft_route_work_from` and of `direct_conv_work_from`, fitted by least
434+
# squares on log(seconds) to 18 single-thread timings (`curvature_route_seconds_from`):
435+
#
436+
# - Aris 512x512 mask (91,507 masked pixels, extent 341x341, Delaunay, nnz 273,661), whole
437+
# NumPy likelihood call (the curvature build is 96-98 % of it), PyAutoArray#628:
438+
# FFT 28.3 / 44.1 / 97.3 s and direct_conv 38.4 / 34.2 / 41.2 s at S = 1,459 / 2,577 / 5,746.
439+
# - autolens_profiling#226 synthetic bake-off (`bakeoff_v2026.8.17.1.json`, curvature build
440+
# only, `rfft2_numpy` = the same algorithm as the NumPy FFT route): sma / alma / alma_high
441+
# extents 70^2 / 140^2 / 280^2, Delaunay S = 1,500 (P = 3) and rectangular S = 784 (P = 4).
442+
#
443+
# Only the ratio of the two constants (0.430) decides a route. Measured / predicted
444+
# FFT-to-direct_conv time ratio over the nine paired points lies in 0.68-1.47 (log-rms 0.26),
445+
# and the predicted faster route matches the measured one at all nine:
446+
#
447+
# point nnz/col measured t_fft/t_dc predicted
448+
# bake-off sma Delaunay 7.7 5.93 6.62
449+
# bake-off sma rect 19.7 2.83 2.59
450+
# bake-off alma Delaunay 30.8 1.54 1.88
451+
# bake-off alma rect 78.5 0.99 0.74
452+
# bake-off alma_high Del. 123.1 0.63 0.47
453+
# bake-off alma_high rect 314.1 0.27 0.18
454+
# Aris 512, S = 1,459 187.6 0.74 0.89
455+
# Aris 512, S = 2,577 106.2 1.29 1.57
456+
# Aris 512, S = 5,746 47.6 2.36 3.49
457+
#
458+
# The absolute residuals are larger (log-rms 0.55 FFT, 0.52 direct_conv): the two campaigns ran
459+
# on the same i9-10885H ~2.5x apart in clock state (the bake-off documents the drift), which
460+
# moves both routes together and cancels in the ratio. Implied crossovers in non-zeros per
461+
# source column: ~51 / 58 / 58 at the 70^2 / 140^2 / 280^2 bake-off extents (measured ~60 Delaunay,
462+
# ~77 rectangular), ~166 at the Aris 512 extent 341^2 (measured between 106 and 188) and ~309 at
463+
# the Aris 125 extent 83^2, whose FFT length 166 = 2 * 83 is pocketfft's Bluestein case.
464+
#
465+
# The FFT work form was checked (not fitted) with a single-thread `rfft2` + `irfft2`
466+
# micro-benchmark at padded lengths 140 / 166 / 240 / 280 / 560 / 662 / 682: seconds per
467+
# `fft_route_work_from` unit varied 1.8x across them, against 5.8x for seconds per `N log2 N`
468+
# -- the awkward lengths 166, 662 and 682 cost 2-5x more per element than smooth ones.
469+
FFT_ROUTE_SECONDS_PER_UNIT = 7.62e-10
470+
DIRECT_CONV_SECONDS_PER_UNIT = 1.77e-09
471+
472+
473+
def curvature_route_seconds_from(source_pixels: int, nnz: int, extent_shape):
474+
"""
475+
Predicted single-thread seconds of the FFT route and of the `direct_conv` route.
476+
477+
Both are computed from quantities known before any curvature work is done -- the
478+
number of source pixels `S = mapper.params`, the mapping operator's non-zeros
479+
`nnz = mapper.pix_sizes_for_sub_slim_index.sum()` and the extent rectangle
480+
`(ny, nx) = mask.shape_native_masked_pixels` (the sparse operator's
481+
`(y_shape, x_shape)`) -- so the factory can route on them.
482+
483+
The model is
484+
485+
t_fft = c_fft · S · [rfft2 + irfft2 work on (2ny, 2nx)]
486+
t_dc = c_dc · nnz · ny · nx
487+
488+
The FFT route scales with the number of source *columns* and with the padded extent's
489+
transform cost, including its factorisation; `direct_conv` scales with the number of
490+
non-zeros times the extent area and is flat in `S` at fixed `nnz` (34-41 s from 2,577
491+
to 5,746 vertices on the Aris 512x512 mask). A fixed non-zeros-per-column gate is the
492+
special case of this model at one extent, which is why it could be right for the
493+
125x125 reference geometry (crossover ~60) and wrong for the 512x512 one (crossover
494+
~130-150).
495+
496+
Returns
497+
-------
498+
`(t_fft, t_direct_conv)` in seconds on the calibration machine.
499+
"""
500+
t_fft = FFT_ROUTE_SECONDS_PER_UNIT * fft_route_work_from(
501+
source_pixels=source_pixels, extent_shape=extent_shape
502+
)
503+
t_direct_conv = DIRECT_CONV_SECONDS_PER_UNIT * direct_conv_work_from(
504+
nnz=nnz, extent_shape=extent_shape
505+
)
506+
507+
return t_fft, t_direct_conv

0 commit comments

Comments
 (0)