|
| 1 | +""" |
| 2 | +Configurable driver for the jaxnnls primal-dual interior-point NNLS solver. |
| 3 | +
|
| 4 | +jaxnnls hard-codes its convergence tolerance (``n * eps * 5e3``, capped at |
| 5 | +1e-2) and iteration cap (``MAX_ITER = 50``) inside ``pdip.solve_nnls``, and |
| 6 | +neither is exposed through ``solve_nnls_primal``. This module re-implements |
| 7 | +only that ``while_loop`` driver with both knobs as arguments, reusing every |
| 8 | +jaxnnls building block (``initialize``, ``pdip_pc_step``, |
| 9 | +``solve_relaxed_nnls``, ``diff_nnls``) and the same custom-vjp relaxed-KKT |
| 10 | +backward pass. With the knobs at their defaults the solve and its gradients |
| 11 | +are identical to upstream jaxnnls. |
| 12 | +
|
| 13 | +The knobs are read from ``general.yaml -> inversion -> nnls_solver_tol / |
| 14 | +nnls_max_iter`` by ``inversion_util.reconstruction_positive_only_from``. |
| 15 | +Measured motivation (PyAutoArray#369, real HST pixelization+MGE systems): |
| 16 | +each PDIP iteration is a fresh dense Cholesky of the (n, n) KKT system, so |
| 17 | +iterations are the whole cost; ``solver_tol=1e-6`` saves ~15-20% of solve |
| 18 | +time with a log-evidence shift of order 1e-8. Under ``vmap`` the while_loop |
| 19 | +runs until the slowest lane converges, so ``max_iter`` also caps the |
| 20 | +worst-case batched cost. |
| 21 | +
|
| 22 | +JAX is imported inside functions, never at module level (see |
| 23 | +``docs/agents/jax_and_decorators.md``); this module must only be imported |
| 24 | +on the ``xp=jnp`` path. The solver knobs are static closure parameters — |
| 25 | +the ``lru_cache`` returns the same function object for repeated settings so |
| 26 | +``jax.jit`` tracing caches hit (a fresh closure per call would cache-bust). |
| 27 | +""" |
| 28 | + |
| 29 | +from functools import lru_cache |
| 30 | + |
| 31 | + |
| 32 | +def solve_nnls(Q, q, solver_tol=None, max_iter=50): |
| 33 | + """ |
| 34 | + Solve the non-negative least squares problem with the jaxnnls PDIP |
| 35 | + algorithm, with configurable convergence tolerance and iteration cap. |
| 36 | +
|
| 37 | + Mirrors ``jaxnnls.pdip.solve_nnls`` exactly at the default settings. |
| 38 | +
|
| 39 | + Parameters |
| 40 | + ---------- |
| 41 | + Q |
| 42 | + The (n, n) positive definite matrix (the curvature-regularization |
| 43 | + matrix of an inversion). |
| 44 | + q |
| 45 | + The (n,) vector (the data vector of an inversion). |
| 46 | + solver_tol |
| 47 | + Infinity-norm KKT residual below which the solve is converged. |
| 48 | + ``None`` (default) reproduces jaxnnls's own tolerance |
| 49 | + ``min(n * eps * 5e3, 1e-2)``. |
| 50 | + max_iter |
| 51 | + Maximum number of PDIP iterations (jaxnnls hard-codes 50). |
| 52 | +
|
| 53 | + Returns |
| 54 | + ------- |
| 55 | + The tuple (x, s, z, converged, pdip_iter) of the primal solution, slack |
| 56 | + and dual variables, convergence flag and iteration count. |
| 57 | + """ |
| 58 | + import jax |
| 59 | + import jax.numpy as jnp |
| 60 | + from jaxnnls.pdip import EPSILON, initialize, pdip_pc_step |
| 61 | + |
| 62 | + x, s, z = initialize(Q, q) |
| 63 | + |
| 64 | + if solver_tol is None: |
| 65 | + solver_tol = jax.lax.min(Q.shape[0] * EPSILON, 1e-2) |
| 66 | + solver_tol = jnp.asarray(solver_tol, dtype=q.dtype) |
| 67 | + |
| 68 | + def converged_check(inputs): |
| 69 | + _, _, _, _, _, _, converged, pdip_iter = inputs |
| 70 | + return jnp.logical_and(pdip_iter < max_iter, converged == 0) |
| 71 | + |
| 72 | + init_inputs = (Q, q, x, s, z, solver_tol, 0, 0) |
| 73 | + outputs = jax.lax.while_loop(converged_check, pdip_pc_step, init_inputs) |
| 74 | + _, _, x, s, z, _, converged, pdip_iter = outputs |
| 75 | + return x, s, z, converged, pdip_iter |
| 76 | + |
| 77 | + |
| 78 | +@lru_cache(maxsize=None) |
| 79 | +def _solve_nnls_primal_with(target_kappa, solver_tol, max_iter): |
| 80 | + """ |
| 81 | + Build (and cache) the differentiable primal solver for one static |
| 82 | + setting of the knobs. The returned function takes only (Q, q), so the |
| 83 | + custom-vjp backward pass returns exactly (dQ, dq). |
| 84 | + """ |
| 85 | + import jax |
| 86 | + from jaxnnls.diff_qp import diff_nnls |
| 87 | + from jaxnnls.pdip_relaxed import solve_relaxed_nnls |
| 88 | + |
| 89 | + def primal(Q, q): |
| 90 | + return solve_nnls(Q, q, solver_tol=solver_tol, max_iter=max_iter)[0] |
| 91 | + |
| 92 | + def forward(Q, q): |
| 93 | + x, s, z, _, _ = solve_nnls(Q, q, solver_tol=solver_tol, max_iter=max_iter) |
| 94 | + # Relax the solution with vanilla Newton steps on the relaxed KKT |
| 95 | + # conditions; only the backward pass consumes the relaxed variables. |
| 96 | + xr, sr, zr, _, _ = solve_relaxed_nnls( |
| 97 | + Q, q, x, s, z, target_kappa=target_kappa |
| 98 | + ) |
| 99 | + return x, (Q, xr, sr, zr) |
| 100 | + |
| 101 | + def backward(res, input_grad): |
| 102 | + Q, xr, sr, zr = res |
| 103 | + return diff_nnls(Q, xr, sr, zr, input_grad) |
| 104 | + |
| 105 | + primal = jax.custom_vjp(primal) |
| 106 | + primal.defvjp(forward, backward) |
| 107 | + return primal |
| 108 | + |
| 109 | + |
| 110 | +def solve_nnls_primal(Q, q, target_kappa=1e-3, solver_tol=None, max_iter=50): |
| 111 | + """ |
| 112 | + Solve the non-negative least squares problem, differentiable via the |
| 113 | + relaxed-KKT implicit backward pass. |
| 114 | +
|
| 115 | + Drop-in replacement for ``jaxnnls.solve_nnls_primal`` with two extra |
| 116 | + knobs; at their defaults (``solver_tol=None``, ``max_iter=50``) the |
| 117 | + forward solve and gradients are identical to upstream. |
| 118 | + """ |
| 119 | + return _solve_nnls_primal_with(target_kappa, solver_tol, max_iter)(Q, q) |
0 commit comments