From d8a1c84ac6fc6de772cb8359d58cbe4e029eaef0 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Thu, 9 Jul 2026 22:18:03 +0100 Subject: [PATCH 1/2] =?UTF-8?q?feat:=20configurable=20NNLS=20solver=20knob?= =?UTF-8?q?s=20(nnls=5Fsolver=5Ftol=20/=20nnls=5Fmax=5Fiter)=20=E2=80=94?= =?UTF-8?q?=20PARKED,=20not=20for=20merge=20(#369)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Vendored jaxnnls while_loop driver with configurable tolerance and iteration cap; defaults bit-identical to upstream (solutions and gradients, validated on real production systems). Parked at user request 2026-07-09 — the measured win (~15-20% of solve) was judged not worth the new solver code path; this branch preserves the validated implementation should that change. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01QgtjXWS2iJegXMTDU4GHth --- autoarray/config/general.yaml | 2 + .../inversion/inversion/inversion_util.py | 32 ++++- autoarray/util/jax_nnls.py | 119 ++++++++++++++++++ test_autoarray/util/test_jax_nnls.py | 58 +++++++++ 4 files changed, 207 insertions(+), 4 deletions(-) create mode 100644 autoarray/util/jax_nnls.py create mode 100644 test_autoarray/util/test_jax_nnls.py diff --git a/autoarray/config/general.yaml b/autoarray/config/general.yaml index 53f42b9e..c4c38508 100644 --- a/autoarray/config/general.yaml +++ b/autoarray/config/general.yaml @@ -8,6 +8,8 @@ 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 diff --git a/autoarray/inversion/inversion/inversion_util.py b/autoarray/inversion/inversion/inversion_util.py index f3ee923e..a7d8b274 100644 --- a/autoarray/inversion/inversion/inversion_util.py +++ b/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,19 @@ 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. + 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 +326,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: diff --git a/autoarray/util/jax_nnls.py b/autoarray/util/jax_nnls.py new file mode 100644 index 00000000..6d238696 --- /dev/null +++ b/autoarray/util/jax_nnls.py @@ -0,0 +1,119 @@ +""" +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 read from ``general.yaml -> inversion -> nnls_solver_tol / +nnls_max_iter`` 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) diff --git a/test_autoarray/util/test_jax_nnls.py b/test_autoarray/util/test_jax_nnls.py new file mode 100644 index 00000000..636650c6 --- /dev/null +++ b/test_autoarray/util/test_jax_nnls.py @@ -0,0 +1,58 @@ +import importlib +import sys +from pathlib import Path + +import numpy as np +import pytest +import yaml + +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__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"] + + assert inversion["nnls_solver_tol"] is None + assert inversion["nnls_max_iter"] == 50 + + +def test__reconstruction_positive_only_from__numpy_path(): + # 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. + 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) From 8771f29b22dbb8e7c3d68abc85b87de876519654 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Thu, 9 Jul 2026 22:38:25 +0100 Subject: [PATCH 2/2] feat: NNLS solver knobs via Settings (nnls_solver_tol / nnls_max_iter), default off (#369) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Rework of the parked d8a1c84a: 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 Claude-Session: https://claude.ai/code/session_01QgtjXWS2iJegXMTDU4GHth --- autoarray/config/general.yaml | 2 - .../inversion/inversion/inversion_util.py | 17 +++---- autoarray/settings.py | 17 +++++++ autoarray/util/jax_nnls.py | 5 +- test_autoarray/util/test_jax_nnls.py | 50 ++++++++++--------- 5 files changed, 52 insertions(+), 39 deletions(-) diff --git a/autoarray/config/general.yaml b/autoarray/config/general.yaml index c4c38508..53f42b9e 100644 --- a/autoarray/config/general.yaml +++ b/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 diff --git a/autoarray/inversion/inversion/inversion_util.py b/autoarray/inversion/inversion/inversion_util.py index a7d8b274..d02c66d9 100644 --- a/autoarray/inversion/inversion/inversion_util.py +++ b/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: diff --git a/autoarray/settings.py b/autoarray/settings.py index 102e8a0c..76b8f20e 100644 --- a/autoarray/settings.py +++ b/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 diff --git a/autoarray/util/jax_nnls.py b/autoarray/util/jax_nnls.py index 6d238696..e7abd179 100644 --- a/autoarray/util/jax_nnls.py +++ b/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 diff --git a/test_autoarray/util/test_jax_nnls.py b/test_autoarray/util/test_jax_nnls.py index 636650c6..796bd503 100644 --- a/test_autoarray/util/test_jax_nnls.py +++ b/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)