Repository: SpikeInterface/spikeinterface
Pull request #552 (2022-03-15 17:29:07)

Support rel z in nwb
Based on https://github.com/SpikeInterface/spikeinterface/pull/536

- Added tests
- Added `planarize()` function in `BaseRecording`
- Improved `apply_dimensions_to_positions()` function to handle any combination of x-y-z
- Fixed `set_dummy_probe_from_locations` when locations are 3d by adding mock `plane_axes`


@magland I was not able to push directly on your fork for some reason. So let me know if this looks good for you and @lfrank


Patch:
diff --git a/requirements.txt b/requirements.txt
index d3fb31aa9..4ed20b6b2 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -2,4 +2,4 @@ numpy
 neo>=0.10.2
 joblib
 tqdm
-probeinterface>=0.2.7
+probeinterface>=0.2.8
diff --git a/spikeinterface/core/baserecording.py b/spikeinterface/core/baserecording.py
index a861edd7b..9e8df43a0 100644
--- a/spikeinterface/core/baserecording.py
+++ b/spikeinterface/core/baserecording.py
@@ -4,7 +4,7 @@ import warnings
 
 import numpy as np
 
-from probeinterface import Probe, ProbeGroup, write_probeinterface, read_probeinterface
+from probeinterface import Probe, ProbeGroup, write_probeinterface, read_probeinterface, select_axes
 
 from .base import BaseExtractor, BaseSegment
 from .core_tools import write_binary_recording, write_memory_recording
@@ -415,10 +415,30 @@ class BaseRecording(BaseExtractor):
                     probe.set_planar_contour(contour)
         return probegroup
 
-    def set_dummy_probe_from_locations(self, locations, shape="circle", shape_params={"radius": 1}):
-        probe = Probe()
+    def set_dummy_probe_from_locations(self, locations, shape="circle", shape_params={"radius": 1},
+                                       axes="xy"):
+        """
+        Sets a 'dummy' probe based on locations.
+
+        Parameters
+        ----------
+        locations : np.array
+            Array with channel locations (num_channels, ndim) [ndim can be 2 or 3]
+        shape : str, optional
+            Electrode shapes, by default "circle"
+        shape_params : dict, optional
+            Shape parameters, by default {"radius": 1}
+        axes : str, optional
+            If ndim is 3, indicates the axes that define the plane of the electrodes, by default "xy"
+        """
+        ndim = locations.shape[1]
+        probe = Probe(ndim=ndim)
         probe.set_contacts(locations, shapes=shape, shape_params=shape_params)
         probe.set_device_channel_indices(np.arange(self.get_num_channels()))
+
+        if ndim == 3:
+            probe = probe.to_3d(axes=axes)
+
         self.set_probe(probe, in_place=True)
 
     def set_channel_locations(self, locations, channel_ids=None):
@@ -426,7 +446,7 @@ class BaseRecording(BaseExtractor):
             raise ValueError('set_channel_locations(..) destroy the probe description, prefer set_probes(..)')
         self.set_property('location', locations, ids=channel_ids)
 
-    def get_channel_locations(self, channel_ids=None, locations_2d=True):
+    def get_channel_locations(self, channel_ids=None, axes: str = 'xy'):
         if channel_ids is None:
             channel_ids = self.get_channel_ids()
         channel_indices = self.ids_to_indices(channel_ids)
@@ -454,14 +474,17 @@ class BaseRecording(BaseExtractor):
                                             for cp in probe_j.contact_positions])):
                             raise Exception("Probes are overlapping! Retrieve locations of single probes separately")
                 all_positions = np.vstack([probe.contact_positions for probe in all_probes])
-                positions = all_positions[channel_indices]  
-            return positions
+                positions = all_positions[channel_indices]
+            return select_axes(positions, axes)
         else:
             locations = self.get_property('location')
             if locations is None:
                 raise Exception('There are no channel locations')
             locations = np.asarray(locations)[channel_indices]
-            return locations
+            return select_axes(locations, axes)
+
+    def has_3d_locations(self):
+        return self.get_property('location').shape[1] == 3
 
     def clear_channel_locations(self, channel_ids=None):
         if channel_ids is None:
@@ -541,6 +564,29 @@ class BaseRecording(BaseExtractor):
                 recordings[value] = subrec
         return recordings
 
