diff --git a/autoarray/inversion/inversion/inversion_util.py b/autoarray/inversion/inversion/inversion_util.py index f3ee923e..d02c66d9 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,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: 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 new file mode 100644 index 00000000..e7abd179 --- /dev/null +++ b/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) diff --git a/test_autoarray/util/test_jax_nnls.py b/test_autoarray/util/test_jax_nnls.py new file mode 100644 index 00000000..796bd503 --- /dev/null +++ b/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)