Create benchmark.py

benchmark.py

@@ -0,0 +1,66 @@
+"""
+Benchmark SpikeCascade against an offline baseline and emit the
+calibration contract required by the paper.
+
+Usage:
+    python benchmark.py --data path/to/recording.npy \
+                        --baseline path/to/kilosort_output.npz
+
+The output is a single JSON object with the four calibration fields:
+    activity_fraction, stage3_fraction, latency_ms, correlation
+plus a 'speedup_vs_full' field computed by the theorem
+    S = (1 + alpha) / (1/k + alpha),  alpha = (1 - T_s/T) / (T_s/T).
+"""
+
+import argparse
+import json
+import time
+
+import numpy as np
+
+from modeling_spikecascade import SpikeCascadeForSorting
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--data", required=True, help="(n_channels, n_samples) .npy")
+    parser.add_argument("--baseline", required=True, help=".npz with 'unit_ids' and 'spike_times_ms'")
+    parser.add_argument("--model", default="your-org/SpikeCascade")
+    args = parser.parse_args()
+
+    voltage = np.load(args.data).astype(np.float32)
+    baseline = np.load(args.baseline)
+    baseline_counts = {}
+    for uid in baseline["unit_ids"]:
+        baseline_counts[int(uid)] = baseline_counts.get(int(uid), 0) + 1
+
+    model = SpikeCascadeForSorting.from_pretrained(args.model)
+
+    # Warm up.
+    _ = model.infer(voltage[:, :3000])
+
+    t0 = time.perf_counter()
+    report = model.infer(voltage, baseline_spike_counts=baseline_counts)
+    wall_ms = (time.perf_counter() - t0) * 1000
+
+    # The theorem: S = (1 + alpha) / (1/k + alpha), with
+    # alpha = (1 - f) / f and f = stage3_fraction.
+    f = report.stage3_fraction
+    alpha = (1.0 - f) / f if f > 0 else float("inf")
+    k = 20.0  # conservative cost ratio (full clustering / template match)
+    speedup = (1.0 + alpha) / (1.0 / k + alpha) if f > 0 else 1.0
+
+    out = {
+        "activity_fraction": report.activity_fraction,
+        "stage3_fraction": report.stage3_fraction,
+        "latency_ms": report.latency_ms,
+        "wall_ms": wall_ms,
+        "correlation": report.correlation,
+        "speedup_vs_full": speedup,
+        "warning": report.warning,
+    }
+    print(json.dumps(out, indent=2))
+
+
+if __name__ == "__main__":
+    main()