From f258e2ea8149b9836cd1b8068872a6622eb95147 Mon Sep 17 00:00:00 2001 From: Samuel Garcia Date: Wed, 17 Dec 2025 21:21:14 +0100 Subject: [PATCH 1/2] Implement the by_side grouping with probes --- .../core/baserecordingsnippets.py | 22 +++++++++++++++++-- 1 file changed, 20 insertions(+), 2 deletions(-) diff --git a/src/spikeinterface/core/baserecordingsnippets.py b/src/spikeinterface/core/baserecordingsnippets.py index 56f558bce7..526aa9bbc6 100644 --- a/src/spikeinterface/core/baserecordingsnippets.py +++ b/src/spikeinterface/core/baserecordingsnippets.py @@ -105,7 +105,7 @@ def _set_probes(self, probe_or_probegroup, group_mode="by_probe", in_place=False The probe(s) to be attached to the recording group_mode: "by_probe" | "by_shank", default: "by_probe" "by_probe" or "by_shank". Adds grouping property to the recording based on the probes ("by_probe") - or shanks ("by_shank") + or shanks ("by_shank") or by probe side ("by_side") in_place: bool False by default. Useful internally when extractor do self.set_probegroup(probe) @@ -115,7 +115,7 @@ def _set_probes(self, probe_or_probegroup, group_mode="by_probe", in_place=False sub_recording: BaseRecording A view of the recording (ChannelSlice or clone or itself) """ - assert group_mode in ("by_probe", "by_shank"), "'group_mode' can be 'by_probe' or 'by_shank'" + assert group_mode in ("by_probe", "by_shank", "by_side"), "'group_mode' can be 'by_probe' or 'by_shank' or 'by_side'" # handle several input possibilities if isinstance(probe_or_probegroup, Probe): @@ -213,6 +213,24 @@ def _set_probes(self, probe_or_probegroup, group_mode="by_probe", in_place=False probe_as_numpy_array["shank_ids"] == a["shank_ids"] ) groups[mask] = group + elif group_mode == "by_side": + assert all( + probe.contact_sides is not None for probe in probegroup.probes + ), "contact_sides is None in probe, you cannot group by side" + if probe.shank_ids is None: + for group, a in enumerate(np.unique(probe_as_numpy_array[["probe_index", "contact_sides"]])): + mask = (probe_as_numpy_array["probe_index"] == a["probe_index"]) & ( + probe_as_numpy_array["contact_sides"] == a["contact_sides"] + ) + groups[mask] = group + else: + for group, a in enumerate(np.unique(probe_as_numpy_array[["probe_index", "shank_ids", "contact_sides"]])): + mask = (probe_as_numpy_array["probe_index"] == a["probe_index"]) & ( + probe_as_numpy_array["shank_ids"] == a["shank_ids"]) & ( + probe_as_numpy_array["contact_sides"] == a["contact_sides"] + ) + groups[mask] = group + sub_recording.set_property("group", groups, ids=None) # add probe annotations to recording From d5f15d2a2b5dd00b66a31a0620e1374f02d96b98 Mon Sep 17 00:00:00 2001 From: Samuel Garcia Date: Fri, 19 Dec 2025 12:04:15 +0100 Subject: [PATCH 2/2] better handling of grouping when set_probe() add "auto" mode --- .../core/baserecordingsnippets.py | 72 +++++++++---------- .../core/tests/test_baserecording.py | 12 ++-- 2 files changed, 41 insertions(+), 43 deletions(-) diff --git a/src/spikeinterface/core/baserecordingsnippets.py b/src/spikeinterface/core/baserecordingsnippets.py index 526aa9bbc6..adfd493446 100644 --- a/src/spikeinterface/core/baserecordingsnippets.py +++ b/src/spikeinterface/core/baserecordingsnippets.py @@ -61,7 +61,7 @@ def is_filtered(self): # the is_filtered is handle with annotation return self._annotations.get("is_filtered", False) - def set_probe(self, probe, group_mode="by_probe", in_place=False): + def set_probe(self, probe, group_mode="auto", in_place=False): """ Attach a list of Probe object to a recording. @@ -69,9 +69,9 @@ def set_probe(self, probe, group_mode="by_probe", in_place=False): ---------- probe_or_probegroup: Probe, list of Probe, or ProbeGroup The probe(s) to be attached to the recording - group_mode: "by_probe" | "by_shank", default: "by_probe - "by_probe" or "by_shank". Adds grouping property to the recording based on the probes ("by_probe") - or shanks ("by_shanks") + group_mode: "auto" | "by_probe" | "by_shank" | "by_side", default: "auto" + How to add the "group" property. + "auto" is the best splitting possible that can be all at once when multiple, probe with multiple shanks and 2 sides. in_place: bool False by default. Useful internally when extractor do self.set_probegroup(probe) @@ -86,10 +86,10 @@ def set_probe(self, probe, group_mode="by_probe", in_place=False): probegroup.add_probe(probe) return self._set_probes(probegroup, group_mode=group_mode, in_place=in_place) - def set_probegroup(self, probegroup, group_mode="by_probe", in_place=False): + def set_probegroup(self, probegroup, group_mode="auto", in_place=False): return self._set_probes(probegroup, group_mode=group_mode, in_place=in_place) - def _set_probes(self, probe_or_probegroup, group_mode="by_probe", in_place=False): + def _set_probes(self, probe_or_probegroup, group_mode="auto", in_place=False): """ Attach a list of Probe objects to a recording. For this Probe.device_channel_indices is used to link contacts to recording channels. @@ -103,9 +103,9 @@ def _set_probes(self, probe_or_probegroup, group_mode="by_probe", in_place=False ---------- probe_or_probegroup: Probe, list of Probe, or ProbeGroup The probe(s) to be attached to the recording - group_mode: "by_probe" | "by_shank", default: "by_probe" - "by_probe" or "by_shank". Adds grouping property to the recording based on the probes ("by_probe") - or shanks ("by_shank") or by probe side ("by_side") + group_mode: "auto" | "by_probe" | "by_shank" | "by_side", default: "auto" + How to add the "group" property. + "auto" is the best splitting possible that can be all at once when multiple, probe with multiple shanks and 2 sides. in_place: bool False by default. Useful internally when extractor do self.set_probegroup(probe) @@ -115,7 +115,7 @@ def _set_probes(self, probe_or_probegroup, group_mode="by_probe", in_place=False sub_recording: BaseRecording A view of the recording (ChannelSlice or clone or itself) """ - assert group_mode in ("by_probe", "by_shank", "by_side"), "'group_mode' can be 'by_probe' or 'by_shank' or 'by_side'" + assert group_mode in ("auto", "by_probe", "by_shank", "by_side"), "'group_mode' can be 'auto' 'by_probe' 'by_shank' or 'by_side'" # handle several input possibilities if isinstance(probe_or_probegroup, Probe): @@ -199,38 +199,32 @@ def _set_probes(self, probe_or_probegroup, group_mode="by_probe", in_place=False sub_recording.set_property("location", locations, ids=None) # handle groups - groups = np.zeros(probe_as_numpy_array.size, dtype="int64") - if group_mode == "by_probe": - for group, probe_index in enumerate(np.unique(probe_as_numpy_array["probe_index"])): - mask = probe_as_numpy_array["probe_index"] == probe_index - groups[mask] = group + all_has_shank_id = all(probe.shank_ids is not None for probe in probegroup.probes) + all_has_contact_side = all(probe.contact_sides is not None for probe in probegroup.probes) + if group_mode == "auto": + group_keys = ["probe_index"] + if all_has_shank_id: + group_keys += ["shank_ids"] + if all_has_contact_side: + group_keys += ["contact_sides"] + elif group_mode == "by_probe": + group_keys = ["probe_index"] elif group_mode == "by_shank": - assert all( - probe.shank_ids is not None for probe in probegroup.probes - ), "shank_ids is None in probe, you cannot group by shank" - for group, a in enumerate(np.unique(probe_as_numpy_array[["probe_index", "shank_ids"]])): - mask = (probe_as_numpy_array["probe_index"] == a["probe_index"]) & ( - probe_as_numpy_array["shank_ids"] == a["shank_ids"] - ) - groups[mask] = group + assert all_has_shank_id, "shank_ids is None in probe, you cannot group by shank" + group_keys = ["probe_index", "shank_ids"] elif group_mode == "by_side": - assert all( - probe.contact_sides is not None for probe in probegroup.probes - ), "contact_sides is None in probe, you cannot group by side" - if probe.shank_ids is None: - for group, a in enumerate(np.unique(probe_as_numpy_array[["probe_index", "contact_sides"]])): - mask = (probe_as_numpy_array["probe_index"] == a["probe_index"]) & ( - probe_as_numpy_array["contact_sides"] == a["contact_sides"] - ) - groups[mask] = group + assert all_has_contact_side, "contact_sides is None in probe, you cannot group by side" + if all_has_shank_id: + group_keys = ["probe_index", "shank_ids", "contact_sides"] else: - for group, a in enumerate(np.unique(probe_as_numpy_array[["probe_index", "shank_ids", "contact_sides"]])): - mask = (probe_as_numpy_array["probe_index"] == a["probe_index"]) & ( - probe_as_numpy_array["shank_ids"] == a["shank_ids"]) & ( - probe_as_numpy_array["contact_sides"] == a["contact_sides"] - ) - groups[mask] = group - + group_keys = ["probe_index", "contact_sides"] + groups = np.zeros(probe_as_numpy_array.size, dtype="int64") + unique_keys = np.unique(probe_as_numpy_array[group_keys]) + for group, a in enumerate(unique_keys): + mask = np.ones(probe_as_numpy_array.size, dtype=bool) + for k in group_keys: + mask &= (probe_as_numpy_array[k] == a[k]) + groups[mask] = group sub_recording.set_property("group", groups, ids=None) # add probe annotations to recording diff --git a/src/spikeinterface/core/tests/test_baserecording.py b/src/spikeinterface/core/tests/test_baserecording.py index 9de800b33d..d4fa26f5f8 100644 --- a/src/spikeinterface/core/tests/test_baserecording.py +++ b/src/spikeinterface/core/tests/test_baserecording.py @@ -179,11 +179,15 @@ def test_BaseRecording(create_cache_folder): # set/get Probe only 2 channels probe = Probe(ndim=2) - positions = [[0.0, 0.0], [0.0, 15.0], [0, 30.0]] - probe.set_contacts(positions=positions, shapes="circle", shape_params={"radius": 5}) - probe.set_device_channel_indices([2, -1, 0]) + positions = [[0.0, 0.0], [0.0, 15.0], [0, 30.0], + [100.0, 0.0], [100.0, 15.0], [100.0, 30.0], + ] + probe.set_contacts(positions=positions, shapes="circle", shape_params={"radius": 5}, shank_ids=["a"]*3 + ["b"]*3) + probe.set_device_channel_indices([2, -1, 0, -1, -1, -1 ], ) probe.create_auto_shape() + print("ici", probe.shank_ids) + rec_p = rec.set_probe(probe, group_mode="by_shank") rec_p = rec.set_probe(probe, group_mode="by_probe") positions2 = rec_p.get_channel_locations() @@ -216,7 +220,7 @@ def test_BaseRecording(create_cache_folder): # set unconnected probe probe = Probe(ndim=2) positions = [[0.0, 0.0], [0.0, 15.0], [0, 30.0]] - probe.set_contacts(positions=positions, shapes="circle", shape_params={"radius": 5}) + probe.set_contacts(positions=positions, shapes="circle", shape_params={"radius": 5}, shank_ids=["a", "a", "a"]) probe.set_device_channel_indices([-1, -1, -1]) probe.create_auto_shape()