Two side probe integration.

Author: samuelgarcia

src/probeinterface/plotting.py

@@ -102,6 +102,7 @@ def plot_probe(
     ylims: tuple | None = None,
     zlims: tuple | None = None,
     show_channel_on_click: bool = False,
+    side=None,
 ):
     """Plot a Probe object.
     Generates a 2D or 3D axis, depending on Probe.ndim
@@ -138,6 +139,8 @@ def plot_probe(
         Limits for z dimension
     show_channel_on_click : bool, default: False
         If True, the channel information is shown upon click
+    side : None | "front" | "back
+        If the probe is two side, then the side must be given otherwise this raises an error.
 
     Returns
     -------
@@ -148,6 +151,14 @@ def plot_probe(
     """
     import matplotlib.pyplot as plt
 
+    if probe.contact_sides is not None:
+        if side is None or side not in ('front', 'back'):
+            raise ValueError("The probe has two side, you must give which one to plot. plot_probe(probe, side='front'|'back')")
+        mask = probe.contact_sides == side
+        probe = probe.get_slice(mask)
+        probe._contact_sides = None
+
+
     if ax is None:
         if probe.ndim == 2:
             fig, ax = plt.subplots()

src/probeinterface/probe.py

@@ -101,6 +101,7 @@ def __init__(
         self.probe_planar_contour = None
 
         # This handles the shank id per contact
+        # If None then one shank only
         self._shank_ids = None
 
         # This handles the wiring to device : channel index on device side.
@@ -112,6 +113,10 @@ def __init__(
         #  This must be unique at Probe AND ProbeGroup level
         self._contact_ids = None
 
+        # Handle contact side for double face probes
+        # If None then one face only
+        self._contact_sides = None
+
         # annotation:  a dict that contains all meta information about
         # the probe (name, manufacturor, date of production, ...)
         self.annotations = dict()
@@ -153,6 +158,10 @@ def contact_ids(self):
     def shank_ids(self):
         return self._shank_ids
 
+    @property
+    def contact_sides(self):
+        return self._contact_sides
+
     @property
     def name(self):
         return self.annotations.get("name", None)
@@ -237,6 +246,8 @@ def get_title(self) -> str:
             if self.shank_ids is not None:
                 num_shank = self.get_shank_count()
                 txt += f" - {num_shank}shanks"
+            if self._contact_sides is not None:
+                txt += f" - 2 sides"
         return txt
 
     def __repr__(self):
@@ -291,7 +302,7 @@ def get_shank_count(self) -> int:
         return n
 
     def set_contacts(
-        self, positions, shapes="circle", shape_params={"radius": 10}, plane_axes=None, contact_ids=None, shank_ids=None
+        self, positions, shapes="circle", shape_params={"radius": 10}, plane_axes=None, contact_ids=None, shank_ids=None, contact_sides=None
     ):
         """Sets contacts to a Probe.
 
@@ -320,16 +331,29 @@ def set_contacts(
         shank_ids : array[str] | None, default: None
             Defines the shank ids for the contacts. If None, then
             these are assigned to a unique Shank.
+        contact_sides : array[str] | None, default: None
+            If probe is double sided, defines sides by a vector of ['front' | 'back']
         """
         positions = np.array(positions)
         if positions.shape[1] != self.ndim:
             raise ValueError(f"positions.shape[1]: {positions.shape[1]} and ndim: {self.ndim} do not match!")
 
-        # Check for duplicate positions
-        unique_positions = np.unique(positions, axis=0)
-        positions_are_not_unique = unique_positions.shape[0] != positions.shape[0]
-        if positions_are_not_unique:
-            _raise_non_unique_positions_error(positions)
+        
+        if contact_sides is None:
+            # Check for duplicate positions
+            unique_positions = np.unique(positions, axis=0)
+            positions_are_not_unique = unique_positions.shape[0] != positions.shape[0]
+            if positions_are_not_unique:
+                _raise_non_unique_positions_error(positions)
+        else:
+            # Check for duplicate positions side by side
+            contact_sides = np.asarray(contact_sides).astype(str)
+            for side in ("front", "back"):
+                mask = contact_sides == "font"
+                unique_positions = np.unique(positions[mask], axis=0)
+                positions_are_not_unique = unique_positions.shape[0] != positions[mask].shape[0]
+                if positions_are_not_unique:
+                    _raise_non_unique_positions_error(positions[mask])                
 
         self._contact_positions = positions
         n = positions.shape[0]
@@ -355,6 +379,15 @@ def set_contacts(
             self._shank_ids = np.asarray(shank_ids).astype(str)
             if self.shank_ids.size != n:
                 raise ValueError(f"shank_ids have wrong size: {self.shanks.ids.size} != {n}")
+        
+        if contact_sides is None:
+            self._contact_sides = contact_sides
+        else:
+            self._contact_sides = contact_sides
+            if self._contact_sides.size != n:
+                raise ValueError(f"contact_sides have wrong size: {self._contact_sides.ids.size} != {n}")
+            if not np.all(np.isin(self._contact_sides, ["front", "back"])):
+                raise ValueError(f"contact_sides must 'front' or 'back'")
 
         # shape
         if isinstance(shapes, str):
@@ -592,6 +625,13 @@ def __eq__(self, other):
         ):
             return False
 
+        if self._contact_sides is None:
+            if other._contact_sides is not None:
+                return False
+        else:
+            if not np.array_equal(self._contact_sides, other._contact_sides):
+                return False
+
         # Compare contact_annotations dictionaries
         if self.contact_annotations.keys() != other.contact_annotations.keys():
             return False
@@ -842,6 +882,7 @@ def rotate_contacts(self, thetas: float | np.array[float] | list[float]):
         "device_channel_indices",
         "_contact_ids",
         "_shank_ids",
+        "_contact_sides",
     ]
 
     def to_dict(self, array_as_list: bool = False) -> dict:
@@ -895,6 +936,9 @@ def from_dict(d: dict) -> "Probe":
             plane_axes=d["contact_plane_axes"],
             shapes=d["contact_shapes"],
             shape_params=d["contact_shape_params"],
+            contact_ids=d.get("contact_ids", None),
+            shank_ids=d.get("shank_ids", None),
+            contact_sides=d.get("contact_sides", None),
         )
 
         v = d.get("probe_planar_contour", None)
