Enhance probegroup API

examples/ex_03_generate_probe_group.py

@@ -46,4 +46,64 @@
 
 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_contacts()` returns a new `ProbeGroup` with a sub-selection
+# of contacts. The selection can be done by ``contact_ids``, by ``probe_ids``,
+# or by both at the same time.
+#
+# Selecting by ``probe_ids`` alone keeps every contact of the matching probes,
+# which is a convenient way to grab a whole hemisphere:
+
+left_hemisphere = probegroup.select_contacts(probe_ids=["left_hemisphere"])
+print("contacts in the left hemisphere:", left_hemisphere.get_contact_count())
+
+##############################################################################
+# We can also select by ``contact_ids``. Note that if ``contact_ids`` are not 
+# unique across probes, the selection will be ambiguous and an error will be 
+# raised. In this case, providing ``probe_ids`` disambiguates the selection:
+
+# 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_contacts = probegroup.select_contacts(contact_ids=["0", "1", "2"], probe_ids=["left_hemisphere"])
+print("contacts selected from the left hemisphere:", left_contacts.get_contact_count())
+
+left_and_right_contacts = probegroup.select_contacts(
+    contact_ids=["0", "1", "2"],
+    probe_ids=["left_hemisphere", "right_hemisphere"]
+)
+print("contacts selected from the left and right hemispheres:", left_and_right_contacts.get_contact_count())
+
+# 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()

examples/ex_05_device_channel_indices.py

@@ -87,4 +87,46 @@
 fig, ax = plt.subplots()
 plot_probegroup(probegroup, with_contact_id=True, same_axes=True, ax=ax)
 
+##############################################################################
+# Reordering contacts with a global contact order
+# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+#
+# By default the contact order of a `ProbeGroup` is the "natural" one: the
+# contacts of each probe are stacked one probe after the other. But sometimes
+# the contacts of the different probes are *interleaved* in the recording file
+# (e.g. the acquisition system alternates between probes sample by sample).
+#
+# `ProbeGroup.set_global_contact_order()` lets us store this external ordering.
+# The order is an array of indices into the natural (stacked) order, and it is
+# applied whenever the group is exported with `to_numpy()` / `to_dataframe()`.
+
+probegroup = ProbeGroup()
+probegroup.add_probe(probe0.copy())
+probegroup.add_probe(probe1.copy())
+
+n = probegroup.get_contact_count()
+print("default global contact order:", probegroup._global_contact_order)
+
+# interleave probe0 and probe1 contacts as they appear in the recording file
+global_contact_order = np.zeros(n, dtype="int64")
+global_contact_order[0::2] = np.arange(0, n // 2)        # probe0 contacts
+global_contact_order[1::2] = np.arange(n // 2, n)        # probe1 contacts
+probegroup.set_global_contact_order(global_contact_order)
+
+##############################################################################
+# Now `to_numpy()` returns the contacts in the interleaved order: the
+# ``probe_index`` column alternates between the two probes.
+
+contact_vector = probegroup.to_numpy()
+print("probe_index in global order:", contact_vector["probe_index"][:8])
+
+##############################################################################
+# The global order interacts with `set_global_device_channel_indices()`: the
+# ``device_channel_indices`` you pass are interpreted in the (reordered) order
+# returned by `to_numpy()`, so they map directly onto the acquisition channels.
+
+probegroup.set_global_device_channel_indices(np.arange(n))
+print("device_channel_indices (global order):",
+      probegroup.to_numpy(complete=True)["device_channel_indices"][:8])
+
 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 beof 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 = []