feat: NNLS solver knobs via Settings (nnls_solver_tol / nnls_max_iter), default off (#369)

Rework of the parked d8a1c84: the knobs move from general.yaml config keys
to per-fit Settings attributes (defaults None = jaxnnls's own tolerance
formula and 50-iteration cap — behaviour identical when unset). The vendored
autoarray/util/jax_nnls.py driver is unchanged: reuses all jaxnnls building
blocks + the relaxed-KKT custom-vjp backward; defaults validated bit-identical
to upstream in solutions and gradients on real production systems.

Measured on the real HST pixelization+MGE systems (PyAutoArray#369):
Settings(nnls_solver_tol=1e-6) saves ~15-20% of solve time with rel
delta-objective 3.8e-13 (delta log-evidence ~1e-8); nnls_max_iter also caps
the vmap worst-case lane. Full ledger: autolens_profiling
results/notes/nnls_solver_ledger.md.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01QgtjXWS2iJegXMTDU4GHth

Author: 2 people

autoarray/config/general.yaml

@@ -8,8 +8,6 @@ inversion:
   use_border_relocator: false          # If True, by default a pixelization's border is used to relocate all pixels outside its border to the border.
   nnls_jacobi_preconditioning: true   # If True (default), the curvature matrix passed to jaxnnls.solve_nnls_primal is Jacobi-preconditioned (D Q D y = D q, x = D y). Fixes NaN backward-pass gradients on ill-conditioned Q and roughly halves forward solve time. Set False to restore the raw unpreconditioned solve.
   nnls_target_kappa: 1.0e-11          # Central-path relaxation parameter passed to jaxnnls.solve_nnls_primal. Larger values smooth the relaxed-KKT backward pass and prevent NaN gradients on ill-conditioned Q; smaller values tighten the primal solve. Verified finite gradients across all MGE/rectangular/delaunay pipelines (imaging + interferometer) with scale invariance over 5 orders of magnitude in noise. jaxnnls's own default (1e-3) is too aggressive for the backward pass.
-  nnls_solver_tol: null               # Convergence tolerance (infinity-norm KKT residual) of the JAX NNLS interior-point solve. null reproduces jaxnnls's own tolerance min(n * eps * 5e3, 1e-2) ~ 1.7e-9 at n=1500 fp64. Each solver iteration is a fresh dense Cholesky of the (n, n) system, so looser tolerances buy real speed: 1e-6 saves ~15-20% of solve time with a log-evidence shift of order 1e-8 (PyAutoArray#369).
-  nnls_max_iter: 50                   # Iteration cap of the JAX NNLS interior-point solve (jaxnnls hard-codes 50; production HST pixelization+MGE systems converge in ~19-21 iterations). Under vmap the solve runs until the slowest lane in the batch converges, so this also caps the worst-case batched cost.
   reconstruction_vmax_factor: 0.5     # Plots of an Inversion's reconstruction use the reconstructed data's bright value multiplied by this factor.
 numba:
   use_numba: true

autoarray/inversion/inversion/inversion_util.py

@@ -302,17 +302,12 @@ def reconstruction_positive_only_from(
             # magnitude in noise.
             target_kappa = 1.0e-11
 
-        try:
-            solver_tol = conf.instance["general"]["inversion"]["nnls_solver_tol"]
-        except KeyError:
-            # Workspaces ship their own general.yaml that shadows autoarray's;
-            # None reproduces jaxnnls's own tolerance min(n * eps * 5e3, 1e-2).
-            solver_tol = None
-
-        try:
-            max_iter = conf.instance["general"]["inversion"]["nnls_max_iter"]
-        except KeyError:
-            # jaxnnls's own hard-coded iteration cap.
+        # Per-fit solver knobs from the Settings class; the defaults (None)
+        # reproduce jaxnnls's own tolerance min(n * eps * 5e3, 1e-2) and its
+        # hard-coded 50-iteration cap exactly.
+        solver_tol = settings.nnls_solver_tol if settings is not None else None
+        max_iter = settings.nnls_max_iter if settings is not None else None
+        if max_iter is None:
             max_iter = 50
 
         if use_jacobi:

autoarray/settings.py

@@ -15,6 +15,8 @@ def __init__(
         use_edge_zeroed_pixels: Optional[bool] = None,
         use_border_relocator: Optional[bool] = None,
         no_regularization_add_to_curvature_diag_value: float = None,
+        nnls_solver_tol: Optional[float] = None,
+        nnls_max_iter: Optional[int] = None,
     ):
         """
         The settings of an Inversion, customizing how a linear set of equations are solved for.
@@ -80,8 +82,23 @@ def __init__(
         no_regularization_add_to_curvature_diag_value
             If a linear func object does not have a corresponding regularization, this value is added to its
             diagonal entries of the curvature regularization matrix to ensure the matrix is positive-definite.
+        nnls_solver_tol
+            Convergence tolerance (infinity-norm KKT residual) of the JAX positive-only (NNLS) interior-point
+            solve. `None` (default) uses jaxnnls's own tolerance ``min(n * eps * 5e3, 1e-2)`` (~1.7e-9 at
+            n=1500 fp64) — behaviour is identical to not having this setting. Each solver iteration is a fresh
+            dense Cholesky of the (n, n) system, so looser tolerances buy real speed: ``1e-6`` saves ~15-20% of
+            solve time with a log-evidence shift of order 1e-8 on production HST pixelization+MGE fits
+            (measured in https://github.com/PyAutoLabs/PyAutoArray/issues/369). Only the JAX (`xp=jnp`) path
+            honors this; the NumPy fnnls path is unaffected.
+        nnls_max_iter
+            Iteration cap of the JAX positive-only (NNLS) interior-point solve. `None` (default) uses
+            jaxnnls's own hard-coded cap of 50 (production HST pixelization+MGE systems converge in ~19-21
+            iterations). Under `vmap` the solve runs until the slowest lane in the batch converges, so this
+            also caps the worst-case batched cost. Only the JAX (`xp=jnp`) path honors this.
         """
         self.use_mixed_precision = use_mixed_precision
+        self.nnls_solver_tol = nnls_solver_tol
+        self.nnls_max_iter = nnls_max_iter
         self._use_positive_only_solver = use_positive_only_solver
         self._use_edge_zeroed_pixels = use_edge_zeroed_pixels
         self._use_border_relocator = use_border_relocator

autoarray/util/jax_nnls.py

@@ -10,8 +10,9 @@
 backward pass. With the knobs at their defaults the solve and its gradients
 are identical to upstream jaxnnls.
 
-The knobs are read from ``general.yaml -> inversion -> nnls_solver_tol /
-nnls_max_iter`` by ``inversion_util.reconstruction_positive_only_from``.
+The knobs are exposed per-fit through the ``Settings`` class
+(``nnls_solver_tol`` / ``nnls_max_iter``, defaults ``None`` = upstream
+behaviour) and read by ``inversion_util.reconstruction_positive_only_from``.
 Measured motivation (PyAutoArray#369, real HST pixelization+MGE systems):
 each PDIP iteration is a fresh dense Cholesky of the (n, n) KKT system, so
 iterations are the whole cost; ``solver_tol=1e-6`` saves ~15-20% of solve

test_autoarray/util/test_jax_nnls.py

@@ -1,10 +1,8 @@
 import importlib
 import sys
-from pathlib import Path
 
 import numpy as np
 import pytest
-import yaml
 
 import autoarray as aa
 
@@ -23,36 +21,40 @@ def test__jax_nnls_module_never_imports_jax_at_module_level(monkeypatch):
     assert hasattr(module, "solve_nnls_primal")
 
 
-def test__nnls_solver_config_keys_have_upstream_default_values():
-    # The packaged general.yaml defaults must reproduce jaxnnls's own
-    # hard-coded behaviour (tolerance formula and MAX_ITER = 50) so that
-    # installing this change alone does not alter any solve. The packaged
-    # file is parsed directly because the test suite (like workspaces)
-    # pushes its own shadowing config which omits the nnls keys — the
-    # solver reads missing keys via try/except fallbacks to these values.
-    packaged = Path(aa.__file__).parent / "config" / "general.yaml"
-    inversion = yaml.safe_load(packaged.read_text())["inversion"]
+def test__settings_nnls_knobs_default_off():
+    # The Settings defaults must reproduce jaxnnls's own hard-coded
+    # behaviour (tolerance formula, MAX_ITER = 50) so a default Settings
+    # object — or none at all — does not alter any solve.
+    settings = aa.Settings()
 
-    assert inversion["nnls_solver_tol"] is None
-    assert inversion["nnls_max_iter"] == 50
+    assert settings.nnls_solver_tol is None
+    assert settings.nnls_max_iter is None
 
+    settings = aa.Settings(nnls_solver_tol=1e-6, nnls_max_iter=30)
 
-def test__reconstruction_positive_only_from__numpy_path():
+    assert settings.nnls_solver_tol == 1e-6
+    assert settings.nnls_max_iter == 30
+
+
+def test__reconstruction_positive_only_from__numpy_path_ignores_knobs():
     # The xp=np path (fnnls_cholesky) is untouched by the JAX solver knobs;
     # solve a small system whose unconstrained solution has a negative
-    # component and check the non-negative solution.
+    # component and check the non-negative solution, with and without
+    # knob-carrying settings.
     data_vector = np.array([1.0, 1.0, 2.0])
 
     curvature_reg_matrix = np.array(
         [[2.0, 1.0, 0.0], [1.0, 3.0, 1.0], [0.0, 1.0, 1.0]]
     )
 
-    reconstruction = aa.util.inversion.reconstruction_positive_only_from(
-        data_vector=data_vector,
-        curvature_reg_matrix=curvature_reg_matrix,
-    )
-
-    assert np.all(reconstruction >= 0.0)
-    # Unconstrained solution is [1, -1, 3]; the NNLS solution zeroes the
-    # negative component and re-solves the free ones.
-    assert reconstruction == pytest.approx(np.array([0.5, 0.0, 2.0]), 1.0e-4)
+    for settings in [None, aa.Settings(nnls_solver_tol=1e-6, nnls_max_iter=30)]:
+        reconstruction = aa.util.inversion.reconstruction_positive_only_from(
+            data_vector=data_vector,
+            curvature_reg_matrix=curvature_reg_matrix,
+            settings=settings,
+        )
+
+        assert np.all(reconstruction >= 0.0)
+        # Unconstrained solution is [1, -1, 3]; the NNLS solution zeroes the
+        # negative component and re-solves the free ones.
+        assert reconstruction == pytest.approx(np.array([0.5, 0.0, 2.0]), 1.0e-4)