Skip to content

Commit 7a89e19

Browse files
author
James Nightingale
authored
Merge pull request #595 from PyAutoLabs/feature/raw-pdip-forward-polish
fix: return the polished raw-forward PDIP iterate as the forward value (#594)
2 parents bd03e09 + 42c5235 commit 7a89e19

7 files changed

Lines changed: 412 additions & 45 deletions

File tree

‎autoarray/config/general.yaml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ inversion:
88
use_border_relocator: false # If True, by default a pixelization's border is used to relocate all pixels outside its border to the border.
99
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.
1010
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. The relaxed solve must start from an iterate whose complementarity s*z is not far above this value; the "raw" no-mapper mode polishes its forward iterate to ensure that (PyAutoArray#573).
11-
nnls_preconditioning_no_mapper: raw # How the JAX positive-only PDIP solve scales inversions with NO mapper (linear light profiles / MGE only). "raw" (default) runs the forward solve on the un-preconditioned system with a data-scaled KKT tolerance (1e-2 * n * eps * max(1, max|data_vector|)) and keeps the Jacobi-space relaxed-KKT gradient, whose relaxed solve starts from the forward iterate polished by <= 10 tight warm-started PDIP iterations on the Jacobi system (without the polish the loose forward tolerance leaves s*z far above nnls_target_kappa and the relaxed solve diverged to NaN gradients on 4/16 jax_grad/mge.py points, PyAutoArray#573); "jacobi" uses the Jacobi-preconditioned solve. Jacobi scaling of signal-free MGE columns (diagonal = the no-regularization floor) made the PDIP dual diverge on 14/48 SLaM source_lp[1] points (PyAutoArray#571). Inversions with a mapper always use jacobi; the NumPy path is unaffected.
11+
nnls_preconditioning_no_mapper: raw # How the JAX positive-only PDIP solve scales inversions with NO mapper (linear light profiles / MGE only). "raw" (default) runs the forward solve on the un-preconditioned system with a data-scaled KKT tolerance (1e-2 * n * eps * max(1, max|data_vector|)), polishes that iterate with <= 10 tight warm-started PDIP iterations on the Jacobi system and returns the polished iterate (the loose tolerance also judges complementarity, so unpolished it left 11.5 % of the fnnls flux on inactive columns of the euclid vis_lp system; polished <= 3.3e-4 over 81 corpus systems, ~5 extra iterations, PyAutoArray#594), and keeps the Jacobi-space relaxed-KKT gradient started from the same polished point (without the polish the relaxed solve diverged to NaN gradients on 4/16 jax_grad/mge.py points, PyAutoArray#573); "jacobi" uses the Jacobi-preconditioned solve. Jacobi scaling of signal-free MGE columns (diagonal = the no-regularization floor) made the PDIP dual diverge on 14/48 SLaM source_lp[1] points (PyAutoArray#571). Inversions with a mapper always use jacobi; the NumPy path is unaffected.
1212
nnls_warm_start_memo: true # If True (default), the NumPy/numba positive-only (fnnls) solve warm-starts its active set from the previous likelihood evaluation's passive set, cutting active-set iterations on successive sampler evaluations. The NNLS optimum is unique so the reconstruction is unchanged. On by default as of PyAutoArray#498, measured on the euclid+hst Delaunay-1250 fiducial (9.9x / 4.0x fewer active-set iterations on successive evaluations, reconstruction unchanged). Set false, or AUTOARRAY_NNLS_WARM_START=0, to disable. JAX path unaffected.
1313
nnls_warm_start_error_tolerance: 1.5 # Relative quality guard on a warm-start memo seed. Each memo entry remembers the error fraction of the most recent dense-sign-started solve for that key; a seeded solve whose own error fraction exceeds this multiple of that reference is dropped, so the next solve restarts from the dense-sign start and refreshes the reference. Default 1.5 sits above the worst seed/dense error-fraction ratio seen in the PyAutoArray#498 32-cell robustness matrix (1.42), so it is protective against unmeasured regimes rather than flapping. Any non-finite or non-positive value (e.g. .inf) disables the guard. NumPy/numba fnnls path only.
1414
positive_only_solver: pdip # Which solver the JAX (xp=jnp) positive-only reconstruction uses. "pdip" (default) is the jaxnnls interior-point solve; "certified" is the certified active-set solve (budgeted masked-Cholesky passes that stop once the KKT conditions certify, exact implicit gradient, PDIP fallback), measured 1.2-2.6x faster on source-only inversions (PyAutoArray#566). Applied only to mapper-only JAX inversions (MGE / linear light profiles keep PDIP); the NumPy path always uses fnnls. Opt-in until the batched (vmap) policy is measured.

‎autoarray/inversion/inversion/inversion_util.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -368,7 +368,9 @@ def reconstruction_positive_only_from(
368368
byte-identical to before this option existed): the Jacobi-preconditioned solve ``(D Q D) y = D q``
369369
governed by the ``nnls_jacobi_preconditioning`` config key. ``"raw"``: the forward PDIP solve runs on
370370
the un-preconditioned ``(Q, q)`` with the data-scaled tolerance
371-
:func:`autoarray.util.jax_nnls.data_scaled_solver_tol` (or ``settings.nnls_solver_tol`` if set), and the
371+
:func:`autoarray.util.jax_nnls.data_scaled_solver_tol` (or ``settings.nnls_solver_tol`` if set), its iterate
372+
is polished by a few tight PDIP iterations on the Jacobi-scaled system and the polished iterate is returned
373+
(PyAutoArray#594; ``stats["converged"]`` / ``stats["iterations"]`` describe the raw forward solve), and the
372374
gradient is the Jacobi-space relaxed-KKT pass as in ``"jacobi"`` -- see
373375
:func:`autoarray.util.jax_nnls.solve_nnls_primal_raw_forward`. Jacobi scaling makes the signal-free
374376
columns of linear-object-only (MGE) inversions, whose diagonal is only the

‎autoarray/settings.py‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -230,7 +230,11 @@ def __init__(
230230
231231
- ``"raw"`` (default) -- the forward PDIP solve runs on the un-preconditioned system with a
232232
data-scaled KKT tolerance (``1e-2 * n * eps_pdip * max(1, max|data_vector|)``, or
233-
`nnls_solver_tol` if set); the gradient is the same Jacobi-space relaxed-KKT pass as ``"jacobi"``.
233+
`nnls_solver_tol` if set), polished by at most 10 tight warm-started PDIP iterations on the
234+
Jacobi-scaled system; the polished iterate is the returned reconstruction (unpolished, the loose
235+
tolerance left 11.5 % of the fnnls flux on inactive columns of the euclid vis_lp system,
236+
PyAutoArray#594) and the start of the gradient, the same Jacobi-space relaxed-KKT pass as
237+
``"jacobi"``.
234238
On the SLaM `source_lp[1]` MGE model (2 x 20 lens + 20 source Gaussians) Jacobi scaling made 14/48
235239
near-truth points hit the 50-iteration cap with wrong log-likelihoods (signal-free Gaussian columns,
236240
whose diagonal is only `no_regularization_add_to_curvature_diag_value`, become degenerate

‎autoarray/util/jax_nnls.py‎

Lines changed: 83 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -27,11 +27,18 @@
2727
solve on the un-preconditioned system with a data-scaled tolerance
2828
(:func:`data_scaled_solver_tol`) and keeps the Jacobi-space backward pass. It
2929
exists because Jacobi scaling of signal-free MGE columns (diagonal = the
30-
no-regularization floor) makes the PDIP dual diverge. Its backward pass polishes
31-
the mapped forward iterate with a few tight, warm-started PDIP iterations on the
32-
Jacobi system before the relaxed-KKT solve, which otherwise diverges to NaN from
33-
the loose forward tolerance (PyAutoArray#573); :func:`raw_forward_backward_status`
34-
reports that pass's convergence.
30+
no-regularization floor) makes the PDIP dual diverge. The raw iterate is then
31+
*polished* with a few tight, warm-started PDIP iterations on the Jacobi system, and
32+
the polished iterate is both the forward value and the start of the backward
33+
pass's relaxed-KKT solve. The data-scaled stop judges complementarity ``s * z``
34+
against the same loose threshold, so a column fnnls holds at zero can keep
35+
``x ~ tol / z`` while "converged": on the euclid vis_lp system the unpolished
36+
iterate left 11.5 % of the fnnls flux on reference-inactive columns (a +5.8 %
37+
latent source flux), invisible to the objective and the log likelihood; polished,
38+
that is <= 3.3e-4 across 81 corpus systems for ~5 extra iterations
39+
(PyAutoArray#594). The polish first went into the backward pass only, where the
40+
relaxed-KKT solve otherwise diverges to NaN from the loose forward tolerance
41+
(PyAutoArray#573); :func:`raw_forward_backward_status` reports its convergence.
3542
3643
JAX is imported inside functions, never at module level (see
3744
``docs/agents/jax_and_decorators.md``); this module must only be imported
@@ -187,51 +194,75 @@ def solve_nnls_primal(Q, q, target_kappa=1e-3, solver_tol=None, max_iter=50):
187194
)[0]
188195

189196

190-
# The backward pass of the ``"raw"`` mode first polishes the mapped raw-forward iterate with at most this many
191-
# PDIP iterations on the Jacobi-scaled system at jaxnnls's own tight tolerance (PyAutoArray#573). Measured on the
192-
# SLaM MGE fixture, the 48 SLaM ``source_lp[1]`` systems and the jax_grad/mge.py points: 4-6 iterations.
193-
RAW_BACKWARD_POLISH_MAX_ITER = 10
197+
# The ``"raw"`` mode polishes the mapped raw-forward iterate with at most this many PDIP iterations on the
198+
# Jacobi-scaled system at jaxnnls's own tight tolerance; the polished iterate is the forward value (PyAutoArray#594)
199+
# and the backward pass's starting point (PyAutoArray#573). Measured on the SLaM MGE fixture, the 48 SLaM
200+
# ``source_lp[1]`` systems and the jax_grad/mge.py points: 4-6 iterations.
201+
RAW_POLISH_MAX_ITER = 10
202+
# The name under which the polish cap was introduced (#573, backward pass only); kept for external importers.
203+
RAW_BACKWARD_POLISH_MAX_ITER = RAW_POLISH_MAX_ITER
194204

195205

196-
def _raw_forward_backward_point(
197-
Q_pc, q_pc, Q, q, D, target_kappa, solver_tol, max_iter
198-
):
206+
def _raw_forward_polished(Q_pc, q_pc, Q, q, D, solver_tol, max_iter):
199207
"""
200-
The forward solve and the relaxed-KKT point of the ``"raw"`` mode (shared by
201-
the custom-vjp forward pass and :func:`raw_forward_backward_status`).
202-
203-
Returns ``(y, converged, pdip_iter)`` of the raw forward solve (mapped to the
204-
Jacobi coordinates), the relaxed point ``(yr, sr, zr)`` the backward pass
205-
differentiates at, and the status ``(relaxed_converged, relaxed_iter,
206-
polish_converged, polish_iter)``.
208+
The forward value of the ``"raw"`` mode: the raw PDIP solve, mapped to the
209+
Jacobi coordinates and polished there (shared by the custom-vjp primal and
210+
forward passes, so every call path returns the same ``y`` bit-for-bit).
211+
212+
Returns ``(y_out, converged, pdip_iter, (yp, sp, zp), ok, polish_iter)``:
213+
``y_out`` the forward value (the polished iterate if ``ok``, else the mapped
214+
raw iterate), ``converged`` / ``pdip_iter`` the raw forward solve's flag and
215+
iteration count, ``(yp, sp, zp)`` the point the backward pass starts from
216+
(selected by the same ``ok``), ``ok`` whether the polish was accepted and
217+
``polish_iter`` its iteration count.
207218
"""
208219
import jax.numpy as jnp
209-
from jaxnnls.pdip_relaxed import solve_relaxed_nnls
210220

211221
tol = data_scaled_solver_tol(q) if solver_tol is None else solver_tol
212222
x, s, z, converged, pdip_iter = solve_nnls(Q, q, solver_tol=tol, max_iter=max_iter)
213223
y, sy, zy = x / D, s / D, z * D
214224

215-
# Polish (PyAutoArray#573): the data-scaled tolerance leaves s * z ~ 1e-10 .. 1e-9, far above
216-
# ``target_kappa``, so the relaxed solve below would have to push toward the boundary from z / s ~ 1e13
217-
# and its fixed 50-iteration while_loop overshoots to NaN. A few tight PDIP iterations on the scaled
218-
# system, warm-started from the mapped iterate, bring s * z down to the jaxnnls tolerance first. If the
219-
# polish does not converge (the scaled dual is what diverges on #571's systems from a cold start), the
220-
# mapped iterate is kept, i.e. the pre-polish behaviour.
225+
# Polish (PyAutoArray#573, #594): the data-scaled tolerance also judges complementarity, so it leaves
226+
# s * z ~ 1e-10 .. 1e-9 -- a column fnnls holds at zero can keep x ~ tol / z (11.5 % of the reference flux on
227+
# inactive columns of the euclid vis_lp system), and the relaxed solve of the backward pass would have to push
228+
# toward the boundary from z / s ~ 1e13, where its fixed 50-iteration while_loop overshoots to NaN. A few
229+
# tight PDIP iterations on the scaled system, warm-started from the mapped iterate, bring s * z down to the
230+
# jaxnnls tolerance. If the polish does not converge (the scaled dual is what diverges on #571's systems from
231+
# a cold start), the mapped iterate is kept, i.e. the pre-polish behaviour.
221232
yp, sp, zp, polish_converged, polish_iter = solve_nnls(
222-
Q_pc, q_pc, max_iter=RAW_BACKWARD_POLISH_MAX_ITER, init=(y, sy, zy)
233+
Q_pc, q_pc, max_iter=RAW_POLISH_MAX_ITER, init=(y, sy, zy)
223234
)
224235
ok = jnp.logical_and(
225236
polish_converged == 1,
226237
jnp.all(jnp.isfinite(yp)) & jnp.all(sp > 0) & jnp.all(zp > 0),
227238
)
228239
yp, sp, zp = (jnp.where(ok, a, b) for a, b in ((yp, y), (sp, sy), (zp, zy)))
240+
return yp, converged, pdip_iter, (yp, sp, zp), ok, polish_iter
241+
242+
243+
def _raw_forward_backward_point(
244+
Q_pc, q_pc, Q, q, D, target_kappa, solver_tol, max_iter
245+
):
246+
"""
247+
The forward value and the relaxed-KKT point of the ``"raw"`` mode (shared by
248+
the custom-vjp forward pass and :func:`raw_forward_backward_status`).
229249
250+
Returns ``(y, converged, pdip_iter)`` with ``y`` the polished forward value
251+
(:func:`_raw_forward_polished`) and ``converged`` / ``pdip_iter`` those of the
252+
raw forward solve, the relaxed point ``(yr, sr, zr)`` the backward pass
253+
differentiates at, and the status ``(relaxed_converged, relaxed_iter,
254+
polish_converged, polish_iter)``.
255+
"""
256+
from jaxnnls.pdip_relaxed import solve_relaxed_nnls
257+
258+
y_out, converged, pdip_iter, (yp, sp, zp), ok, polish_iter = _raw_forward_polished(
259+
Q_pc, q_pc, Q, q, D, solver_tol, max_iter
260+
)
230261
yr, sr, zr, relaxed_converged, relaxed_iter = solve_relaxed_nnls(
231262
Q_pc, q_pc, yp, sp, zp, target_kappa=target_kappa
232263
)
233264
status = (relaxed_converged, relaxed_iter, ok.astype(int), polish_iter)
234-
return (y, converged, pdip_iter), (yr, sr, zr), status
265+
return (y_out, converged, pdip_iter), (yr, sr, zr), status
235266

236267

237268
@lru_cache(maxsize=None)
@@ -251,16 +282,25 @@ def _solve_nnls_raw_forward_with(target_kappa, solver_tol, max_iter):
251282
linear-object-only (MGE) systems, Jacobi scaling turns signal-free columns
252283
whose diagonal is only the no-regularization floor into degenerate
253284
coordinates that make the PDIP dual diverge; the raw solve does not.
285+
The mapped iterate is then *polished*: at most
286+
:data:`RAW_POLISH_MAX_ITER` PDIP iterations on ``(Q_pc, q_pc)`` at
287+
jaxnnls's tight tolerance, warm-started from it (kept only if the polish
288+
converges with a finite, strictly interior point). The polished iterate is
289+
the returned ``y`` (PyAutoArray#594): the data-scaled stop judges
290+
complementarity ``s * z`` against the loose tolerance, so an unpolished
291+
column fnnls holds at zero can keep ``x ~ tol / z`` -- 11.5 % of the fnnls
292+
flux on inactive columns of the euclid vis_lp system, invisible to the
293+
objective and log likelihood; polished, <= 3.3e-4 over 81 corpus systems
294+
for ~5 extra iterations. The custom-vjp primal and forward rule compute it
295+
through the same :func:`_raw_forward_polished`, so plain, jitted and
296+
differentiated calls return the same value. ``converged`` and
297+
``pdip_iter`` are those of the raw forward solve.
254298
- **Backward:** the relaxed-KKT implicit derivative on ``Q_pc`` (as the
255-
Jacobi mode), started from the mapped iterate after a *polish*: at most
256-
:data:`RAW_BACKWARD_POLISH_MAX_ITER` PDIP iterations on ``(Q_pc, q_pc)``
257-
at jaxnnls's tight tolerance, warm-started from the mapped iterate
258-
(kept only if it converges). Without it the loose forward tolerance
259-
leaves complementarity ``s * z`` orders of magnitude above
299+
Jacobi mode), started from the polished point. Without the polish the
300+
loose forward tolerance leaves ``s * z`` orders of magnitude above
260301
``target_kappa`` and the relaxed solve diverges to NaN on a fraction of
261302
points (PyAutoArray#573); with it the relaxed solve converges in about
262-
one iteration. The primal ``y`` is the unpolished forward solution, so
263-
the forward value is unchanged. The relaxed-KKT pass on the raw,
303+
one iteration. The relaxed-KKT pass on the raw,
264304
ill-conditioned ``Q`` produces NaN gradients, which is why Jacobi scaling
265305
was introduced. ``(Q, q, D)`` get zero cotangents: ``y`` depends only on
266306
``(Q_pc, q_pc)``, and the caller's autodiff carries the dependence of
@@ -274,11 +314,10 @@ def _solve_nnls_raw_forward_with(target_kappa, solver_tol, max_iter):
274314
from jaxnnls.diff_qp import diff_nnls
275315

276316
def primal(Q_pc, q_pc, Q, q, D):
277-
tol = data_scaled_solver_tol(q) if solver_tol is None else solver_tol
278-
x, _, _, converged, pdip_iter = solve_nnls(
279-
Q, q, solver_tol=tol, max_iter=max_iter
317+
y_out, converged, pdip_iter, _, _, _ = _raw_forward_polished(
318+
Q_pc, q_pc, Q, q, D, solver_tol, max_iter
280319
)
281-
return x / D, converged, pdip_iter
320+
return y_out, converged, pdip_iter
282321

283322
def forward(Q_pc, q_pc, Q, q, D):
284323
out, (yr, sr, zr), status = _raw_forward_backward_point(
@@ -305,8 +344,9 @@ def raw_forward_backward_status(
305344
``(relaxed_converged, relaxed_iter, polish_converged, polish_iter)``.
306345
307346
``relaxed_*`` describe the relaxed-KKT solve whose point the gradient is
308-
taken at; ``polish_*`` the tight warm-started PDIP polish before it
309-
(``polish_converged == 0`` means the mapped iterate was used unpolished).
347+
taken at; ``polish_*`` the tight warm-started PDIP polish before it,
348+
whose iterate is also the forward value (``polish_converged == 0`` means the
349+
mapped raw iterate was used unpolished, for both).
310350
Arguments are those of :func:`solve_nnls_primal_raw_forward`.
311351
"""
312352
return _raw_forward_backward_point(
@@ -319,7 +359,8 @@ def solve_nnls_primal_raw_forward(
319359
):
320360
"""
321361
The ``"raw"`` positive-only mode: forward PDIP on the raw system with a
322-
data-scaled tolerance, backward pass on the Jacobi-scaled system. Returns
362+
data-scaled tolerance, polished on the Jacobi-scaled system, backward pass
363+
on the Jacobi-scaled system. Returns
323364
``(y, converged, pdip_iter)`` with ``x = D * y``; see
324365
:func:`_solve_nnls_raw_forward_with`.
325366
"""

‎test_autoarray/inversion/inversion/files/README.md‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,3 +15,13 @@
1515
Captured 2026-09-25 on CPU fp64 via a `jax.debug.callback` on `reconstruction_positive_only_from`, with
1616
autolens_workspace_test 5ec64413d2, PyAutoArray 3de624b5b9, PyAutoGalaxy 70a61e26cd, PyAutoLens 86054bbc19,
1717
PyAutoFit dd9fbe0aab, jax 0.10.2, numpy 2.5.3.
18+
- `mge_solver_reference_systems.npz` — 9 positive-only MGE systems (60x60 fp64) with their fnnls reference
19+
solutions, for the raw-forward PDIP amplitude regression (PyAutoArray#594): keys `Q_<name>` / `q_<name>` /
20+
`x_ref_<name>` for names `k0`..`k7` (the 8 #571 SLaM `source_lp[1]` systems, identical to
21+
`mge_slam_nnls_systems.npz`) and `euclid_vis_lp` (the euclid pipeline vis_lp system behind
22+
`test_latent_euclid_variables_traces_under_jax_jit`, captured with PyAutoArray d4298445). `meta` is a JSON
23+
string with per-system `group`, `source_column_index_list`, `max_abs_q`, `cond_Q`, `category` and reference
24+
active-column counts, plus provenance. Copied verbatim (no recomputation) from the autolens_profiling solver
25+
corpus `results/lens/solver/corpus/{slam_fixture_571,euclid_vis_lp}.npz` + `manifest.json` at
26+
autolens_profiling 3ad68af (phase-1 record `complete/2026/09/linear-solver-accuracy-study.md`) by a one-off
27+
script outside the repo that checks symmetry of `Q`, `x_ref >= 0` and `max|q|` against the manifest.
Binary file not shown.

0 commit comments

Comments
 (0)