Optimizations: faster waveforms/templates and seeding iterative_isosplit properly (#4402)

Co-authored-by: Samuel Garcia <sam.garcia.die@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: 3 people

src/spikeinterface/core/waveform_tools.py

@@ -381,14 +381,19 @@ def _worker_distribute_buffers(segment_index, start_frame, end_frame, worker_dic
     l1 = i1 + s0
 
     if l1 > l0:
-        start = spikes[l0]["sample_index"] - nbefore
-        end = spikes[l1 - 1]["sample_index"] + nafter
+
+        sub_spikes = in_seg_spikes[i0:i1]
+        start = sub_spikes[0]["sample_index"] - nbefore
+        end = sub_spikes[-1]["sample_index"] + nafter
 
         # load trace in memory
         traces = recording.get_traces(
             start_frame=start, end_frame=end, segment_index=segment_index, return_in_uV=return_in_uV
         )
 
+        onset = start + nbefore
+        offset = nbefore + nafter
+
         for unit_ind, unit_id in enumerate(unit_ids):
             # find pos
             inds = inds_by_unit[unit_id]
@@ -404,8 +409,8 @@ def _worker_distribute_buffers(segment_index, start_frame, end_frame, worker_dic
                 wfs = worker_dict["waveforms_by_units"][unit_id]
 
             for pos in in_chunk_pos:
-                sample_index = spikes[inds[pos]]["sample_index"]
-                wf = traces[sample_index - start - nbefore : sample_index - start + nafter, :]
+                sample_index = spikes["sample_index"][inds[pos]] - onset
+                wf = traces[sample_index : sample_index + offset, :]
 
                 if sparsity_mask is None:
                     wfs[pos, :, :] = wf
@@ -639,23 +644,24 @@ def _worker_distribute_single_buffer(segment_index, start_frame, end_frame, work
         in_seg_spikes["sample_index"], [max(start_frame, nbefore), min(end_frame, seg_size - nafter)]
     )
 
-    # slice in absolut in spikes vector
-    l0 = i0 + s0
-    l1 = i1 + s0
-
-    if l1 > l0:
-        start = spikes[l0]["sample_index"] - nbefore
-        end = spikes[l1 - 1]["sample_index"] + nafter
+    if i1 > i0:
+        sub_spikes = in_seg_spikes[i0:i1]
+        start = sub_spikes[0]["sample_index"] - nbefore
+        end = sub_spikes[-1]["sample_index"] + nafter
 
         # load trace in memory
         traces = recording.get_traces(
             start_frame=start, end_frame=end, segment_index=segment_index, return_in_uV=return_in_uV
         )
 
-        for spike_index in range(l0, l1):
-            sample_index = spikes[spike_index]["sample_index"]
-            unit_index = spikes[spike_index]["unit_index"]
-            wf = traces[sample_index - start - nbefore : sample_index - start + nafter, :]
+        onset = start + nbefore
+        offset = nbefore + nafter
+        sample_indices = sub_spikes["sample_index"] - onset
+        unit_indices = sub_spikes["unit_index"]
+        spike_indices = s0 + np.arange(i0, i1)
+
+        for sample_index, unit_index, spike_index in zip(sample_indices, unit_indices, spike_indices):
+            wf = traces[sample_index : sample_index + offset, :]
 
             if sparsity_mask is None:
                 all_waveforms[spike_index, :, :] = wf
@@ -1055,23 +1061,25 @@ def _worker_estimate_templates(segment_index, start_frame, end_frame, worker_dic
         in_seg_spikes["sample_index"], [max(start_frame, nbefore), min(end_frame, seg_size - nafter)]
     )
 
-    # slice in absolut in spikes vector
-    l0 = i0 + s0
-    l1 = i1 + s0
+    if i1 > i0:
+        sub_spikes = in_seg_spikes[i0:i1]
 
-    if l1 > l0:
-        start = spikes[l0]["sample_index"] - nbefore
-        end = spikes[l1 - 1]["sample_index"] + nafter
+        start = sub_spikes[0]["sample_index"] - nbefore
+        end = sub_spikes[-1]["sample_index"] + nafter
 
         # load trace in memory
         traces = recording.get_traces(
             start_frame=start, end_frame=end, segment_index=segment_index, return_in_uV=return_in_uV
         )
 
-        for spike_index in range(l0, l1):
-            sample_index = spikes[spike_index]["sample_index"]
-            unit_index = spikes[spike_index]["unit_index"]
-            wf = traces[sample_index - start - nbefore : sample_index - start + nafter, :]
+        onset = start + nbefore
+        offset = nbefore + nafter
+        sample_indices = sub_spikes["sample_index"] - onset
+        unit_indices = sub_spikes["unit_index"]
+
+        for sample_index, unit_index in zip(sample_indices, unit_indices):
+
+            wf = traces[sample_index : sample_index + offset, :]
 
             if sparsity_mask is None:
                 waveform_accumulator_per_worker[worker_index, unit_index, :, :] += wf

src/spikeinterface/sortingcomponents/clustering/iterative_isosplit.py

@@ -116,7 +116,8 @@ def main_function(cls, recording, peaks, params, job_kwargs=dict()):
         debug_folder = params["debug_folder"]
 
         params_peak_svd = params["peaks_svd"].copy()
-        params_peak_svd["seed"] = params["seed"]
+        seed = params["seed"]
+        params_peak_svd["seed"] = seed
         motion = params_peak_svd["motion"]
         motion_aware = motion is not None
 
@@ -136,9 +137,14 @@ def main_function(cls, recording, peaks, params, job_kwargs=dict()):
         else:
             peaks_svd, sparse_mask, svd_model = outs
 
+        # Clustering: channel index > split > merge
         # Clustering: channel index > split > merge
         split_params = params["split"].copy()
 
+        if seed is not None:
+            params_peak_svd.update(seed=seed)
+            split_params["method_kwargs"].update(seed=seed)
+
         split_radius_um = split_params.pop("split_radius_um")
         neighbours_mask = get_channel_distances(recording) <= split_radius_um
         split_params["method_kwargs"]["neighbours_mask"] = neighbours_mask