Merge pull request #3670 from chrishalcrow/fix-concatenate-mode

Allow for concatenated mode in `get_some_projections` and `compute_pc_metrics`

Author: alejoe91

src/spikeinterface/postprocessing/principal_component.py

@@ -260,7 +260,11 @@ def get_some_projections(self, channel_ids=None, unit_ids=None):
         spike_unit_indices = some_spikes["unit_index"][selected_inds]
 
         if sparsity is None:
-            some_projections = all_projections[selected_inds, :, :][:, :, channel_indices]
+            if self.params["mode"] == "concatenated":
+                some_projections = all_projections[selected_inds, :]
+            else:
+                some_projections = all_projections[selected_inds, :, :][:, :, channel_indices]
+
         else:
             # need re-alignement
             some_projections = np.zeros((selected_inds.size, num_components, channel_indices.size), dtype=dtype)

src/spikeinterface/postprocessing/tests/test_principal_component.py

@@ -50,6 +50,13 @@ def test_mode_concatenated(self):
         assert pca.ndim == 2
         assert pca.shape[1] == n_components
 
+        ext_rand = sorting_analyzer.get_extension("random_spikes")
+        num_rand_spikes = len(ext_rand.get_data())
+
+        some_projections = ext.get_some_projections()
+        assert some_projections[0].shape[0] == num_rand_spikes
+        assert some_projections[0].shape[1] == n_components
+
     @pytest.mark.parametrize("sparse", [True, False])
     def test_get_projections(self, sparse):
         """

src/spikeinterface/qualitymetrics/pca_metrics.py

@@ -157,7 +157,10 @@ def compute_pc_metrics(
         neighbor_channel_indices = sorting_analyzer.channel_ids_to_indices(neighbor_channel_ids)
 
         labels = all_labels[np.isin(all_labels, neighbor_unit_ids)]
-        pcs = dense_projections[np.isin(all_labels, neighbor_unit_ids)][:, :, neighbor_channel_indices]
+        if pca_ext.params["mode"] == "concatenated":
+            pcs = dense_projections[np.isin(all_labels, neighbor_unit_ids)]
+        else:
+            pcs = dense_projections[np.isin(all_labels, neighbor_unit_ids)][:, :, neighbor_channel_indices]
         pcs_flat = pcs.reshape(pcs.shape[0], -1)
 
         func_args = (pcs_flat, labels, non_nn_metrics, unit_id, unit_ids, metric_params, max_threads_per_worker)