import argparse
import csv
import os
from pathlib import Path
def main():
parser = argparse.ArgumentParser(description="Plot FFN low-rank sweep results")
parser.add_argument("--csv", default="target/low_rank_ffn.csv", help="input CSV path")
parser.add_argument("--png", default="target/low_rank_ffn.png", help="output PNG path")
args = parser.parse_args()
try:
import matplotlib.pyplot as plt
except ImportError as exc:
raise SystemExit("matplotlib is required: pip install matplotlib") from exc
ranks = []
low_rank_params = []
logit_l2 = []
logit_mse = []
top1 = []
residual_l2 = []
residual_top1 = []
with open(args.csv, newline="") as f:
reader = csv.DictReader(f)
for row in reader:
ranks.append(int(row["rank"]))
low_rank_params.append(int(row["low_rank_params"]))
logit_l2.append(float(row["logit_l2"]))
logit_mse.append(float(row["logit_mse"]))
top1.append(float(row["top1_agreement"]) * 100.0)
residual_l2.append(float(row["residual_logit_l2"]))
residual_top1.append(float(row["residual_top1_agreement"]) * 100.0)
fig, axes = plt.subplots(2, 2, figsize=(12, 10))
fig.suptitle("FFN low-rank + residual prototype")
ax = axes[0, 0]
ax.plot(ranks, logit_l2, "o-", label="low-rank")
ax.plot(ranks, residual_l2, "s--", label="low-rank + dense residual")
ax.set_xlabel("rank")
ax.set_ylabel("mean logit L2 drift")
ax.set_xscale("log", base=2)
ax.set_yscale("log")
ax.legend()
ax.grid(True, which="both", ls="--", alpha=0.5)
ax = axes[0, 1]
ax.plot(low_rank_params, logit_l2, "o-", label="low-rank")
ax.plot(low_rank_params, residual_l2, "s--", label="low-rank + dense residual")
ax.set_xlabel("low-rank parameters")
ax.set_ylabel("mean logit L2 drift")
ax.set_xscale("log")
ax.set_yscale("log")
ax.legend()
ax.grid(True, which="both", ls="--", alpha=0.5)
ax = axes[1, 0]
ax.plot(ranks, logit_mse, "o-", label="low-rank MSE")
ax.set_xlabel("rank")
ax.set_ylabel("mean logit MSE")
ax.set_xscale("log", base=2)
ax.set_yscale("log")
ax.legend()
ax.grid(True, which="both", ls="--", alpha=0.5)
ax = axes[1, 1]
ax.plot(ranks, top1, "o-", label="low-rank top-1")
ax.plot(ranks, residual_top1, "s--", label="low-rank + residual top-1")
ax.set_xlabel("rank")
ax.set_ylabel("top-1 agreement (%)")
ax.set_xscale("log", base=2)
ax.set_ylim(-5, 105)
ax.legend()
ax.grid(True, which="both", ls="--", alpha=0.5)
fig.tight_layout()
out = Path(args.png)
out.parent.mkdir(parents=True, exist_ok=True)
fig.savefig(out, dpi=150)
print(f"Saved {out}")
if __name__ == "__main__":
main()