Build one probe per electrode group in NwbRecordingExtractor

src/spikeinterface/extractors/nwbextractors.py

@@ -5,6 +5,7 @@
 import importlib.util
 
 import numpy as np
+from probeinterface import Probe, ProbeGroup
 
 from spikeinterface import get_global_tmp_folder
 from spikeinterface.core import BaseRecording, BaseRecordingSegment, BaseSorting, BaseSortingSegment
@@ -1040,9 +1041,21 @@ def __init__(
         self.set_channel_gains(gains_to_uV)
         self.set_channel_offsets(offsets_to_uV)
         locations = self._reader.locations()
-        if locations is not None:
-            self.set_channel_locations(locations)
         groups = self._reader.groups()
+        if locations is not None:
+            # rel_x, rel_y and rel_z are relative to each electrode group, so each group is its own
+            # probe. Two groups can then share positions (e.g. two identical probes in one
+            # ElectricalSeries), so the overlap check across probes is skipped.
+            probegroup = ProbeGroup()
+            for group_name in np.unique(groups):
+                channel_indices = np.flatnonzero(groups == group_name)
+                probe = Probe(ndim=2)
+                probe.set_contacts(locations[channel_indices, :2], shapes="circle", shape_params={"radius": 1})
+                probe.set_device_channel_indices(channel_indices)
+                if locations.shape[1] == 3:
+                    probe = probe.to_3d(axes="xy")
+                probegroup.add_probe(probe)
+            self.set_probegroup(probegroup, check_overlap=False)
         if groups is not None:
             self.set_channel_groups(groups)
 

src/spikeinterface/extractors/tests/test_nwbextractors.py

@@ -311,6 +311,48 @@ def test_nwb_extractor_electrodes_region_out_of_order(tmp_path, use_pynwb):
     assert np.array_equal(recording.get_channel_offsets(), np.array([float(i) for i in region]) * 1e6)
 
 
+@pytest.mark.parametrize("use_pynwb", [True, False])
+def test_nwb_extractor_one_probe_per_electrode_group(tmp_path, use_pynwb):
+    """rel_x and rel_y are relative to each electrode group, so two identical probes written to one
+    ElectricalSeries share positions. Building a single probe from all channels failed on the
+    duplicate positions (GH-3999)."""
+    from pynwb import NWBHDF5IO
+    from pynwb.testing.mock.file import mock_NWBFile
+    from pynwb.testing.mock.device import mock_Device
+    from pynwb.testing.mock.ecephys import mock_ElectrodeGroup, mock_ElectricalSeries
+
+    nwbfile = mock_NWBFile()
+    device = mock_Device(nwbfile=nwbfile)
+
+    # Both groups get the same contact positions, so the two probes overlap.
+    group_names = ["probe_A", "probe_B"]
+    probe_locations = np.array([[0.0, 0.0], [0.0, 20.0], [0.0, 40.0], [0.0, 60.0]])
+    for group_name in group_names:
+        electrode_group = mock_ElectrodeGroup(name=group_name, device=device, nwbfile=nwbfile)
+        for rel_x, rel_y in probe_locations:
+            nwbfile.add_electrode(group=electrode_group, location="brain", rel_x=rel_x, rel_y=rel_y)
+
+    channels_per_group = len(probe_locations)
+    num_electrodes = len(group_names) * channels_per_group
+    electrode_region = nwbfile.create_electrode_table_region(region=list(range(num_electrodes)), description="all")
+    mock_ElectricalSeries(
+        name="ElectricalSeries", data=np.ones((10, num_electrodes)), electrodes=electrode_region, nwbfile=nwbfile
+    )
+
+    nwbfile_path = tmp_path / "two_electrode_groups.nwb"
+    with NWBHDF5IO(str(nwbfile_path), mode="w") as io:
+        io.write(nwbfile)
+
+    recording = NwbRecordingExtractor(
+        nwbfile_path, electrical_series_path="acquisition/ElectricalSeries", use_pynwb=use_pynwb
+    )
+
+    assert len(recording.get_probegroup().probes) == len(group_names)
+    recordings_by_group = recording.split_by("group")
+    for group_name in group_names:
+        assert np.array_equal(recordings_by_group[group_name].get_channel_locations(), probe_locations)
+
+
 @pytest.mark.parametrize("use_pynwb", [True, False])
 def test_nwb_extractor_offset_from_series(generate_nwbfile, use_pynwb):
     """Test that the offset is retrieved from the ElectricalSeries if it is present."""