Create modeling_spikecascade.py

modeling_spikecascade.py

@@ -0,0 +1,259 @@
+"""
+SpikeCascade: three-stage fidelity cascade for real-time spike sorting.
+
+Stage 1: amplitude threshold (cheap, ~1 us).
+Stage 2: CNN template match (medium, ~50 us).
+Stage 3: full KMeans clustering (expensive, ~1 ms, ~1% of spikes).
+
+The CascadeReport returned by `infer` implements the calibration
+contract from the paper: it reports activity_fraction, stage3_fraction,
+latency_ms, and correlation (NaN if the baseline is empty).
+"""
+
+from __future__ import annotations
+
+import math
+import time
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel
+
+from configuration_spikecascade import SpikeCascadeConfig
+from preprocessing import bandpass_filter, extract_snippets
+
+
+# ─────────────────────────────────────────────────────────────────────
+# Stage 2: CNN template matcher
+# ─────────────────────────────────────────────────────────────────────
+
+class TemplateCNN(nn.Module):
+    """
+    A small 1D CNN that maps a multi-channel spike snippet to a
+    probability distribution over a bank of learned templates.
+
+    Input:  (batch, n_channels, template_len)
+    Output: (batch, template_bank_size) logits, (batch, hidden_dim)
+    """
+
+    def __init__(self, cfg: SpikeCascadeConfig):
+        super().__init__()
+        c1, c2, c3 = cfg.cnn_channels
+        self.conv1 = nn.Conv1d(cfg.n_channels, c1, kernel_size=7, padding=3)
+        self.conv2 = nn.Conv1d(c1, c2, kernel_size=5, padding=2)
+        self.conv3 = nn.Conv1d(c2, c3, kernel_size=3, padding=1)
+        self.pool = nn.AdaptiveAvgPool1d(1)
+        self.fc = nn.Linear(c3, cfg.hidden_dim)
+        self.head = nn.Linear(cfg.hidden_dim, cfg.template_bank_size)
+
+    def forward(self, x: torch.Tensor):
+        h = F.relu(self.conv1(x))
+        h = F.relu(self.conv2(h))
+        h = F.relu(self.conv3(h))
+        h = self.pool(h).squeeze(-1)
+        h = F.relu(self.fc(h))
+        return self.head(h), h
+
+
+# ─────────────────────────────────────────────────────────────────────
+# Cascade report
+# ─────────────────────────────────────────────────────────────────────
+
+@dataclass
+class CascadeReport:
+    units: Dict[int, List[float]] = field(default_factory=dict)
+    activity_fraction: float = 0.0
+    stage3_fraction: float = 0.0
+    latency_ms: float = 0.0
+    correlation: float = float("nan")
+    warning: Optional[str] = None
+
+
+# ─────────────────────────────────────────────────────────────────────
+# The model
+# ─────────────────────────────────────────────────────────────────────
+
+class SpikeCascadeForSorting(PreTrainedModel):
+    config_class = SpikeCascadeConfig
+
+    def __init__(self, config: SpikeCascadeConfig):
+        super().__init__(config)
+        self.config = config
+        self.template_cnn = TemplateCNN(config)
+
+        # Template bank: (template_bank_size, n_channels, template_len).
+        # Populated by `set_template_bank` or loaded from the checkpoint.
+        self.register_buffer(
+            "template_bank",
+            torch.zeros(
+                config.template_bank_size,
+                config.n_channels,
+                config.template_len_samples,
+            ),
+        )
+        # Per-template confidence calibration (set during training).
+        self.register_buffer(
+            "template_prior", torch.ones(config.template_bank_size) / config.template_bank_size
+        )
+
+        # Stage 3 fallback: KMeans centroids. Populated at inference time.
+        self._stage3_centroids: Optional[np.ndarray] = None
+
+    # ── template bank management ───────────────────────────────────
+    def set_template_bank(self, templates: torch.Tensor, priors: Optional[torch.Tensor] = None):
+        """
+        templates: (template_bank_size, n_channels, template_len)
+        priors:    (template_bank_size,) optional
+        """
+        assert templates.shape == self.template_bank.shape
+        self.template_bank.copy_(templates)
+        if priors is not None:
+            self.template_prior.copy_(priors / priors.sum())
+
+    # ── stage 1: threshold ─────────────────────────────────────────
+    def _stage1_detect(self, spike_band: np.ndarray) -> np.ndarray:
+        """
+        Returns indices (sample positions) where any channel crosses
+        |x| > threshold_sigma * robust_sigma.
+        """
+        sigma = np.median(np.abs(spike_band), axis=1, keepdims=True) / 0.6745
+        sigma = np.maximum(sigma, 1e-6)
+        crossings = (np.abs(spike_band) > self.config.threshold_sigma * sigma).any(axis=0)
+        return np.flatnonzero(crossings)
+
+    # ── stage 2: CNN template match ────────────────────────────────
+    @torch.no_grad()
+    def _stage2_match(self, snippets: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+        """
+        snippets: (n_spikes, n_channels, template_len)
+        Returns: (unit_ids, confidences)
+        """
+        x = torch.from_numpy(snippets).float()
+        logits, _ = self.template_cnn(x)
+        probs = F.softmax(logits, dim=-1).numpy()
+        unit_ids = probs.argmax(axis=1)
+        confidences = probs.max(axis=1)
+        return unit_ids, confidences
+
+    # ── stage 3: KMeans fallback ───────────────────────────────────
+    def _stage3_cluster(self, snippets: np.ndarray) -> np.ndarray:
+        """
+        Full clustering on the ambiguous snippets. Fits KMeans on the
+        fly if no centroids are cached, then assigns.
+        """
+        from sklearn.cluster import KMeans
+
+        flat = snippets.reshape(snippets.shape[0], -1)
+        if self._stage3_centroids is None or self._stage3_centroids.shape[1] != flat.shape[1]:
+            k = min(self.config.stage3_kmeans_k, max(2, flat.shape[0] // 4))
+            km = KMeans(n_clusters=k, max_iter=self.config.stage3_max_iter, n_init=1)
+            labels = km.fit_predict(flat)
+            self._stage3_centroids = km.cluster_centers_
+            return labels + self.config.template_bank_size  # offset so IDs do not collide
+        # Assign to nearest centroid
+        d = ((flat[:, None, :] - self._stage3_centroids[None, :, :]) ** 2).sum(axis=2)
+        return d.argmin(axis=1) + self.config.template_bank_size
+
+    # ── the cascade ────────────────────────────────────────────────
+    def infer(
+        self,
+        voltage: np.ndarray,
+        sample_rate: Optional[int] = None,
+        baseline_spike_counts: Optional[Dict[int, int]] = None,
+    ) -> CascadeReport:
+        """
+        voltage: (n_channels, n_samples), float32, raw or spike-band.
+        sample_rate: defaults to config.sample_rate.
+        baseline_spike_counts: optional {unit_id: count} from an offline
+            sort, used to compute the correlation. If None, correlation
+            is left as NaN and a warning is issued.
+        """
+        t0 = time.perf_counter()
+        sample_rate = sample_rate or self.config.sample_rate
+        n_channels, n_samples = voltage.shape
+        if n_channels != self.config.n_channels:
+            raise ValueError(
+                f"expected {self.config.n_channels} channels, got {n_channels}"
+            )
+
+        # Preprocess: bandpass to the spike band.
+        spike_band = bandpass_filter(
+            voltage, sample_rate, self.config.spike_band_hz[0], self.config.spike_band_hz[1]
+        )
+
+        # Stage 1.
+        spike_positions = self._stage1_detect(spike_band)
+        n_windows = max(1, n_samples - self.config.template_len_samples)
+        activity_fraction = len(spike_positions) / n_windows
+
+        if len(spike_positions) == 0:
+            # Zero-spike branch: report NaN, not 0.0.
+            latency_ms = (time.perf_counter() - t0) * 1000
+            return CascadeReport(
+                units={},
+                activity_fraction=0.0,
+                stage3_fraction=0.0,
+                latency_ms=latency_ms,
+                correlation=float("nan"),
+                warning=(
+                    "zero spikes detected; correlation is NaN, not 0. "
+                    "The speedup over an empty baseline is undefined."
+                ),
+            )
+
+        # Extract snippets centred on the crossings.
+        snippets = extract_snippets(
+            spike_band, spike_positions, self.config.template_len_samples
+        )  # (n_spikes, n_channels, template_len)
+
+        # Stage 2.
+        unit_ids_s2, conf_s2 = self._stage2_match(snippets)
+
+        # Stage 3: only the low-confidence spikes.
+        ambiguous = conf_s2 < self.config.template_confidence_threshold
+        stage3_fraction = ambiguous.mean()
+        if ambiguous.any():
+            unit_ids_s3 = self._stage3_cluster(snippets[ambiguous])
+            unit_ids = unit_ids_s2.copy()
+            unit_ids[ambiguous] = unit_ids_s3
+        else:
+            unit_ids = unit_ids_s2
+
+        # Convert positions to milliseconds.
+        spike_times_ms = spike_positions / sample_rate * 1000.0
+
+        # Build the per-unit spike-time dictionary.
+        units: Dict[int, List[float]] = {}
+        for uid, t in zip(unit_ids.tolist(), spike_times_ms.tolist()):
+            units.setdefault(int(uid), []).append(float(t))
+
+        # Correlation against the offline baseline (if provided).
+        correlation = float("nan")
+        if baseline_spike_counts is not None:
+            correlation = self._correlation(units, baseline_spike_counts)
+
+        latency_ms = (time.perf_counter() - t0) * 1000
+        return CascadeReport(
+            units=units,
+            activity_fraction=activity_fraction,
+            stage3_fraction=stage3_fraction,
+            latency_ms=latency_ms,
+            correlation=correlation,
+        )
+
+    @staticmethod
+    def _correlation(
+        units: Dict[int, List[float]], baseline: Dict[int, int]
+    ) -> float:
+        all_ids = sorted(set(units) | set(baseline))
+        if not all_ids:
+            return float("nan")
+        a = np.array([len(units.get(u, [])) for u in all_ids], dtype=float)
+        b = np.array([baseline.get(u, 0) for u in all_ids], dtype=float)
+        if a.std() == 0.0 or b.std() == 0.0:
+            return float("nan")
+        return float(np.corrcoef(a, b)[0, 1])