diff --git a/src/spikeinterface/postprocessing/principal_component.py b/src/spikeinterface/postprocessing/principal_component.py index 225b37c9f0..bf54a3e644 100644 --- a/src/spikeinterface/postprocessing/principal_component.py +++ b/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 diff --git a/src/spikeinterface/postprocessing/tests/test_principal_component.py b/src/spikeinterface/postprocessing/tests/test_principal_component.py index 879f4ff055..8dc7c7ecbe 100644 --- a/src/spikeinterface/postprocessing/tests/test_principal_component.py +++ b/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.