Repository navigation
Extract margin tapering and make chunk ownership explicit #4887
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
c0f7197
9f02608
bfd5258
bb5ff6a
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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): | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. These can be parameterized on the test rather than loops. The benefit is that each test gets a line in the pytest output (and its slightly more readable). To avoid repeated test data set up could be served by a fixture with
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I would rename |
||
| 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) | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. These are good tests of the ownership but could also could consider explicitly testing the taper is applied correctly with an independent calculation
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. In general, it would be nice to confirm that these functions are indeed generating the correct chunk, windowed correctly, by directly indexing into parent_data here (and independently applying the window) |
||
| legacy, legacy_left, legacy_right = fetch( | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. maybe legacy could be renamed as it sounds a bit like its a legacy function, maybe |
||
| 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) | ||
| Original file line number | Diff line number | Diff line change | ||||
|---|---|---|---|---|---|---|
|
|
@@ -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)) | ||||||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I might be out of date here, do these functions handle arbitrary dimension (e.g. for the new imaging stuff?). If its just for ephys and traces is always 2D, so |
||||||
| 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: | ||||||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Suggested change
just for readability |
||||||
| 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 | ||||||
|
|
||||||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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 | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. could we create a helper function somewhere get_cosine_window to avoid the duplicate computation here and in |
||
|
|
||
| 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, | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. maybe best to set explicitly |
||
| 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, :] | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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(): | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. could we have a quick docstring here explain what this test does and why? (comes in very handy when tests fail) |
||
| 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) | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This docstring could be expanded more just to describe the purpose of the test (I find that comment a little mysterious)