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

autoarray/config/general.yaml

@@ -18,7 +18,8 @@ inversion:
   reconstruction_vmax_factor: 0.5     # Plots of an Inversion's reconstruction use the reconstructed data's bright value multiplied by this factor.
   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.
   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.
-  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.
+  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.
+  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.
 numba:
   use_numba: true
   cache: true

autoarray/inversion/inversion/factory.py

@@ -1,5 +1,5 @@
 import numpy as np
-from typing import List, Union
+from typing import List, Optional, Tuple, Union
 
 from autoarray.dataset.imaging.dataset import Imaging
 from autoarray.dataset.interferometer.dataset import Interferometer
@@ -214,6 +214,10 @@ def inversion_interferometer_from(
             linear_obj_list=linear_obj_list,
             settings=settings,
             xp=xp,
+            extent_shape=(
+                dataset.sparse_operator.y_shape,
+                dataset.sparse_operator.x_shape,
+            ),
         ):
             return InversionInterferometerSparseNumba(
                 dataset=dataset,
@@ -243,6 +247,7 @@ def _use_interferometer_numba(
     linear_obj_list: List[LinearObj],
     settings: Settings = None,
     xp=np,
+    extent_shape: Optional[Tuple[int, int]] = None,
 ) -> bool:
     """
     Whether an interferometer inversion is routed to the numba `direct_conv` curvature
@@ -258,12 +263,25 @@ def _use_interferometer_numba(
     The conditions are, in order of cost to evaluate:
 
     - `xp is np` -- the kernel is numba, with no JAX path.
-    - `settings.interferometer_numba_nnz_per_source_max > 0` -- `0` is the kill switch.
+    - `settings.interferometer_route` -- `"fft"` never routes to numba; `"numba"` forces it
+      whenever the kernel can represent the inversion (the conditions below that are not
+      about speed), bypassing the geometry gate and the kill switch. For benchmarking.
+    - `settings.interferometer_numba_nnz_per_source_max` is not `0` -- `0` is the kill
+      switch.
     - exactly one `Mapper` and no `AbstractLinearObjFuncList` -- the kernel builds a single
       mapper-mapper block and has no off-diagonal or function blocks.
     - no over-sampling (`sub_fraction == 1`) -- the kernel uses the mapper's weights as-is.
-    - the mapper's mean non-zeros per source column is at or below the gate -- above it the
-      FFT route is faster (see `Settings.interferometer_numba_nnz_per_source_max`).
+    - the geometry gate, which picks the faster route:
+
+      - `interferometer_numba_nnz_per_source_max` positive -- a hard gate: numba at or
+        below that many mean non-zeros per source column (the pre-#628 behaviour, kept
+        so explicit A/B settings are honoured).
+      - `interferometer_numba_nnz_per_source_max` `None` (the packaged default) -- the
+        cost model `inversion_interferometer_numba_util.curvature_route_seconds_from`,
+        which predicts both routes' time from the source pixels, the mapping operator's
+        non-zeros and the extent rectangle `extent_shape`, and routes to numba when
+        `direct_conv` is predicted no slower. Without an `extent_shape` the model cannot
+        be evaluated and the FFT route is used.
     - `import numba` succeeds.
 
     NaN mapping weights (an overflowing or degenerate `DelaunayNN` Sibson row) are
@@ -277,19 +295,29 @@ def _use_interferometer_numba(
     linear_obj_list
         The linear objects reconstructing the data.
     settings
-        The inversion settings, whose `interferometer_numba_nnz_per_source_max` is the
-        geometry gate. `None` uses the packaged defaults.
+        The inversion settings, whose `interferometer_route` and
+        `interferometer_numba_nnz_per_source_max` select the route. `None` uses the
+        packaged defaults.
     xp
         The array module the inversion runs on.
+    extent_shape
+        The `(ny, nx)` extent rectangle the `W~` operator lives on, i.e. the dataset's
+        `sparse_operator.y_shape, sparse_operator.x_shape`
+        (`= mask.shape_native_masked_pixels`). Needed only by the default cost model.
     """
     if xp is not np:
         return False
 
     settings = settings if settings is not None else Settings()
 
+    route = settings.interferometer_route
+
+    if route == "fft":
+        return False
+
     nnz_max = settings.interferometer_numba_nnz_per_source_max
 
-    if nnz_max is None or nnz_max <= 0:
+    if route is None and nnz_max is not None and nnz_max <= 0:
         return False
 
     if any(
@@ -312,12 +340,31 @@ def _use_interferometer_numba(
     if not np.all(sub_fraction == 1.0):
         return False
 
-    nnz_per_source_column = (
-        inversion_interferometer_numba_util.nnz_per_source_column_from(mapper=mapper)
-    )
+    if route is None:
+        if nnz_max is not None:
+            nnz_per_source_column = (
+                inversion_interferometer_numba_util.nnz_per_source_column_from(
+                    mapper=mapper
+                )
+            )
 
-    if nnz_per_source_column > nnz_max:
-        return False
+            if nnz_per_source_column > nnz_max:
+                return False
+
+        else:
+            if extent_shape is None:
+                return False
+
+            t_fft, t_direct_conv = (
+                inversion_interferometer_numba_util.curvature_route_seconds_from(
+                    source_pixels=int(mapper.params),
+                    nnz=int(np.asarray(mapper.pix_sizes_for_sub_slim_index).sum()),
+                    extent_shape=extent_shape,
+                )
+            )
+
+            if t_direct_conv > t_fft:
+                return False
 
     try:
         import numba  # noqa: F401

autoarray/inversion/inversion/interferometer_numba/inversion_interferometer_numba_util.py

@@ -33,9 +33,9 @@ def curvature_direct_conv(
     Cost `O(nnz·M + S·nnz)` with `M = ny·nx`, versus the historic pair loop's `O(N² P²)`:
     the convolution replaces the `N²` pixel-pair space with the `M`-cell extent rectangle,
     which is the whole point of the extent-grid form. It beats the FFT route (the JAX and
-    NumPy `InterferometerSparseOperator` paths) while the source columns stay sparse --
-    measured crossover ~60 non-zeros per source column on Delaunay meshes, ~77 on
-    rectangular ones (autolens_profiling issue #226).
+    NumPy `InterferometerSparseOperator` paths) while the source columns stay sparse; where
+    the two cross depends on the extent as well as the column density, and the factory
+    decides with :func:`curvature_route_seconds_from` (PyAutoArray#628).
 
     The result is a complete, symmetric `[pix_pixels, pix_pixels]` matrix: the kernel
     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:
     """
     The mean number of non-zero mapping weights per source column of `A`.
 
-    This is the geometry the `direct_conv` kernel's cost scales with, and therefore the
-    quantity the factory gates on: the kernel's convolution step costs `O(nnz · M)`,
-    while the FFT route's cost is set by the number of source *columns* rather than their
-    density, so the two cross at a roughly fixed non-zeros-per-column value (~60 on
-    Delaunay meshes, ~77 on rectangular ones).
+    The kernel's convolution step costs `O(nnz · M)` while the FFT route's cost is set by
+    the number of source *columns* rather than their density, so at a fixed extent the two
+    cross at a fixed non-zeros-per-column value (~60 on Delaunay meshes, ~77 on
+    rectangular ones at the autolens_profiling#226 geometries). That value moves with the
+    extent, so the default routing uses :func:`curvature_route_seconds_from`; this ratio is
+    what an explicit `Settings.interferometer_numba_nnz_per_source_max` hard gate compares
+    against.
 
     `nnz = pix_sizes_for_sub_slim_index.sum()` is the total number of valid triplets and
     `mapper.params` the number of source pixels, so this is simply their ratio.
@@ -300,3 +302,208 @@ def nnz_per_source_column_from(mapper) -> float:
         return 0.0
 
     return nnz / params
+
+
+def _fft_length_cost(n: int) -> float:
+    """
+    The work of one length-`n` complex FFT, as `scipy.fft`'s pocketfft backend estimates it.
+
+    This mirrors pocketfft's own planner heuristic (`util::cost_guess`): `n` times the sum
+    of its prime factors, with a 1.1 penalty on factors above 5, so a smooth length such as
+    `560 = 2^4 * 5 * 7` costs `560 * 20` while an awkward one such as
+    `682 = 2 * 11 * 31` costs `682 * 48.2` -- 2.4x more per element although
+    `log2(682)` and `log2(560)` differ by 3 %. That factorisation penalty is what makes the
+    FFT route relatively slower on the Aris masks (extents 83 and 341, so FFT lengths
+    `166 = 2 * 83` and `682`) than `n log n` predicts, and is the reason the route gate is
+    a cost model rather than a fixed non-zeros-per-column constant.
+
+    For a length with a large prime factor (`p^2 > n`, `n >= 50`) pocketfft switches to
+    Bluestein's algorithm when that is cheaper -- two FFTs of the smallest 11-smooth length
+    `>= 2n - 1`, with pocketfft's 1.5 fudge factor -- and so does this estimate.
+    """
+    n = int(n)
+
+    def cost_guess(m: int) -> float:
+        m_in = m
+        result = 0.0
+
+        while m % 4 == 0:
+            result += 2.0
+            m //= 4
+
+        while m % 2 == 0:
+            result += 2.0
+            m //= 2
+
+        x = 3
+        while x * x <= m:
+            while m % x == 0:
+                result += float(x) if x <= 5 else 1.1 * float(x)
+                m //= x
+            x += 2
+
+        if m > 1:
+            result += float(m) if m <= 5 else 1.1 * float(m)
+
+        return result * float(m_in)
+
+    def largest_prime_factor(m: int) -> int:
+        largest = 1
+        x = 2
+        while x * x <= m:
+            while m % x == 0:
+                largest = x
+                m //= x
+            x += 1
+        return m if m > 1 else largest
+
+    def good_size(m: int) -> int:
+        """The smallest `2^a 3^b 5^c 7^d 11^e >= m` (pocketfft's `good_size_cmplx`)."""
+        best = 2 * m
+        f11 = 1
+        while f11 < best:
+            f7 = f11
+            while f7 < best:
+                f5 = f7
+                while f5 < best:
+                    f3 = f5
+                    while f3 < best:
+                        f2 = f3
+                        while f2 < m:
+                            f2 *= 2
+                        best = min(best, f2)
+                        f3 *= 3
+                    f5 *= 5
+                f7 *= 7
+            f11 *= 11
+        return best
+
+    direct = cost_guess(n)
+
+    if n < 50:
+        return direct
+
+    p = largest_prime_factor(n)
+
+    if p * p <= n:
+        return direct
+
+    bluestein = 1.5 * 2.0 * cost_guess(good_size(2 * n - 1))
+
+    return min(direct, bluestein)
+
+
+def fft_route_work_from(source_pixels: int, extent_shape) -> float:
+    """
+    The FFT route's curvature work, in pocketfft cost units.
+
+    `InterferometerSparseOperator.curvature_matrix_diag_from` (NumPy branch) applies `W~`
+    to every source column once: an `rfft2` and an `irfft2` of the `(2ny, 2nx)` padded
+    extent, in batches of `batch_size` columns. Batching changes the memory traffic, not
+    the number of transforms, so the work is `S` times one forward plus one inverse 2D
+    real transform. A 2D real transform of shape `(a, b)` is `a` real transforms of length
+    `b` (half the cost of a complex one) followed by `b // 2 + 1` complex transforms of
+    length `a`.
+
+    The sparse `Aᵀ G` projection (`S · nnz` multiply-adds) is left out: the `direct_conv`
+    kernel does the identical gather, so it cancels in the comparison.
+    """
+    ny, nx = int(extent_shape[0]), int(extent_shape[1])
+    a, b = 2 * ny, 2 * nx
+
+    per_transform = 0.5 * a * _fft_length_cost(b) + (b // 2 + 1) * _fft_length_cost(a)
+
+    return float(source_pixels) * 2.0 * per_transform
+
+
+def direct_conv_work_from(nnz: int, extent_shape) -> float:
+    """
+    The `direct_conv` kernel's curvature work, in multiply-adds.
+
+    :func:`curvature_direct_conv` scatters a shifted copy of the `W~` kernel over the whole
+    `(ny, nx)` extent rectangle for every non-zero of the mapping operator, so the
+    convolution step is `nnz · ny · nx` multiply-adds whatever the number of source pixels
+    (its `S · nnz` gather is shared with the FFT route and left out of both).
+    """
+    ny, nx = int(extent_shape[0]), int(extent_shape[1])
+
+    return float(nnz) * float(ny) * float(nx)
+
+
+# Seconds per unit of `fft_route_work_from` and of `direct_conv_work_from`, fitted by least
+# squares on log(seconds) to 18 single-thread timings (`curvature_route_seconds_from`):
+#
+# - Aris 512x512 mask (91,507 masked pixels, extent 341x341, Delaunay, nnz 273,661), whole
+#   NumPy likelihood call (the curvature build is 96-98 % of it), PyAutoArray#628:
+#   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.
+# - autolens_profiling#226 synthetic bake-off (`bakeoff_v2026.8.17.1.json`, curvature build
+#   only, `rfft2_numpy` = the same algorithm as the NumPy FFT route): sma / alma / alma_high
+#   extents 70^2 / 140^2 / 280^2, Delaunay S = 1,500 (P = 3) and rectangular S = 784 (P = 4).
+#
+# Only the ratio of the two constants (0.430) decides a route. Measured / predicted
+# FFT-to-direct_conv time ratio over the nine paired points lies in 0.68-1.47 (log-rms 0.26),
+# and the predicted faster route matches the measured one at all nine:
+#
+#   point                    nnz/col  measured t_fft/t_dc  predicted
+#   bake-off sma Delaunay       7.7          5.93            6.62
+#   bake-off sma rect          19.7          2.83            2.59
+#   bake-off alma Delaunay     30.8          1.54            1.88
+#   bake-off alma rect         78.5          0.99            0.74
+#   bake-off alma_high Del.   123.1          0.63            0.47
+#   bake-off alma_high rect   314.1          0.27            0.18
+#   Aris 512, S = 1,459       187.6          0.74            0.89
+#   Aris 512, S = 2,577       106.2          1.29            1.57
+#   Aris 512, S = 5,746        47.6          2.36            3.49
+#
+# The absolute residuals are larger (log-rms 0.56 FFT, 0.52 direct_conv): the two campaigns ran
+# on the same i9-10885H ~2.5x apart in clock state (the bake-off documents the drift), which
+# moves both routes together and cancels in the ratio. Implied crossovers in non-zeros per
+# source column: ~51 / 58 / 58 at the 70^2 / 140^2 / 280^2 bake-off extents (measured ~60 Delaunay,
+# ~77 rectangular), ~166 at the Aris 512 extent 341^2 (measured between 106 and 188) and ~309 at
+# the Aris 125 extent 83^2, whose FFT length 166 = 2 * 83 is pocketfft's Bluestein case.
+#
+# The FFT work form was checked (not fitted) with a single-thread `rfft2` + `irfft2`
+# micro-benchmark at padded lengths 140 / 166 / 240 / 280 / 560 / 662 / 682: seconds per
+# `fft_route_work_from` unit varied 1.8x across them, against 5.8x for seconds per `N log2 N`
+# -- the awkward lengths 166, 662 and 682 cost 2-5x more per element than smooth ones.
+FFT_ROUTE_SECONDS_PER_UNIT = 7.62e-10
+DIRECT_CONV_SECONDS_PER_UNIT = 1.77e-09
+
+
+def curvature_route_seconds_from(source_pixels: int, nnz: int, extent_shape):
+    """
+    Predicted single-thread seconds of the FFT route and of the `direct_conv` route.
+
+    Both are computed from quantities known before any curvature work is done -- the
+    number of source pixels `S = mapper.params`, the mapping operator's non-zeros
+    `nnz = mapper.pix_sizes_for_sub_slim_index.sum()` and the extent rectangle
+    `(ny, nx) = mask.shape_native_masked_pixels` (the sparse operator's
+    `(y_shape, x_shape)`) -- so the factory can route on them.
+
+    The model is
+
+        t_fft = c_fft · S · [rfft2 + irfft2 work on (2ny, 2nx)]
+        t_dc  = c_dc  · nnz · ny · nx
+
+    The FFT route scales with the number of source *columns* and with the padded extent's
+    transform cost, including its factorisation; `direct_conv` scales with the number of
+    non-zeros times the extent area, so at a large extent it is ~flat in `S` at fixed
+    `nnz` (34-41 s from 2,577 to 5,746 vertices on the Aris 512x512 mask). Both routes
+    also pay `S`-dependent costs the model leaves out because they are common to both
+    (the `S · nnz` gather, the dense solve). A fixed non-zeros-per-column gate is the
+    special case of this model at one extent, which is why it could be right at the
+    autolens_profiling#226 extents (70^2-280^2, crossover ~60) and wrong on the Aris
+    512x512 mask (model crossover ~166, measured between 106 and 188).
+
+    Returns
+    -------
+    `(t_fft, t_direct_conv)` in seconds on the calibration machine.
+    """
+    t_fft = FFT_ROUTE_SECONDS_PER_UNIT * fft_route_work_from(
+        source_pixels=source_pixels, extent_shape=extent_shape
+    )
+    t_direct_conv = DIRECT_CONV_SECONDS_PER_UNIT * direct_conv_work_from(
+        nnz=nnz, extent_shape=extent_shape
+    )
+
+    return t_fft, t_direct_conv