feat(contextual): pipeline part 3 — drift analysis + eval

analysis/contextual_drift.py

@@ -0,0 +1,173 @@
+"""Contextual drift vs the anchor year, from finalized per-year centroids.
+
+Loads the centered contextual centroids (`models/contextual/{year}_centered.npy`)
+into one coordinate system (frozen encoder => no alignment needed), computes
+cosine drift of each word vs the anchor year, gates every (word, year) on a
+minimum observation count so noisy low-frequency centroids don't pollute the
+ranking, and emits a long table + a per-word summary that carries a dispersion
+(polysemy) column.
+
+Reuses `analysis/drift.py::compute_drift_from_base`; the summary aggregation is
+local because that drift is keyed by a single `year` column (vs the consecutive
+`year_a`/`year_b` pairs that `compute_drift_summary` expects).
+"""
+from __future__ import annotations
+
+import warnings
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+
+from analysis.drift import compute_drift_from_base
+from config import ANCHOR_YEAR, CONTEXTUAL_MIN_COUNT
+
+
+def load_centroids(
+    contextual_dir: Path, years: list[int] | None = None
+) -> tuple[dict[int, np.ndarray], dict[int, np.ndarray], dict[int, np.ndarray]]:
+    """Return (centered centroids, counts, dispersion) keyed by year, for years present on disk."""
+    contextual_dir = Path(contextual_dir)
+    aligned: dict[int, np.ndarray] = {}
+    counts: dict[int, np.ndarray] = {}
+    disp: dict[int, np.ndarray] = {}
+
+    candidate_years = years
+    if candidate_years is None:
+        candidate_years = sorted(
+            int(p.stem.replace("_centered", ""))
+            for p in contextual_dir.glob("*_centered.npy")
+        )
+
+    for y in candidate_years:
+        cen = contextual_dir / f"{y}_centered.npy"
+        if not cen.exists():
+            continue
+        aligned[y] = np.load(cen)
+        counts[y] = np.load(contextual_dir / f"{y}_count.npy")
+        disp[y] = np.load(contextual_dir / f"{y}_dispersion.npy")
+    return aligned, counts, disp
+
+
+def _mean_dispersion(disp: dict[int, np.ndarray], counts: dict[int, np.ndarray], min_count: int) -> np.ndarray:
+    """Per-word mean dispersion across years where the word is trusted (count >= min_count)."""
+    years = sorted(disp.keys())
+    disp_stack = np.stack([disp[y] for y in years])      # [Y, V]
+    cnt_stack = np.stack([counts[y] for y in years])     # [Y, V]
+    masked = np.where(cnt_stack >= min_count, disp_stack, np.nan)
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", category=RuntimeWarning)  # all-nan slices -> nan
+        return np.nanmean(masked, axis=0)                # [V]; nan where never trusted
+
+
+def compute_contextual_drift(
+    contextual_dir: Path,
+    vocab: dict[str, int],
+    *,
+    base_year: int = ANCHOR_YEAR,
+    min_count: int = CONTEXTUAL_MIN_COUNT,
+    years: list[int] | None = None,
+) -> tuple[pd.DataFrame, pd.DataFrame]:
+    """Return (long_df, summary_df). long_df is per (word, year) drift vs base_year,
+    gated to (word,year) pairs where both that year and the base year have
+    count >= min_count. summary_df is per-word, sorted by total drift, with a
+    dispersion column."""
+    aligned, counts, disp = load_centroids(contextual_dir, years)
+    if base_year not in aligned:
+        raise FileNotFoundError(
+            f"Anchor year {base_year} not found in {contextual_dir} (have {sorted(aligned)})."
+        )
+
+    long_df = compute_drift_from_base(aligned, vocab, base_year)
+    if long_df.empty:
+        return long_df, long_df
+
+    long_df["wid"] = long_df["word"].map(vocab).astype(int)
+    base_cnt = counts[base_year]
+    long_df["count_base"] = base_cnt[long_df["wid"].values]
+    long_df["count_year"] = 0.0
+    for y in aligned:
+        if y == base_year:
+            continue
+        m = long_df["year"] == y
+        long_df.loc[m, "count_year"] = counts[y][long_df.loc[m, "wid"].values]
+
+    valid = (long_df["count_year"] >= min_count) & (long_df["count_base"] >= min_count)
+    long_df = long_df[valid].reset_index(drop=True)
+    if long_df.empty:
+        return long_df, long_df
+
+    # Noise-normalized drift (a t-like statistic): how many sampling-error units
+    # the two yearly centroids sit apart. A centroid is a sample mean of `count`
+    # contextual vectors, so its squared standard error is dispersion/count
+    # (dispersion = E||h - c||^2 = trace of the within-word covariance); for the
+    # difference of two independent years the SE^2 adds. z = ||Δcentroid|| / SE
+    # therefore de-ranks high-variance / low-count words whose centroids merely
+    # wobble, while leaving genuine sense shifts (||Δ|| >> SE) high. Δ is taken
+    # on the *centered* centroids (anisotropy/global drift removed).
+    base_vec = aligned[base_year]
+    long_df["disp_base"] = disp[base_year][long_df["wid"].values]
+    long_df["disp_year"] = 0.0
+    long_df["euclid"] = 0.0
+    for y in aligned:
+        if y == base_year:
+            continue
+        m = (long_df["year"] == y).values
+        wids = long_df.loc[m, "wid"].values
+        long_df.loc[m, "disp_year"] = disp[y][wids]
+        long_df.loc[m, "euclid"] = np.linalg.norm(aligned[y][wids] - base_vec[wids], axis=1)
+    se = np.sqrt(long_df["disp_base"] / long_df["count_base"] + long_df["disp_year"] / long_df["count_year"])
+    long_df["z_score"] = (long_df["euclid"] / se.replace(0.0, np.nan)).fillna(0.0)
+
+    summary = (
+        long_df.groupby("word")
+        .agg(
+            total_z=("z_score", "sum"),
+            max_z=("z_score", "max"),
+            total_drift=("cosine_distance", "sum"),
+            mean_drift=("cosine_distance", "mean"),
+            max_drift=("cosine_distance", "max"),
+            n_years=("year", "count"),
+        )
+        .reset_index()
+    )
+    peak_rows = long_df.loc[long_df.groupby("word")["z_score"].idxmax()][
+        ["word", "year"]
+    ].rename(columns={"year": "peak_year"})
+    summary = summary.merge(peak_rows, on="word")
+
+    mean_disp = _mean_dispersion(disp, counts, min_count)
+    summary["dispersion"] = mean_disp[summary["word"].map(vocab).astype(int).values]
+
+    # Rank by the noise-normalized score so high-variance words don't dominate.
+    summary = summary.sort_values("total_z", ascending=False).reset_index(drop=True)
+    return long_df, summary
+
+
+def run_drift(
+    contextual_dir: Path,
+    vocab: dict[str, int],
+    *,
+    base_year: int = ANCHOR_YEAR,
+    min_count: int = CONTEXTUAL_MIN_COUNT,
+    years: list[int] | None = None,
+) -> pd.DataFrame:
+    """Compute drift and write drift_vs_{base}.parquet + drift_summary.parquet."""
+    contextual_dir = Path(contextual_dir)
+    long_df, summary = compute_contextual_drift(
+        contextual_dir, vocab, base_year=base_year, min_count=min_count, years=years
+    )
+    if summary.empty:
+        print(f"  [drift] no (word,year) pairs passed min_count={min_count}; nothing written", flush=True)
+        return summary
+
+    long_path = contextual_dir / f"drift_vs_{base_year}.parquet"
+    summary_path = contextual_dir / "drift_summary.parquet"
+    long_df.to_parquet(long_path)
+    summary.to_parquet(summary_path)
+    print(
+        f"  [drift] {len(long_df):,} (word,year) rows, {len(summary):,} words "
+        f"(min_count={min_count}) -> {long_path.name}, {summary_path.name}",
+        flush=True,
+    )
+    return summary

