diff --git a/src/spikeinterface/preprocessing/phase_shift.py b/src/spikeinterface/preprocessing/phase_shift.py index a6beb78e1b..273e828da3 100644 --- a/src/spikeinterface/preprocessing/phase_shift.py +++ b/src/spikeinterface/preprocessing/phase_shift.py @@ -1,12 +1,18 @@ +import importlib.util + import numpy as np from spikeinterface.core.core_tools import define_function_handling_dict_from_class from spikeinterface.core import get_chunk_with_margin -from spikeinterface.core.time_series_tools import _taper_margins from .basepreprocessor import BasePreprocessor, BasePreprocessorSegment +HAVE_NUMBA = importlib.util.find_spec("numba") is not None + +# 16 taps keep the response flat across the Neuropixels passband; fewer start to cut into it. +_DEFAULT_FIR_TAPS = 16 + class PhaseShiftRecording(BasePreprocessor): """ @@ -22,14 +28,21 @@ class PhaseShiftRecording(BasePreprocessor): recording : Recording The recording. It need to have "inter_sample_shift" in properties. margin_ms : float, default: 40.0 - Margin in ms for computation. - 40ms ensure a very small error when doing chunk processing + Margin in ms read on each side of a chunk with ``method="fft"``; 40 ms keeps the error at chunk edges small. + Ignored with ``method="fir"``. inter_sample_shift : None or numpy array, default: None If "inter_sample_shift" is not in recording properties, we can externally provide one. dtype : None | str | dtype, default: None - Dtype of input and output `recording` objects. - + Dtype of the output. None keeps the parent's. + method : "fft" | "fir", default: "fft" + - ``"fft"``: shifts each chunk in the frequency domain. Exact up to Nyquist. + - ``"fir"``: convolves each channel with a short windowed-sinc filter. Faster, and with the default + ``n_taps`` flat up to about a third of the sampling rate, which covers the Neuropixels passband. For + signal above that, raise ``n_taps`` or use ``"fft"``. + n_taps : int, default: 16 + Filter length with ``method="fir"``; must be even. Longer filters stay flat to higher frequencies and run + slower. Ignored with ``method="fft"``. Returns ------- @@ -37,7 +50,21 @@ class PhaseShiftRecording(BasePreprocessor): The phase shifted recording object """ - def __init__(self, recording, margin_ms=40.0, inter_sample_shift=None, dtype=None): + def __init__( + self, + recording, + margin_ms=40.0, + inter_sample_shift=None, + dtype=None, + method="fft", + n_taps=_DEFAULT_FIR_TAPS, + ): + if method not in ("fft", "fir"): + raise ValueError(f"method must be 'fft' or 'fir', got {method!r}") + if method == "fir": + if n_taps < 2 or n_taps % 2 != 0: + raise ValueError(f"n_taps must be a positive even integer, got {n_taps}") + if inter_sample_shift is None: assert "inter_sample_shift" in recording.get_property_keys(), "'inter_sample_shift' is not a property!" sample_shifts = recording.get_property("inter_sample_shift") @@ -47,20 +74,23 @@ def __init__(self, recording, margin_ms=40.0, inter_sample_shift=None, dtype=Non ), "the 'inter_sample_shift' must be same size at the num_channels " sample_shifts = np.asarray(inter_sample_shift) - margin = int(margin_ms * recording.get_sampling_frequency() / 1000.0) + dtype = np.dtype(recording.get_dtype() if dtype is None else dtype) - if dtype is None: - dtype = recording.get_dtype() - # the "apply_shift" function returns a float64 buffer. In case the dtype is different - # than float64, we need a temporary casting and to force the buffer back to the original dtype - if str(dtype) != "float64": - tmp_dtype = np.dtype("float64") + if method == "fft": + margin = int(margin_ms * recording.get_sampling_frequency() / 1000.0) else: - tmp_dtype = None + margin = n_taps // 2 BasePreprocessor.__init__(self, recording, dtype=dtype) for parent_segment in recording.segments: - rec_segment = PhaseShiftRecordingSegment(parent_segment, sample_shifts, margin, dtype, tmp_dtype) + rec_segment = PhaseShiftRecordingSegment( + parent_segment, + sample_shifts, + margin, + dtype, + method=method, + n_taps=n_taps, + ) self.add_recording_segment(rec_segment) # for dumpability @@ -70,44 +100,110 @@ def __init__(self, recording, margin_ms=40.0, inter_sample_shift=None, dtype=Non recording=recording, margin_ms=float(margin_ms), inter_sample_shift=inter_sample_shift, - dtype=np.dtype(dtype).str, + dtype=dtype.str, + method=method, + n_taps=int(n_taps), ) class PhaseShiftRecordingSegment(BasePreprocessorSegment): - def __init__(self, parent_recording_segment, sample_shifts, margin, dtype, tmp_dtype): + def __init__( + self, + parent_recording_segment, + sample_shifts, + margin, + dtype, + method="fft", + n_taps=_DEFAULT_FIR_TAPS, + ): BasePreprocessorSegment.__init__(self, parent_recording_segment) - self.sample_shifts = sample_shifts + self.sample_shifts = np.asarray(sample_shifts) self.margin = margin - self.dtype = dtype - self.tmp_dtype = tmp_dtype - self._margin_window = (1 - np.cos(np.arange(margin) / margin * np.pi)) / 2 + self.dtype = np.dtype(dtype) + self.method = method + self.n_taps = int(n_taps) + self._fir_kernels_kc = None + if method == "fir": + # The numba kernel indexes taps forward in time, the reverse of convolution order. Store them as (K, C) + # so its innermost loop, over channels, reads adjacent memory. + taps = _build_fir_taps(self.sample_shifts, self.n_taps) + self._fir_kernels_kc = np.ascontiguousarray(taps[:, ::-1].T, dtype=np.float32) def get_traces(self, start_frame, end_frame, channel_indices): if channel_indices is None: channel_indices = slice(None) + if self.method == "fir": + return self._get_traces_fir(start_frame, end_frame, channel_indices) + return self._get_traces_fft(start_frame, end_frame, channel_indices) - # The in-place taper requires an owned buffer, including for float64 parents. + def _get_traces_fft(self, start_frame, end_frame, channel_indices): + # The shift computes in float64 for every parent and output dtype. The chunk must be float in any case, + # since its margins are tapered in place. traces_chunk, left_margin, right_margin = get_chunk_with_margin( self.parent_recording_segment, start_frame, end_frame, channel_indices, self.margin, - dtype=self.tmp_dtype, + dtype=np.float64, add_zeros=True, - copy=True, + window_on_margin=True, ) - _taper_margins(traces_chunk, self._margin_window) traces_shift = apply_frequency_shift(traces_chunk, self.sample_shifts[channel_indices], axis=0) + return _cast_owned_traces(traces_shift[left_margin:-right_margin, :], self.dtype) + + def _get_traces_fir(self, start_frame, end_frame, channel_indices): + """Shift each channel by filtering it with its own windowed-sinc kernel. + + The kernel reaches ``n_taps // 2`` samples to each side, so the chunk is read with that margin and padded + with zeros past the ends of the recording. The kernel computes in float32: int16 and float32 traces go in as + they are and other dtypes are cast to float32 first. Without numba, a numpy version gives the same result + more slowly. + """ + half = self.n_taps // 2 + parent = self.parent_recording_segment + length = int(parent.get_num_samples()) + fetch_start = max(0, start_frame - half) + fetch_end = min(length, end_frame + half) + left_pad = half - (start_frame - fetch_start) + right_pad = half - (fetch_end - end_frame) + + traces = parent.get_traces(start_frame=fetch_start, end_frame=fetch_end, channel_indices=channel_indices) + + if left_pad > 0 or right_pad > 0: + full_len = (end_frame - start_frame) + 2 * half + padded = np.zeros((full_len, traces.shape[1]), dtype=traces.dtype) + padded[left_pad : left_pad + traces.shape[0], :] = traces + traces = padded + + if isinstance(channel_indices, slice) and channel_indices == slice(None): + kernels_kc = self._fir_kernels_kc + else: + kernels_kc = np.ascontiguousarray(self._fir_kernels_kc[:, channel_indices]) + + if not HAVE_NUMBA: + shifted = _sinc_fir_numpy_tc(traces, kernels_kc) + else: + if traces.dtype not in (np.int16, np.float32): + traces = traces.astype(np.float32) + shifted = _sinc_fir_kernel_tc(np.ascontiguousarray(traces), kernels_kc) + + return _cast_owned_traces(shifted[half : half + (end_frame - start_frame), :], self.dtype) - traces_shift = traces_shift[left_margin:-right_margin, :] - if self.tmp_dtype is not None: - if np.issubdtype(self.dtype, np.integer): - traces_shift = traces_shift.round() - traces_shift = traces_shift.astype(self.dtype) - return traces_shift +def _cast_owned_traces(traces, dtype): + """Cast filtered float traces to ``dtype``, rounding and clipping in place for an integer ``dtype``. + + A filter can push a saturated signal past the integer range, and an unclipped cast would wrap it to the opposite + sign. ``traces`` must be a buffer the caller owns. + """ + if traces.dtype == dtype: + return traces + if np.issubdtype(dtype, np.integer): + dtype_info = np.iinfo(dtype) + np.round(traces, out=traces) + np.clip(traces, dtype_info.min, dtype_info.max, out=traces) + return traces.astype(dtype) # function for API @@ -216,3 +312,110 @@ def apply_fshift_ibl(w, s, axis=0, ns=None): W = np.real(irfft(W, ns, axis=axis)) W = W.astype(w.dtype) return W + + +def _build_fir_taps(shift_samples, n_taps, beta=None): + """Build one windowed-sinc kernel per channel that delays the channel by a fraction of a sample. + + An ideal delay of ``d`` samples convolves the signal with ``sinc(t - d)``, which never ends. Each kernel keeps + ``n_taps`` samples of it around the centre and tapers them with a Kaiser window, the standard windowed-sinc + fractional delay (Laakso et al., 1996). Each kernel is then scaled to sum to 1, so a constant signal passes + unchanged; the windowed sinc's sum varies with ``d``, which would otherwise scale each ADC group differently. + + Parameters + ---------- + shift_samples : array of shape (C,) + Delay of each channel in samples. A positive delay gives ``y[n] = x[n - d]``, as in ``apply_frequency_shift``. + n_taps : int + Kernel length ``K``. Longer kernels stay accurate to higher frequencies. + beta : float | None, default: None + Kaiser window shape. None uses ``(K + 1) / 2``, which minimises the error in 0.3-10 kHz at 30 kHz sampling + for each ``K``; a fixed value suits only one length. + + Returns + ------- + taps : array of shape (C, K) + Taps in convolution order: ``y[n] = sum_j taps[c, j] * x[n + center - j]`` with ``center = (K - 1) // 2``. + + References + ---------- + T. I. Laakso, V. Välimäki, M. Karjalainen and U. K. Laine, "Splitting the unit delay", IEEE Signal Processing + Magazine 13(1), 30-60, 1996. + """ + # Tap j weighs x[n + center - j], an offset of t = center - j samples, so the ideal delay gives sinc(t + shift). + offsets = (n_taps - 1) // 2 - np.arange(n_taps, dtype=np.float64) + shifts = np.asarray(shift_samples, dtype=np.float64)[:, np.newaxis] + if beta is None: + beta = (n_taps + 1) / 2 + # The window reaches half a tap past each end tap. np.kaiser(K) would end exactly on them, where it is near zero, + # and waste two taps. + window_positions = np.arange(n_taps) - (n_taps - 1) / 2 + window = np.i0(beta * np.sqrt(1 - (window_positions / (n_taps / 2)) ** 2)) / np.i0(beta) + taps = np.sinc(offsets + shifts) * window + taps /= taps.sum(axis=1, keepdims=True) + return taps + + +if HAVE_NUMBA: + import numba + + # Single-threaded: SpikeInterface already runs this inside parallel chunk workers, and numba's parallel=True + # crashes when called from worker threads or forked worker processes. + @numba.njit(nogil=True, cache=False, boundscheck=False) + def _sinc_fir_kernel_tc(signal_tc, kernels_kc): + """Apply each channel's FIR kernel to a ``(T, C)`` signal, treating samples outside it as zero. + + Computes ``out[n, c] = sum_k kernels_kc[k, c] * signal_tc[n + k - K // 2, c]``, the same as + ``scipy.ndimage.correlate1d(signal_tc[:, c], kernels_kc[:, c], mode="constant")`` for each channel. With a + few taps, summing directly in the time domain costs less than an FFT-based filter such as + ``scipy.signal.oaconvolve``. ``scipy.ndimage`` sums directly but takes one kernel per call, so it would loop + over channels; this covers all channels in one pass, with the innermost loop over adjacent channels. + + Numba compiles a version per input dtype, so int16 and float32 signals need no cast; both return float32. + """ + T, C = signal_tc.shape + K = kernels_kc.shape[0] + half = K // 2 + interior_start = half + interior_end = T - (K - 1 - half) + # np.zeros would defer zeroing to page faults inside the loop; zeroing each row just before use keeps it in cache + out = np.empty((T, C), dtype=np.float32) + for n in range(T): + out[n] = 0.0 + # Away from the ends every tap lands inside the signal, so the bounds check is skipped. + if interior_start <= n < interior_end: + base = n - half + for k in range(K): + for c in range(C): + # Casting each sample keeps int16 input from promoting the sum to float64. + out[n, c] += kernels_kc[k, c] * np.float32(signal_tc[base + k, c]) + else: + for k in range(K): + idx = n + k - half + if 0 <= idx < T: + for c in range(C): + out[n, c] += kernels_kc[k, c] * np.float32(signal_tc[idx, c]) + return out + + +def _sinc_fir_numpy_tc(signal_tc, kernels_kc): + """Numpy equivalent of the numba kernel, used when numba is not installed. + + Accumulates one tap at a time over the whole ``(T, C)`` block, in the same tap order as the numba kernel, + so the two agree to float32 rounding. + """ + signal = np.asarray(signal_tc, dtype=np.float32) + T = signal.shape[0] + K = kernels_kc.shape[0] + half = K // 2 + out = np.zeros(signal.shape, dtype=np.float32) + product = np.empty(signal.shape, dtype=np.float32) + for k in range(K): + offset = k - half # out[n] += h[k] * x[n + offset], zero outside [0, T) + n0 = max(0, -offset) + n1 = min(T, T - offset) + if n1 <= n0: + continue + np.multiply(signal[n0 + offset : n1 + offset], kernels_kc[k], out=product[n0:n1]) + out[n0:n1] += product[n0:n1] + return out diff --git a/src/spikeinterface/preprocessing/tests/test_phase_shift.py b/src/spikeinterface/preprocessing/tests/test_phase_shift.py index 6570991789..bd8a5b209c 100644 --- a/src/spikeinterface/preprocessing/tests/test_phase_shift.py +++ b/src/spikeinterface/preprocessing/tests/test_phase_shift.py @@ -1,4 +1,5 @@ import numpy as np +import pytest from spikeinterface import NumpyRecording, load from spikeinterface.core.testing import check_recordings_equal @@ -96,7 +97,8 @@ def test_phase_shift(): # ~ plt.show() -def test_phase_shift_recovers_delayed_signal(): +@pytest.mark.parametrize("method", ["fft", "fir"]) +def test_phase_shift_recovers_delayed_signal(method): sampling_frequency = 1000.0 times = np.arange(5000) / sampling_frequency inter_sample_shift = np.tile([0.0, 0.25, 0.5, 0.75], 8) @@ -108,7 +110,7 @@ def signal(time): recording = NumpyRecording([traces], sampling_frequency) recording.set_property("inter_sample_shift", inter_sample_shift) - shifted_traces = phase_shift(recording).get_traces() + shifted_traces = phase_shift(recording, method=method).get_traces() interior = slice(100, -100) expected = np.broadcast_to(signal(times)[:, np.newaxis], traces.shape) @@ -116,6 +118,11 @@ def signal(time): def test_phase_shift_taper_preserves_parent(): + """FFT phase shift matches tapering and shifting each chunk by hand, and never writes into the parent's traces. + + The parent is read-only, so any write into it raises. Chunks at both ends and in the middle cover the zero + padding and the read of real margins. + """ rng = np.random.default_rng(0) shifts = np.array([0.0, 0.25, 0.5, 0.75]) for dtype in ("int16", "float32", "float64"): @@ -156,5 +163,112 @@ def test_phase_shift_dtype_survives_reload(): check_recordings_equal(shifted, reloaded, return_in_uV=False) +def test_phase_shift_fir_matches_fft_in_spike_band(): + """PhaseShift(method="fir") output must track method="fft" within ~1% spike-band RMS.""" + import scipy.signal + + traces, sampling_frequency, inter_sample_shift = create_shifted_channel() + rec = NumpyRecording([traces.astype("float32")], sampling_frequency) + rec.set_property("inter_sample_shift", inter_sample_shift) + + fft_rec = phase_shift(rec, method="fft") + fir_rec = phase_shift(rec, method="fir") + + fft_out = fft_rec.get_traces().astype(np.float64) + fir_out = fir_rec.get_traces().astype(np.float64) + assert fft_out.shape == fir_out.shape + + # Signal-band RMS error vs the FFT reference. The fixture has tones at + # 2.5 Hz and 8.5 Hz (fs=1 kHz), so we band-limit to [1, 30] Hz to isolate + # them while excluding the raised-cosine edge artifacts. + sos = scipy.signal.butter(4, [1.0, 30.0], btype="bandpass", fs=sampling_frequency, output="sos") + edge = int(0.05 * sampling_frequency) + ref = scipy.signal.sosfiltfilt(sos, fft_out, axis=0)[edge:-edge, :] + out = scipy.signal.sosfiltfilt(sos, fir_out, axis=0)[edge:-edge, :] + sig_rms = float(np.sqrt(np.mean(ref**2))) + err_rms = float(np.sqrt(np.mean((out - ref) ** 2))) + assert err_rms / sig_rms < 0.01, f"FIR spike-band RMS error {err_rms / sig_rms:.2%} > 1%" + + +@pytest.mark.parametrize("method, dtype", [("fir", "float32"), ("fft", "float64")]) +def test_phase_shift_int16_to_float(method, dtype): + """An int16 parent gives float output of the requested dtype.""" + traces, sampling_frequency, inter_sample_shift = create_shifted_channel() + rec = NumpyRecording([traces.astype("int16")], sampling_frequency) + rec.set_property("inter_sample_shift", inter_sample_shift) + shifted_rec = phase_shift(rec, method=method, dtype=dtype) + assert shifted_rec.get_dtype() == np.dtype(dtype) + out = shifted_rec.get_traces(start_frame=0, end_frame=100) + assert out.dtype == np.dtype(dtype) + + +@pytest.mark.parametrize("method", ["fft", "fir"]) +def test_phase_shift_clips_int16_output(method): + """A step between the int16 rails overshoots after filtering, and the output clips rather than wrapping.""" + traces = np.full((2000, 2), np.iinfo("int16").max, dtype="int16") + traces[:1000] = np.iinfo("int16").min + rec = NumpyRecording([traces], 30000.0) + shifted_traces = phase_shift(rec, method=method, inter_sample_shift=[0.0, 0.5]).get_traces() + # Wrapping would put large values of the opposite sign on each side of the step. The half-sample shift puts a + # sample near zero at the step itself, so skip the samples next to it. + assert shifted_traces[900:998].max() < 0 + assert shifted_traces[1002:1100].min() > 0 + + +@pytest.mark.parametrize("parent_dtype", ["int16", "float32"]) +@pytest.mark.parametrize("dtype", [None, "float32"]) +def test_phase_shift_fir_numpy_fallback_matches_numba(monkeypatch, parent_dtype, dtype): + """Without numba, method='fir' runs a numpy path that matches the numba kernels.""" + import importlib + + # the package re-exports the phase_shift function under the module's name + phase_shift_module = importlib.import_module("spikeinterface.preprocessing.phase_shift") + if not phase_shift_module.HAVE_NUMBA: + pytest.skip("needs numba for the reference output") + rng = np.random.default_rng(0) + num_channels = 24 + traces = (rng.standard_normal((5000, num_channels)) * 200).astype(parent_dtype) + rec = NumpyRecording([traces], 30000.0) + rec.set_property("inter_sample_shift", (np.arange(num_channels) // 2 % 13) / 13.0) + fir_rec = phase_shift(rec, method="fir", dtype=dtype) + channel_ids = rec.channel_ids[3:11] + frame_ranges = [dict(), dict(start_frame=0, end_frame=1000), dict(start_frame=2000, end_frame=4990)] + + # Read the whole recording, a chunk at each end and a subset of channels, first with numba as the reference and + # then with the numpy fallback. + def read_all(): + return [fir_rec.get_traces(**frames) for frames in frame_ranges] + [fir_rec.get_traces(channel_ids=channel_ids)] + + numba_outputs = read_all() + monkeypatch.setattr(phase_shift_module, "HAVE_NUMBA", False) + numpy_outputs = read_all() + + for numpy_traces, numba_traces in zip(numpy_outputs, numba_outputs): + assert numpy_traces.dtype == numba_traces.dtype + if np.issubdtype(numba_traces.dtype, np.integer): + # values near .5 may round either way after float32 summation-order differences + assert np.abs(numpy_traces.astype(np.int64) - numba_traces).max() <= 1 + else: + np.testing.assert_allclose(numpy_traces, numba_traces, rtol=0, atol=1e-5 * np.abs(numba_traces).max()) + + +def test_phase_shift_fir_taps_are_in_convolution_order(): + """Each channel's taps, convolved with the input and centred on (K - 1) // 2, give the FIR phase shift.""" + import importlib + + phase_shift_module = importlib.import_module("spikeinterface.preprocessing.phase_shift") + rng = np.random.default_rng(0) + num_channels, n_taps = 8, 16 + traces = rng.standard_normal((2000, num_channels)).astype("float32") + shifts = np.linspace(-0.5, 0.9, num_channels) + rec = NumpyRecording([traces], 30000.0) + fir_rec = phase_shift(rec, method="fir", inter_sample_shift=shifts, n_taps=n_taps) + taps = phase_shift_module._build_fir_taps(shifts, n_taps) + center = (n_taps - 1) // 2 + expected = np.stack([np.convolve(traces[:, c], taps[c])[center : center + 2000] for c in range(num_channels)], 1) + np.testing.assert_allclose(fir_rec.get_traces(), expected, rtol=0, atol=1e-5) + + if __name__ == "__main__": test_phase_shift() + test_phase_shift_fir_matches_fft_in_spike_band()