+    def planarize(self, axes: str = "xy"):
+        """
+        Returns a Recording with a 2D probe from one with a 3D probe
+
+        Parameters
+        ----------
+        axes : str, optional
+            The axes to keep, by default "xy"
+
+        Returns
+        -------
+        BaseRecording
+            The recording with 2D positions
+        """
+        assert self.has_3d_locations, "The 'planarize' function needs a recording with 3d locations"
+        assert len(axes) == 2, "You need to specify 2 dimensions (e.g. 'xy', 'zy')"
+
+        probe2d = self.get_probe().to_2d(axes=axes)
+        recording2d = self.clone()
+        recording2d.set_probe(probe2d, in_place=True)
+
+        return recording2d
+
 
 class BaseRecordingSegment(BaseSegment):
     """
diff --git a/spikeinterface/core/core_tools.py b/spikeinterface/core/core_tools.py
index 9c590ad32..86454e87f 100644
--- a/spikeinterface/core/core_tools.py
+++ b/spikeinterface/core/core_tools.py
@@ -1,5 +1,6 @@
 from pathlib import Path
 import os
+import sys
 import datetime
 
 import numpy as np
diff --git a/spikeinterface/extractors/nwbextractors.py b/spikeinterface/extractors/nwbextractors.py
index c598255c8..6360f7d27 100644
--- a/spikeinterface/extractors/nwbextractors.py
+++ b/spikeinterface/extractors/nwbextractors.py
@@ -142,11 +142,17 @@ class NwbRecordingExtractor(BaseRecording):
         properties = dict()
         for es_ind, (channel_id, electrode_table_index) in enumerate(zip(channel_ids, self._es.electrodes.data)):
             if 'rel_x' in self._nwbfile.electrodes:
+                ndim = 2  # assume 2 dimensions
+                if 'rel_z' in self._nwbfile.electrodes:
+                    ndim = 3  # if we have rel_z, it is 3 dimensions
+
                 if 'location' not in properties:
-                    properties['location'] = np.zeros((self.get_num_channels(), 2), dtype=float)
+                    properties['location'] = np.zeros((self.get_num_channels(), ndim), dtype=float)
                 properties['location'][es_ind, 0] = self._nwbfile.electrodes['rel_x'][electrode_table_index]
                 if 'rel_y' in self._nwbfile.electrodes:
                     properties['location'][es_ind, 1] = self._nwbfile.electrodes['rel_y'][electrode_table_index]
+                if 'rel_z' in self._nwbfile.electrodes:
+                    properties['location'][es_ind, 2] = self._nwbfile.electrodes['rel_z'][electrode_table_index]
 
             for col in self._nwbfile.electrodes.colnames:
                 if isinstance(self._nwbfile.electrodes[col][electrode_table_index], ElectrodeGroup):
@@ -167,7 +173,7 @@ class NwbRecordingExtractor(BaseRecording):
                     if 'offset' not in properties:
                         properties['offset'] = np.zeros(self.get_num_channels(), dtype=type(offset))
                     properties['offset'][es_ind] = offset
-                elif col in ['x', 'y', 'z', 'rel_x', 'rel_y']:
+                elif col in ['x', 'y', 'z', 'rel_x', 'rel_y', 'rel_z']:
                     continue
                 else:
                     val = self._nwbfile.electrodes[col][electrode_table_index]


Test patch:
diff --git a/spikeinterface/core/testing_tools.py b/spikeinterface/core/testing_tools.py
index f66821acf..77167a9a3 100644
--- a/spikeinterface/core/testing_tools.py
+++ b/spikeinterface/core/testing_tools.py
@@ -9,7 +9,8 @@ def generate_recording(
         num_channels=2,
         sampling_frequency=30000.,  # in Hz
         durations=[10.325, 3.5],  #  in s for 2 segments
-        set_probe=True
+        set_probe=True,
+        ndim=2
 ):
     num_segments = len(durations)
     num_timepoints = [int(sampling_frequency * d) for d in durations]
@@ -24,9 +25,11 @@ def generate_recording(
 
     if set_probe:
         probe = generate_linear_probe(num_elec=num_channels)
+        if ndim == 3:
+            probe = probe.to_3d()
         probe.set_device_channel_indices(np.arange(num_channels))
         recording.set_probe(probe, in_place=True)
-
+        probe = generate_linear_probe(num_elec=num_channels)
     return recording
 
 
diff --git a/spikeinterface/core/tests/test_baserecording.py b/spikeinterface/core/tests/test_baserecording.py
index 6a3b4d575..5f92d032f 100644
--- a/spikeinterface/core/tests/test_baserecording.py
+++ b/spikeinterface/core/tests/test_baserecording.py
@@ -13,10 +13,13 @@ from probeinterface import Probe
 from spikeinterface.core import BinaryRecordingExtractor, NumpyRecording, load_extractor
 from spikeinterface.core.base import BaseExtractor
 
+from spikeinterface.core.testing_tools import generate_recording
+
 if hasattr(pytest, "global_test_folder"):
     cache_folder = pytest.global_test_folder / "core"
 else:
     cache_folder = Path("cache_folder") / "core"
+    cache_folder.mkdir(exist_ok=True, parents=True)
 
 
 def test_BaseRecording():
@@ -198,6 +201,25 @@ def test_BaseRecording():
     rec3 = load_extractor(folder)
     assert np.allclose(times1, rec3.get_times(1))
 
+    # test 3d probe
+    rec_3d = generate_recording(ndim=3, num_channels=30)
+    locations_3d = rec_3d.get_property("location")
+
+    locations_xy = rec_3d.get_channel_locations(axes="xy")
+    assert np.allclose(locations_xy, locations_3d[:, [0, 1]])
+
+    locations_xz = rec_3d.get_channel_locations(axes="xz")
+    assert np.allclose(locations_xz, locations_3d[:, [0, 2]])
+
+    locations_zy = rec_3d.get_channel_locations(axes="zy")
+    assert np.allclose(locations_zy, locations_3d[:, [2, 1]])
+
+    locations_xzy = rec_3d.get_channel_locations(axes="xzy")
+    assert np.allclose(locations_xzy, locations_3d[:, [0, 2, 1]])
+
+    rec_2d = rec_3d.planarize(axes="zy")
+    assert np.allclose(rec_2d.get_channel_locations(), locations_3d[:, [2, 1]])
+
 
 if __name__ == '__main__':
     test_BaseRecording()