scripts/contextual_drift.py

@@ -5,14 +5,14 @@
             finalize that year (writes models/contextual/{year}*.npy). Resumable.
   finalize  re-derive {year}*.npy from existing partial accumulators (cheap way
             to re-tune centering / PCA-removal without re-streaming).
-  all       stream (which finalizes each year on completion).
-
-The drift stage (parquet + summary + eval) lands in a follow-up PR alongside
-analysis/contextual_drift.py.
+  all       stream (finalizing each year on completion) + drift.
+  drift     compute cosine drift vs the anchor year from finalized centroids,
+            writing drift_vs_{anchor}.parquet + drift_summary.parquet.
 
     uv run python scripts/contextual_drift.py --smoke              # ~2 min code-path + resume check
     uv run python scripts/contextual_drift.py --stage stream       # full run (background, ~1.7 days)
     uv run python scripts/contextual_drift.py --stage finalize     # re-finalize from partials
+    uv run python scripts/contextual_drift.py --stage drift        # (re)compute drift artifacts
 """
 import argparse
 import hashlib
@@ -25,8 +25,10 @@
 sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
 
 from config import (
+    ANCHOR_YEAR,
     CONTEXTUAL_DIR,
     CONTEXTUAL_LAYER,
+    CONTEXTUAL_MIN_COUNT,
     CONTEXTUAL_MODEL,
     CONTEXTUAL_POOLING,
     CONTEXTUAL_MAX_LEN,
@@ -37,6 +39,7 @@
     VOCAB_DIR,
     YEARS,
 )
+from analysis.contextual_drift import run_drift
 from pipeline.contextual_aggregate import ContextualAccumulator
 from pipeline.contextual_finalize import finalize_year, load_partial
 from pipeline.contextual_model import ContextualEncoder
@@ -67,7 +70,7 @@ def main() -> None:
     ap.add_argument("--device", default="cuda")
     ap.add_argument("--model", default=CONTEXTUAL_MODEL)
     ap.add_argument("--years", type=int, nargs="+", default=None)
-    ap.add_argument("--stage", choices=["stream", "finalize", "all"], default="all")
+    ap.add_argument("--stage", choices=["stream", "finalize", "drift", "all"], default="all")
     ap.add_argument("--target-tokens", type=int, default=None)
     a = ap.parse_args()
 
@@ -104,7 +107,7 @@ def snapshots_for(y: int) -> list[str] | None:
         encoder = ContextualEncoder(model_name=a.model, device=a.device)
         print(f"  encoder hidden_dim={encoder.hidden_dim}", flush=True)
 
-    for year in years:
+    for year in years if a.stage != "drift" else []:
         final_path = out_dir / f"{year}.npy"
 
         if a.stage == "finalize":
@@ -141,6 +144,9 @@ def snapshots_for(y: int) -> list[str] | None:
         if a.device == "cuda":
             torch.cuda.empty_cache()
 
+    if a.stage in ("drift", "all"):
+        run_drift(out_dir, vocab, base_year=ANCHOR_YEAR, min_count=CONTEXTUAL_MIN_COUNT, years=years)
+
     print("DONE", flush=True)
 
 

scripts/contextual_eval.py

@@ -0,0 +1,142 @@
+"""Local evaluation of the contextual-drift artifacts (plan section Verification).
+
+Runs the sanity checks that tell us the contextual signal is real:
+  - anisotropy:   random-pair cosine high on raw centroids, ~0 after centering.
+  - stable words: house/table/water move little between first and last year.
+  - drifters:     known sense-shifters (covid/zoom/mask/remote/delve) rank high;
+                  function-ish words sit low (top-N from the drift summary).
+  - dispersion:   polysemous words (bank/apple/python) > monosemous (house/water).
+  - noise floor:  median per-word total drift vs the drifters' drift (a coarse
+                  SNR; only meaningful with the full multi-year run).
+
+Works with whatever years are present, so it is useful on the 2-year smoke too.
+
+    uv run python scripts/contextual_eval.py --smoke
+    uv run python scripts/contextual_eval.py                 # models/contextual
+"""
+import argparse
+import sys
+from pathlib import Path
+
+import numpy as np
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+
+from analysis.contextual_drift import compute_contextual_drift, load_centroids
+from analysis.drift import cosine_similarity_rows
+from config import ANCHOR_YEAR, CONTEXTUAL_DIR, CONTEXTUAL_MIN_COUNT, VOCAB_DIR
+from pipeline.vocab import load_vocab
+
+STABLE = ["house", "table", "water", "mother", "river", "stone"]
+DRIFTERS = ["covid", "zoom", "mask", "remote", "delve", "lockdown", "vaccine", "pandemic"]
+POLYSEMOUS = ["bank", "apple", "python", "mouse", "spring", "crane"]
+
+
+def _mean_pair_cos(mat: np.ndarray, ids: np.ndarray, n: int = 4000, seed: int = 0) -> float:
+    rng = np.random.default_rng(seed)
+    a, b = rng.choice(ids, n), rng.choice(ids, n)
+    va, vb = mat[a], mat[b]
+    na, nb = np.linalg.norm(va, axis=1), np.linalg.norm(vb, axis=1)
+    ok = (na > 0) & (nb > 0)
+    return float((np.sum(va[ok] * vb[ok], axis=1) / (na[ok] * nb[ok])).mean())
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--smoke", action="store_true")
+    ap.add_argument("--dir", default=None)
+    ap.add_argument("--min-count", type=int, default=CONTEXTUAL_MIN_COUNT)
+    a = ap.parse_args()
+
+    contextual_dir = Path(a.dir) if a.dir else (Path("models/contextual_smoke") if a.smoke else CONTEXTUAL_DIR)
+    vocab = load_vocab(VOCAB_DIR / "vocab.json")
+    aligned, counts, disp = load_centroids(contextual_dir)
+    years = sorted(aligned)
+    if not years:
+        print(f"No centered centroids in {contextual_dir}; run the pipeline first.")
+        return
+    print(f"dir={contextual_dir} years={years} min_count={a.min_count}\n")
+
+    # 1) Anisotropy (per year): raw {year}.npy vs centered {year}_centered.npy.
+    print("== anisotropy (mean random-pair cosine; raw high, centered ~0) ==")
+    for y in years:
+        raw = np.load(contextual_dir / f"{y}.npy")
+        cen = aligned[y]
+        trusted = np.where(counts[y] >= a.min_count)[0]
+        trusted = trusted[trusted != 0]
+        if len(trusted) < 50:
+            print(f"  {y}: only {len(trusted)} trusted words, skipping")
+            continue
+        print(f"  {y}: raw {_mean_pair_cos(raw, trusted):+.4f}   centered {_mean_pair_cos(cen, trusted):+.4f}")
+
+    # 2) Stable words: similarity between first and last available year.
+    y0, y1 = years[0], years[-1]
+    print(f"\n== stable words: cos(centered {y0}, centered {y1}) [expect high] ==")
+    a0, a1 = aligned[y0], aligned[y1]
+    for w in STABLE:
+        wid = vocab.get(w)
+        if wid is None or counts[y0][wid] < a.min_count or counts[y1][wid] < a.min_count:
+            print(f"  {w:9s} (insufficient obs)")
+            continue
+        sim = float(cosine_similarity_rows(a0[wid:wid+1], a1[wid:wid+1])[0])
+        print(f"  {w:9s} cos={sim:+.3f}  drift={1-sim:.3f}  (n={counts[y0][wid]:.0f}/{counts[y1][wid]:.0f})")
+
+    # 3) Dispersion: polysemous vs monosemous (year-invariant polysemy signal).
+    print("\n== dispersion (polysemous should exceed monosemous) ==")
+    for label, words in [("polysemous", POLYSEMOUS), ("monosemous", STABLE)]:
+        vals = []
+        for w in words:
+            wid = vocab.get(w)
+            ds = [disp[y][wid] for y in years if wid is not None and counts[y][wid] >= a.min_count]
+            if ds:
+                vals.append(np.mean(ds))
+        if vals:
+            print(f"  {label:11s} mean dispersion = {np.mean(vals):.2f}  (over {len(vals)} words)")
+
+    # 4) Drift ranking + coarse SNR.
+    print(f"\n== drift vs {ANCHOR_YEAR}: ranking + noise floor ==")
+    if ANCHOR_YEAR not in aligned:
+        print(f"  anchor {ANCHOR_YEAR} absent; skipping drift ranking")
+        return
+    long_df, summary = compute_contextual_drift(contextual_dir, vocab, base_year=ANCHOR_YEAR, min_count=a.min_count)
+    if summary.empty:
+        print("  no words passed gating")
+        return
+    noise = float(summary["total_drift"].median())
+    print(f"  words ranked by noise-normalized z: {len(summary):,}   "
+          f"median total_drift (cosine noise floor) = {noise:.3f}")
+    print("  top 15 drifters (z-ranked):")
+    for _, r in summary.head(15).iterrows():
+        print(f"    {r['word']:15s} z_sum={r['total_z']:6.1f} peak@{int(r['peak_year'])} "
+              f"cos_total={r['total_drift']:.3f} disp={r['dispersion']:.1f}")
+    present = summary[summary["word"].isin(DRIFTERS)].copy()
+    if not present.empty:
+        print("  known drifters present:")
+        for _, r in present.sort_values("total_z", ascending=False).iterrows():
+            snr = r["total_drift"] / noise if noise else float("nan")
+            print(f"    {r['word']:15s} z_sum={r['total_z']:6.1f} peak@{int(r['peak_year'])} "
+                  f"cos_total={r['total_drift']:.3f} (cos SNR {snr:.1f}x)")
+    else:
+        print("  (no known drifters in gated set — expected on the 2-year smoke; "
+              "covid/zoom emerge only with 2020+ years in the full run)")
+
+    # Per-year trajectories: the clean, interpretable view. Drifters should be
+    # flat then spike at their event year; stable controls flat throughout.
+    drift_years = sorted(long_df["year"].unique())
+    print(f"\n== per-year drift vs {ANCHOR_YEAR} (cosine distance) — trajectories ==")
+    print("  " + "word".ljust(13) + "".join(f"{y:>8}" for y in drift_years))
+    for label, words in [("drifters", DRIFTERS), ("stable", STABLE)]:
+        print(f"  -- {label} --")
+        for w in words:
+            sub = long_df[long_df["word"] == w].set_index("year")
+            if sub.empty:
+                continue
+            cells = "".join(
+                f"{sub.loc[y, 'cosine_distance']:8.3f}" if y in sub.index else f"{'-':>8}"
+                for y in drift_years
+            )
+            print(f"  {w:13s}{cells}")
+
+
+if __name__ == "__main__":
+    main()