feat: numba CPU interferometer curvature path — direct_conv, geometry-gated (#543)

autoarray/config/general.yaml

@@ -13,6 +13,7 @@ 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.
 numba:
   use_numba: true
   cache: true

autoarray/inversion/inversion/factory.py

@@ -11,6 +11,13 @@
 from autoarray.inversion.inversion.interferometer.sparse import (
     InversionInterferometerSparse,
 )
+from autoarray.inversion.inversion.interferometer_numba.sparse import (
+    InversionInterferometerSparseNumba,
+)
+from autoarray.inversion.inversion.interferometer_numba import (
+    inversion_interferometer_numba_util,
+)
+from autoarray.inversion.mappers.abstract import Mapper
 from autoarray.inversion.inversion.dataset_interface import DatasetInterface
 from autoarray.inversion.linear_obj.linear_obj import LinearObj
 from autoarray.inversion.linear_obj.func_list import AbstractLinearObjFuncList
@@ -202,6 +209,19 @@ def inversion_interferometer_from(
 
     if dataset.sparse_operator is not None and use_sparse_operator:
 
+        if _use_interferometer_numba(
+            linear_obj_list=linear_obj_list,
+            settings=settings,
+            xp=xp,
+        ):
+            return InversionInterferometerSparseNumba(
+                dataset=dataset,
+                linear_obj_list=linear_obj_list,
+                settings=settings,
+                xp=xp,
+                preloads=preloads,
+            )
+
         return InversionInterferometerSparse(
             dataset=dataset,
             linear_obj_list=linear_obj_list,
@@ -216,3 +236,85 @@ def inversion_interferometer_from(
         settings=settings,
         xp=xp,
     )
+
+
+def _use_interferometer_numba(
+    linear_obj_list: List[LinearObj],
+    settings: Settings = None,
+    xp=np,
+) -> bool:
+    """
+    Whether an interferometer inversion is routed to the numba `direct_conv` curvature
+    path (`InversionInterferometerSparseNumba`) rather than the FFT one
+    (`InversionInterferometerSparse`).
+
+    Every condition below is a routing decision, not an error: a model the kernel cannot
+    represent, or a geometry where the FFT route is faster, simply falls through to the
+    sparse path silently. Constructing `InversionInterferometerSparseNumba` directly with
+    such inputs still raises -- the class checks the same preconditions itself, so the two
+    cannot drift apart in meaning, only in whether they are fatal.
+
+    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.
+    - 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`).
+    - `import numba` succeeds.
+
+    Parameters
+    ----------
+    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.
+    xp
+        The array module the inversion runs on.
+    """
+    if xp is not np:
+        return False
+
+    settings = settings if settings is not None else Settings()
+
+    nnz_max = settings.interferometer_numba_nnz_per_source_max
+
+    if nnz_max is None or nnz_max <= 0:
+        return False
+
+    if any(
+        isinstance(linear_obj, AbstractLinearObjFuncList)
+        for linear_obj in linear_obj_list
+    ):
+        return False
+
+    mapper_list = [
+        linear_obj for linear_obj in linear_obj_list if isinstance(linear_obj, Mapper)
+    ]
+
+    if len(mapper_list) != 1 or len(mapper_list) != len(linear_obj_list):
+        return False
+
+    mapper = mapper_list[0]
+
+    sub_fraction = np.asarray(mapper.over_sampler.sub_fraction.array)
+
+    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 nnz_per_source_column > nnz_max:
+        return False
+
+    try:
+        import numba  # noqa: F401
+    except ModuleNotFoundError:
+        return False
+
+    return True