perf: symmetric row-restricted direct_conv curvature kernel

autoarray/config/general.yaml

@@ -20,6 +20,7 @@ inversion:
   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_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
   cache: true

autoarray/inversion/inversion/factory.py

@@ -11,9 +11,15 @@
 from autoarray.inversion.inversion.interferometer.sparse import (
     InversionInterferometerSparse,
 )
+from autoarray.inversion.inversion.interferometer.abstract import (
+    real_space_mask_from,
+)
 from autoarray.inversion.inversion.interferometer_numba.sparse import (
     InversionInterferometerSparseNumba,
 )
+from autoarray.inversion.inversion.interferometer_numba import (
+    sparse as interferometer_numba_sparse,
+)
 from autoarray.inversion.inversion.interferometer_numba import (
     inversion_interferometer_numba_util,
 )
@@ -218,6 +224,10 @@ def inversion_interferometer_from(
                 dataset.sparse_operator.y_shape,
                 dataset.sparse_operator.x_shape,
             ),
+            # Only the NumPy cost model reads the mask, so the JAX path never resolves it.
+            real_space_mask=(
+                real_space_mask_from(dataset=dataset) if xp is np else None
+            ),
         ):
             return InversionInterferometerSparseNumba(
                 dataset=dataset,
@@ -248,6 +258,7 @@ def _use_interferometer_numba(
     settings: Settings = None,
     xp=np,
     extent_shape: Optional[Tuple[int, int]] = None,
+    real_space_mask=None,
 ) -> bool:
     """
     Whether an interferometer inversion is routed to the numba `direct_conv` curvature
@@ -281,7 +292,14 @@ def _use_interferometer_numba(
         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.
+        be evaluated and the FFT route is used. When numba is serial and
+        `interferometer_direct_conv_symmetric_max_work_fraction` is positive, the
+        symmetric kernel's work fraction is computed from the mapper's triplets and
+        `real_space_mask` (`symmetric_work_fraction_from`), and if it is at or below that
+        threshold -- i.e. `InversionInterferometerSparseNumba` will run the cheaper
+        symmetric kernel -- it is passed to the cost model, which scales the predicted
+        `direct_conv` time by it. Without a `real_space_mask` whose extent matches
+        `extent_shape`, the full kernel's cost is predicted.
     - `import numba` succeeds.
 
     NaN mapping weights (an overflowing or degenerate `DelaunayNN` Sibson row) are
@@ -304,6 +322,9 @@ def _use_interferometer_numba(
         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.
+    real_space_mask
+        The real-space mask the inversion reconstructs on (`real_space_mask_from`), used
+        only by the default cost model for the symmetric kernel's work fraction.
     """
     if xp is not np:
         return False
@@ -360,6 +381,12 @@ def _use_interferometer_numba(
                     source_pixels=int(mapper.params),
                     nnz=int(np.asarray(mapper.pix_sizes_for_sub_slim_index).sum()),
                     extent_shape=extent_shape,
+                    work_fraction=_symmetric_work_fraction_for_route(
+                        mapper=mapper,
+                        settings=settings,
+                        extent_shape=extent_shape,
+                        real_space_mask=real_space_mask,
+                    ),
                 )
             )
 
@@ -372,3 +399,56 @@ def _use_interferometer_numba(
         return False
 
     return True
+
+
+def _symmetric_work_fraction_for_route(
+    mapper, settings: Settings, extent_shape, real_space_mask
+) -> Optional[float]:
+    """
+    The work fraction to give the route cost model: the symmetric kernel's
+    `work_fraction` when `InversionInterferometerSparseNumba` would run it (serial numba,
+    the gate positive, finite weights and the fraction at or below it), else `None`
+    (predict the full kernel).
+
+    Computed from the mapper's triplets and the mask's extent indices, with the same
+    statistics the inversion's `symmetric_restriction_inputs` uses, so the gate here and
+    the kernel choice there cannot disagree.
+    """
+    threshold = settings.interferometer_direct_conv_symmetric_max_work_fraction
+
+    if (
+        threshold <= 0.0
+        or real_space_mask is None
+        or interferometer_numba_sparse._numba_parallel()
+    ):
+        return None
+
+    if tuple(int(n) for n in real_space_mask.shape_native_masked_pixels) != tuple(
+        int(n) for n in extent_shape
+    ):
+        return None
+
+    # Non-finite weights make the inversion fall back to the full kernel.
+    sizes = np.asarray(mapper.pix_sizes_for_sub_slim_index)
+    weights = np.asarray(mapper.pix_weights_for_sub_slim_index)
+    valid = np.arange(weights.shape[1])[None, :] < sizes[:, None]
+
+    if not np.all(np.isfinite(weights[valid])):
+        return None
+
+    extent_index_for_masked_pixel = np.asarray(
+        real_space_mask.extent_index_for_masked_pixel
+    )[np.asarray(mapper.slim_index_for_sub_slim_index)]
+
+    work_fraction = inversion_interferometer_numba_util.symmetric_work_fraction_from(
+        pix_indexes_for_sub_slim_index=mapper.pix_indexes_for_sub_slim_index,
+        pix_sizes_for_sub_slim_index=mapper.pix_sizes_for_sub_slim_index,
+        extent_index_for_masked_pixel=extent_index_for_masked_pixel,
+        extent_shape=extent_shape,
+        pix_pixels=int(mapper.params),
+    )
+
+    if work_fraction > threshold:
+        return None
+
+    return work_fraction

autoarray/inversion/inversion/interferometer/abstract.py

@@ -15,6 +15,28 @@
 from autoarray.inversion.inversion import inversion_util
 
 
+def real_space_mask_from(dataset) -> Mask2D:
+    """
+    The real-space mask an interferometer inversion of `dataset` reconstructs on: an
+    `Interferometer`'s `real_space_mask`, else a transformer's `real_space_mask`, else the
+    dataset's own `mask` (a `DatasetInterface`'s `grids.lp.mask`).
+
+    A module function so the inversion factory can read the same mask before it builds an
+    inversion (`AbstractInversionInterferometer.mask` returns this).
+    """
+    real_space_mask = getattr(dataset, "real_space_mask", None)
+
+    if real_space_mask is not None:
+        return real_space_mask
+
+    transformer = getattr(dataset, "transformer", None)
+
+    if transformer is not None:
+        return transformer.real_space_mask
+
+    return dataset.mask
+
+
 class AbstractInversionInterferometer(AbstractInversion):
     def __init__(
         self,
@@ -70,17 +92,7 @@ def mask(self) -> Mask2D:
         `real_space_mask`, else the dataset's own `mask` (a `DatasetInterface`'s
         `grids.lp.mask`).
         """
-        real_space_mask = getattr(self.dataset, "real_space_mask", None)
-
-        if real_space_mask is not None:
-            return real_space_mask
-
-        transformer = getattr(self.dataset, "transformer", None)
-
-        if transformer is not None:
-            return transformer.real_space_mask
-
-        return self.dataset.mask
+        return real_space_mask_from(dataset=self.dataset)
 
     @cached_property
     def operated_mapping_matrix_list(self) -> List[np.ndarray]: