Create preprocessing.py

preprocessing.py

@@ -0,0 +1,35 @@
+import numpy as np
+from scipy.signal import butter, sosfiltfilt
+
+
+def bandpass_filter(
+    voltage: np.ndarray, sample_rate: int, low_hz: float, high_hz: float
+) -> np.ndarray:
+    """
+    Zero-phase Butterworth bandpass along the sample axis.
+    voltage: (n_channels, n_samples)
+    """
+    sos = butter(
+        4,
+        [low_hz / (sample_rate / 2), high_hz / (sample_rate / 2)],
+        btype="band",
+        output="sos",
+    )
+    return sosfiltfilt(sos, voltage, axis=1).astype(np.float32)
+
+
+def extract_snippets(
+    spike_band: np.ndarray, positions: np.ndarray, length: int
+) -> np.ndarray:
+    """
+    Extract fixed-length snippets centred on the given positions.
+    Pads with zeros at the edges.
+    Returns: (n_spikes, n_channels, length)
+    """
+    half = length // 2
+    n_channels, n_samples = spike_band.shape
+    padded = np.pad(spike_band, ((0, 0), (half, half)), mode="constant")
+    out = np.empty((len(positions), n_channels, length), dtype=np.float32)
+    for i, p in enumerate(positions):
+        out[i] = padded[:, p : p + length]
+    return out