Merge pull request #391 from samuelgarcia/shank_id_none

First attemps to make shank_ids optional (None) instead of array of [""]

Author: alejoe91

src/probeinterface/probe.py

@@ -297,8 +297,11 @@ def get_shank_count(self) -> int:
         """
         Return the number of shanks for this probe.
         """
-        assert self.shank_ids is not None
-        n = len(np.unique(self.shank_ids))
+        # assert self.shank_ids is not None
+        if self.shank_ids is None:
+            n = 1
+        else:
+            n = len(np.unique(self.shank_ids))
         return n
 
     def set_contacts(
@@ -380,7 +383,8 @@ def set_contacts(
             self.set_contact_ids(contact_ids)
 
         if shank_ids is None:
-            self._shank_ids = np.zeros(n, dtype=str)
+            # self._shank_ids = np.zeros(n, dtype=str)
+            self._shank_ids = None
         else:
             self._shank_ids = np.asarray(shank_ids).astype(str)
             if self.shank_ids.size != n:
@@ -601,11 +605,15 @@ def get_shanks(self):
         """
         Return the list of Shank objects for this Probe
         """
-        assert self.shank_ids is not None, "Can only get shanks if `shank_ids` exist"
-        shanks = []
-        for shank_id in np.unique(self.shank_ids):
-            shank = Shank(probe=self, shank_id=shank_id)
-            shanks.append(shank)
+        # assert self.shank_ids is not None, "Can only get shanks if `shank_ids` exist"
+        if self.shank_ids is None:
+            # has a unique shank
+            shanks = [Shank(probe=self, shank_id=None)]
+        else:
+            shanks = []
+            for shank_id in np.unique(self.shank_ids):
+                shank = Shank(probe=self, shank_id=shank_id)
+                shanks.append(shank)
         return shanks
 
     def __eq__(self, other):
@@ -937,13 +945,18 @@ def from_dict(d: dict) -> "Probe":
         """
         probe = Probe(ndim=d["ndim"], si_units=d["si_units"])
 
+        shank_ids = d.get("shank_ids", None)
+        if shank_ids is not None and np.all(shank_ids == ""):
+            # backward compatible hack with previous version
+            shank_ids = None
+
         probe.set_contacts(
             positions=d["contact_positions"],
             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),
+            shank_ids=shank_ids,
             contact_sides=d.get("contact_sides", None),
         )
 
@@ -1032,7 +1045,11 @@ def to_numpy(self, complete: bool = False) -> np.array:
                     param_shape.append(k)
         for k in param_shape:
             dtype += [(k, "float64")]
-        dtype += [("shank_ids", "U64"), ("contact_ids", "U64")]
+
+        if self._shank_ids is not None:
+            dtype += [("shank_ids", "U64")]
+
+        dtype += [("contact_ids", "U64")]
 
         if self._contact_sides is not None:
             dtype += [
@@ -1060,7 +1077,8 @@ def to_numpy(self, complete: bool = False) -> np.array:
             for k, v in p.items():
                 arr[k][i] = v
 
-        arr["shank_ids"] = self.shank_ids
+        if self._shank_ids is not None:
+            arr["shank_ids"] = self.shank_ids
 
         if self._contact_sides is not None:
             arr["contact_sides"] = self.contact_sides

src/probeinterface/shank.py

@@ -15,7 +15,10 @@ def __init__(self, probe, shank_id):
         self.shank_id = shank_id
 
     def get_indices(self):
-        (inds,) = np.nonzero(self.probe.shank_ids == self.shank_id)
+        if self.probe.shank_ids is None:
+            inds = np.arange(self.probe.get_contact_count(), dtype=int)
+        else:
+            inds = np.flatnonzero(self.probe.shank_ids == self.shank_id)
         return inds
 
     def get_contact_count(self):

tests/test_io/test_3brain.py

@@ -43,3 +43,7 @@ def test_3brain():
         assert np.all(np.isclose(np.diff(unique_rows), contact_pitch)), file
         unique_cols = np.unique(probe.contact_positions[:, 0])
         assert np.all(np.isclose(np.diff(unique_cols), contact_pitch))
+
+
+if __name__ == "__main__":
+    test_3brain()

tests/test_io/test_io.py

@@ -83,7 +83,8 @@ def test_BIDS_format(tmp_path):
         probe.set_contact_ids(probe_el_ids)
 
         # switch to more generic dtype for shank_ids
-        probe.set_shank_ids(probe.shank_ids.astype(str))
+        if probe.shank_ids is not None:
+            probe.set_shank_ids(probe.shank_ids.astype(str))
 
     write_BIDS_probe(folder_path, probegroup)
 
@@ -103,7 +104,8 @@ def test_BIDS_format(tmp_path):
         t = np.array([list(probe_read.contact_ids).index(elid) for elid in probe_orig.contact_ids])
 
         assert all(probe_orig.contact_ids == probe_read.contact_ids[t])
-        assert all(probe_orig.shank_ids == probe_read.shank_ids[t])
+        if probe_orig.shank_ids is not None:
+            assert all(probe_orig.shank_ids == probe_read.shank_ids[t])
         assert all(probe_orig.contact_shapes == probe_read.contact_shapes[t])
         assert probe_orig.ndim == probe_read.ndim
         assert probe_orig.si_units == probe_read.si_units
@@ -206,8 +208,12 @@ def test_prb(tmp_path):
 
 
 if __name__ == "__main__":
-    # test_probeinterface_format()
-    # test_BIDS_format()
+    import tempfile
+
+    tmp_path = Path(tempfile.mkdtemp())
+
+    # test_probeinterface_format(tmp_path)
+    test_BIDS_format(tmp_path)
     # test_BIDS_format_empty()
     # test_BIDS_format_minimal()
     pass

tests/test_probe.py

@@ -229,10 +229,10 @@ def test_double_side_probe():
 
 
 if __name__ == "__main__":
-    test_probe()
+    import tempfile
 
-    tmp_path = Path("tmp")
-    tmp_path.mkdir(exist_ok=True)
-    test_save_to_zarr(tmp_path)
+    tmp_path = Path(tempfile.mkdtemp())
 
+    test_probe()
+    test_save_to_zarr(tmp_path)
     test_double_side_probe()