benchmark.py

"""
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()