Handle waveform extraction when max_spikes_per_unit_is None
@@ -138,11 +138,13 @@ class WaveformExtractor:
if dtype is None:
dtype = self.recording.get_dtype()
+ if max_spikes_per_unit is not None:
+ max_spikes_per_unit = int(max_spikes_per_unit)
self._params = dict(
ms_before=float(ms_before),
ms_after=float(ms_after),
- max_spikes_per_unit=int(max_spikes_per_unit),
+ max_spikes_per_unit=max_spikes_per_unit,
dtype=dtype.str)
(self.folder / 'params.json').write_text(
@@ -327,7 +329,8 @@ class WaveformExtractor:
else:
init_args = (self.recording.to_dict(), self.sorting.to_dict(),)
init_args = init_args + (wfs_memmap, selected_spikes, selected_spike_times, nbefore, nafter)
- processor = ChunkRecordingExecutor(self.recording, func, init_func, init_args, job_name='extract waveforms', **job_kwargs)
+ processor = ChunkRecordingExecutor(self.recording, func, init_func, init_args, job_name='extract waveforms',
+ **job_kwargs)
processor.run()
@@ -346,11 +349,14 @@ def select_random_spikes_uniformly(recording, sorting, max_spikes_per_unit, nbef
n_per_segment = [sorting.get_unit_spike_train(unit_id, segment_index=i).size for i in range(num_seg)]
cum_sum = [0] + np.cumsum(n_per_segment).tolist()
total = np.sum(n_per_segment)
+ if max_spikes_per_unit is not None:
if total > max_spikes_per_unit:
global_inds = np.random.choice(total, size=max_spikes_per_unit, replace=False)
global_inds = np.sort(global_inds)
else:
global_inds = np.arange(total)
+ else:
+ global_inds = np.arange(total)
sel_spikes = []
for segment_index in range(num_seg):
in_segment = (global_inds >= cum_sum[segment_index]) & (global_inds < cum_sum[segment_index + 1])
@@ -441,7 +447,6 @@ def _waveform_extractor_chunk(segment_index, start_frame, end_frame, worker_ctx)
wfs[pos, :, :] = traces[st - start - nbefore:st - start + nafter, :]
-
def extract_waveforms(recording, sorting, folder,
load_if_exists=False,
ms_before=3., ms_after=4., max_spikes_per_unit=500, dtype=None, **job_kwargs):
@@ -458,7 +463,3 @@ def extract_waveforms(recording, sorting, folder,
we.run(**job_kwargs)
return we
-
-
-
-
\ No newline at end of file
