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

autoarray/inversion/inversion/inversion_util.py

@@ -274,9 +274,10 @@ def reconstruction_positive_only_from(
     """
     if xp.__name__.startswith("jax"):
 
-        import jaxnnls
         from autoconf import conf
 
+        from autoarray.util.jax_nnls import solve_nnls_primal
+
         try:
             use_jacobi = conf.instance["general"]["inversion"][
                 "nnls_jacobi_preconditioning"
@@ -301,6 +302,14 @@ def reconstruction_positive_only_from(
             # magnitude in noise.
             target_kappa = 1.0e-11
 
+        # 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:
             # Ill-conditioned Q makes jaxnnls's relaxed-KKT backward pass
             # produce NaN gradients. Rescale Q so its diagonal is unit:
@@ -312,12 +321,22 @@ def reconstruction_positive_only_from(
             Q_pc = (curvature_reg_matrix * D[:, None]) * D[None, :]
             q_pc = data_vector * D
             return (
-                jaxnnls.solve_nnls_primal(Q_pc, q_pc, target_kappa=target_kappa)
+                solve_nnls_primal(
+                    Q_pc,
+                    q_pc,
+                    target_kappa=target_kappa,
+                    solver_tol=solver_tol,
+                    max_iter=max_iter,
+                )
                 * D
             )
 
-        return jaxnnls.solve_nnls_primal(
-            curvature_reg_matrix, data_vector, target_kappa=target_kappa
+        return solve_nnls_primal(
+            curvature_reg_matrix,
+            data_vector,
+            target_kappa=target_kappa,
+            solver_tol=solver_tol,
+            max_iter=max_iter,
         )
 
     try:

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

@@ -0,0 +1,120 @@
+"""
+Configurable driver for the jaxnnls primal-dual interior-point NNLS solver.
+
+jaxnnls hard-codes its convergence tolerance (``n * eps * 5e3``, capped at
+1e-2) and iteration cap (``MAX_ITER = 50``) inside ``pdip.solve_nnls``, and
+neither is exposed through ``solve_nnls_primal``. This module re-implements
+only that ``while_loop`` driver with both knobs as arguments, reusing every
+jaxnnls building block (``initialize``, ``pdip_pc_step``,
+``solve_relaxed_nnls``, ``diff_nnls``) and the same custom-vjp relaxed-KKT
+backward pass. With the knobs at their defaults the solve and its gradients
+are identical to upstream jaxnnls.
+
+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
+time with a log-evidence shift of order 1e-8. Under ``vmap`` the while_loop
+runs until the slowest lane converges, so ``max_iter`` also caps the
+worst-case batched cost.
+
+JAX is imported inside functions, never at module level (see
+``docs/agents/jax_and_decorators.md``); this module must only be imported
+on the ``xp=jnp`` path. The solver knobs are static closure parameters —
+the ``lru_cache`` returns the same function object for repeated settings so
+``jax.jit`` tracing caches hit (a fresh closure per call would cache-bust).
+"""
+
+from functools import lru_cache
+
+
+def solve_nnls(Q, q, solver_tol=None, max_iter=50):
+    """
+    Solve the non-negative least squares problem with the jaxnnls PDIP
+    algorithm, with configurable convergence tolerance and iteration cap.
+
+    Mirrors ``jaxnnls.pdip.solve_nnls`` exactly at the default settings.
+
+    Parameters
+    ----------
+    Q
+        The (n, n) positive definite matrix (the curvature-regularization
+        matrix of an inversion).
+    q
+        The (n,) vector (the data vector of an inversion).
+    solver_tol
+        Infinity-norm KKT residual below which the solve is converged.
+        ``None`` (default) reproduces jaxnnls's own tolerance
+        ``min(n * eps * 5e3, 1e-2)``.
+    max_iter
+        Maximum number of PDIP iterations (jaxnnls hard-codes 50).
+
+    Returns
+    -------
+    The tuple (x, s, z, converged, pdip_iter) of the primal solution, slack
+    and dual variables, convergence flag and iteration count.
+    """
+    import jax
+    import jax.numpy as jnp
+    from jaxnnls.pdip import EPSILON, initialize, pdip_pc_step
+
+    x, s, z = initialize(Q, q)
+
+    if solver_tol is None:
+        solver_tol = jax.lax.min(Q.shape[0] * EPSILON, 1e-2)
+    solver_tol = jnp.asarray(solver_tol, dtype=q.dtype)
+
+    def converged_check(inputs):
+        _, _, _, _, _, _, converged, pdip_iter = inputs
+        return jnp.logical_and(pdip_iter < max_iter, converged == 0)
+
+    init_inputs = (Q, q, x, s, z, solver_tol, 0, 0)
+    outputs = jax.lax.while_loop(converged_check, pdip_pc_step, init_inputs)
+    _, _, x, s, z, _, converged, pdip_iter = outputs
+    return x, s, z, converged, pdip_iter
+
+
+@lru_cache(maxsize=None)
+def _solve_nnls_primal_with(target_kappa, solver_tol, max_iter):
+    """
+    Build (and cache) the differentiable primal solver for one static
+    setting of the knobs. The returned function takes only (Q, q), so the
+    custom-vjp backward pass returns exactly (dQ, dq).
+    """
+    import jax
+    from jaxnnls.diff_qp import diff_nnls
+    from jaxnnls.pdip_relaxed import solve_relaxed_nnls
+
+    def primal(Q, q):
+        return solve_nnls(Q, q, solver_tol=solver_tol, max_iter=max_iter)[0]
+
+    def forward(Q, q):
+        x, s, z, _, _ = solve_nnls(Q, q, solver_tol=solver_tol, max_iter=max_iter)
+        # Relax the solution with vanilla Newton steps on the relaxed KKT
+        # conditions; only the backward pass consumes the relaxed variables.
+        xr, sr, zr, _, _ = solve_relaxed_nnls(
+            Q, q, x, s, z, target_kappa=target_kappa
+        )
+        return x, (Q, xr, sr, zr)
+
+    def backward(res, input_grad):
+        Q, xr, sr, zr = res
+        return diff_nnls(Q, xr, sr, zr, input_grad)
+
+    primal = jax.custom_vjp(primal)
+    primal.defvjp(forward, backward)
+    return primal
+
+
+def solve_nnls_primal(Q, q, target_kappa=1e-3, solver_tol=None, max_iter=50):
+    """
+    Solve the non-negative least squares problem, differentiable via the
+    relaxed-KKT implicit backward pass.
+
+    Drop-in replacement for ``jaxnnls.solve_nnls_primal`` with two extra
+    knobs; at their defaults (``solver_tol=None``, ``max_iter=50``) the
+    forward solve and gradients are identical to upstream.
+    """
+    return _solve_nnls_primal_with(target_kappa, solver_tol, max_iter)(Q, q)

test_autoarray/util/test_jax_nnls.py

@@ -0,0 +1,60 @@
+import importlib
+import sys
+
+import numpy as np
+import pytest
+
+import autoarray as aa
+
+
+def test__jax_nnls_module_never_imports_jax_at_module_level(monkeypatch):
+    # Library unit tests are NumPy-only: importing the module must succeed
+    # even when jax / jaxnnls are unimportable. A None entry in sys.modules
+    # makes any `import jax` raise ImportError, so a module-level import
+    # would fail this reload.
+    monkeypatch.setitem(sys.modules, "jax", None)
+    monkeypatch.setitem(sys.modules, "jaxnnls", None)
+
+    module = importlib.reload(importlib.import_module("autoarray.util.jax_nnls"))
+
+    assert hasattr(module, "solve_nnls")
+    assert hasattr(module, "solve_nnls_primal")
+
+
+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 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)
+
+    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, 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]]
+    )
+
+    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)