Enhance probegroup API

examples/ex_03_generate_probe_group.py

@@ -46,4 +46,65 @@
 
 plot_probegroup(probegroup, same_axes=False, with_contact_id=True)
 
+##############################################################################
+# Identifying probes with a ``probe_id``
+# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+#
+# Each probe in a `ProbeGroup` can be given a human-readable ``probe_id`` when
+# it is added. This is handy to keep track of which probe targets which brain
+# area or hemisphere. If no ``probe_id`` is given, a default one
+# (``"probe_1"``, ``"probe_2"``, ...) is generated automatically.
+
+probe0 = generate_dummy_probe(elec_shapes='square')
+probe1 = generate_dummy_probe(elec_shapes='circle')
+probe1.move([250, -90])
+
+probegroup = ProbeGroup()
+probegroup.add_probe(probe0, probe_id="left_hemisphere")
+probegroup.add_probe(probe1, probe_id="right_hemisphere")
+
+print(probegroup)
+print("probe_ids:", probegroup.probe_ids)
+
+##############################################################################
+# `ProbeGroup.select_probes()` returns a new `ProbeGroup` with a sub-selection
+# of probes given by probe_ids.
+
+left_hemisphere_probe = probegroup.select_probes(probe_ids=["left_hemisphere"])
+print(left_hemisphere_probe)
+
+##############################################################################
+# We can also select by specific contacts from a probegroup with the
+# ``select_contacts`` function. Note that if ``contact_ids`` are not
+# unique across probes, you need to disambiguate the selection by specifying the
+# probe_ids as well. Otherwise, a ValueError is raised.
+
+# check if any contact_id is not unique across probes
+contact_ids = probegroup.get_global_contact_ids()
+if len(contact_ids) != len(set(contact_ids)):
+    print("contact_ids are not unique across probes, you should provide probe_ids to disambiguate")
+
+##############################################################################
+# Because the contact ids are not unique across probes, combining ``contact_ids``
+# with ``probe_ids`` lets us pull specific contacts from a single hemisphere:
+
+left_probegroup = probegroup.select_contacts(
+    contact_ids=["0", "1", "2"],
+    probe_ids=["left_hemisphere", "left_hemisphere", "left_hemisphere"]
+)
+print(left_probegroup)
+
+# Now select contacts from both hemispheres by providing the corresponding probe_ids for each contact_id:
+left_and_right_probegroup = probegroup.select_contacts(
+    contact_ids=["0", "1", "2"],
+    probe_ids=["left_hemisphere", "right_hemisphere", "left_hemisphere"]
+)
+print(left_and_right_probegroup)
+
+# Without providing probe_ids, the selection is ambiguous and an error is raised:
+try:
+    ambiguous_selection = probegroup.select_contacts(contact_ids=["0", "1", "2"])
+except ValueError as e:
+    print("Error raised for ambiguous selection:", e)
+
 plt.show()

src/probeinterface/__init__.py

@@ -52,4 +52,8 @@
     cache_full_library,
     clear_cache,
 )
-from .wiring import get_available_pathways
+from .wiring import (
+    get_available_pathways,
+    get_pathway,
+    wire_probe
+)

src/probeinterface/io.py

@@ -203,10 +203,9 @@ def read_BIDS_probe(folder: str | Path, prefix: str | None = None) -> ProbeGroup
 
         # create probe object and register with probegroup
         probe = Probe.from_dataframe(df=df_probe)
-        probe.annotate(probe_id=probe_id)
 
         probes[str(probe_id)] = probe
