feat: matrix-free interferometer route via Settings (#638 phase 5a)

autoarray/__init__.py

@@ -51,6 +51,9 @@
 from .inversion.inversion.interferometer.sparse import (
     InversionInterferometerSparse,
 )
+from .inversion.inversion.interferometer.matrix_free import (
+    InversionInterferometerMatrixFree,
+)
 from .inversion.inversion.interferometer.mapping import InversionInterferometerMapping
 from .inversion.inversion.interferometer.inversion_interferometer_util import (
     InterferometerSparseOperator,

autoarray/config/general.yaml

@@ -19,7 +19,7 @@ inversion:
   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: 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.
+  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. "matrix_free" uses InversionInterferometerMatrixFree, which never forms the dense curvature or regularization matrix (PCG / interior-point solve, stochastic Lanczos log-det); it needs NumPy, one RectangularUniform mapper with MaternSparse regularization, no linear functions and (positive-only) use_edge_zeroed_pixels false, raises on anything else, and has no dense attributes such as reconstruction_noise_map (PyAutoArray#657).
   interferometer_direct_conv_symmetric_max_work_fraction: 0.7 # When the serial (numba parallel: false) direct_conv curvature route runs, use the symmetric row-restricted kernel (curvature_direct_conv_symmetric: source columns ranked by their top extent row, each convolving only the extent rows later-ranked columns span and only each row's mask x-range, then a mirror pass) when its geometry work fraction (its convolution work / the full kernel's) is at or below this value. Same F to round-off (computed entries bit-identical). Synthetic Aris-512 geometry (compact columns): f ~0.41, 1.81-1.90x single-thread on RAL; Aris's real 512 Delaunay mesh (multiply imaged, tall columns): f ~0.51, ~1.46x measured (PyAutoArray#635). The route cost model scales t_dc by min(1, 1.32 x work fraction), so the kernel wins below f ~0.76; 0.7 keeps a predicted gain >= 1.08x. Set 0 to always use the full kernel. The parallel kernel is unaffected.
 numba:
   use_numba: true # If True, use numba-accelerated mapper and Sibson interpolation paths when eligible.

autoarray/inversion/inversion/factory.py

@@ -1,6 +1,7 @@
 import numpy as np
 from typing import List, Optional, Tuple, Union
 
+from autoarray import exc
 from autoarray.dataset.imaging.dataset import Imaging
 from autoarray.dataset.interferometer.dataset import Interferometer
 from autoarray.inversion.inversion.imaging.mapping import InversionImagingMapping
@@ -17,6 +18,9 @@
 from autoarray.inversion.inversion.interferometer_numba.sparse import (
     InversionInterferometerSparseNumba,
 )
+from autoarray.inversion.inversion.interferometer.matrix_free import (
+    InversionInterferometerMatrixFree,
+)
 from autoarray.inversion.inversion.interferometer_numba import (
     sparse as interferometer_numba_sparse,
 )
@@ -224,8 +228,25 @@ def inversion_interferometer_from(
     # `transform_mapping_matrix` of shape (N_vis, S), which dominates run time and memory at large N_vis, and
     # `InversionInterferometerSparse` handles zero mappers (the func-func curvature blocks and the data vector
     # need no mapper). `_use_interferometer_numba` returns False whenever a func-list is present.
+    route_is_matrix_free = (
+        settings is not None and settings.interferometer_route == "matrix_free"
+    )
+
     if dataset.sparse_operator is not None:
 
+        # The matrix-free route is opt-in only. `InversionInterferometerMatrixFree` raises on
+        # every configuration it cannot represent (JAX, linear functions, several mappers, a
+        # non-rectangular mesh, a regularization without a sparse form, ...); there is no silent
+        # fallback to the dense routes.
+        if route_is_matrix_free:
+            return InversionInterferometerMatrixFree(
+                dataset=dataset,
+                linear_obj_list=linear_obj_list,
+                settings=settings,
+                xp=xp,
+                preloads=preloads,
+            )
+
         # The numba `direct_conv` kernel uses the mapper weights as-is and cannot apply a primary
         # beam, so a dataset with a non-trivial one (any value other than 1) always takes the FFT
         # sparse route. `P = ones` keeps the numba route, so it is bit-for-bit `P = None`.
@@ -260,6 +281,14 @@ def inversion_interferometer_from(
             preloads=preloads,
         )
 
+    if route_is_matrix_free:
+        raise exc.InversionException(
+            '`Settings(interferometer_route="matrix_free")` was set but the interferometer '
+            "dataset has no `sparse_operator`. The matrix-free route reads the w-tilde kernel "
+            "and dirty image from it; call `dataset.apply_sparse_operator(...)` first, or use "
+            '`interferometer_route="fft"` (or leave it unset).'
+        )
+
     return InversionInterferometerMapping(
         dataset=dataset,
         linear_obj_list=linear_obj_list,

autoarray/inversion/inversion/interferometer/inversion_interferometer_matrix_free_util.py

@@ -30,8 +30,8 @@
 
 Phase 1 landed the operator data term and the solver as library utilities, phase 2 the natively sparse
 Matérn precision (`aa.reg.MaternSparse`), phase 3 the evidence log-determinants (below) and phase 4 the
-positivity constraint (below), each with parity tests against the dense sparse-path quantities. None of
-it is wired into fits and there is no `Settings` flag.
+positivity constraint (below), each with parity tests against the dense sparse-path quantities. Phase 5a
+wires them into fits (`InversionInterferometerMatrixFree`, `Settings(interferometer_route="matrix_free")`).
 
 Positivity (phase 4)
 --------------------
@@ -46,7 +46,8 @@
 
 Remaining work (phase 5)
 ------------------------
-- wiring as an inversion class or a `Settings` option, so fits can select the matrix-free path;
+- wiring into fits: landed in phase 5a (#657) as `InversionInterferometerMatrixFree`
+  (`interferometer/matrix_free.py`), selected by `Settings(interferometer_route="matrix_free")`;
 - a profiling note, dense vs matrix-free at `50^2`, `64^2`, `128^2` and `192^2` mesh pixels, ingested
   by PyAutoPulse;
 - a workspace example;