File size: 5,867 Bytes
8f46582 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 | #!/usr/bin/env python3
"""Regenerate intervention figures for the README from saved JSON."""
from __future__ import annotations
import json
from pathlib import Path
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
ROOT = Path(__file__).resolve().parents[1]
IN = ROOT / "figs" / "interventions"
OUT = IN
COLORS = {"CSO": "#c0392b", "BT W=5": "#e67e22", "BT W=2": "#2980b9"}
ORDER = ["CSO", "BT W=5", "BT W=2"]
def load_methods():
out = {}
for name, fname in [
("CSO", "L20_CSO.json"),
("BT W=5", "L20_BT_W5.json"),
("BT W=2", "L20_BT_W2.json"),
]:
out[name] = json.load(open(IN / fname))
return out
def plot_all_methods(methods):
fig, axes = plt.subplots(2, 2, figsize=(12.5, 8.5))
fig.suptitle(
"L20 interventions: CSO vs BT W=5 vs BT W=2 (128 graphs)",
fontsize=13,
fontweight="bold",
)
# A: pinlast-k
ax = axes[0, 0]
for name in ORDER:
pk = methods[name]["pinlastk"]
ks = [0] + list(pk["ks"])
ys = [pk["clean"]] + list(pk["rand_mean"])
ax.plot(ks, ys, "o-", color=COLORS[name], label=name, lw=2, ms=5)
ax.set_xlabel("k earlier latents corrupted (last pinned)")
ax.set_ylabel("leaf accuracy")
ax.set_ylim(0, 1.05)
ax.set_title("A. Pin-last: corrupt k earlier thoughts")
ax.grid(True, alpha=0.3)
ax.legend(frameon=False)
# B: aggregate bars
ax = axes[0, 1]
keys = [
("clean", "clean"),
("earlier L1..L19 -> noise (last pinned)", "earlier→noise\n(last pin)"),
("earlier L1..L19 -> other-graph (last pinned)", "earlier→donor\n(last pin)"),
("LAST thought L20 -> other-graph", "last→donor"),
("PROPAGATED: corrupt L10, recompute rest", "prop L10"),
("PROPAGATED: corrupt L19, recompute rest", "prop L19"),
]
x = np.arange(len(keys))
width = 0.25
for i, name in enumerate(ORDER):
rows = {r["label"]: r["acc"] for r in methods[name]["aggregate"]["rows"]}
ys = [rows[k] for k, _ in keys]
ax.bar(x + (i - 1) * width, ys, width, color=COLORS[name], label=name)
ax.set_xticks(x)
ax.set_xticklabels([lab for _, lab in keys], fontsize=8)
ax.set_ylabel("leaf accuracy")
ax.set_ylim(0, 1.05)
ax.set_title("B. Aggregate interventions")
ax.legend(frameon=False, fontsize=8)
ax.grid(True, axis="y", alpha=0.3)
# C: per-slot pin-last
ax = axes[1, 0]
for name in ORDER:
ys = methods[name]["perslot"]["pin_accs"]
xs = np.arange(1, len(ys) + 1)
ax.plot(xs, ys, "o-", color=COLORS[name], label=name, lw=2, ms=4)
ax.axvline(20, color="gray", ls="--", alpha=0.5)
ax.set_xlabel("corrupted latent slot (last pinned unless slot=20)")
ax.set_ylabel("leaf accuracy")
ax.set_ylim(0, 1.05)
ax.set_title("C. Per-slot pin-last")
ax.grid(True, alpha=0.3)
ax.legend(frameon=False)
# D: per-slot propagate
ax = axes[1, 1]
for name in ORDER:
ys = methods[name]["perslot"]["prop_accs"]
xs = np.arange(1, len(ys) + 1)
ax.plot(xs, ys, "o-", color=COLORS[name], label=name, lw=2, ms=4)
ax.set_xlabel("corrupted latent slot (then recompute rest)")
ax.set_ylabel("leaf accuracy")
ax.set_ylim(0, 1.05)
ax.set_title("D. Per-slot propagate")
ax.grid(True, alpha=0.3)
ax.legend(frameon=False)
fig.tight_layout(rect=[0, 0, 1, 0.96])
path = OUT / "L20_all_methods.png"
fig.savefig(path, dpi=160)
plt.close(fig)
print(f"saved {path}")
def plot_pinlast_k(methods):
fig, ax = plt.subplots(figsize=(7.2, 4.2))
for name in ORDER:
pk = methods[name]["pinlastk"]
ks = [0] + list(pk["ks"])
ys = [pk["clean"]] + list(pk["rand_mean"])
ax.plot(ks, ys, "o-", color=COLORS[name], label=name, lw=2.2, ms=6)
ax.set_title("L20: pin-last intervention")
ax.set_xlabel("k earlier latents corrupted (last pinned)")
ax.set_ylabel("leaf accuracy")
ax.set_ylim(0, 1.05)
ax.grid(True, alpha=0.3)
ax.legend(frameon=False)
fig.tight_layout()
path = OUT / "pinlast_k_L20.png"
fig.savefig(path, dpi=160)
plt.close(fig)
print(f"saved {path}")
def plot_pinlast_table(methods):
ks = methods["BT W=2"]["pinlastk"]["ks"]
rows = ["clean"] + [str(k) for k in ks] + ["all 19"]
col_labels = ["k"] + ORDER
cell = []
for name in ORDER:
pk = methods[name]["pinlastk"]
vals = [pk["clean"]] + list(pk["rand_mean"]) + [pk["all"]]
cell.append([f"{v:.3f}" for v in vals])
data = [[rows[i]] + [cell[j][i] for j in range(3)] for i in range(len(rows))]
fig, ax = plt.subplots(figsize=(7.5, 4.8))
ax.axis("off")
ax.set_title(
"L20 pin-last: k earlier thoughts corrupted (last pinned)",
fontsize=11,
pad=12,
)
table = ax.table(
cellText=data,
colLabels=col_labels,
loc="center",
cellLoc="center",
)
table.auto_set_font_size(False)
table.set_fontsize(9)
table.scale(1.15, 1.35)
for (r, c), cell_obj in table.get_celld().items():
if r == 0:
cell_obj.set_facecolor("#ecf0f1")
cell_obj.set_text_props(fontweight="bold")
if c == 0:
cell_obj.set_text_props(fontweight="bold")
ckpts = ", ".join(
f"{n}={Path(methods[n]['ckpt']).name.split('_')[-1]}" for n in ORDER
)
fig.text(0.5, 0.02, f"ckpts: {ckpts}", ha="center", fontsize=8, color="#555")
fig.tight_layout(rect=[0, 0.05, 1, 1])
path = OUT / "pinlast_k_L20_table.png"
fig.savefig(path, dpi=160)
plt.close(fig)
print(f"saved {path}")
def main():
methods = load_methods()
plot_all_methods(methods)
plot_pinlast_k(methods)
plot_pinlast_table(methods)
if __name__ == "__main__":
main()
|