Fix run_for_all_spikes units and segments after the first

src/spikeinterface/postprocessing/principal_component.py

@@ -399,7 +399,6 @@ def run_for_all_spikes(self, file_path=None, verbose=False, **job_kwargs):
         ), "To compute PCA projections for all spikes, the sorting analyzer needs the recording"
         recording = sorting_analyzer.recording
 
-        # assert sorting.get_num_segments() == 1
         assert p["mode"] in ("by_channel_local", "by_channel_global")
 
         assert file_path is not None
@@ -438,6 +437,7 @@ def run_for_all_spikes(self, file_path=None, verbose=False, **job_kwargs):
             waveforms_ext.nafter,
             unit_channels,
             pca_model,
+            sorting_analyzer.return_in_uV,
         )
         processor = TimeSeriesChunkExecutor(
             recording, func, init_func, init_args, job_name="extract PCs", verbose=verbose, **job_kwargs
@@ -639,10 +639,13 @@ def _all_pc_extractor_chunk(segment_index, start_frame, end_frame, worker_ctx):
     nafter = worker_ctx["nafter"]
     unit_channels = worker_ctx["unit_channels"]
     pca_model = worker_ctx["pca_model"]
+    return_in_uV = worker_ctx["return_in_uV"]
 
     seg_size = recording.get_num_samples(segment_index=segment_index)
 
-    i0, i1 = np.searchsorted(spike_times, [start_frame, end_frame])
+    # i0 and i1 index the spike vector of all segments, like all_pcs
+    s0, s1 = worker_ctx["segment_slices"][segment_index]
+    i0, i1 = s0 + np.searchsorted(spike_times[s0:s1], [start_frame, end_frame])
 
     if i0 != i1:
         # protect from spikes on border :  spike_time<0 or spike_time>seg_size
@@ -660,7 +663,9 @@ def _all_pc_extractor_chunk(segment_index, start_frame, end_frame, worker_ctx):
     # all waveforms in one go without worrying about borders.
     start = int(spike_times[i0] - nbefore)
     end = int(spike_times[i1 - 1] + nafter)
-    traces = recording.get_traces(start_frame=start, end_frame=end, segment_index=segment_index)
+    traces = recording.get_traces(
+        start_frame=start, end_frame=end, segment_index=segment_index, return_in_uV=return_in_uV
+    )  # IMPT: in the same units (uV) as the waveforms the PCA model was fitted on!
 
     nsamples = nbefore + nafter
 
@@ -713,24 +718,24 @@ def _all_pc_extractor_chunk(segment_index, start_frame, end_frame, worker_ctx):
             pass
 
 
-def _init_work_all_pc_extractor(recording, sorting, all_pcs_args, nbefore, nafter, unit_channels, pca_model):
+def _init_work_all_pc_extractor(
+    recording, sorting, all_pcs_args, nbefore, nafter, unit_channels, pca_model, return_in_uV
+):
     worker_ctx = {}
     worker_ctx["recording"] = recording
     worker_ctx["sorting"] = sorting
 
-    spikes = sorting.to_spike_vector(concatenated=False)
-    # This is the first segment only
-    spikes = spikes[0]
-    spike_times = spikes["sample_index"]
-    spike_labels = spikes["unit_index"]
+    spikes, segment_slices = sorting.to_spike_vector(return_slices=True)
 
     worker_ctx["all_pcs"] = np.lib.format.open_memmap(**all_pcs_args)
-    worker_ctx["spike_times"] = spike_times
-    worker_ctx["spike_labels"] = spike_labels
+    worker_ctx["spike_times"] = spikes["sample_index"]
+    worker_ctx["spike_labels"] = spikes["unit_index"]
+    worker_ctx["segment_slices"] = segment_slices
     worker_ctx["nbefore"] = nbefore
     worker_ctx["nafter"] = nafter
     worker_ctx["unit_channels"] = unit_channels
     worker_ctx["pca_model"] = pca_model
+    worker_ctx["return_in_uV"] = return_in_uV
 
     return worker_ctx
 

src/spikeinterface/postprocessing/tests/test_principal_component.py

@@ -144,11 +144,18 @@ def test_compute_for_all_spikes(self, sparse):
         Compute the principal component scores, checking the shape
         matches the number of spikes as expected. This is re-run
         with n_jobs=2 and output projection score matrices
-        checked against n_jobs=1.
+        checked against n_jobs=1. The projections of the random spikes
+        must equal the stored ones, on both segments of a recording
+        whose gain is not 1.
         """
-        sorting_analyzer = self._prepare_sorting_analyzer(
-            format="memory", sparse=sparse, extension_class=ComputePrincipalComponents
-        )
+        recording = self.recording.clone()
+        recording.set_channel_gains(0.195)  # so traces in uV differ from raw traces
+        recording.set_channel_offsets(-100.0)
+        sparsity = self.sparsity if sparse else None
+        sorting_analyzer = self.get_sorting_analyzer(recording, self.sorting, sparsity=sparsity, name="pca_all_spikes")
+        # waveforms zero pad spikes at segment borders but run_for_all_spikes skips them
+        sorting_analyzer.compute("random_spikes", max_spikes_per_unit=20, margin_size=100, seed=2205)
+        sorting_analyzer.compute("waveforms")
 
         num_spikes = sorting_analyzer.sorting.to_spike_vector().size
 
@@ -167,6 +174,9 @@ def test_compute_for_all_spikes(self, sparse):
 
         np.testing.assert_almost_equal(all_pc1, all_pc2, decimal=3)
 
+        random_spikes_indices = sorting_analyzer.get_extension("random_spikes").get_data()
+        np.testing.assert_allclose(all_pc1[random_spikes_indices], ext.data["pca_projection"], rtol=1e-5, atol=1e-3)
+
     def test_project_new(self):
         """
         `project_new` projects new (unseen) waveforms onto the PCA components.