| 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 |
