Skip to content

Commit cfd7f80

Browse files
Jammy2211claude
andcommitted
perf: relative error-tolerance guard on the NNLS warm-start memo (#498)
Each memo entry now also carries the error fraction of the most recent dense-sign-started solve for its key; a seeded solve whose own error fraction exceeds Settings.nnls_warm_start_error_tolerance (config default 1.5) times that reference is dropped, so the next solve restarts from the dense-sign start and refreshes the reference. Relative, not absolute: the 32-cell robustness matrix showed the absolute seed error fraction does not separate seeds that save iterations from seeds that cost them (0.048-0.138 overlap), while the seed/dense ratio does (helpful <= 0.89, worst 1.42). Non-finite or non-positive tolerance disables the guard. reconstruction_positive_only_from records seed_source and warm_start_fallback in the stats dict. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01XySQZP6Npj4adZMkYZ96Bo
1 parent 105f6ea commit cfd7f80

7 files changed

Lines changed: 369 additions & 38 deletions

File tree

‎autoarray/config/general.yaml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ inversion:
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.
1111
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.
12+
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.
1213
reconstruction_vmax_factor: 0.5 # Plots of an Inversion's reconstruction use the reconstructed data's bright value multiplied by this factor.
1314
log_det_method: cholesky # How the Bayesian-evidence log-determinant terms are computed. "cholesky" (default) is the historical 2*sum(log(diag(cholesky(M)))); "slogdet" uses logabsdet of slogdet(M), which is identical where M is positive-definite but finite (not NaN) where the Cholesky fails, for gradient-based searches (opt-in, non-default; does not change the default evidence). Under "slogdet" the kernel regularization schemes (Matern/Gaussian/Exponential) also compute the regularization log-det analytically from a Cholesky of their covariance instead of factorizing the formed inverse. See PyAutoArray#391.
1415
regularization_term_method: matmul # How the Bayesian-evidence regularization term s^T H s is computed. "matmul" (default) is the historical s @ (H @ s) against the explicitly formed regularization matrix; "cho_solve" evaluates coefficient * s^T C^-1 s for the kernel schemes (Matern/Gaussian/Exponential/MaternAdapt) via one Cholesky solve of their covariance C, avoiding the explicit inverse whose round-off is amplified by cond(C) (~1e9 on clustered traced mesh vertices). Opt-in, non-default; does not change the default evidence. Schemes with no such factorization fall back to the formed matrix.

‎autoarray/inversion/inversion/inversion_util.py‎

Lines changed: 57 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -301,6 +301,16 @@ def reconstruction_positive_only_from(
301301
Identifies the index space this solve's passive set lives in, enabling the cross-evaluation warm-start
302302
memo (`Settings.nnls_warm_start_memo`) on the NumPy path. `None` disables the memo for this call.
303303
304+
Notes
305+
-----
306+
On the NumPy path this function writes two keys into the `stats` dict it passes to `fnnls_cholesky`
307+
that `fnnls_cholesky` itself knows nothing about: ``seed_source`` (``"memo"`` if the solve that
308+
produced the returned reconstruction started from a memo seed, ``"dense"`` if it started from the
309+
sign of the unconstrained dense solve) and ``warm_start_fallback`` (`True` if a memo seed breached
310+
`Settings.nnls_warm_start_error_tolerance` and its entry was dropped, so the next solve for that key
311+
restarts dense). They are set after `fnnls_cholesky` returns, so a diagnostic wrapping the solver
312+
must read the dict it handed in *after* the evaluation, not at the point the solver returns.
313+
304314
Returns
305315
-------
306316
Non-negative S that minimizes the Eq.(2) of https://arxiv.org/pdf/astro-ph/0302587.pdf.
@@ -392,20 +402,22 @@ def reconstruction_positive_only_from(
392402
else None
393403
)
394404

395-
warm_start = (
396-
nnls_memo.passive_set_get(key=key, n=data_vector.shape[0])
397-
if use_memo
398-
else None
399-
)
405+
n = data_vector.shape[0]
406+
407+
entry = nnls_memo.passive_set_get(key=key, n=n) if use_memo else None
400408

401-
if warm_start is not None:
409+
stats["seed_source"] = "dense"
410+
stats["warm_start_fallback"] = False
411+
412+
if entry is not None:
402413
try:
403414
reconstruction = fnnls_cholesky(
404415
curvature_reg_matrix,
405416
(data_vector).T,
406-
P_initial=warm_start,
417+
P_initial=entry.passive_set,
407418
stats=stats,
408419
)
420+
stats["seed_source"] = "memo"
409421
except (RuntimeError, np.linalg.LinAlgError, ValueError):
410422
# A seed from a previous evaluation is a guess about a
411423
# different matrix, so it can factorise badly where the
@@ -423,9 +435,46 @@ def reconstruction_positive_only_from(
423435
P_initial=np.linalg.solve(curvature_reg_matrix, data_vector) > 0,
424436
stats=stats,
425437
)
438+
stats["seed_source"] = "dense"
426439

427440
if use_memo:
428-
nnls_memo.passive_set_put(key=key, passive_set=stats["passive_set"])
441+
error_fraction = stats["warm_start_errors"] / max(n, 1)
442+
443+
if stats["seed_source"] == "memo":
444+
# A memo seed is judged against the dense-sign start it
445+
# replaced, not against an absolute error count: the absolute
446+
# fraction does not separate seeds that save iterations from
447+
# seeds that cost them, the ratio to the dense-sign reference
448+
# does. A seed that breaches the tolerance is discarded, so the
449+
# next solve for this key restarts dense and refreshes the
450+
# reference -- no stale seed can be dragged through a run in a
451+
# regime the reference was never measured in.
452+
tolerance = settings.nnls_warm_start_error_tolerance
453+
454+
guard_active = (
455+
tolerance is not None and np.isfinite(tolerance) and tolerance > 0.0
456+
)
457+
458+
if (
459+
guard_active
460+
and error_fraction > tolerance * entry.dense_error_fraction
461+
):
462+
nnls_memo.memo_drop(key=key)
463+
stats["warm_start_fallback"] = True
464+
else:
465+
# The reference describes the dense-sign start, so it is
466+
# carried forward unchanged; only a dense solve refreshes it.
467+
nnls_memo.passive_set_put(
468+
key=key,
469+
passive_set=stats["passive_set"],
470+
dense_error_fraction=entry.dense_error_fraction,
471+
)
472+
else:
473+
nnls_memo.passive_set_put(
474+
key=key,
475+
passive_set=stats["passive_set"],
476+
dense_error_fraction=error_fraction,
477+
)
429478

430479
return reconstruction
431480

‎autoarray/inversion/inversion/nnls_memo.py‎

Lines changed: 55 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import os
22

33
import numpy as np
4-
from typing import Dict, Optional
4+
from typing import Dict, NamedTuple, Optional
55

66
# Cross-evaluation memo for the positive-only (fnnls) solve's passive set.
77
#
@@ -22,8 +22,36 @@
2222
# stale entry costs iterations, not correctness. Forked pool workers inherit a
2323
# copy at fork and diverge from there, which is fine for the same reason.
2424
#
25+
# An entry therefore carries TWO things: the passive set to seed from, and
26+
# `dense_error_fraction` -- the error fraction of the most recent solve for
27+
# that key that started from the DENSE-SIGN guess. That number is the
28+
# reference the fallback guard in `reconstruction_positive_only_from` measures
29+
# a seed against: the PyAutoArray#498 robustness matrix showed the absolute
30+
# error fraction of a seed does NOT separate seeds that save iterations from
31+
# seeds that cost them (helpful and unhelpful cells overlap at 0.048-0.138),
32+
# but the ratio of the seed's fraction to the dense-sign start's does (helpful
33+
# cells never exceed 0.89, the worst seed reaches 1.42). The reference is
34+
# per-key and self-calibrating, so a solve regime far outside anything the
35+
# matrix probed cannot drag a stale seed through a whole run: once a seed is
36+
# that much worse than the dense-sign start, the entry is dropped
37+
# (`memo_drop`) and the next solve for that key restarts dense, refreshing the
38+
# reference.
39+
#
2540
# Disable with AUTOARRAY_NNLS_WARM_START=0.
26-
_nnls_passive_set_memo: Dict[str, np.ndarray] = {}
41+
42+
43+
class MemoEntry(NamedTuple):
44+
"""
45+
One memoized solve: the passive set to seed the next solve for this key
46+
from, and the error fraction of the most recent dense-sign-started solve
47+
for the same key (the reference the fallback guard compares a seed to).
48+
"""
49+
50+
passive_set: np.ndarray
51+
dense_error_fraction: float
52+
53+
54+
_nnls_passive_set_memo: Dict[str, MemoEntry] = {}
2755

2856
_NNLS_PASSIVE_SET_MEMO_MAX_ENTRIES = 8
2957

@@ -43,31 +71,35 @@ def memo_key(n: int, fingerprint) -> str:
4371
return f"{n}:{fingerprint}"
4472

4573

46-
def passive_set_get(key: str, n: int) -> Optional[np.ndarray]:
74+
def passive_set_get(key: str, n: int) -> Optional[MemoEntry]:
4775
"""
48-
The memoized passive set for `key`, or None on a miss.
76+
The memoized entry for `key` -- its passive set and dense-sign reference
77+
error fraction -- or None on a miss.
4978
5079
An entry whose indices do not all fit a size-`n` solve is a miss, not a
5180
hit: `n` is already part of the key, so this only fires if a caller
5281
fingerprints two different index spaces identically, and a miss is always
5382
a safe outcome.
5483
"""
55-
passive_set = _nnls_passive_set_memo.get(key)
84+
entry = _nnls_passive_set_memo.get(key)
5685

57-
if passive_set is None:
86+
if entry is None:
5887
return None
5988

60-
if passive_set.size and passive_set.max() >= n:
89+
if entry.passive_set.size and entry.passive_set.max() >= n:
6190
return None
6291

63-
return passive_set
92+
return entry
6493

6594

66-
def passive_set_put(key: str, passive_set: np.ndarray) -> None:
95+
def passive_set_put(
96+
key: str, passive_set: np.ndarray, dense_error_fraction: float
97+
) -> None:
6798
"""
68-
Store a solve's final passive set, evicting the oldest entry once the memo
69-
is full (FIFO; the memo tracks one inversion's recent history, not a
70-
working set worth ranking).
99+
Store a solve's final passive set alongside the dense-sign reference error
100+
fraction to carry forward, evicting the oldest entry once the memo is full
101+
(FIFO; the memo tracks one inversion's recent history, not a working set
102+
worth ranking).
71103
"""
72104
stored = np.asarray(passive_set, dtype=int).copy()
73105
stored.setflags(write=False)
@@ -78,4 +110,14 @@ def passive_set_put(key: str, passive_set: np.ndarray) -> None:
78110
):
79111
_nnls_passive_set_memo.pop(next(iter(_nnls_passive_set_memo)))
80112

81-
_nnls_passive_set_memo[key] = stored
113+
_nnls_passive_set_memo[key] = MemoEntry(
114+
passive_set=stored, dense_error_fraction=float(dense_error_fraction)
115+
)
116+
117+
118+
def memo_drop(key: str) -> None:
119+
"""
120+
Forget `key`, so the next solve for it restarts from the dense-sign guess
121+
and refreshes the reference error fraction. A no-op if the key is absent.
122+
"""
123+
_nnls_passive_set_memo.pop(key, None)

‎autoarray/settings.py‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ def __init__(
1818
nnls_solver_tol: Optional[float] = None,
1919
nnls_max_iter: Optional[int] = None,
2020
nnls_warm_start_memo: Optional[bool] = None,
21+
nnls_warm_start_error_tolerance: Optional[float] = None,
2122
log_det_method: Optional[str] = None,
2223
regularization_term_method: Optional[str] = None,
2324
):
@@ -110,6 +111,21 @@ def __init__(
110111
(`true`); setting ``AUTOARRAY_NNLS_WARM_START=0`` is the process-wide kill-switch that disables the
111112
memo whatever this setting says. Only the NumPy (`xp=np`) fnnls path honors this; the JAX path
112113
ignores it.
114+
nnls_warm_start_error_tolerance
115+
The relative guard on a memo seed's quality. Each memo entry carries a reference: the error
116+
fraction (warm-start errors / solve size) of the most recent solve for that key which started
117+
from the dense-sign guess. A memo-seeded solve whose own error fraction exceeds
118+
``tolerance * reference`` is judged worse than simply starting dense, so its entry is discarded
119+
and the next solve for that key restarts from the dense-sign start, refreshing the reference.
120+
The guard is relative because the absolute error fraction does not separate helpful from
121+
unhelpful seeds (the two populations overlap at 0.048-0.138 in the PyAutoArray#498 32-cell
122+
robustness matrix) whereas the ratio to the dense-sign start does (helpful cells top out at
123+
0.89, the worst seed measured reaches 1.42). `None` (default) reads the packaged value
124+
(``1.5``), chosen above that worst observed ratio so the guard is protective against regimes
125+
far outside the matrix rather than flapping inside it. The guard is **disabled** by any value
126+
that is not finite and positive -- ``float("inf")`` is the idiomatic choice, and ``0`` or a
127+
negative value disables it too rather than meaning "drop everything". Only the NumPy
128+
(`xp=np`) fnnls path honors this; the JAX path ignores it.
113129
log_det_method
114130
Which computation is used for the two Bayesian-evidence log-determinant terms
115131
(``log_det_curvature_reg_matrix_term`` and ``log_det_regularization_matrix_term``). `None`
@@ -163,6 +179,7 @@ def __init__(
163179
self.nnls_solver_tol = nnls_solver_tol
164180
self.nnls_max_iter = nnls_max_iter
165181
self._nnls_warm_start_memo = nnls_warm_start_memo
182+
self._nnls_warm_start_error_tolerance = nnls_warm_start_error_tolerance
166183
self._use_positive_only_solver = use_positive_only_solver
167184
self._use_edge_zeroed_pixels = use_edge_zeroed_pixels
168185
self._use_border_relocator = use_border_relocator
@@ -229,6 +246,28 @@ def nnls_warm_start_memo(self) -> bool:
229246

230247
return self._nnls_warm_start_memo
231248

249+
@property
250+
def nnls_warm_start_error_tolerance(self) -> float:
251+
"""
252+
How much worse than the dense-sign start a memo seed may be before its entry is dropped.
253+
254+
A seeded solve whose error fraction exceeds this multiple of the entry's dense-sign reference
255+
fraction is discarded, so the next solve for that key restarts dense. Any value that is not
256+
finite and positive disables the guard.
257+
"""
258+
if self._nnls_warm_start_error_tolerance is None:
259+
try:
260+
return conf.instance["general"]["inversion"][
261+
"nnls_warm_start_error_tolerance"
262+
]
263+
except KeyError:
264+
# Workspaces ship their own general.yaml that shadows autoarray's
265+
# and will not carry this key, so in practice this fallback *is*
266+
# the production default, matching the packaged value.
267+
return 1.5
268+
269+
return self._nnls_warm_start_error_tolerance
270+
232271
@property
233272
def log_det_method(self):
234273
if self._log_det_method is None:

‎test_autoarray/inversion/inversion/test_nnls_memo.py‎

Lines changed: 28 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
from autoarray.inversion.inversion.nnls_memo import (
1919
_NNLS_PASSIVE_SET_MEMO_MAX_ENTRIES,
2020
_nnls_passive_set_memo,
21+
memo_drop,
2122
memo_key,
2223
passive_set_get,
2324
passive_set_put,
@@ -53,7 +54,9 @@ def solve_ids_to_keep(self):
5354
return self._ids_to_keep
5455

5556

56-
def _inversion_from(curvature_reg_matrix, data_vector, memo, ids_to_keep=None):
57+
def _inversion_from(
58+
curvature_reg_matrix, data_vector, memo, ids_to_keep=None, tolerance=None
59+
):
5760
n = data_vector.shape[0]
5861

5962
kwargs = dict(
@@ -66,6 +69,7 @@ def _inversion_from(curvature_reg_matrix, data_vector, memo, ids_to_keep=None):
6669
use_positive_only_solver=True,
6770
use_edge_zeroed_pixels=False,
6871
nnls_warm_start_memo=memo,
72+
nnls_warm_start_error_tolerance=tolerance,
6973
),
7074
)
7175

@@ -115,7 +119,7 @@ def test__memoized_reconstruction__subset_branch__matches_unmemoized(seed):
115119
# the subset -- not the full parameter vector.
116120
(key,) = _nnls_passive_set_memo
117121
assert key.startswith(f"{len(ids_to_keep)}:")
118-
assert np.all(_nnls_passive_set_memo[key] < len(ids_to_keep))
122+
assert np.all(_nnls_passive_set_memo[key].passive_set < len(ids_to_keep))
119123

120124

121125
def test__fingerprint__changes_with_ids_to_keep():
@@ -141,36 +145,51 @@ def test__fingerprint__changes_with_ids_to_keep():
141145

142146
def test__passive_set_put__evicts_the_oldest_entry_when_full():
143147
for i in range(_NNLS_PASSIVE_SET_MEMO_MAX_ENTRIES + 2):
144-
passive_set_put(key=f"key_{i}", passive_set=np.array([i]))
148+
passive_set_put(
149+
key=f"key_{i}", passive_set=np.array([i]), dense_error_fraction=0.1
150+
)
145151

146152
assert len(_nnls_passive_set_memo) == _NNLS_PASSIVE_SET_MEMO_MAX_ENTRIES
147153
assert "key_0" not in _nnls_passive_set_memo
148154
assert "key_1" not in _nnls_passive_set_memo
149155
assert f"key_{_NNLS_PASSIVE_SET_MEMO_MAX_ENTRIES + 1}" in _nnls_passive_set_memo
150156

151157

152-
def test__passive_set_put__stores_a_read_only_copy():
158+
def test__passive_set_put__stores_a_read_only_copy_and_the_reference_fraction():
153159
passive_set = np.array([0, 3, 4])
154160

155-
passive_set_put(key="key", passive_set=passive_set)
161+
passive_set_put(key="key", passive_set=passive_set, dense_error_fraction=0.25)
156162

157163
passive_set[0] = 99
158164

159-
stored = passive_set_get(key="key", n=5)
165+
entry = passive_set_get(key="key", n=5)
160166

161-
assert np.array_equal(stored, np.array([0, 3, 4]))
167+
assert np.array_equal(entry.passive_set, np.array([0, 3, 4]))
168+
assert entry.dense_error_fraction == 0.25
162169
with pytest.raises(ValueError):
163-
stored[0] = 1
170+
entry.passive_set[0] = 1
164171

165172

166173
def test__passive_set_get__miss_on_out_of_range_indices_and_unknown_key():
167-
passive_set_put(key="key", passive_set=np.array([0, 3, 4]))
174+
passive_set_put(
175+
key="key", passive_set=np.array([0, 3, 4]), dense_error_fraction=0.0
176+
)
168177

169178
assert passive_set_get(key="key", n=5) is not None
170179
assert passive_set_get(key="key", n=4) is None
171180
assert passive_set_get(key="other", n=5) is None
172181

173182

183+
def test__memo_drop__forgets_the_key_and_is_a_no_op_when_absent():
184+
passive_set_put(key="key", passive_set=np.array([0, 2]), dense_error_fraction=0.1)
185+
186+
memo_drop(key="key")
187+
188+
assert passive_set_get(key="key", n=5) is None
189+
190+
memo_drop(key="key")
191+
192+
174193
def test__memo_enabled__reads_the_environment(monkeypatch):
175194
monkeypatch.delenv("AUTOARRAY_NNLS_WARM_START", raising=False)
176195
assert nnls_memo.memo_enabled() is True

0 commit comments

Comments
 (0)