|
| 1 | +import os |
| 2 | + |
| 3 | +import numpy as np |
| 4 | +from typing import Dict, Optional |
| 5 | + |
| 6 | +# Cross-evaluation memo for the positive-only (fnnls) solve's passive set. |
| 7 | +# |
| 8 | +# The Bro & De Jong active set iteration is warm-started from a guess at which |
| 9 | +# reconstruction entries are non-zero. The production guess -- the sign of the |
| 10 | +# unconstrained dense solve -- gets ~150 of ~1560 entries wrong on a Delaunay |
| 11 | +# euclid fit, and each wrong entry costs an active-set iteration in a solve |
| 12 | +# that is ~70% of the whole numba CPU likelihood evaluation. Successive |
| 13 | +# sampler evaluations sit close together in parameter space and share nearly |
| 14 | +# all of their passive set, so seeding from the previous evaluation's FINAL |
| 15 | +# passive set is a much better guess. |
| 16 | +# |
| 17 | +# Entries are keyed by the index space the passive set lives in (mesh/data |
| 18 | +# shapes, plus the edge-zeroed subset when one is in use) -- never by the |
| 19 | +# matrix values, since the whole point is to hit across nearby parameter |
| 20 | +# points. A wrong seed cannot corrupt the answer: the NNLS optimum is unique |
| 21 | +# and the solver adds and removes entries until the KKT conditions hold, so a |
| 22 | +# stale entry costs iterations, not correctness. Forked pool workers inherit a |
| 23 | +# copy at fork and diverge from there, which is fine for the same reason. |
| 24 | +# |
| 25 | +# Disable with AUTOARRAY_NNLS_WARM_START=0. |
| 26 | +_nnls_passive_set_memo: Dict[str, np.ndarray] = {} |
| 27 | + |
| 28 | +_NNLS_PASSIVE_SET_MEMO_MAX_ENTRIES = 8 |
| 29 | + |
| 30 | + |
| 31 | +def memo_enabled() -> bool: |
| 32 | + """ |
| 33 | + Whether the passive-set memo is active in this process. |
| 34 | + """ |
| 35 | + return os.environ.get("AUTOARRAY_NNLS_WARM_START", "1") != "0" |
| 36 | + |
| 37 | + |
| 38 | +def memo_key(n: int, fingerprint) -> str: |
| 39 | + """ |
| 40 | + The memo key for a solve of size `n` whose index space is described by |
| 41 | + `fingerprint` (see `AbstractInversion._nnls_warm_start_fingerprint`). |
| 42 | + """ |
| 43 | + return f"{n}:{fingerprint}" |
| 44 | + |
| 45 | + |
| 46 | +def passive_set_get(key: str, n: int) -> Optional[np.ndarray]: |
| 47 | + """ |
| 48 | + The memoized passive set for `key`, or None on a miss. |
| 49 | +
|
| 50 | + An entry whose indices do not all fit a size-`n` solve is a miss, not a |
| 51 | + hit: `n` is already part of the key, so this only fires if a caller |
| 52 | + fingerprints two different index spaces identically, and a miss is always |
| 53 | + a safe outcome. |
| 54 | + """ |
| 55 | + passive_set = _nnls_passive_set_memo.get(key) |
| 56 | + |
| 57 | + if passive_set is None: |
| 58 | + return None |
| 59 | + |
| 60 | + if passive_set.size and passive_set.max() >= n: |
| 61 | + return None |
| 62 | + |
| 63 | + return passive_set |
| 64 | + |
| 65 | + |
| 66 | +def passive_set_put(key: str, passive_set: np.ndarray) -> None: |
| 67 | + """ |
| 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). |
| 71 | + """ |
| 72 | + stored = np.asarray(passive_set, dtype=int).copy() |
| 73 | + stored.setflags(write=False) |
| 74 | + |
| 75 | + if ( |
| 76 | + key not in _nnls_passive_set_memo |
| 77 | + and len(_nnls_passive_set_memo) >= _NNLS_PASSIVE_SET_MEMO_MAX_ENTRIES |
| 78 | + ): |
| 79 | + _nnls_passive_set_memo.pop(next(iter(_nnls_passive_set_memo))) |
| 80 | + |
| 81 | + _nnls_passive_set_memo[key] = stored |
0 commit comments