3737# (`memo_drop`) and the next solve for that key restarts dense, refreshing the
3838# reference.
3939#
40+ # That guard judges a seed only AFTER the seeded solve has run. On a
41+ # scattered evaluation stream (iid draws, e.g. a nested sampler's early live
42+ # points) every seed is bad, so the stream alternates dense solve -> bad
43+ # seeded solve -> drop -> dense solve ...: half the solves pay for a bad seed.
44+ # autolens_profiling#332 measured the alma interferometer Delaunay solve 2.17x
45+ # slower memo-on than memo-off on such a stream. The per-key back-off below
46+ # stops that: after `_NNLS_BACKOFF_AFTER_FALLBACKS` seeded solves for a key
47+ # have fallen back in a row, the next solve(s) for it skip the seed and start
48+ # dense (still refreshing the entry), for 1, 2, 4, ... up to
49+ # `_NNLS_BACKOFF_MAX_SKIP` solves, then the seed is probed again. One accepted
50+ # seed resets the streak, and a stream that never falls back twice in a row (a
51+ # local walk) never meets the back-off at all.
52+ #
53+ # A backed-off solve also judges, for free, the seed it skipped: the dense
54+ # solve's final passive set is the unique optimum's, so the seed's error count
55+ # on this system is just the size of its disagreement with that set
56+ # (`seed_error_fraction`) -- exactly what `warm_start_errors` would have
57+ # reported had the seed been used. If the skipped seed would have passed the
58+ # fallback guard against this dense solve, the remaining skips are cancelled
59+ # and the next solve is seeded again: a stream that turns local regains the
60+ # memo after one solve, without waiting out the schedule. The streak is kept
61+ # (only an accepted REAL seeded solve resets it), so a shadow check that passes
62+ # by chance on a scattered stream costs one probe and lengthens the next skip.
63+ #
4064# Disable with AUTOARRAY_NNLS_WARM_START=0.
4165
4266
@@ -55,6 +79,31 @@ class MemoEntry(NamedTuple):
5579
5680_NNLS_PASSIVE_SET_MEMO_MAX_ENTRIES = 8
5781
82+ # Consecutive seeded-solve fallbacks for one key before the back-off engages.
83+ # 2, not 1: a single fallback already costs only one bad solve (the next solve
84+ # restarts dense anyway), and a local walk that crosses one sharp change must
85+ # keep the memo on the very next probe.
86+ _NNLS_BACKOFF_AFTER_FALLBACKS = 2
87+
88+ # Cap on the number of solves skipped between probes of a backed-off key. On a
89+ # stream where every seed is bad, one probe in (cap + 1) solves still pays for
90+ # a bad seed; on a stream that turns local the memo is back within cap solves.
91+ _NNLS_BACKOFF_MAX_SKIP = 32
92+
93+
94+ class BackoffState (NamedTuple ):
95+ """
96+ Per-key back-off bookkeeping: how many seeded solves in a row fell back
97+ (`fallback_streak`) and how many upcoming solves still skip the seed
98+ (`skip_remaining`).
99+ """
100+
101+ fallback_streak : int
102+ skip_remaining : int
103+
104+
105+ _nnls_backoff : Dict [str , BackoffState ] = {}
106+
58107
59108def memo_enabled () -> bool :
60109 """
@@ -121,3 +170,89 @@ def memo_drop(key: str) -> None:
121170 and refreshes the reference error fraction. A no-op if the key is absent.
122171 """
123172 _nnls_passive_set_memo .pop (key , None )
173+
174+
175+ def memo_clear () -> None :
176+ """
177+ Forget every memo entry AND every back-off state, returning the process to
178+ a cold start. Harnesses that compare memo arms must use this rather than
179+ clearing `_nnls_passive_set_memo` alone: back-off state left by one arm
180+ (e.g. an iid stream) would otherwise carry into the next.
181+ """
182+ _nnls_passive_set_memo .clear ()
183+ _nnls_backoff .clear ()
184+
185+
186+ def backoff_should_skip (key : str ) -> bool :
187+ """
188+ Whether the solve for `key` about to run should skip the memo seed and
189+ start dense, consuming one skip if so. False for a key with no back-off.
190+ """
191+ state = _nnls_backoff .get (key )
192+
193+ if state is None or state .skip_remaining <= 0 :
194+ return False
195+
196+ _nnls_backoff [key ] = state ._replace (skip_remaining = state .skip_remaining - 1 )
197+
198+ return True
199+
200+
201+ def backoff_record_fallback (key : str ) -> None :
202+ """
203+ Record that a seeded solve for `key` breached the fallback guard. Once
204+ `_NNLS_BACKOFF_AFTER_FALLBACKS` have happened in a row, schedule the next
205+ 1, 2, 4, ... (capped at `_NNLS_BACKOFF_MAX_SKIP`) solves to skip the seed.
206+ """
207+ state = _nnls_backoff .get (key , BackoffState (0 , 0 ))
208+
209+ streak = state .fallback_streak + 1
210+
211+ excess = streak - _NNLS_BACKOFF_AFTER_FALLBACKS
212+
213+ skip = 0 if excess < 0 else min (2 ** excess , _NNLS_BACKOFF_MAX_SKIP )
214+
215+ if (
216+ key not in _nnls_backoff
217+ and len (_nnls_backoff ) >= _NNLS_PASSIVE_SET_MEMO_MAX_ENTRIES
218+ ):
219+ _nnls_backoff .pop (next (iter (_nnls_backoff )))
220+
221+ _nnls_backoff [key ] = BackoffState (fallback_streak = streak , skip_remaining = skip )
222+
223+
224+ def backoff_end_skip (key : str ) -> None :
225+ """
226+ Cancel the remaining skips for `key` -- the next solve is seeded again --
227+ while keeping its fallback streak, so a failed probe resumes the back-off
228+ at the escalated length. A no-op if the key has no back-off.
229+ """
230+ state = _nnls_backoff .get (key )
231+
232+ if state is not None :
233+ _nnls_backoff [key ] = state ._replace (skip_remaining = 0 )
234+
235+
236+ def backoff_record_accept (key : str ) -> None :
237+ """
238+ Record that a seeded solve for `key` passed the fallback guard, resetting
239+ its back-off entirely. A no-op if the key has no back-off.
240+ """
241+ _nnls_backoff .pop (key , None )
242+
243+
244+ def seed_error_fraction (seed_passive_set : np .ndarray , passive_set , n : int ) -> float :
245+ """
246+ The fraction of a size-`n` solve's entries that the warm-start passive set
247+ `seed_passive_set` gets wrong relative to the solve's final `passive_set`
248+ -- the quantity `fnnls_cholesky` reports as `warm_start_errors / n` when it
249+ is seeded from `seed_passive_set`. Used to judge a seed the back-off skipped
250+ against the dense solve that ran instead.
251+ """
252+ seed_mask = np .zeros (n , dtype = bool )
253+ seed_mask [np .asarray (seed_passive_set , dtype = int )] = True
254+
255+ final_mask = np .zeros (n , dtype = bool )
256+ final_mask [np .asarray (passive_set , dtype = int )] = True
257+
258+ return np .count_nonzero (seed_mask != final_mask ) / max (n , 1 )
0 commit comments