fig_regimes.py: reproducible cross-model depth-regimes figure (E2B+12B filled, 3 scans pending)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-15 22:05:29 +02:00
parent 8187b79b50
commit dc43769118
2 changed files with 68 additions and 0 deletions
@@ -0,0 +1,68 @@
+"""Cross-model depth-regimes figure: sensor / motor / persistence curves by
+layer, one panel per lens-mapped model, workspace band shaded (from
+results/REGIMES.json). Replaces the ad-hoc single-model regimes.png.
+
+Panel sources: per-model regimes*.pt diagnostics (exp4_regimes.py output:
+dict with per-layer 'sensor', 'motor', 'persist' fractions). Pending models
+render as placeholders and fill in on re-render.
+"""
+import json
+from pathlib import Path
+
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import torch
+
+ROOT = Path(__file__).resolve().parent.parent
+REG = json.load(open(ROOT / "results/REGIMES.json"))
+
+PANELS = [
+    ("gemma-4-E2B (2B eff.)", "results/regimes.pt", "google/gemma-4-E2B-it"),
+    ("gemma-4-E4B (4B eff., E2B ⊂ E4B)", "results/regimes_e4b.pt",
+     "google/gemma-4-E4B-it"),
+    ("gemma-4-12B", "results-12b/regimes.pt", "google/gemma-4-12B-it"),
+    ("gemma-4-26B-A4B (MoE)", "results/regimes_26b_a4b.pt",
+     "google/gemma-4-26B-A4B-it"),
+    ("gemma-4-31B", "results/regimes_31b.pt", "google/gemma-4-31B-it"),
+]
+
+fig, axes = plt.subplots(1, len(PANELS), figsize=(3.4 * len(PANELS), 3.4),
+                         sharey=True)
+for ax, (title, ptfile, key) in zip(axes, PANELS):
+    ax.set_title(title, fontsize=9)
+    ax.set_xlabel("layer", fontsize=8)
+    ax.grid(True, color="#ececec", lw=0.6)
+    ax.set_axisbelow(True)
+    for s in ("top", "right"):
+        ax.spines[s].set_visible(False)
+    p = ROOT / ptfile
+    if not p.exists():
+        ax.text(0.5, 0.5, "scan pending", transform=ax.transAxes,
+                ha="center", color="#999", fontsize=10)
+        continue
+    d = torch.load(p, map_location="cpu")
+    sensor = torch.as_tensor(d["sensor"]).float()
+    motor = torch.as_tensor(d["motor"]).float()
+    persist = torch.as_tensor(d["persist"]).float()
+    L = len(sensor)
+    xs = range(L)
+    band = REG.get(key, {}).get("workspace_band")
+    if band:
+        ax.axvspan(band[0], band[1], color="#2f855a", alpha=0.10)
+    kv = REG.get(key, {}).get("kv_share_start")
+    if isinstance(kv, int):
+        ax.axvline(kv, color="#c53030", lw=0.9, ls=":")
+    ax.plot(xs, sensor, color="#2b6cb0", lw=1.6, label="sensor")
+    ax.plot(xs, motor, color="#c53030", lw=1.6, label="motor")
+    ax.plot(xs, persist, color="#2f855a", lw=1.4, ls="--", label="persistence")
+axes[0].set_ylabel("fraction", fontsize=9)
+axes[0].legend(fontsize=7.5, frameon=False, loc="upper left")
+fig.suptitle("Depth regimes under the J-lens across gemma-4 scales — green "
+             "shade: workspace band (REGIMES.json); red dotted: KV-share "
+             "boundary (entrances above are structurally null)",
+             fontsize=10.5, y=1.04)
+fig.tight_layout()
+out = ROOT / "results/regimes_all_models.png"
+fig.savefig(out, dpi=140, facecolor="white", bbox_inches="tight")
+print("wrote", out)