"""Exogenous noise for counterfactual replay. Every random decision a step makes is drawn through a :class:`NoiseHandle`. Draws are *addressed*, not streamed: the value for ``(world, step, key, element)`` is a deterministic hash of that address. Consequences: * Swapping step A's implementation never shifts the random numbers seen by step B (no "RNG stream drift"), so every coalition of fixes is evaluated under the same exogenous noise -- common random numbers (CRN). * Categorical draws are keyed by option *label*, so a counterfactual call that scores the same options with new logits reuses the same Gumbel noise per option. That is the Gumbel-max structural causal model: ``argmax(logits + g)`` with ``g`` fixed is the counterfactual sample. A :class:`NoiseSource` ("world") may carry *overrides*: explicit noise values for some addresses. They are used to replay recorded noise exactly (:meth:`NoiseSource.from_trace`) or to plug in noise abducted from an observed trace (:func:`abduct`), when the original noise is unknown. """ from __future__ import annotations import hashlib import math from dataclasses import dataclass, field from typing import Any, Hashable, Mapping, Sequence import numpy as np __all__ = [ "Draw", "NoiseHandle", "NoiseSource", "abduct", "gumbel_max", "posterior_gumbels", "truncated_gumbel", ] # --------------------------------------------------------------------------- # Gumbel-max helpers # --------------------------------------------------------------------------- def _u01(*parts: Any) -> float: """Deterministic uniform in the open interval (0, 1) from a hashed address.""" h = hashlib.blake2b(digest_size=8) for p in parts: h.update(repr(p).encode()) h.update(b"\x1f") x = int.from_bytes(h.digest(), "little") return (x + 0.5) / 2.0**64 def _gumbel_from_u(u: float) -> float: return -math.log(-math.log(u)) def gumbel_max(logits: Sequence[float], gumbels: Sequence[float]) -> int: """Index of ``argmax(logits + gumbels)``. If ``gumbels`` are i.i.d. standard Gumbel, the result is distributed as ``softmax(logits)``. Keeping ``gumbels`` fixed while changing ``logits`` gives the Gumbel-max counterfactual sample. """ best, idx = -math.inf, 0 for j, (l, g) in enumerate(zip(logits, gumbels)): # pure Python: options lists are short if l + g > best: best, idx = l + g, j return idx def truncated_gumbel(rng: np.random.Generator, loc: float, upper: float) -> float: """Sample from Gumbel(loc) conditioned on being <= ``upper``. Uses the inverse-CDF identity: if ``z ~ Gumbel(loc)`` then ``-log(exp(-upper) + exp(-z))`` has the truncated distribution. """ z = loc + rng.gumbel() # -log(exp(-upper) + exp(-z)), computed stably m = min(upper, z) return m - math.log1p(math.exp(-abs(upper - z))) def posterior_gumbels(logits: Sequence[float], observed: int, rng: np.random.Generator) -> np.ndarray: """Sample Gumbel noise ``g`` from its posterior given ``argmax(logits+g) == observed``. Top-down construction (Maddison et al. 2014; used for Gumbel-max SCM counterfactuals by Oberst & Sontag 2019): the max of the perturbed values is Gumbel(logsumexp(logits)) and is independent of the argmax; it is assigned to the observed index, and every other perturbed value is a Gumbel(logit_j) truncated below that max. Returns the *noise* (perturbed value minus logit), so ``gumbel_max(logits, g) == observed`` always. Options with ``-inf`` logit get unconditioned prior noise. """ logits = [float(x) for x in logits] if not math.isfinite(logits[observed]): raise ValueError("observed option has -inf logit (probability zero)") fin = [x for x in logits if math.isfinite(x)] mx = max(fin) lse = mx + math.log(sum(math.exp(x - mx) for x in fin)) g = rng.gumbel(size=len(logits) + 1) # one batch call: cheaper than many top = lse + g[-1] noise = np.empty(len(logits)) for j, lj in enumerate(logits): if j == observed: noise[j] = top - lj elif math.isfinite(lj): z = lj + g[j] # z ~ Gumbel(lj), truncate below top m = min(top, z) noise[j] = m - math.log1p(math.exp(-abs(top - z))) - lj else: noise[j] = g[j] return noise # --------------------------------------------------------------------------- # Noise sources and per-step handles # --------------------------------------------------------------------------- @dataclass class Draw: """One recorded random draw made by a step.""" step: str key: str # address within the step, e.g. "choice#0" kind: str # "categorical" | "uniform" labels: list # element ids (option labels, or 0..n-1 for uniforms) values: list # raw noise: Gumbel per label, or uniforms logits: list | None = None chosen: Any = None class NoiseSource: """A "world": one realization of all exogenous noise for one test case. Parameters ---------- seed: Any hashable; with no overrides the world is fully determined by it. overrides: ``{(step, key): {element_id: value}}`` explicit noise values that take precedence over hashed values (used for exact replay and abduction). name: Unique identifier used in step-level caching. Defaults to ``seed=...``. Two sources with the same name must produce the same noise. """ def __init__(self, seed: Hashable, overrides: Mapping | None = None, name: str | None = None): self.seed = seed self.overrides = dict(overrides or {}) self.world_id = name if name is not None else f"seed={seed!r}" if self.overrides and name is None: # Overrides change the noise, so they must change the cache identity. self.world_id += f"/ov{abs(hash(repr(sorted(self.overrides.items(), key=repr)))):x}" def handle(self, step: str) -> "NoiseHandle": return NoiseHandle(self, step) @classmethod def from_trace(cls, trace, name: str | None = None) -> "NoiseSource": """A world that replays exactly the noise recorded in ``trace``. Useful when the original noise came from somewhere non-reproducible (e.g. it was logged in production); addresses not in the trace fall back to hashed noise of the trace's seed. """ ov = {(d.step, d.key): dict(zip(d.labels, d.values)) for d in trace.draws} return cls(trace.world_seed, ov, name or f"replay({trace.world_id})") def __repr__(self) -> str: return f"NoiseSource({self.world_id})" class NoiseHandle: """What a step receives as ``noise``: addressed random draws for one step call. Repeated calls with the same ``key`` inside one step execution get distinct addresses (``key#0``, ``key#1``, ...), aligned by call order. """ def __init__(self, source: NoiseSource, step: str): self.source = source self.step = step self.draws: list[Draw] = [] self._counts: dict[str, int] = {} def _address(self, key: str) -> str: n = self._counts.get(key, 0) self._counts[key] = n + 1 return f"{key}#{n}" def _raw(self, addr: str, element: Hashable, kind: str) -> float: ov = self.source.overrides.get((self.step, addr)) if ov is not None and element in ov: return float(ov[element]) u = _u01(self.source.seed, self.step, addr, element) return _gumbel_from_u(u) if kind == "gumbel" else u # -- public draw API ---------------------------------------------------- def uniform(self, key: str = "u") -> float: """One U(0,1) draw.""" return float(self.uniforms(1, key)[0]) def uniforms(self, n: int, key: str = "u") -> np.ndarray: addr = self._address(key) vals = [self._raw(addr, i, "uniform") for i in range(n)] self.draws.append(Draw(self.step, addr, "uniform", list(range(n)), vals)) return np.array(vals) def bernoulli(self, p: float, key: str = "bern") -> bool: """Monotone coupling: ``u < p``. Raising p never flips True to False.""" return self.uniform(key) < p def gumbels(self, labels: Sequence[Hashable], key: str = "choice") -> np.ndarray: """Standard Gumbel noise, one per label (does not record a choice).""" addr = self._address(key) vals = [self._raw(addr, lab, "gumbel") for lab in labels] self.draws.append(Draw(self.step, addr, "gumbel", list(labels), vals)) return np.array(vals) def categorical(self, logits: Mapping[Hashable, float] | Sequence[float], key: str = "choice") -> Any: """Sample one option via Gumbel-max and record it. ``logits`` is either ``{label: logit}`` (returns a label) or a sequence (returns an index). Noise is keyed by label, so counterfactual calls with different logits over the same labels share noise. """ if isinstance(logits, Mapping): labels = list(logits.keys()) lv = [float(logits[k]) for k in labels] as_label = True else: lv = [float(x) for x in logits] labels = list(range(len(lv))) as_label = False addr = self._address(key) g = [self._raw(addr, lab, "gumbel") for lab in labels] idx = gumbel_max(lv, g) chosen = labels[idx] if as_label else idx self.draws.append(Draw(self.step, addr, "categorical", labels, g, lv, chosen)) return chosen def abduct(trace, seed: int, name: str | None = None) -> NoiseSource: """Posterior world given only the *observed choices* in a factual trace. For each categorical draw we sample Gumbel noise from its posterior given the observed option (:func:`posterior_gumbels`). Because the logits of a step are a deterministic function of its observed inputs, the posterior factorizes over draws. Non-categorical draws cannot be abducted from their outcome alone and are re-sampled from the prior (hashed from ``seed``). """ rng = np.random.default_rng(seed) ov = {} for d in trace.draws: if d.kind != "categorical": continue g = posterior_gumbels(d.logits, d.labels.index(d.chosen), rng) ov[(d.step, d.key)] = dict(zip(d.labels, g.tolist())) return NoiseSource(("posterior", trace.world_id, seed), ov, name or f"posterior({trace.world_id},{seed})")