diff --git a/src/spikeinterface/core/numpyextractors.py b/src/spikeinterface/core/numpyextractors.py index 7f5f5321b9..45622fceb3 100644 --- a/src/spikeinterface/core/numpyextractors.py +++ b/src/spikeinterface/core/numpyextractors.py @@ -114,7 +114,9 @@ def from_recording(source_recording, with_metadata=True, with_time_vector=False, class NumpyRecordingSegment(BaseRecordingSegment): def __init__(self, traces, sampling_frequency, t_start): BaseRecordingSegment.__init__(self, sampling_frequency=sampling_frequency, t_start=t_start) - self._traces = traces + # Read-only, so writing into get_traces() output raises instead of changing the recording. + self._traces = traces.view() + self._traces.flags.writeable = False self.num_samples = traces.shape[0] def get_num_samples(self) -> int: diff --git a/src/spikeinterface/core/recording_tools.py b/src/spikeinterface/core/recording_tools.py index 09d0ce0deb..1bce36e05d 100644 --- a/src/spikeinterface/core/recording_tools.py +++ b/src/spikeinterface/core/recording_tools.py @@ -834,6 +834,8 @@ def get_chunk_with_margin( add_reflect_padding=False, window_on_margin=False, dtype=None, + *, + copy=False, ): """ Helper to get chunk with margin @@ -843,6 +845,10 @@ def get_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". + + With ``copy=True``, the returned chunk is a writable copy that does not + share memory with the segment, even when no padding or casting is needed. + With ``copy=False`` (the default), the chunk may be a view of parent data. """ # wrapper to keep backward compatibility with the previous naming and signature @@ -857,4 +863,5 @@ def get_chunk_with_margin( add_reflect_padding=add_reflect_padding, window_on_margin=window_on_margin, dtype=dtype, + copy=copy, ) diff --git a/src/spikeinterface/core/tests/test_numpy_extractors.py b/src/spikeinterface/core/tests/test_numpy_extractors.py index d0a8c5be70..938fec8d71 100644 --- a/src/spikeinterface/core/tests/test_numpy_extractors.py +++ b/src/spikeinterface/core/tests/test_numpy_extractors.py @@ -34,6 +34,15 @@ def setup_NumpyRecording(tmp_path_factory): return cache_folder +def test_NumpyRecording_traces_read_only(): + traces = np.zeros((100, 4), dtype="float32") + rec = NumpyRecording(traces, 30000.0) + with pytest.raises(ValueError): + rec.get_traces()[:] = 1 + assert np.all(rec.get_traces() == 0) + assert traces.flags.writeable # the caller's array is untouched + + def test_SharedMemoryRecording(): rec0 = generate_recording(num_channels=2, durations=[4.0, 3.0]) # print(rec0) diff --git a/src/spikeinterface/core/tests/test_time_series_tools.py b/src/spikeinterface/core/tests/test_time_series_tools.py index 50d91734c3..b41315e535 100644 --- a/src/spikeinterface/core/tests/test_time_series_tools.py +++ b/src/spikeinterface/core/tests/test_time_series_tools.py @@ -1,6 +1,6 @@ import numpy as np -from spikeinterface.core import generate_recording, MockRecording +from spikeinterface.core import generate_recording, MockRecording, NumpyRecording, get_chunk_with_margin from spikeinterface.core.binaryrecordingextractor import BinaryRecordingExtractor @@ -10,6 +10,8 @@ write_memory, get_random_sample_slices, get_chunks, + get_time_series_chunk_with_margin, + _taper_margins, ) @@ -171,3 +173,34 @@ def test_get_chunks(): rec = generate_recording(num_channels=1, sampling_frequency=1000.0, durations=[10.0, 20.0]) chunks = get_chunks(rec, num_chunks_per_segment=50, chunk_size=500, seed=0) assert chunks.shape == (50000, 1) + + +def test_chunk_with_margin_copy_and_taper(): + # A read-only float64 parent catches accidental in-place writes even when no cast is needed. + original = np.arange(400, dtype="float64").reshape(200, 2) + parent_data = original.copy() + parent_data.flags.writeable = False + segment = NumpyRecording(parent_data, 1000.0).segments[0] + margin = 10 + window = (1 - np.cos(np.arange(margin) / margin * np.pi)) / 2 + + for fetch in (get_chunk_with_margin, get_time_series_chunk_with_margin): + for padding in ({}, {"add_zeros": True}, {"add_reflect_padding": True}): + for start, end in ((0, 30), (50, 80), (170, 200)): + view, _, _ = fetch(segment, start, end, None, margin, **padding) + owned, left, right = fetch(segment, start, end, None, margin, copy=True, **padding) + np.testing.assert_array_equal(owned, view) + assert not np.shares_memory(owned, parent_data) + assert owned.flags.writeable + if start == 50: + assert np.shares_memory(view, parent_data) + if padding == {"add_zeros": True}: + _taper_margins(owned, window) + legacy, legacy_left, legacy_right = fetch( + segment, start, end, None, margin, window_on_margin=True, **padding + ) + np.testing.assert_array_equal(owned, legacy) + assert (left, right) == (legacy_left, legacy_right) + else: + owned[:] = 0 + np.testing.assert_array_equal(parent_data, original) diff --git a/src/spikeinterface/core/time_series_tools.py b/src/spikeinterface/core/time_series_tools.py index 1ebf0374df..4c4abc3492 100644 --- a/src/spikeinterface/core/time_series_tools.py +++ b/src/spikeinterface/core/time_series_tools.py @@ -595,6 +595,19 @@ def get_chunks(time_series: TimeSeries, concatenated=True, get_data_kwargs=None, return chunk_list +def _taper_margins(traces, window): + """Taper traces in place using a one-dimensional window and its reverse. + + The caller must own a writable floating-point buffer. When reading from a + segment, use ``get_time_series_chunk_with_margin(..., copy=True)`` first. + """ + if window.size == 0: + return + taper = window.reshape((window.size,) + (1,) * (traces.ndim - 1)) + traces[: window.size] *= taper + traces[-window.size :] *= taper[::-1] + + def get_time_series_chunk_with_margin( chunkable_segment: TimeSeriesSegment, start_frame, @@ -605,6 +618,8 @@ def get_time_series_chunk_with_margin( add_reflect_padding=False, window_on_margin=False, dtype=None, + *, + copy=False, ): """ Helper to get chunk with margin @@ -614,6 +629,10 @@ 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". + + With ``copy=True``, the returned chunk is a writable copy that does not + share memory with the segment, even when no padding or casting is needed. + With ``copy=False`` (the default), the chunk may be a view of parent data. """ length = int(chunkable_segment.get_num_samples()) @@ -645,6 +664,8 @@ def get_time_series_chunk_with_margin( end_frame + right_margin, last_dimension_indices, ) + if copy: + data_chunk = data_chunk.copy() else: # either add_zeros or reflect_padding @@ -672,7 +693,7 @@ def get_time_series_chunk_with_margin( data_chunk = chunkable_segment.get_data(start_frame2, end_frame2, last_dimension_indices) - if dtype is not None or window_on_margin or left_pad > 0 or right_pad > 0: + if copy or dtype is not None or window_on_margin or left_pad > 0 or right_pad > 0: need_copy = True else: need_copy = False @@ -696,11 +717,8 @@ 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] + window = (1 - np.cos(np.arange(margin) / margin * np.pi)) / 2 + _taper_margins(data_chunk2, window) data_chunk = data_chunk2 elif add_reflect_padding: # in this case, we don't want to taper diff --git a/src/spikeinterface/preprocessing/astype.py b/src/spikeinterface/preprocessing/astype.py index 26d45dd711..ebc1194ebb 100644 --- a/src/spikeinterface/preprocessing/astype.py +++ b/src/spikeinterface/preprocessing/astype.py @@ -73,7 +73,7 @@ def get_traces(self, start_frame, end_frame, channel_indices): channel_indices = slice(None) traces = self.parent_recording_segment.get_traces(start_frame, end_frame, channel_indices) if self.round: - np.round(traces, out=traces) + traces = np.round(traces) return traces.astype(self.dtype, copy=False) diff --git a/src/spikeinterface/preprocessing/phase_shift.py b/src/spikeinterface/preprocessing/phase_shift.py index 4713707c56..a6beb78e1b 100644 --- a/src/spikeinterface/preprocessing/phase_shift.py +++ b/src/spikeinterface/preprocessing/phase_shift.py @@ -3,6 +3,7 @@ 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 @@ -80,12 +81,13 @@ def __init__(self, parent_recording_segment, sample_shifts, margin, dtype, tmp_d self.margin = margin self.dtype = dtype self.tmp_dtype = tmp_dtype + self._margin_window = (1 - np.cos(np.arange(margin) / margin * np.pi)) / 2 def get_traces(self, start_frame, end_frame, channel_indices): if channel_indices is None: channel_indices = slice(None) - # this return a copy with margin + taper on border always + # The in-place taper requires an owned buffer, including for float64 parents. traces_chunk, left_margin, right_margin = get_chunk_with_margin( self.parent_recording_segment, start_frame, @@ -94,8 +96,9 @@ def get_traces(self, start_frame, end_frame, channel_indices): self.margin, dtype=self.tmp_dtype, add_zeros=True, - window_on_margin=True, + copy=True, ) + _taper_margins(traces_chunk, self._margin_window) traces_shift = apply_frequency_shift(traces_chunk, self.sample_shifts[channel_indices], axis=0) traces_shift = traces_shift[left_margin:-right_margin, :] diff --git a/src/spikeinterface/preprocessing/tests/test_phase_shift.py b/src/spikeinterface/preprocessing/tests/test_phase_shift.py index 5d4b74a567..6570991789 100644 --- a/src/spikeinterface/preprocessing/tests/test_phase_shift.py +++ b/src/spikeinterface/preprocessing/tests/test_phase_shift.py @@ -3,6 +3,8 @@ from spikeinterface.core.testing import check_recordings_equal from spikeinterface.preprocessing import phase_shift +from spikeinterface.preprocessing.phase_shift import apply_frequency_shift +from spikeinterface.core import get_chunk_with_margin def create_shifted_channel(): @@ -113,6 +115,36 @@ def signal(time): np.testing.assert_allclose(shifted_traces[interior], expected[interior], atol=1e-3) +def test_phase_shift_taper_preserves_parent(): + rng = np.random.default_rng(0) + shifts = np.array([0.0, 0.25, 0.5, 0.75]) + for dtype in ("int16", "float32", "float64"): + original = (rng.standard_normal((500, 4)) * 100).astype(dtype) + data = original.copy() + data.flags.writeable = False + rec = NumpyRecording(data, 1000.0) + shifted = phase_shift(rec, inter_sample_shift=shifts, margin_ms=10) + for start, end in ((0, 100), (200, 300), (400, 500)): + chunk, left, right = get_chunk_with_margin( + rec.segments[0], + start, + end, + None, + 10, + add_zeros=True, + window_on_margin=True, + dtype=None if dtype == "float64" else "float64", + ) + expected = apply_frequency_shift(chunk, shifts)[left:-right] + if dtype == "int16": + expected = expected.round() + expected = expected.astype(dtype) + actual = shifted.get_traces(start_frame=start, end_frame=end) + np.testing.assert_array_equal(actual, expected) + np.testing.assert_array_equal(shifted.get_traces(start_frame=start, end_frame=end), actual) + np.testing.assert_array_equal(data, original) + + def test_phase_shift_dtype_survives_reload(): traces, sampling_frequency, inter_sample_shift = create_shifted_channel() recording = NumpyRecording([traces.astype("int16")], sampling_frequency)