Repository navigation
Shift subsets of
channels at a time in phase_shift to reduce peak memory - #4904
Open
h-mayorquin wants to merge 1 commit into
Open
h-mayorquin wants to merge 1 commit into
h-mayorquin wants to merge 1 commit into
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR continues #4840 (Filter subsets of channels at a time to reduce peak memory) and #4903 (Filter subsets of channels at a time in
gaussian_filterto reduce peak memory), now forphase_shift. It is the step that sets the peak memory of the usual Neuropixels chain, see the overview in #4837 (Process channels in banks in preprocessors to reduce RAM usage).PhaseShiftRecordingSegment.get_tracescasts the whole chunk to float64 and takes its FFT across all the channels at once. So an int16 chunk peaks at about 17 times its size. This PR shifts 16 channels at a time (reusingNUM_CHANNELS_PER_BLOCKfromfilter.py) and writes each block into a preallocated output. The peak goes down to 2 times the chunk for the phase shift and from 18 to 4 times for the whole chain. The output is identical to main for int16, float32 and float64 traces, for int16 traces withdtype="float32", at the edges of the recording and on a subset of channels.With 1 s chunks it is also about 10% faster. With 10 s chunks it is about 25% slower, and about half of that comes from computing the rotations once per block instead of once per chunk.
phase_shift, 1 s, mainphase_shift, 1 s, this PRphase_shift, 10 s, mainphase_shift, 10 s, this PRNeuropixels 2.0 recording, 384 channels at 30 kHz, int16. The chain is
phase_shift,bandpass_filterto float32 andcommon_referencewith the global median. The peak is fromtracemallocand includes the output array, the time is the best of 3 runs.I am adding this now because I am doing some benchmarks in #4833 (Add
num_chunk_per_batchto job kwargs and use it for zarr v3 shard writes), but it can come after #4563 (FIR phase-shift + extractapply_raised_cosine_taperfromget_chunk_with_margin) and I am fine with fixing the conflicts.