"Allow" for 3D channel locations for compute unit locs and template metrics (#4572)

Co-authored-by: Alessio Buccino <alejoe9187@gmail.com>

Author: chrishalcrow

src/spikeinterface/metrics/template/template_metrics.py

@@ -201,9 +201,10 @@ def _set_params(
         if include_multi_channel_metrics or (
             metric_names is not None and any([m in get_multi_channel_template_metric_names() for m in metric_names])
         ):
-            assert (
-                self.sorting_analyzer.get_channel_locations().shape[1] == 2
-            ), "If multi-channel metrics are computed, channel locations must be 2D."
+            if self.sorting_analyzer.get_channel_locations().shape[1] == 3:
+                warnings.warn(
+                    "Multi-channel metrics assume 2D channel locations. We will assume the first two dimensions are the physically relevant ones"
+                )
 
         if metric_names is None:
             metric_names = get_single_channel_template_metric_names()
@@ -260,7 +261,9 @@ def _prepare_data(self, sorting_analyzer, unit_ids):
         )
         all_templates = get_dense_templates_array(sorting_analyzer, return_in_uV=True, operator=operator)
 
-        channel_locations = sorting_analyzer.get_channel_locations()
+        analyzer_channel_locations = sorting_analyzer.get_channel_locations()
+        # the template metrics only work for 2D probes. We warn users with 3D locations above.
+        channel_locations = analyzer_channel_locations[:, :2]
 
         main_channel_templates = []
         peaks_info = []

src/spikeinterface/postprocessing/localization_tools.py

@@ -104,7 +104,7 @@ def compute_monopolar_triangulation(
     unit_location = np.zeros((unit_ids.size, 4), dtype="float64")
     for i, unit_id in enumerate(unit_ids):
         chan_inds = sparsity.unit_id_to_channel_indices[unit_id]
-        local_contact_locations = contact_locations[chan_inds, :]
+        local_contact_locations = contact_locations[chan_inds, :2]
 
         # wf is (nsample, nchan) - chann is only nieghboor
         wf = templates[i, :, :][:, chan_inds]
@@ -182,7 +182,7 @@ def compute_center_of_mass(
     unit_location = np.zeros((len(unit_ids), 2), dtype="float64")
     for i, unit_id in enumerate(unit_ids):
         chan_inds = sparsity.unit_id_to_channel_indices[unit_id]
-        local_contact_locations = contact_locations[chan_inds, :]
+        local_contact_locations = contact_locations[chan_inds, :2]
 
         wf = templates[i, :, :]
 
@@ -247,7 +247,7 @@ def compute_grid_convolution(
     unit_location: np.array
     """
 
-    contact_locations = sorting_analyzer_or_templates.get_channel_locations()
+    contact_locations = sorting_analyzer_or_templates.get_channel_locations()[:, :2]
 
     templates = get_dense_templates_array(
         sorting_analyzer_or_templates, return_in_uV=get_return_in_uV(sorting_analyzer_or_templates)
@@ -693,7 +693,7 @@ def compute_location_max_channel(
     extremum_channels_index = get_template_extremum_channel(
         templates_or_sorting_analyzer, peak_sign=peak_sign, mode=mode, outputs="index", operator=operator
     )
-    contact_locations = templates_or_sorting_analyzer.get_channel_locations()
+    contact_locations = templates_or_sorting_analyzer.get_channel_locations()[:, :2]
     if unit_ids is None:
         unit_ids = templates_or_sorting_analyzer.unit_ids
     else:

src/spikeinterface/postprocessing/tests/test_unit_locations.py

@@ -1,6 +1,9 @@
 from spikeinterface.postprocessing.tests.common_extension_tests import AnalyzerExtensionCommonTestSuite
 from spikeinterface.postprocessing import ComputeUnitLocations
 import pytest
+from probeinterface import Probe
+from spikeinterface.core import create_sorting_analyzer, generate_ground_truth_recording
+import numpy as np
 
 
 class TestUnitLocationsExtension(AnalyzerExtensionCommonTestSuite):
@@ -18,3 +21,46 @@ class TestUnitLocationsExtension(AnalyzerExtensionCommonTestSuite):
     )
     def test_extension(self, params):
         self.run_extension_tests(ComputeUnitLocations, params=params)
+
+
+def test_2d_and_3d_unit_localization():
+    """
+    Our localization tools do not use the 3rd dimension of contact position.
+    Hence if we pass the same data with a 2D probe and a 3D probe (with the
+    same 2D positions), we should get the same result for all methods.
+
+    Also serves as an integration test of all the localization methods for
+    2D and 3D channel locations.
+    """
+
+    # make a 2D synthetic recording
+    positions_2D = [[0, 0], [0, 1], [1, 0], [1, 1]]
+
+    probe = Probe(ndim=2, si_units="um")
+    probe.set_contacts(positions=positions_2D)
+    probe.set_device_channel_indices(np.arange(4))
+
+    recording, sorting = generate_ground_truth_recording(num_channels=4, num_units=2, probe=probe, seed=1205)
+    analyzer_2D = create_sorting_analyzer(sorting, recording, sparse=False)
+    analyzer_2D.compute(["random_spikes", "templates"])
+
+    # make a 3D synthetic recording
+    positions_3D = [[0, 0, 10], [0, 1, 15], [1, 0, 20], [1, 1, 100]]
+
+    probe = Probe(ndim=3, si_units="um")
+    probe.set_contacts(positions=positions_3D, plane_axes=[[[1, 0, 0], [0, 1, 0]] * 4])
+
+    probe.set_device_channel_indices(np.arange(4))
+
+    recording_3D, sorting_3D = generate_ground_truth_recording(num_channels=4, num_units=2, probe=probe, seed=1205)
+    analyzer_3D = create_sorting_analyzer(sorting_3D, recording_3D, sparse=False)
+    analyzer_3D.compute(["random_spikes", "templates"])
+
+    for method in ["center_of_mass", "grid_convolution", "monopolar_triangulation", "max_channel"]:
+
+        analyzer_2D.compute("unit_locations", method=method)
+        analyzer_3D.compute("unit_locations", method=method)
+        unit_locations_2D = analyzer_2D.get_extension("unit_locations").get_data()
+        unit_locations_3D = analyzer_3D.get_extension("unit_locations").get_data()
+
+        assert np.all(unit_locations_2D[:, :2] == unit_locations_3D[:, :2])