File size: 2,436 Bytes
769c9b8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 | """
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),
} |