diff --git a/src/spikeinterface/core/__init__.py b/src/spikeinterface/core/__init__.py index 72222753f7..bb4b2a2468 100644 --- a/src/spikeinterface/core/__init__.py +++ b/src/spikeinterface/core/__init__.py @@ -110,6 +110,7 @@ get_closest_channels, get_noise_levels, get_chunk_with_margin, + apply_raised_cosine_taper, order_channels_by_depth, ) from .sorting_tools import ( diff --git a/src/spikeinterface/core/recording_tools.py b/src/spikeinterface/core/recording_tools.py index 09d0ce0deb..1b4a034d75 100644 --- a/src/spikeinterface/core/recording_tools.py +++ b/src/spikeinterface/core/recording_tools.py @@ -18,7 +18,12 @@ split_job_kwargs, ) -from .time_series_tools import get_random_sample_slices, get_chunks, get_time_series_chunk_with_margin +from .time_series_tools import ( + get_random_sample_slices, + get_chunks, + get_time_series_chunk_with_margin, + apply_raised_cosine_taper, +) from .time_series_tools import write_binary as _write_binary from .time_series_tools import write_memory as _write_memory from .time_series_tools import _write_time_series_to_zarr diff --git a/src/spikeinterface/core/time_series_tools.py b/src/spikeinterface/core/time_series_tools.py index 758482b965..2eaa66b494 100644 --- a/src/spikeinterface/core/time_series_tools.py +++ b/src/spikeinterface/core/time_series_tools.py @@ -594,6 +594,35 @@ def get_chunks(time_series: TimeSeries, concatenated=True, get_data_kwargs=None, return chunk_list +def apply_raised_cosine_taper(data, margin, *, inplace=True): + """Taper the first and last `margin` rows of `data` with a raised cosine that is zero at the ends. + + Use it on a zero-padded chunk before an FFT, to limit spectral leakage from the padded edges. + + Parameters + ---------- + data : ndarray + Array of shape ``(num_samples, ...)``; must be floating point when ``inplace=True``. + margin : int + Number of rows to taper at each end. + inplace : bool, default: True + If True, modify `data` and return it; otherwise return a tapered copy. + + Returns + ------- + ndarray + The tapered array. + """ + if margin <= 0: + return data + taper = (1 - np.cos(np.arange(margin) / margin * np.pi)) / 2 + taper = taper[:, np.newaxis] + out = data if inplace else np.array(data, copy=True) + out[:margin] *= taper + out[-margin:] *= taper[::-1] + return out + + def get_time_series_chunk_with_margin( chunkable_segment: TimeSeriesSegment, start_frame, @@ -613,6 +642,8 @@ def get_time_series_chunk_with_margin( of `add_zeros` or `add_reflect_padding` is True. In the first case zero padding is used, in the second case np.pad is called with mod="reflect". + + `window_on_margin=True` is deprecated: call :func:`apply_raised_cosine_taper` on the returned chunk instead. """ length = int(chunkable_segment.get_num_samples()) @@ -695,11 +726,13 @@ def get_time_series_chunk_with_margin( i1 = left_pad + data_chunk.shape[0] data_chunk2[i0:i1, :] = data_chunk if window_on_margin: - # apply inplace taper on border - taper = (1 - np.cos(np.arange(margin) / margin * np.pi)) / 2 - taper = taper[:, np.newaxis] - data_chunk2[:margin] *= taper - data_chunk2[-margin:] *= taper[::-1] + warnings.warn( + "window_on_margin=True is deprecated and will be removed in a future release; " + "call apply_raised_cosine_taper() on the returned chunk instead.", + DeprecationWarning, + stacklevel=2, + ) + apply_raised_cosine_taper(data_chunk2, margin, inplace=True) data_chunk = data_chunk2 elif add_reflect_padding: # in this case, we don't want to taper diff --git a/src/spikeinterface/preprocessing/phase_shift.py b/src/spikeinterface/preprocessing/phase_shift.py index c01a9acf59..fa67ce256c 100644 --- a/src/spikeinterface/preprocessing/phase_shift.py +++ b/src/spikeinterface/preprocessing/phase_shift.py @@ -2,10 +2,13 @@ from spikeinterface.core.core_tools import define_function_handling_dict_from_class -from spikeinterface.core import get_chunk_with_margin +from spikeinterface.core import get_chunk_with_margin, apply_raised_cosine_taper from .basepreprocessor import BasePreprocessor, BasePreprocessorSegment +# 16 taps keep the response flat across the Neuropixels passband; fewer start to cut into it. +_DEFAULT_FIR_TAPS = 16 + class PhaseShiftRecording(BasePreprocessor): """ @@ -21,14 +24,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 ------- @@ -36,7 +46,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") @@ -46,61 +70,140 @@ 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) - 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) + # apply_frequency_shift computes in float64; other dtypes are cast back after it. + tmp_dtype = np.dtype("float64") if str(dtype) != "float64" else None else: + margin = n_taps // 2 tmp_dtype = None - BasePreprocessor.__init__(self, recording, dtype=dtype) + advertised_dtype = np.dtype(dtype) + + BasePreprocessor.__init__(self, recording, dtype=advertised_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, + advertised_dtype, + tmp_dtype, + method=method, + n_taps=n_taps, + ) self.add_recording_segment(rec_segment) # for dumpability if inter_sample_shift is not None: inter_sample_shift = list(inter_sample_shift) - self._kwargs = dict(recording=recording, margin_ms=float(margin_ms), inter_sample_shift=inter_sample_shift) + self._kwargs = dict( + recording=recording, + margin_ms=float(margin_ms), + inter_sample_shift=inter_sample_shift, + dtype=advertised_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, + tmp_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.dtype = np.dtype(dtype) self.tmp_dtype = tmp_dtype + self.method = method + self.n_taps = int(n_taps) + self._fir_kernels_kc = None + if method == "fir": + # The kernels correlate, out[n] += kernels_kc[k] * x[n - half + k], in (K, C) so the channel loop is + # contiguous. + 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) - # this return a copy with margin + taper on border always + def _get_traces_fft(self, start_frame, end_frame, channel_indices): + # A float64 parent can return a view of its own data, so pin the dtype to get a copy the taper can modify. + compute_dtype = self.tmp_dtype if self.tmp_dtype is not None else np.dtype("float64") 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=compute_dtype, add_zeros=True, - window_on_margin=True, + window_on_margin=False, ) + apply_raised_cosine_taper(traces_chunk, self.margin, inplace=True) traces_shift = apply_frequency_shift(traces_chunk, self.sample_shifts[channel_indices], axis=0) 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() + if np.issubdtype(self.dtype, np.integer): + traces_shift = traces_shift.round() + if traces_shift.dtype != self.dtype: traces_shift = traces_shift.astype(self.dtype) return traces_shift + def _get_traces_fir(self, start_frame, end_frame, channel_indices): + 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 _HAS_NUMBA: + shifted = _sinc_fir_numpy_tc(traces, kernels_kc) + elif traces.dtype == np.int16: + traces = np.ascontiguousarray(traces) + shifted = _sinc_fir_kernel_int16_tc(traces, kernels_kc) + else: + sig_f32 = np.ascontiguousarray(traces, dtype=np.float32) + shifted = _sinc_fir_kernel_tc(sig_f32, kernels_kc) + + out = shifted[half : half + (end_frame - start_frame), :] + + if out.dtype == self.dtype: + return out + if np.issubdtype(self.dtype, np.integer): + return out.round().astype(self.dtype) + return out.astype(self.dtype) + # function for API phase_shift = define_function_handling_dict_from_class(source_class=PhaseShiftRecording, name="phase_shift") @@ -205,3 +308,104 @@ 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=8.6): + """Per-channel Kaiser-windowed sinc kernels that delay each channel by its fractional shift, as ``(C, K)``. + + Rows are in ``fir_filter``'s tap order, ``y[n] = sum_j taps[j] * x[n + center - j]`` with its default + ``center = (K - 1) // 2``. A positive shift delays, ``y[n] = x[n - shift]``, as in ``apply_frequency_shift``. + ``beta=8.6`` gives about -80 dB sidelobes. + """ + # 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] + taps = np.sinc(offsets + shifts) * np.kaiser(n_taps, beta=beta) + # Unit DC gain: the windowed sinc's sum varies with the shift, which would scale each ADC group differently. + taps /= taps.sum(axis=1, keepdims=True) + return taps + + +try: + import numba + + _HAS_NUMBA = True +except ImportError: # pragma: no cover - the FIR path falls back to _sinc_fir_numpy_tc + _HAS_NUMBA = False + + +if _HAS_NUMBA: + # Single-threaded: chunk workers supply the parallelism, and parallel=True crashes under numba's workqueue + # layer with the thread engine and under its omp layer with forked processes. + @numba.njit(nogil=True, cache=True, boundscheck=False) + def _sinc_fir_kernel_tc(signal_tc, kernels_kc): + """Per-channel FIR on a ``(T, C)`` float32 signal, treating samples outside it as zero.""" + 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 + if interior_start <= n < interior_end: + base = n - half + for k in range(K): + for c in range(C): + out[n, c] += kernels_kc[k, c] * 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] * signal_tc[idx, c] + return out + + @numba.njit(nogil=True, cache=True, boundscheck=False) + def _sinc_fir_kernel_int16_tc(signal_tc, kernels_kc): + """int16-input variant: reads int16 and writes float32, so an int16 parent needs no cast first.""" + 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 + if interior_start <= n < interior_end: + base = n - half + for k in range(K): + for c in range(C): + 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 kernels, 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 kernels, + 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 23718f6ac6..1c0890bffa 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 from spikeinterface.preprocessing import phase_shift @@ -112,5 +113,112 @@ def signal(time): np.testing.assert_allclose(shifted_traces[interior], expected[interior], atol=1e-3) +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%" + + +def test_phase_shift_fir_int16_to_float32(): + """method='fir' with dtype=float32 on an int16 parent yields float32 output.""" + 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) + fir_rec = phase_shift(rec, method="fir", dtype=np.float32) + assert fir_rec.get_dtype() == np.dtype("float32") + out = fir_rec.get_traces(start_frame=0, end_frame=100) + assert out.dtype == np.dtype("float32") + + +def test_phase_shift_fft_still_matches_stock_after_taper_refactor(): + """Regression guard: FFT path must still behave the same after get_chunk_with_margin's + taper was split into apply_raised_cosine_taper. This reruns the core chunked-vs-full + identity check on one representative combo. + """ + 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) + rec2 = phase_shift(rec, margin_ms=30.0, method="fft") + rec3 = rec2.save(format="memory", chunk_size=500, n_jobs=1, progress_bar=False) + t2 = rec2.get_traces() + t3 = rec3.get_traces() + rms = float(np.sqrt(np.mean(traces**2))) + err = float(np.sqrt(np.mean((t2 - t3) ** 2))) + assert err / rms < 0.001 + + +@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._HAS_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] + windows = [dict(), dict(start_frame=0, end_frame=1000), dict(start_frame=2000, end_frame=4990)] + + expected = [fir_rec.get_traces(**w) for w in windows] + [fir_rec.get_traces(channel_ids=channel_ids)] + monkeypatch.setattr(phase_shift_module, "_HAS_NUMBA", False) + actual = [fir_rec.get_traces(**w) for w in windows] + [fir_rec.get_traces(channel_ids=channel_ids)] + + for a, e in zip(actual, expected): + assert a.dtype == e.dtype + if np.issubdtype(e.dtype, np.integer): + # values near .5 may round either way after float32 summation-order differences + assert np.abs(a.astype(np.int64) - e).max() <= 1 + else: + np.testing.assert_allclose(a, e, rtol=0, atol=1e-5 * np.abs(e).max()) + + if __name__ == "__main__": test_phase_shift() + test_phase_shift_fir_matches_fft_in_spike_band() + test_phase_shift_fir_int16_advertises_float32() + test_phase_shift_fft_still_matches_stock_after_taper_refactor() + + +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)