Repository navigation
Shift subsets of
channels at a time in phase_shift to reduce peak memory - #4904
h-mayorquin wants to merge 4 commits into
Conversation
alejoe91
added
preprocessing
|
Good point. I took a look and the blocks were computing the rotations again in every block. Each block of 16 channels has 8 different shifts on a Neuropixels 2.0 probe. So it was 192 rotations per chunk instead of 16. I changed it so the unique rotations are computed once per chunk and each block looks up the ones of its channels. I prefer this to one block per shift as it keeps the two concerns separate: the blocks are there to bound the memory and computing each rotation once is there for the speed. One block per shift was also a bit slower in my tests.
So the speed is back to I merged main into this to fix the conflicts, this shoudl be ready to go. |
| add_zeros=True, | ||
| window_on_margin=True, | ||
| ) | ||
| traces_shift = apply_frequency_shift(traces_chunk, self.sample_shifts[channel_indices], axis=0) |
There was a problem hiding this comment.
So now apply_frequency_shift is not used anymore, right?
There was a problem hiding this comment.
Not anymore, but is a public method.
| self.margin, | ||
| dtype=np.float64, | ||
| add_zeros=True, | ||
| window_on_margin=True, |
There was a problem hiding this comment.
why do we need to move _taper_margins out? I think tapering all trace_chunk at once should have very very minimal impact on performance...
There was a problem hiding this comment.
Tapering requires casting integer traces to float because it multiplies by a cosine window. Previously, that conversion happened inside get_chunk_with_margin for all channels at once. I moved tapering into the loop so only one block needs to be converted to float64 at a time. This reduces peak memory.
I considered calling get_chunk_with_margin per channel block to keep tapering inside the helper. But some upstream preprocessors still read or process all channels even when only a subset is requested. That would repeat their full-channel allocations and computation for every block. I kept a single parent read here.
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.