| """ |
| Online wrapper around SpikeCascadeForSorting. Consumes 50 ms chunks |
| from a SpikeGLX-like stream and returns per-chunk CascadeReports. |
| |
| This is the class a closed-loop experiment would use: the loop runs |
| at ~1 kHz (one chunk per millisecond of wall clock), and the latency |
| of each `process_chunk` call is the closed-loop latency. |
| """ |
|
|
| import time |
| from collections import deque |
| from typing import Deque, Optional |
|
|
| import numpy as np |
|
|
| from modeling_spikecascade import SpikeCascadeForSorting, CascadeReport |
|
|
|
|
| class StreamingSorter: |
| def __init__( |
| self, |
| model: SpikeCascadeForSorting, |
| chunk_ms: float = 50.0, |
| overlap_ms: float = 2.0, |
| ): |
| self.model = model |
| self.sample_rate = model.config.sample_rate |
| self.chunk_samples = int(chunk_ms / 1000 * self.sample_rate) |
| self.overlap_samples = int(overlap_ms / 1000 * self.sample_rate) |
| self.buffer: Deque[np.ndarray] = deque() |
| self.reports: list[CascadeReport] = [] |
|
|
| def push(self, chunk: np.ndarray): |
| """chunk: (n_channels, n_samples).""" |
| assert chunk.shape[1] == self.chunk_samples |
| self.buffer.append(chunk) |
|
|
| def step(self) -> Optional[CascadeReport]: |
| """ |
| Pop one chunk, run the cascade, and return the report. |
| Returns None if the buffer is empty. |
| """ |
| if not self.buffer: |
| return None |
| chunk = self.buffer.popleft() |
| report = self.model.infer(chunk, sample_rate=self.sample_rate) |
| self.reports.append(report) |
| return report |
|
|
| def run(self, n_steps: int) -> list[CascadeReport]: |
| out = [] |
| for _ in range(n_steps): |
| r = self.step() |
| if r is None: |
| break |
| out.append(r) |
| return out |
|
|
| def summary(self) -> dict: |
| if not self.reports: |
| return {} |
| return { |
| "n_chunks": len(self.reports), |
| "mean_activity_fraction": float(np.mean([r.activity_fraction for r in self.reports])), |
| "mean_stage3_fraction": float(np.mean([r.stage3_fraction for r in self.reports])), |
| "mean_latency_ms": float(np.mean([r.latency_ms for r in self.reports])), |
| "mean_correlation": float( |
| np.nanmean([r.correlation for r in self.reports]) |
| ), |
| "zero_spike_chunks": sum(1 for r in self.reports if r.activity_fraction == 0.0), |
| } |
