add set_probe method to BaseRecordingExtractorInterface

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ### Features
 * Added Pydantic data models of `BackendConfiguration` for both HDF5 and Zarr datasets (container/mapper of all the `DatasetConfiguration`s for a particular file). [PR #568](https://github.com/catalystneuro/neuroconv/pull/568)
+* Added set_probe() method to BaseRecordingExtractorInterface
 
 
 

src/neuroconv/datainterfaces/ecephys/baserecordingextractorinterface.py

@@ -225,6 +225,29 @@ def set_aligned_segment_starting_times(self, aligned_segment_starting_times: Lis
             ]
             self.set_aligned_segment_timestamps(aligned_segment_timestamps=aligned_segment_timestamps)
 
+    def set_probe(self, probe, *, group_mode: Union[Literal["by_shank"], Literal["by_probe"]]):
+        """
+        Set the probe information via a ProbeInterface object.
+
+        Parameters
+        ----------
+        probe : probeinterface.Probe
+            The probe object.
+        group_mode : {'by_shank', 'by_probe'}
+            How to group the channels. If 'by_shank', channels are grouped by the shank_id column.
+            If 'by_probe', channels are grouped by the probe_id column.
+            This is a required parameter to avoid the pitfall of using the wrong mode.
+        """
+        # Set the probe to the recording extractor
+        self.recording_extractor.set_probe(
+            probe,
+            in_place=True,
+            group_mode=group_mode,
+        )
+        # Spike interface sets the "group" property
+        # Buy neuroconv expects "group_name" to be set
+        self.recording_extractor.set_property("group_name", self.recording_extractor.get_property("group").astype(str))
+
     def align_by_interpolation(
         self,
         unaligned_timestamps: np.ndarray,