Probe refactoring with wiring

pyproject.toml

@@ -220,7 +220,7 @@ docs = [
     "huggingface_hub", # For automated curation
 
     # for release we need pypi, so this needs to be commented
-    "probeinterface @ git+https://github.com/SpikeInterface/probeinterface.git",  # We always build from the latest version
+    "probeinterface @ git+https://github.com/SpikeInterface/probeinterface.git",
     "neo @ git+https://github.com/NeuralEnsemble/python-neo.git",  # We always build from the latest version
 ]
 

src/spikeinterface/core/base.py

@@ -574,6 +574,9 @@ def to_dict(
                 folder_metadata = Path(folder_metadata).resolve().absolute().relative_to(relative_to)
             dump_dict["folder_metadata"] = str(folder_metadata)
 
+        if getattr(self, "_probegroup", None) is not None:
+            dump_dict["probegroup"] = self._probegroup.to_dict(array_as_list=True)
+
         return dump_dict
 
     @staticmethod
@@ -607,12 +610,11 @@ def from_dict(dictionary: dict, base_folder: Path | str | None = None) -> "BaseE
         return extractor
 
     def load_metadata_from_folder(self, folder_metadata):
-        # hack to load probe for recording
         folder_metadata = Path(folder_metadata)
 
-        self._extra_metadata_from_folder(folder_metadata)
-
-        # load properties
+        # load properties first so that `_extra_metadata_from_folder` can see
+        # restored state like the `wiring` property and skip re-running
+        # `set_probegroup` when the mapping is already in place.
         prop_folder = folder_metadata / "properties"
         if prop_folder.is_dir():
             for prop_file in prop_folder.iterdir():
@@ -621,6 +623,8 @@ def load_metadata_from_folder(self, folder_metadata):
                     key = prop_file.stem
                     self.set_property(key, values)
 
+        self._extra_metadata_from_folder(folder_metadata)
+
     def save_metadata_to_folder(self, folder_metadata):
         self._extra_metadata_to_folder(folder_metadata)
 
@@ -1155,9 +1159,53 @@ def _load_extractor_from_dict(dic) -> "BaseExtractor":
     for k, v in dic["properties"].items():
         extractor.set_property(k, v)
 
+    if "probegroup" in dic:
+        from probeinterface import ProbeGroup
+
+        probegroup = ProbeGroup.from_dict(dic["probegroup"])
+        # The `wiring` per-channel property was restored above by the standard
+        # property-load loop; we just attach the probegroup object.
+        extractor._probegroup = probegroup
+    elif "contact_vector" in dic.get("properties", {}):
+        _restore_probegroup_from_legacy_contact_vector(extractor)
+
     return extractor
 
 
+def _restore_probegroup_from_legacy_contact_vector(extractor) -> None:
+    """
+    Reconstruct a `ProbeGroup` from the legacy `contact_vector` property.
+
+    Recordings saved before the probegroup refactor stored the probe as a structured numpy
+    array under the `contact_vector` property, with probe-level annotations under a separate
+    `probes_info` annotation and per-probe planar contours under `probe_{i}_planar_contour`
+    annotations. This function reconstructs a `ProbeGroup` from those legacy fields, attaches
+    it via the canonical `set_probegroup` path, and removes the legacy property so the new
+    and old representations do not coexist on the loaded extractor.
+    """
+    from probeinterface import ProbeGroup
+
+    contact_vector_array = extractor.get_property("contact_vector")
+    probegroup = ProbeGroup.from_numpy(contact_vector_array)
+
+    if "probes_info" in extractor.get_annotation_keys():
+        probes_info = extractor.get_annotation("probes_info")
+        for probe, probe_info in zip(probegroup.probes, probes_info):
+            probe.annotations = probe_info
+
+    for probe_index, probe in enumerate(probegroup.probes):
+        contour = extractor._annotations.get(f"probe_{probe_index}_planar_contour")
+        if contour is not None:
+            probe.set_planar_contour(contour)
+
+    if hasattr(extractor, "set_probegroup"):
+        extractor.set_probegroup(probegroup, in_place=True)
+    else:
+        extractor._probegroup = probegroup
+
+    extractor._properties.pop("contact_vector", None)
+
+
 def _get_class_from_string(class_string):
     class_name = class_string.split(".")[-1]
     module = ".".join(class_string.split(".")[:-1])

src/spikeinterface/core/baserecording.py

@@ -21,6 +21,7 @@ class BaseRecording(BaseRecordingSnippets, ChunkableMixin):
     _main_properties = [
         "group",
         "location",
+        "wiring",
         "gain_to_uV",
         "offset_to_uV",
         "gain_to_physical_unit",
@@ -392,9 +393,14 @@ def _save(self, format="binary", verbose: bool = False, **save_kwargs):
         else:
             raise ValueError(f"format {format} not supported")
 
-        if self.get_property("contact_vector") is not None:
-            probegroup = self.get_probegroup()
-            cached.set_probegroup(probegroup)
+        if self.has_probe() and not cached.has_probe():
+            # Share the probegroup by reference. We deliberately skip
+            # `set_probegroup` (which re-runs _set_probes and validates dci)
+            # because a child of `split_by` references the parent's full
+            # probegroup whose dci values can exceed the child's channel
+            # count. Wiring/location/group properties are carried over by
+            # the caller's `copy_metadata` step.
+            cached._probegroup = self._probegroup
 
         return cached
 
@@ -403,7 +409,14 @@ def _extra_metadata_from_folder(self, folder):
         folder = Path(folder)
         if (folder / "probe.json").is_file():
             probegroup = read_probeinterface(folder / "probe.json")
-            self.set_probegroup(probegroup, in_place=True)
+            if "wiring" in self.get_property_keys():
+                # wiring was restored via the property-load loop; the stored
+                # probegroup's dci refers to the parent's channel space, so
+                # re-running `_set_probes` would fail for sliced children.
+                # Attach the probegroup object directly.
+                self._probegroup = probegroup
+            else:
+                self.set_probegroup(probegroup, in_place=True)
 
         # load time vector if any
         for segment_index, rs in enumerate(self.segments):
@@ -414,7 +427,7 @@ def _extra_metadata_from_folder(self, folder):
 
     def _extra_metadata_to_folder(self, folder):
         # save probe
-        if self.get_property("contact_vector") is not None:
+        if self.has_probe():
             probegroup = self.get_probegroup()
             write_probeinterface(folder / "probe.json", probegroup)