fix(identification): weakref-validated posterior cache for ProxySVAR

src/impulso/identification.py

@@ -1,6 +1,7 @@
 """Identification schemes for structural VAR analysis."""
 
-from typing import TYPE_CHECKING, ClassVar
+import weakref
+from typing import TYPE_CHECKING, Any, ClassVar, Final
 
 import numpy as np
 import pandas as pd
@@ -15,6 +16,132 @@
     from impulso.data import VARData
 
 
+_CACHE_MISS: Final = object()
+"""Sentinel returned by `_PosteriorCache.get` when there is no valid entry.
+
+A dedicated sentinel (rather than `None`) keeps `None` usable as a cached
+value, and makes the hit test at call sites an unambiguous `is` check.
+"""
+
+
+class _PosteriorCache:
+    """Single-slot memo keyed on object identity, validated by weak references.
+
+    Several identification schemes memoise a quantity that depends only on
+    the posterior (plus a few scalars) so that a per-period identification
+    loop — `IdentifiedVAR._identify_per_t` calls `identify()` once per time
+    index with the same posterior — pays the expensive part once instead of
+    `T` times.
+
+    Keying such a memo on `id(posterior)` is unsafe: once the posterior is
+    garbage collected its address can be reused by a new object, and a
+    subsequent lookup with a *different* posterior that happens to land on
+    the recycled address returns a stale value silently (issue #203). This
+    cache stores `weakref.ref(owner)` instead and treats a dead referent —
+    or a live referent that is not the object being looked up — as a miss.
+    The referent identity check (`ref() is owner`) subsumes an `id()`
+    comparison exactly, so no `id()` is kept in the key.
+
+    Usage is a get/compute/set triple:
+
+        cached = self._cache.get(posterior, (n_lags,))
+        if cached is _CACHE_MISS:
+            cached = expensive(posterior, n_lags)
+            self._cache.set(posterior, (n_lags,), cached)
+
+    `owners` may be a single object or a tuple of objects when validity
+    depends on more than one identity (`ProxySVAR` keys on the posterior
+    *and* the `VARData`). Every owner must support weak references —
+    `xr.Dataset` and Pydantic models such as `VARData` both do. If any
+    owner does not (plain tuples, ints and strings do not), `set` declines
+    to cache rather than falling back to an unsafe key, so the only cost is
+    a lost speed-up.
+
+    Entries in `key` are ordinary scalars compared with `==` (lag order, a
+    variable name, a horizon). Do not put arrays there — elementwise
+    comparison would not yield a bool.
+
+    Intended adoption (issue #203). The other identity-keyed memos in this
+    module collapse to the same three lines once their branches merge:
+
+        MaxShare._spectral_cache:
+            self._spectral_cache.get(posterior, (n_lags, target))
+            self._spectral_cache.set(posterior, (n_lags, target), value)
+
+        LongRunRestriction._lr_cache:
+            self._lr_cache.get(posterior, (n_lags,))
+            self._lr_cache.set(posterior, (n_lags,), value)
+
+    Declare the attribute on the (frozen) scheme with
+    `PrivateAttr(default_factory=_PosteriorCache)` so each instance gets
+    its own slot, and mutate it in place — no `object.__setattr__` needed.
+    """
+
+    __slots__ = ("_key", "_refs", "_value")
+
+    def __init__(self) -> None:
+        self._refs: tuple[weakref.ref, ...] | None = None
+        self._key: tuple[Any, ...] = ()
+        self._value: Any = _CACHE_MISS
+
+    @staticmethod
+    def _as_tuple(owners: Any) -> tuple[Any, ...]:
+        """Normalise a single owner or a tuple of owners to a tuple."""
+        return owners if isinstance(owners, tuple) else (owners,)
+
+    def get(self, owners: Any, key: tuple[Any, ...] = ()) -> Any:
+        """Look up the memoised value.
+
+        Args:
+            owners: The object (or tuple of objects) whose identity the
+                cached value is tied to.
+            key: Scalar key tail — everything the value depends on that is
+                not an owner identity.
+
+        Returns:
+            The cached value, or `_CACHE_MISS` if there is no entry, the
+            key tail differs, the owners differ, or any owner has been
+            garbage collected.
+        """
+        if self._refs is None or self._key != key:
+            return _CACHE_MISS
+        owner_tuple = self._as_tuple(owners)
+        if len(owner_tuple) != len(self._refs):
+            return _CACHE_MISS
+        # A dead referent dereferences to None, which can never be an
+        # owner (weakref.ref(None) is not constructible), so the same
+        # check covers both "collected" and "different object".
+        if any(ref() is not owner for ref, owner in zip(self._refs, owner_tuple, strict=True)):
+            return _CACHE_MISS
+        return self._value
+
+    def set(self, owners: Any, key: tuple[Any, ...], value: Any) -> None:
+        """Store `value`, tying its validity to the owners staying alive.
+
+        Args:
+            owners: The object (or tuple of objects) whose identity the
+                value depends on.
+            key: Scalar key tail.
+            value: The value to memoise.
+        """
+        try:
+            refs = tuple(weakref.ref(owner) for owner in self._as_tuple(owners))
+        except TypeError:
+            # An owner that cannot be weakly referenced has no validity
+            # token, so caching it would reintroduce the id()-reuse bug.
+            self.clear()
+            return
+        self._refs = refs
+        self._key = key
+        self._value = value
+
+    def clear(self) -> None:
+        """Drop any stored entry."""
+        self._refs = None
+        self._key = ()
+        self._value = _CACHE_MISS
+
+
 class Cholesky(ImpulsoModel):
     """Cholesky identification scheme.
 
@@ -326,9 +453,10 @@ class ProxySVAR(ImpulsoBaseModel):
     # — not on L — so under time-varying volatility, where the pipeline
     # calls identify() once per period with the same posterior/data, the
     # expensive residual reconstruction runs once instead of T times.
-    # Keyed by object identity: valid while the caller holds the same
-    # posterior/data objects, which is exactly the per-t loop's lifetime.
-    _impact_cache: tuple | None = PrivateAttr(default=None)
+    # Keyed by object identity, with weak references as the validity token:
+    # valid while the caller holds the same posterior/data objects, which
+    # is exactly the per-t loop's lifetime. See `_PosteriorCache`.
+    _impact_cache: _PosteriorCache = PrivateAttr(default_factory=_PosteriorCache)
 
     def identify(
         self,
@@ -370,10 +498,8 @@ def identify(
             raise ValueError(f"policy_variable {self.policy_variable!r} not in var_names {var_names}")
         policy_idx = var_names.index(self.policy_variable)
 
-        cache_key = (id(posterior), id(data), n_lags)
-        if self._impact_cache is not None and self._impact_cache[0] == cache_key:
-            d = self._impact_cache[1]
-        else:
+        d = self._impact_cache.get((posterior, data), (n_lags,))
+        if d is _CACHE_MISS:
             z, u = self._aligned_residuals(posterior, data, n_lags)
 
             # Impact direction: per-draw covariance between instrument and
@@ -391,7 +517,7 @@ def identify(
                 "proxy_first_stage_f_q05": float(np.quantile(f_draws, 0.05)),
                 "proxy_first_stage_f_q95": float(np.quantile(f_draws, 0.95)),
             }
-            self._impact_cache = (cache_key, d)
+            self._impact_cache.set((posterior, data), (n_lags,), d)
             if f_median < 10.0:
                 import warnings
 

tests/test_identification.py

@@ -1,10 +1,14 @@
 """Tests for identification schemes."""
 
+import gc
+import weakref
+
 import numpy as np
 import pytest
+import xarray as xr
 from pydantic import ValidationError
 
-from impulso.identification import Cholesky, SignRestriction
+from impulso.identification import _CACHE_MISS, Cholesky, SignRestriction, _PosteriorCache
 from impulso.protocols import IdentificationScheme
 
 
@@ -365,3 +369,108 @@ def test_identify_with_horizon_uses_posterior_B(self, synthetic_idata_2v):
         # synthetic_idata_2v has B in posterior; should run without raising.
         P = scheme.identify(L, var_names, posterior=synthetic_idata_2v.posterior)
         assert P.shape == L.shape
+
+
+class TestPosteriorCache:
+    """Weakref-validated identity cache shared by the identification schemes (#203)."""
+
+    def test_xr_dataset_supports_weak_references(self):
+        """The whole design rests on this — assert it rather than assume it."""
+        ds = xr.Dataset()
+        assert weakref.ref(ds)() is ds
+
+    def test_empty_cache_misses(self):
+        cache = _PosteriorCache()
+        assert cache.get(xr.Dataset(), (1,)) is _CACHE_MISS
+
+    def test_hit_on_the_same_owner(self):
+        cache = _PosteriorCache()
+        owner = xr.Dataset()
+        sentinel = object()
+        cache.set(owner, (2,), sentinel)
+        assert cache.get(owner, (2,)) is sentinel
+
+    def test_none_is_a_storable_value(self):
+        cache = _PosteriorCache()
+        owner = xr.Dataset()
+        cache.set(owner, (), None)
+        assert cache.get(owner, ()) is None
+
+    def test_live_owner_with_different_identity_misses(self):
+        """Two live, equal-looking owners must not share an entry."""
+        cache = _PosteriorCache()
+        first, second = xr.Dataset(), xr.Dataset()
+        cache.set(first, (1,), "value")
+        assert cache.get(second, (1,)) is _CACHE_MISS
+        assert cache.get(first, (1,)) == "value"
+
+    def test_key_tail_mismatch_misses(self):
+        cache = _PosteriorCache()
+        owner = xr.Dataset()
+        cache.set(owner, (1, "gdp"), "value")
+        assert cache.get(owner, (2, "gdp")) is _CACHE_MISS
+        assert cache.get(owner, (1, "inflation")) is _CACHE_MISS
+        assert cache.get(owner, (1, "gdp")) == "value"
+
+    def test_dead_referent_misses(self):
+        """The defect this cache exists to close: a collected owner's
+        address may be recycled, so a dead referent must read as a miss."""
+        cache = _PosteriorCache()
+        owner = xr.Dataset()
+        cache.set(owner, (1,), "stale")
+        ref = weakref.ref(owner)
+
+        del owner
+        gc.collect()
+        assert ref() is None
+
+        assert cache.get(xr.Dataset(), (1,)) is _CACHE_MISS
+
+    def test_multiple_owners_all_must_match(self):
+        cache = _PosteriorCache()
+        a, b, other = xr.Dataset(), xr.Dataset(), xr.Dataset()
+        cache.set((a, b), (1,), "value")
+        assert cache.get((a, b), (1,)) == "value"
+        assert cache.get((a, other), (1,)) is _CACHE_MISS
+        assert cache.get((other, b), (1,)) is _CACHE_MISS
+
+    def test_owner_arity_mismatch_misses(self):
+        cache = _PosteriorCache()
+        a, b = xr.Dataset(), xr.Dataset()
+        cache.set((a, b), (1,), "value")
+        assert cache.get(a, (1,)) is _CACHE_MISS
+
+    def test_single_owner_and_one_tuple_are_equivalent(self):
+        cache = _PosteriorCache()
+        owner = xr.Dataset()
+        cache.set(owner, (1,), "value")
+        assert cache.get((owner,), (1,)) == "value"
+
+    def test_non_weakrefable_owner_declines_to_cache(self):
+        """No validity token means no cache — never an unsafe fallback key."""
+        cache = _PosteriorCache()
+        owner = xr.Dataset()
+        cache.set((owner, 3), (1,), "value")
+        assert cache.get((owner, 3), (1,)) is _CACHE_MISS
+
+    def test_non_weakrefable_owner_evicts_a_previous_entry(self):
+        cache = _PosteriorCache()
+        a, b = xr.Dataset(), xr.Dataset()
+        cache.set(a, (1,), "value")
+        cache.set((b, "not-weakrefable"), (1,), "other")
+        assert cache.get(a, (1,)) is _CACHE_MISS
+
+    def test_set_overwrites_the_single_slot(self):
+        cache = _PosteriorCache()
+        a, b = xr.Dataset(), xr.Dataset()
+        cache.set(a, (1,), "first")
+        cache.set(b, (1,), "second")
+        assert cache.get(b, (1,)) == "second"
+        assert cache.get(a, (1,)) is _CACHE_MISS
+
+    def test_clear(self):
+        cache = _PosteriorCache()
+        owner = xr.Dataset()
+        cache.set(owner, (1,), "value")
+        cache.clear()
+        assert cache.get(owner, (1,)) is _CACHE_MISS