feat: add parameter-matched EP comparison

1 files changed, 36 insertions, 0 deletions

diff --git a/experiments/analyze_verified.py b/experiments/analyze_verified.py
index 8a1a841..50863b5 100755
--- a/experiments/analyze_verified.py
+++ b/experiments/analyze_verified.py
@@ -124,6 +124,40 @@ def print_local_baselines(root):
               f"{fmt([r['final']['wall_s'] for r in rows], digits=1)} |")
 
 
+def forward_parameter_count(record):
+    args = record["args"]
+    n_in = args.get("n_in") or {"mnist": 784, "fmnist": 784, "cifar10": 3072}[args["dataset"]]
+    sizes = [n_in] + [args["width"]] * args["depth"] + [10]
+    return sum(n_out * n_in_ + n_out for n_in_, n_out in zip(sizes[:-1], sizes[1:]))
+
+
+def best_epoch_accuracy(record):
+    values = [step["test_acc"] for step in record.get("steps", []) if "test_acc" in step]
+    values.append(record["final"]["test_acc"])
+    return max(values)
+
+
+def print_ep_match(root):
+    records = read_many([
+        os.path.join(root, "ep_original_v2_*.json"),
+        os.path.join(root, "ep_match_v1_*.json"),
+    ])
+    groups = {}
+    for record in records:
+        args = record["args"]
+        key = (method_name(record), args["depth"], args["width"],
+               forward_parameter_count(record), args["epochs"])
+        groups.setdefault(key, []).append(record)
+    print("\n\n## EP near-parameter-matched comparison\n")
+    print("| method | depth × width | forward params | epochs | n | last acc (%) | best acc (%) | wall (s) |")
+    print("|:---|:---|---:|---:|---:|---:|---:|---:|")
+    for (method, depth, width, params, epochs), rows in sorted(groups.items()):
+        print(f"| {method} | {depth} × {width} | {params:,} | {epochs} | {len(rows)} | "
+              f"{fmt([r['final']['test_acc'] for r in rows], percent=True)} | "
+              f"{fmt([best_epoch_accuracy(r) for r in rows], percent=True)} | "
+              f"{fmt([r['final']['wall_s'] for r in rows], digits=1)} |")
+
+
 def provenance_audit(root):
     records = read_many([
         os.path.join(root, "scale_v3_*.json"),
@@ -132,6 +166,7 @@ def provenance_audit(root):
         os.path.join(root, "pepita_tuned_v1_*.json"),
         os.path.join(root, "baseline_budget_v1_mnist_ff_*.json"),
         os.path.join(root, "ep_original_v2_*.json"),
+        os.path.join(root, "ep_match_v1_*.json"),
     ])
     dirty = [r["_path"] for r in records if r.get("provenance", {}).get("git_dirty") is not False]
     print(f"\n\nProvenance: {len(records)} files audited; {len(dirty)} dirty or unknown.")
@@ -146,6 +181,7 @@ def main():
     print_scaling(args.results)
     print_nuisance(args.results)
     print_local_baselines(args.results)
+    print_ep_match(args.results)
     provenance_audit(args.results)