@@ -905,14 +949,6 @@ def from_dict(d: dict) -> "Probe":
         if v is not None:
             probe.set_device_channel_indices(v)
 
-        v = d.get("shank_ids", None)
-        if v is not None:
-            probe.set_shank_ids(v)
-
-        v = d.get("contact_ids", None)
-        if v is not None:
-            probe.set_contact_ids(v)
-
         if "annotations" in d:
             probe.annotate(**d["annotations"])
         if "contact_annotations" in d:
@@ -955,6 +991,7 @@ def to_numpy(self, complete: bool = False) -> np.array:
                 ...
                 ('shank_ids', 'U64'),
                 ('contact_ids', 'U64'),
+                ('contact_sides', 'U8'),
 
                 # The rest is added only if `complete=True`
                 ('device_channel_indices', 'int64', optional),
@@ -991,6 +1028,9 @@ def to_numpy(self, complete: bool = False) -> np.array:
             dtype += [(k, "float64")]
         dtype += [("shank_ids", "U64"), ("contact_ids", "U64")]
 
+        if self._contact_sides is not None:
+            dtype += [("contact_sides", "U8"), ]
+
         if complete:
             dtype += [("device_channel_indices", "int64")]
             dtype += [("si_units", "U64")]
@@ -1014,6 +1054,11 @@ def to_numpy(self, complete: bool = False) -> np.array:
 
         arr["shank_ids"] = self.shank_ids
 
+        if self._contact_sides is not None:
+            arr["contact_sides"] = self.contact_sides
+
+        
+
         if self.contact_ids is None:
             arr["contact_ids"] = [""] * self.get_contact_count()
         else:
@@ -1062,6 +1107,7 @@ def from_numpy(arr: np.ndarray) -> "Probe":
             "contact_shapes",
             "shank_ids",
             "contact_ids",
+            "contact_sides",
             "device_channel_indices",
             "radius",
             "width",
@@ -1118,17 +1164,19 @@ def from_numpy(arr: np.ndarray) -> "Probe":
         else:
             plane_axes = None
 
-        probe.set_contacts(positions=positions, plane_axes=plane_axes, shapes=shapes, shape_params=shape_params)
+        
+        shank_ids = arr["shank_ids"] if "shank_ids" in fields else None
+        contact_sides = arr["contact_sides"] if "contact_sides" in fields else None
+
+        probe.set_contacts(positions=positions, plane_axes=plane_axes, shapes=shapes, shape_params=shape_params, shank_ids=shank_ids, contact_sides=contact_sides)
 
         if "device_channel_indices" in fields:
             dev_channel_indices = arr["device_channel_indices"]
             if not np.all(dev_channel_indices == -1):
                 probe.set_device_channel_indices(dev_channel_indices)
-        if "shank_ids" in fields:
-            probe.set_shank_ids(arr["shank_ids"])
         if "contact_ids" in fields:
             probe.set_contact_ids(arr["contact_ids"])
-
+        
         # contact annotations
         for k in contact_annotation_fields:
             probe.annotate_contacts(**{k: arr[k]})

tests/test_plotting.py

@@ -50,7 +50,21 @@ def test_plot_probegroup():
     plot_probegroup(probegroup_3d, same_axes=True)
 
 
+def test_plot_probe_two_side():
+    probe = Probe()
+    probe.set_contacts(
+        positions=np.array([[0, 0], [0, 10], [0, 20],[0, 0], [0, 10], [0, 20],]),
+        shapes="circle",
+        contact_ids=["F1", "F2", "F3", "B1", "B2", "B3"],
+        contact_sides=["front", "front", "front", "back", "back","back"]
+    )
+
+    plot_probe(probe, with_contact_id=True, side="front")
+    plot_probe(probe, with_contact_id=True, side="back")
+
+
 if __name__ == "__main__":
-    test_plot_probe()
+    # test_plot_probe()
     # test_plot_probe_group()
+    test_plot_probe_two_side()
     plt.show()

tests/test_probe.py

@@ -197,9 +197,34 @@ def test_position_uniqueness():
         probe.set_contacts(positions=positions_with_dups, shapes="circle", shape_params={"radius": 5})
 
 
+def test_double_side_probe():
+
+    probe = Probe()
+    probe.set_contacts(
+        positions=np.array([[0, 0], [0, 10], [0, 20],[0, 0], [0, 10], [0, 20],]),
+        shapes="circle",
+        contact_sides=["front", "front", "front", "back", "back","back"]
+    )
+    print(probe)
+
+    assert "contact_sides" in probe.to_dict()
+
+    probe2 = Probe.from_dict(probe.to_dict())
+    assert probe2 == probe 
+
+    probe3 = Probe.from_numpy(probe.to_numpy())
+    assert probe3 == probe 
+
+    probe4 = Probe.from_dataframe(probe.to_dataframe())
+    assert probe4 == probe
+
+
+
 if __name__ == "__main__":
     test_probe()
 
     tmp_path = Path("tmp")
     tmp_path.mkdir(exist_ok=True)
     test_save_to_zarr(tmp_path)
+
+    test_double_side_probe()