diff --git a/autoarray/config/general.yaml b/autoarray/config/general.yaml index 916e85045..e4dc1b6f5 100644 --- a/autoarray/config/general.yaml +++ b/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 diff --git a/autoarray/inversion/inversion/factory.py b/autoarray/inversion/inversion/factory.py index addb1f643..3b908f034 100644 --- a/autoarray/inversion/inversion/factory.py +++ b/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 diff --git a/autoarray/inversion/inversion/interferometer_numba/inversion_interferometer_numba_util.py b/autoarray/inversion/inversion/interferometer_numba/inversion_interferometer_numba_util.py index 034909ae9..215b2ab3b 100644 --- a/autoarray/inversion/inversion/interferometer_numba/inversion_interferometer_numba_util.py +++ b/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 diff --git a/autoarray/settings.py b/autoarray/settings.py index 54bf6360d..a69832582 100644 --- a/autoarray/settings.py +++ b/autoarray/settings.py @@ -22,6 +22,7 @@ def __init__( log_det_method: Optional[str] = None, regularization_term_method: Optional[str] = None, interferometer_numba_nnz_per_source_max: Optional[float] = None, + interferometer_route: Optional[str] = None, positive_only_solver: Optional[str] = None, certified_pass_budget: Optional[int] = None, certified_fallback: Optional[str] = None, @@ -190,12 +191,23 @@ def __init__( ``curvature_reg_matrix`` is a dense ``F + H`` feeding the dense solve for the reconstruction, so ``H`` is still formed there regardless. interferometer_numba_nnz_per_source_max - The geometry gate above which the numba `direct_conv` interferometer curvature - path is not used, in mean non-zeros per source column - (``mapper.pix_sizes_for_sub_slim_index.sum() / mapper.params``). `None` - (default) reads the packaged value (`60.0`); `0` disables the numba path. See - the property of the same name for the measured crossovers and why the constant - is machine-dependent. + How a NumPy single-mapper interferometer inversion chooses between the numba + `direct_conv` curvature path and the FFT one. `None` (default) reads the + packaged value, which is itself `None`: a cost model predicts both routes' time + from the source pixels, the mapping operator's non-zeros and the extent + rectangle, and picks the faster (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. `0` disables the numba path. See the property of the same + name. + interferometer_route + Forces the interferometer curvature route, for benchmarking. `None` (default) + reads the packaged value, which is `None`: route by + `interferometer_numba_nnz_per_source_max`. ``"fft"`` always uses the FFT route + (`InversionInterferometerSparse`); ``"numba"`` uses the numba `direct_conv` route + (`InversionInterferometerSparseNumba`) whenever it can represent the inversion + (NumPy backend, one mapper, no linear functions, no over-sampling, numba + installed), overriding the geometry gate and the `0` kill switch. positive_only_solver Which solver the JAX (`xp=jnp`) positive-only reconstruction uses. `None` (default) reads the packaged value ``"pdip"``. @@ -262,6 +274,7 @@ def __init__( self._interferometer_numba_nnz_per_source_max = ( interferometer_numba_nnz_per_source_max ) + self._interferometer_route = interferometer_route self._positive_only_solver = positive_only_solver self._certified_pass_budget = certified_pass_budget self._certified_fallback = certified_fallback @@ -275,6 +288,8 @@ def __init__( self.certified_fallback if nnls_preconditioning_no_mapper is not None: self.nnls_preconditioning_no_mapper + if interferometer_route is not None: + self.interferometer_route @property def use_positive_only_solver(self): @@ -374,43 +389,69 @@ def regularization_term_method(self): return self._regularization_term_method @property - def interferometer_numba_nnz_per_source_max(self) -> float: + def interferometer_numba_nnz_per_source_max(self) -> Optional[float]: """ - The geometry gate above which the numba `direct_conv` interferometer curvature - path is not used. - - `InversionInterferometerSparseNumba` convolves each source column of the mapping - operator over the extent rectangle, at a cost that scales with the column's - non-zeros; the FFT route (`InversionInterferometerSparse`) costs the same whatever - the density. The two therefore cross at a roughly fixed number of non-zeros per - source column, `mapper.pix_sizes_for_sub_slim_index.sum() / mapper.params`, and - the factory routes to numba only at or below this value. - - The measured crossovers are **~60 non-zeros per source column on Delaunay meshes** - and **~77 on rectangular meshes** (autolens_profiling issue #226 verdict, section - 2), where the numba kernel is 2-7x faster than JAX-CPU well below the crossover. - The default is the conservative of the two. - - This is a **machine-dependent constant**: it is set by the ratio of scalar AXPY - throughput to FFT throughput on the CPU running the fit, so a machine with a very - different cache hierarchy or FFT library will cross somewhere else. Re-measure - before tuning it for a new machine; `0` disables the numba path entirely. + How the factory chooses between the numba `direct_conv` interferometer curvature + path and the FFT one: `None` (packaged default) for the cost model, a positive + value for a hard non-zeros-per-source-column gate, `0` for never numba. + + `InversionInterferometerSparseNumba` convolves each non-zero of the mapping + operator over the extent rectangle, so its cost scales with the total non-zeros + times the extent area and is flat in the number of source pixels; the FFT route + (`InversionInterferometerSparse`) transforms the padded extent once per source + column, so its cost scales with the source pixels times the padded extent's + transform cost. At a fixed extent the two cross at a fixed mean non-zeros per + source column, `mapper.pix_sizes_for_sub_slim_index.sum() / mapper.params` -- but + that crossover moves with the extent: ~60 (Delaunay) / ~77 (rectangular) at the + autolens_profiling#226 geometries (extents 70^2-280^2), between 106 and 188 on + the 512x512 Aris mask (extent 341^2, model crossover ~166; PyAutoArray#628), so + no single constant is right for both. + + - `None` -- the cost model + (`inversion_interferometer_numba_util.curvature_route_seconds_from`, constants + calibrated in that module) predicts both routes' time and the factory routes to + numba when `direct_conv` is predicted no slower. + - positive -- the pre-#628 hard gate: numba at or below this many non-zeros per + source column, FFT above, whatever the extent. Honoured as given, so A/B scripts + that set it keep their meaning. + - `0` -- the kill switch: never numba. + + The cost model's constants are single-thread measurements on one CPU; the ratio + that decides a route depends on the ratio of scalar AXPY throughput to FFT + throughput, so a very different machine may cross elsewhere. `interferometer_route` + forces a route for benchmarking. """ if self._interferometer_numba_nnz_per_source_max is None: - try: - return conf.instance["general"]["inversion"][ - "interferometer_numba_nnz_per_source_max" - ] - except KeyError: - # A workspace `general.yaml` normally omits this key, so autoconf's - # config-path list falls through to autoarray's packaged value (`60.0`). - # This fallback fires only when the workspace config is the sole config - # path (isolated test configs push one dir) and returns that same value, - # so both routes resolve identically. - return 60.0 + # A workspace `general.yaml` normally omits this key, so autoconf's config-path + # list falls through to autoarray's packaged value (`null`, the cost model). The + # fallback fires only when the workspace config is the sole config path + # (isolated test configs push one dir) and returns that same value. + return self._inversion_config_value( + "interferometer_numba_nnz_per_source_max", None + ) return self._interferometer_numba_nnz_per_source_max + @property + def interferometer_route(self) -> Optional[str]: + """ + A forced interferometer curvature route: `None` (route by + `interferometer_numba_nnz_per_source_max`), ``"numba"`` or ``"fft"``. + + See the constructor docstring; ``"numba"`` still requires an inversion the + `direct_conv` kernel can represent. + """ + value = self._interferometer_route + if value is None: + value = self._inversion_config_value("interferometer_route", None) + + if value not in (None, "numba", "fft"): + raise ValueError( + f"interferometer_route={value!r} is invalid; expected None, 'numba' or 'fft'." + ) + + return value + def _inversion_config_value(self, key, default): # A workspace `general.yaml` normally omits the newer keys, so autoconf's config-path # list falls through to autoarray's packaged value. The fallback fires only when the diff --git a/test_autoarray/inversion/inversion/interferometer_numba/test_interferometer_numba.py b/test_autoarray/inversion/inversion/interferometer_numba/test_interferometer_numba.py index c341b5a5c..b6fa0cca7 100644 --- a/test_autoarray/inversion/inversion/interferometer_numba/test_interferometer_numba.py +++ b/test_autoarray/inversion/inversion/interferometer_numba/test_interferometer_numba.py @@ -385,15 +385,22 @@ def test__factory__routes_to_numba_below_the_gate_and_sparse_above_it(): def test__factory__delaunay_nn__routes_to_numba_with_default_settings(): """ Sibson rows carry more non-zeros than Delaunay ones, so the risk is that a - `DelaunayNN` mapper trips the nnz-per-source-column gate onto the FFT route. With the - packaged gate it must not, and the factory-built inversion matches the FFT one. + `DelaunayNN` mapper trips the geometry gate onto the FFT route. With the packaged + routing (the cost model) it must not, and the factory-built inversion matches the FFT + one. """ mask, _, dataset_sparse = _dataset_from() mapper = _delaunay_nn_mapper_from(mask=mask) - assert numba_util.nnz_per_source_column_from( - mapper=mapper - ) <= aa.Settings().interferometer_numba_nnz_per_source_max + assert aa.Settings().interferometer_numba_nnz_per_source_max is None + + t_fft, t_direct_conv = numba_util.curvature_route_seconds_from( + source_pixels=mapper.params, + nnz=int(np.asarray(mapper.pix_sizes_for_sub_slim_index).sum()), + extent_shape=mask.shape_native_masked_pixels, + ) + + assert t_direct_conv <= t_fft inversion = aa.Inversion(dataset=dataset_sparse, linear_obj_list=[mapper]) @@ -448,6 +455,33 @@ def log_evidence_terms_from(inversion): assert np.isnan(log_evidence_terms_from(inversion)) +def test__factory__interferometer_route_forces_the_route(): + mask, _, dataset_sparse = _dataset_from() + mapper = _delaunay_mapper_from(mask=mask) + + inversion = aa.Inversion( + dataset=dataset_sparse, + linear_obj_list=[mapper], + settings=aa.Settings(interferometer_route="fft"), + ) + + assert isinstance(inversion, InversionInterferometerSparse) + assert not isinstance(inversion, InversionInterferometerSparseNumba) + + # "numba" overrides both the hard gate and the kill switch. + for nnz_max in (0, 0.5): + inversion = aa.Inversion( + dataset=dataset_sparse, + linear_obj_list=[mapper], + settings=aa.Settings( + interferometer_route="numba", + interferometer_numba_nnz_per_source_max=nnz_max, + ), + ) + + assert isinstance(inversion, InversionInterferometerSparseNumba) + + def test__factory__gate_of_zero_disables_the_numba_path(): mask, _, dataset_sparse = _dataset_from() mapper = _delaunay_mapper_from(mask=mask) diff --git a/test_autoarray/inversion/inversion/test_factory.py b/test_autoarray/inversion/inversion/test_factory.py index 706ccbcba..0d8a0b011 100644 --- a/test_autoarray/inversion/inversion/test_factory.py +++ b/test_autoarray/inversion/inversion/test_factory.py @@ -796,3 +796,192 @@ def test__data_linear_func_matrix_dict( assert inversion_mapping.data_linear_func_matrix_dict[linear_obj][ 2 ] == pytest.approx([0.06458333, 0.05972222], 1.0e-4) + + +class _GeometryMapper(aa.m.MockMapper): + """ + A mapper carrying only what the interferometer route gate reads: the number of source + pixels and the per-image-pixel non-zero counts, with no over-sampling. + """ + + def __init__(self, source_pixels, nnz, image_pixels): + import types + + super().__init__( + parameters=source_pixels, + over_sampler=types.SimpleNamespace( + sub_fraction=types.SimpleNamespace(array=np.ones(image_pixels)) + ), + ) + + sizes = np.full(image_pixels, nnz // image_pixels, dtype=np.int64) + sizes[: nnz - int(sizes.sum())] += 1 + + self._pix_sizes = sizes + + @property + def pix_sizes_for_sub_slim_index(self): + return self._pix_sizes + + +def _route_is_numba(source_pixels, nnz, image_pixels, extent_shape, settings=None): + from autoarray.inversion.inversion.factory import _use_interferometer_numba + + mapper = _GeometryMapper( + source_pixels=source_pixels, nnz=nnz, image_pixels=image_pixels + ) + + return _use_interferometer_numba( + linear_obj_list=[mapper], + settings=settings, + xp=np, + extent_shape=extent_shape, + ) + + +# Every measured point the cost model was calibrated or checked against, with the route that +# measured faster (single thread, seconds per call / per curvature build): +# +# - Aris 512x512 mask: 91,507 masked pixels, extent 341x341 (`shape_native_masked_pixels`), +# Delaunay P=3, nnz 273,661 (PyAutoArray#628): FFT 28.3 / 44.1 / 97.3 s vs direct_conv +# 38.4 / 34.2 / 41.2 s. +# - Aris 125x125 mask: 5,496 masked pixels, extent 83x83: direct_conv 0.40 / 2.08 / 6.12 s +# (1 thread) vs the JAX FFT route 1.96 / 6.84 / 12.0 s (2.8-3.6 cores). +# - autolens_profiling#226 bake-off (`bakeoff_v2026.8.17.1.json`), rfft2_numpy vs direct_conv. +_MEASURED_ROUTES = [ + # (label, S, nnz, image pixels, extent, faster route) + ("aris_512_1459", 1459, 273661, 91507, (341, 341), "fft"), + ("aris_512_2577", 2577, 273661, 91507, (341, 341), "numba"), + ("aris_512_5746", 5746, 273661, 91507, (341, 341), "numba"), + ("aris_125_1405", 1405, 16438, 5496, (83, 83), "numba"), + ("aris_125_2947", 2947, 16448, 5496, (83, 83), "numba"), + ("aris_125_5526", 5526, 16488, 5496, (83, 83), "numba"), + ("bakeoff_sma_delaunay", 1500, 11556, 3852, (70, 70), "numba"), + ("bakeoff_sma_rect", 784, 15408, 3852, (70, 70), "numba"), + ("bakeoff_alma_delaunay", 1500, 46140, 15380, (140, 140), "numba"), + ("bakeoff_alma_rect", 784, 61520, 15380, (140, 140), "fft"), + ("bakeoff_alma_high_delaunay", 1500, 184716, 61572, (280, 280), "fft"), + ("bakeoff_alma_high_rect", 784, 246288, 61572, (280, 280), "fft"), +] + + +@pytest.mark.parametrize( + "label, source_pixels, nnz, image_pixels, extent_shape, faster", + _MEASURED_ROUTES, + ids=[row[0] for row in _MEASURED_ROUTES], +) +def test__interferometer_route__default_cost_model_picks_the_measured_faster_route( + label, source_pixels, nnz, image_pixels, extent_shape, faster +): + """ + The Witness of PyAutoArray#628: on both Aris masks the default routing picks the + faster of numba `direct_conv` and the NumPy FFT curvature at every measured vertex + count, and it still does at every autolens_profiling#226 bake-off cell. + """ + pytest.importorskip("numba") + + assert aa.Settings().interferometer_numba_nnz_per_source_max is None + + is_numba = _route_is_numba( + source_pixels=source_pixels, + nnz=nnz, + image_pixels=image_pixels, + extent_shape=extent_shape, + ) + + assert is_numba == (faster == "numba") + + +def test__interferometer_route__crossover_moves_with_the_extent(): + """ + The old fixed gate (60 non-zeros per source column) is right at the + autolens_profiling#226 extents and wrong on the Aris 512x512 one; the cost model + crosses near 60 at the former and between the measured 106 (numba) and 188 (FFT) at + the latter. + """ + pytest.importorskip("numba") + + # alma extent (140x140), where the ~60 Delaunay crossover was measured. + assert _route_is_numba(1500, 1500 * 50, 15380, (140, 140)) + assert not _route_is_numba(1500, 1500 * 80, 15380, (140, 140)) + + # Aris 512x512 extent at 2,577 vertices: 106 non-zeros per column is numba by the cost + # model and FFT by the old gate, which measured 1.29x slower. + assert _route_is_numba(2577, 273661, 91507, (341, 341)) + assert not _route_is_numba( + 2577, + 273661, + 91507, + (341, 341), + settings=aa.Settings(interferometer_numba_nnz_per_source_max=60.0), + ) + + +def test__interferometer_route__settings_overrides(): + pytest.importorskip("numba") + + # Aris 512 at 1,459 vertices: the cost model picks FFT. + geometry = dict( + source_pixels=1459, nnz=273661, image_pixels=91507, extent_shape=(341, 341) + ) + # Aris 512 at 5,746 vertices: the cost model picks numba. + geometry_numba = dict( + source_pixels=5746, nnz=273661, image_pixels=91507, extent_shape=(341, 341) + ) + + assert not _route_is_numba(**geometry) + assert _route_is_numba(**geometry_numba) + + # A positive gate is a hard gate, whatever the extent (pre-#628 behaviour). + assert _route_is_numba( + **geometry, + settings=aa.Settings(interferometer_numba_nnz_per_source_max=200.0), + ) + assert not _route_is_numba( + **geometry_numba, + settings=aa.Settings(interferometer_numba_nnz_per_source_max=40.0), + ) + + # 0 is the kill switch. + assert not _route_is_numba( + **geometry_numba, + settings=aa.Settings(interferometer_numba_nnz_per_source_max=0), + ) + + # The forced route wins over the cost model, the hard gate and the kill switch. + assert _route_is_numba(**geometry, settings=aa.Settings(interferometer_route="numba")) + assert _route_is_numba( + **geometry_numba, + settings=aa.Settings( + interferometer_route="numba", interferometer_numba_nnz_per_source_max=0 + ), + ) + assert not _route_is_numba( + **geometry_numba, settings=aa.Settings(interferometer_route="fft") + ) + + # Without the extent the cost model cannot be evaluated and the FFT route is used. + assert not _route_is_numba(**{**geometry_numba, "extent_shape": None}) + + with pytest.raises(ValueError): + aa.Settings(interferometer_route="direct_conv") + + +def test__interferometer_route__fft_length_cost_penalises_awkward_lengths(): + from autoarray.inversion.inversion.interferometer_numba import ( + inversion_interferometer_numba_util as numba_util, + ) + + def per_element(n): + return numba_util._fft_length_cost(n) / n + + # 682 = 2 * 11 * 31 (the Aris 512 padded extent) costs far more per element than the + # smooth 560 = 2^4 * 5 * 7, although log2 of the two differs by 3 %. + assert per_element(682) > 2.0 * per_element(560) + + # 166 = 2 * 83 (the Aris 125 padded extent) is pocketfft's Bluestein case: the + # estimate is capped by the Bluestein cost, below the direct mixed-radix one. + assert numba_util._fft_length_cost(166) < 166 * (2.0 + 1.1 * 83) + + # A power of two is the cheapest per element of its size class. + assert per_element(512) < per_element(510)