results/parse_and_plot.py

import re
import csv
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

BASE = "/Users/javierhuang/du/Apps/AI/ICML_2026_Agent_Repro/results"

acc_re = re.compile(
    r"step (\d+): train acc ([\d.]+), val acc ([\d.]+), test acc ([\d.]+)"
)
loss_re = re.compile(
    r"step (\d+): train loss ([\d.]+), val loss ([\d.]+)"
)


def parse_claim1(path):
    rows = []
    with open(path) as f:
        for line in f:
            m = acc_re.search(line)
            if m:
                rows.append(
                    {
                        "step": int(m.group(1)),
                        "train_acc": float(m.group(2)),
                        "val_acc": float(m.group(3)),
                        "test_acc": float(m.group(4)),
                    }
                )
    return rows


def parse_claim2(path, section=None):
    rows = []
    in_section = section is None
    with open(path) as f:
        for line in f:
            if section == "rope" and "=== RoPE run" in line:
                in_section = True
                continue
            if section == "rope" and "=== PoPE run" in line:
                in_section = False
                continue
            if section == "pope" and "=== PoPE run" in line:
                in_section = True
                continue
            if not in_section:
                continue
            m = loss_re.search(line)
            if m:
                rows.append(
                    {
                        "step": int(m.group(1)),
                        "train_loss": float(m.group(2)),
                        "val_loss": float(m.group(3)),
                    }
                )
    return rows


# ---- Claim 1 ----
c1_rope = parse_claim1(f"{BASE}/claim1/rope_full.log")
c1_pope = parse_claim1(f"{BASE}/claim1/pope_full.log")

with open(f"{BASE}/claim1/rope_accuracy.csv", "w", newline="") as f:
    w = csv.DictWriter(f, fieldnames=["step", "train_acc", "val_acc", "test_acc"])
    w.writeheader()
    w.writerows(c1_rope)

with open(f"{BASE}/claim1/pope_accuracy.csv", "w", newline="") as f:
    w = csv.DictWriter(f, fieldnames=["step", "train_acc", "val_acc", "test_acc"])
    w.writeheader()
    w.writerows(c1_pope)

fig, ax = plt.subplots(figsize=(8, 5))
ax.plot([r["step"] for r in c1_rope], [r["val_acc"] * 100 for r in c1_rope], label="RoPE (val acc)", color="#d62728")
ax.plot([r["step"] for r in c1_pope], [r["val_acc"] * 100 for r in c1_pope], label="PoPE (val acc)", color="#1f77b4")
ax.axhline(11.16, color="#d62728", linestyle="--", alpha=0.5, label="Paper: RoPE ceiling (11.16%)")
ax.axhline(94.82, color="#1f77b4", linestyle="--", alpha=0.5, label="Paper: PoPE (94.82%)")
ax.set_xlabel("Training iteration")
ax.set_ylabel("Validation accuracy (%)")
ax.set_title("Claim 1: Synthetic Indirect Indexing — RoPE vs PoPE\n(our reproduction vs paper's reported endpoints)")
ax.legend(loc="center left", fontsize=8)
ax.set_ylim(0, 100)
fig.tight_layout()
fig.savefig(f"{BASE}/claim1/accuracy_curve.png", dpi=150)
plt.close(fig)

# ---- Claim 2 ----
seeds = {"1337": f"{BASE}/claim2/seed1337_full.log", "42": f"{BASE}/claim2/seed42_full.log"}
all_rows = {}
for seed, path in seeds.items():
    all_rows[(seed, "rope")] = parse_claim2(path, "rope")
    all_rows[(seed, "pope")] = parse_claim2(path, "pope")

with open(f"{BASE}/claim2/all_runs_loss.csv", "w", newline="") as f:
    w = csv.writer(f)
    w.writerow(["seed", "pos_type", "step", "train_loss", "val_loss"])
    for (seed, pos), rows in all_rows.items():
        for r in rows:
            w.writerow([seed, pos, r["step"], r["train_loss"], r["val_loss"]])

fig, axes = plt.subplots(1, 2, figsize=(12, 5), sharey=True)
for ax, seed in zip(axes, ["1337", "42"]):
    rope_rows = all_rows[(seed, "rope")]
    pope_rows = all_rows[(seed, "pope")]
    ax.plot([r["step"] for r in rope_rows], [r["val_loss"] for r in rope_rows], label="RoPE", color="#d62728")
    ax.plot([r["step"] for r in pope_rows], [r["val_loss"] for r in pope_rows], label="PoPE", color="#1f77b4")
    best_rope = min(rope_rows, key=lambda r: r["val_loss"])
    best_pope = min(pope_rows, key=lambda r: r["val_loss"])
    ax.scatter([best_rope["step"]], [best_rope["val_loss"]], color="#d62728", zorder=5, s=50, marker="*")
    ax.scatter([best_pope["step"]], [best_pope["val_loss"]], color="#1f77b4", zorder=5, s=50, marker="*")
    ax.axhline(0.5081, color="#d62728", linestyle="--", alpha=0.4)
    ax.axhline(0.4889, color="#1f77b4", linestyle="--", alpha=0.4)
    ax.set_title(f"Seed {seed}\nbest: RoPE={best_rope['val_loss']:.4f}, PoPE={best_pope['val_loss']:.4f}")
    ax.set_xlabel("Training iteration")
    ax.legend(fontsize=8)
axes[0].set_ylabel("Validation loss (NLL)")
fig.suptitle("Claim 2: JSB Chorales — RoPE vs PoPE across two seeds\n(dashed lines = paper's reported RoPE/PoPE NLL)")
fig.tight_layout()
fig.savefig(f"{BASE}/claim2/loss_curves.png", dpi=150)
plt.close(fig)

print("Claim 1 RoPE final:", c1_rope[-1])
print("Claim 1 PoPE final:", c1_pope[-1])
for seed in ["1337", "42"]:
    best_rope = min(all_rows[(seed, "rope")], key=lambda r: r["val_loss"])
    best_pope = min(all_rows[(seed, "pope")], key=lambda r: r["val_loss"])
    print(f"Seed {seed}: best RoPE val_loss={best_rope['val_loss']:.4f} @ step {best_rope['step']}, best PoPE val_loss={best_pope['val_loss']:.4f} @ step {best_pope['step']}")

print("Done. Wrote CSVs and PNGs to", BASE)