-        probegroup.add_probe(probe)
+        probegroup.add_probe(probe, probe_id=str(probe_id))
 
         ignore_annotations = [
             "probe_ids",
@@ -326,7 +325,7 @@ def write_BIDS_probe(folder: str | Path, probe_or_probegroup: Probe | ProbeGroup
         probegroup = probe_or_probegroup
     else:
         raise TypeError(
-            f"probe_or_probegroup has to be" "of type Probe or ProbeGroup " f"not type: {type(probe_or_probegroup)}"
+            f"probe_or_probegroup has to be of type Probe or ProbeGroup not type: {type(probe_or_probegroup)}"
         )
     folder = Path(folder)
 
@@ -337,22 +336,12 @@ def write_BIDS_probe(folder: str | Path, probe_or_probegroup: Probe | ProbeGroup
     probes = probegroup.probes
 
     # Step 1: GENERATION OF PROBE.TSV
-    # ensure required keys (probe_id, probe_type) are present
-
-    if any("probe_id" not in p.annotations for p in probes):
-        probegroup.auto_generate_probe_ids()
+    # ensure required keys (probe_type) are present
 
     for probe in probes:
-        if "probe_id" not in probe.annotations:
-            raise ValueError(
-                "Export to BIDS probe format requires "
-                "the probe id to be specified as an annotation "
-                "(probe_id). You can do this via "
-                "`probegroup.auto_generate_ids."
-            )
         if "type" not in probe.annotations:
             raise ValueError(
-                "Export to BIDS probe format requires " "the probe type to be specified as an " "annotation (type)"
+                "Export to BIDS probe format requires the probe type to be specified as an annotation (type)"
             )
 
     # extract all used annotation keys
@@ -361,11 +350,12 @@ def write_BIDS_probe(folder: str | Path, probe_or_probegroup: Probe | ProbeGroup
     annotation_keys = np.unique(keys_concatenated)
 
     # generate a tsv table capturing probe information
-    index = range(len([p.annotations["probe_id"] for p in probes]))
+    index = range(len(probes))
     df = pd.DataFrame(index=index)
     for annotation_key in annotation_keys:
         df[annotation_key] = [p.annotations[annotation_key] for p in probes]
     df["n_shanks"] = [len(np.unique(p.shank_ids)) for p in probes]
+    df["probe_id"] = probegroup.probe_ids
 
     # Note: in principle it would also be possible to add the probe width and
     # depth here based on the probe contour information. However this would
@@ -378,8 +368,7 @@ def write_BIDS_probe(folder: str | Path, probe_or_probegroup: Probe | ProbeGroup
 
     # Step 2: GENERATION OF PROBE.JSON
     probes_dict = {}
-    for probe in probes:
-        probe_id = probe.annotations["probe_id"]
+    for probe_id, probe in zip(probegroup.probe_ids, probes):
         probes_dict[probe_id] = {
             "contour": probe.probe_planar_contour.tolist(),
             "units": probe.si_units,
@@ -403,7 +392,7 @@ def write_BIDS_probe(folder: str | Path, probe_or_probegroup: Probe | ProbeGroup
     index = range(sum([p.get_contact_count() for p in probes]))
     df.rename(columns=tsv_label_map_to_BIDS, inplace=True)
 
-    df["probe_id"] = [p.annotations["probe_id"] for p in probes for _ in p.contact_ids]
+    df["probe_id"] = [probe_id for probe_id, probe in zip(probegroup.probe_ids, probes) for _ in probe.contact_ids]
     df["coordinate_system"] = ["relative cartesian"] * len(index)
 
     channel_indices = []

src/probeinterface/probe.py

@@ -534,7 +534,7 @@ def set_device_channel_indices(self, channel_indices: np.ndarray | list):
             )
         self.device_channel_indices = channel_indices
         if self._probe_group is not None:
-            self._probe_group.check_global_device_wiring_and_ids()
+            self._probe_group._check_global_device_wiring_and_ids()
 
     def wiring_to_device(self, pathway: str, channel_offset: int = 0):
         """
@@ -584,7 +584,7 @@ def set_contact_ids(self, contact_ids: np.ndarray | list):
 
         self._contact_ids = contact_ids
         if self._probe_group is not None:
-            self._probe_group.check_global_device_wiring_and_ids()
+            self._probe_group._check_global_device_wiring_and_ids()
 
     def set_shank_ids(self, shank_ids: np.ndarray | list):
         """
@@ -1140,8 +1140,10 @@ def from_numpy(arr: np.ndarray) -> "Probe":
             "plane_axis_y_1",
             "plane_axis_z_0",
             "plane_axis_z_1",
-            "probe_index",
             "si_units",
+            # these two are for ProbeGroup to avoid duplication of fields
+            "probe_index",
+            "probe_id",
         ]
         contact_annotation_fields = [f for f in fields if f not in main_fields]