Implement copy/get_slice/get_global_contact_positions for probegroup

Author: alejoe91

src/probeinterface/probegroup.py

@@ -50,6 +50,22 @@ def _check_compatible(self, probe: Probe):
     def ndim(self):
         return self.probes[0].ndim
 
+    def copy(self) -> "ProbeGroup":
+        """
+        Create a copy of the ProbeGroup
+
+        Returns
+        -------
+        copy: ProbeGroup
+            A copy of the ProbeGroup
+        """
+        copy = ProbeGroup()
+        for probe in self.probes:
+            copy.add_probe(probe.copy())
+        global_device_channel_indices = self.get_global_device_channel_indices()["device_channel_indices"]
+        copy.set_global_device_channel_indices(global_device_channel_indices)
+        return copy
+
     def get_contact_count(self) -> int:
         """
         Total number of channels.
@@ -249,6 +265,79 @@ def get_global_contact_ids(self) -> np.ndarray:
         contact_ids = self.to_numpy(complete=True)["contact_ids"]
         return contact_ids
 
+    def get_global_contact_positions(self) -> np.ndarray:
+        """
+        Gets all contact positions concatenated across probes
+
+        Returns
+        -------
+        contact_positions: np.ndarray
+            An array of the contact positions across all probes
+        """
+        contact_positions = np.vstack([probe.contact_positions for probe in self.probes])
+        return contact_positions
+
+    def get_slice(self, selection: np.ndarray[bool | int]):
+        """
+        Get a copy of the ProbeGroup with a sub selection of contacts.
+
+        Selection can be boolean or by index
+
+        Parameters
+        ----------
+        selection : np.array of bool or int (for index)
+            Either an np.array of bool or for desired selection of contacts
+            or the indices of the desired contacts
+
+        Returns
+        -------
+        sliced_probe_group: ProbeGroup
+            The sliced probe group
+
+        """
+
+        n = self.get_contact_count()
+
+        selection = np.asarray(selection)
+        if selection.dtype == "bool":
+            assert selection.shape == (n,), (
+                f"if array of bool given it must be the same size " "as the number of contacts {selection.shape} != {n}"
+            )
+            selection_indices, = np.nonzero(selection)
+        elif selection.dtype.kind == "i":
+            assert np.unique(selection).size == selection.size
+            if len(selection) > 0:
+                assert 0 <= np.min(selection) < n, f"An index within your selection is out of bounds {np.min(selection)}"
+                assert 0 <= np.max(selection) < n, f"An index within your selection is out of bounds {np.max(selection)}"
+                selection_indices = selection
+            else:
+                selection_indices = []
+        else:
+            raise TypeError(f"selection must be bool array or int array, not of type: {type(selection)}")
+
+        if len(selection_indices) == 0:
+            return ProbeGroup()
+
+        # Map selection to indices of individual probes
+        d = self.to_dict(array_as_list=False)
+        ind = 0
+        sliced_probes = []
+        for probe in self.probes:
+            n = probe.get_contact_count()
+            probe_limits = (ind, ind + n)
+            ind += n
+
+            probe_selection_indices = selection_indices[(selection_indices >= probe_limits[0]) & (selection_indices < probe_limits[1])]
+            if len(probe_selection_indices) == 0:
+                continue
+            sliced_probe = probe.get_slice(probe_selection_indices - probe_limits[0])
+            sliced_probes.append(sliced_probe)
+
+        sliced_probe_group = ProbeGroup()
+        sliced_probe_group.probes = sliced_probes
+
+        return sliced_probe_group
+
     def check_global_device_wiring_and_ids(self):
         # check unique device_channel_indices for !=-1
         chans = self.get_global_device_channel_indices()

tests/test_probegroup.py

@@ -116,6 +116,196 @@ def test_set_contact_ids_rejects_wrong_size():
         probe.set_contact_ids(["a", "b", "c"])
 
 
+def _make_probegroup(n_probes=3):
+    """Helper: build a ProbeGroup with device channel indices set."""
+    probegroup = ProbeGroup()
+    nchan = 0
+    for i in range(n_probes):
+        probe = generate_dummy_probe()
+        probe.move([i * 100, i * 80])
+        n = probe.get_contact_count()
+        probe.set_device_channel_indices(np.arange(n) + nchan)
+        probe.set_contact_ids([f"e{j}" for j in range(nchan, nchan + n)])
+        nchan += n
+    probegroup.add_probe(probe)
+    return probegroup
+
+
+def _make_probegroup_full(n_probes=3):
+    """Helper: build a ProbeGroup where **every** probe is added."""
+    probegroup = ProbeGroup()
+    nchan = 0
+    for i in range(n_probes):
+        probe = generate_dummy_probe()
+        probe.move([i * 100, i * 80])
+        n = probe.get_contact_count()
+        probe.set_device_channel_indices(np.arange(n) + nchan)
+        probe.set_contact_ids([f"e{j}" for j in range(nchan, nchan + n)])
+        probegroup.add_probe(probe)
+        nchan += n
+    return probegroup
+
+
+# ── copy() tests ────────────────────────────────────────────────────────────
+
+
+def test_copy_returns_new_object():
+    pg = _make_probegroup_full(2)
+    pg_copy = pg.copy()
+    assert pg_copy is not pg
+    assert len(pg_copy.probes) == len(pg.probes)
+    for orig, copied in zip(pg.probes, pg_copy.probes):
+        assert orig is not copied
+
+
+def test_copy_preserves_positions():
+    pg = _make_probegroup_full(2)
+    pg_copy = pg.copy()
+    for orig, copied in zip(pg.probes, pg_copy.probes):
+        np.testing.assert_array_equal(orig.contact_positions, copied.contact_positions)
+
+
+def test_copy_preserves_device_channel_indices():
+    pg = _make_probegroup_full(2)
+    pg_copy = pg.copy()
+    np.testing.assert_array_equal(
+        pg.get_global_device_channel_indices(),
+        pg_copy.get_global_device_channel_indices(),
+    )
+
+
+def test_copy_does_not_preserve_contact_ids():
+    """Probe.copy() intentionally does not copy contact_ids."""
+    pg = _make_probegroup_full(2)
+    pg_copy = pg.copy()
+    # All contact_ids should be empty strings after copy
+    assert all(cid == "" for cid in pg_copy.get_global_contact_ids())
+
+
+def test_copy_is_independent():
+    """Mutating the copy must not affect the original."""
+    pg = _make_probegroup_full(2)
+    original_positions = pg.probes[0].contact_positions.copy()
+    pg_copy = pg.copy()
+    pg_copy.probes[0].move([999, 999])
+    np.testing.assert_array_equal(pg.probes[0].contact_positions, original_positions)
+
+
+# ── get_slice() tests ───────────────────────────────────────────────────────
+
+
+def test_get_slice_by_bool():
+    pg = _make_probegroup_full(2)
+    total = pg.get_contact_count()
+    sel = np.zeros(total, dtype=bool)
+    sel[:5] = True  # first 5 contacts from the first probe
+    sliced = pg.get_slice(sel)
+    assert sliced.get_contact_count() == 5
+
+
+def test_get_slice_by_index():
+    pg = _make_probegroup_full(2)
+    indices = np.array([0, 1, 2, 33, 34])  # contacts from both probes
+    sliced = pg.get_slice(indices)
+    assert sliced.get_contact_count() == 5
+
+
+def test_get_slice_preserves_device_channel_indices():
+    pg = _make_probegroup_full(2)
+    indices = np.array([0, 1, 2])
+    sliced = pg.get_slice(indices)
+    orig_chans = pg.get_global_device_channel_indices()["device_channel_indices"][:3]
+    sliced_chans = sliced.get_global_device_channel_indices()["device_channel_indices"]
+    np.testing.assert_array_equal(sliced_chans, orig_chans)
+
+
+def test_get_slice_preserves_positions():
+    pg = _make_probegroup_full(2)
+    indices = np.array([0, 1, 2])
+    sliced = pg.get_slice(indices)
+    expected = pg.get_global_contact_positions()[indices]
+    np.testing.assert_array_equal(sliced.get_global_contact_positions(), expected)
+
+
+def test_get_slice_empty_selection():
+    pg = _make_probegroup_full(2)
+    sliced = pg.get_slice(np.array([], dtype=int))
+    assert sliced.get_contact_count() == 0
+    assert len(sliced.probes) == 0
+
+
+def test_get_slice_wrong_bool_size():
+    pg = _make_probegroup_full(2)
+    with pytest.raises(AssertionError):
+        pg.get_slice(np.array([True, False]))  # wrong size
+
+
+def test_get_slice_out_of_bounds():
+    pg = _make_probegroup_full(2)
+    total = pg.get_contact_count()
+    with pytest.raises(AssertionError):
+        pg.get_slice(np.array([total + 10]))
+
+
+def test_get_slice_all_contacts():
+    """Slicing with all contacts should give an equivalent ProbeGroup."""
+    pg = _make_probegroup_full(2)
+    total = pg.get_contact_count()
+    sliced = pg.get_slice(np.arange(total))
+    assert sliced.get_contact_count() == total
+    np.testing.assert_array_equal(
+        sliced.get_global_contact_positions(),
+        pg.get_global_contact_positions(),
+    )
+
+
+# ── get_global_contact_positions() tests ────────────────────────────────────
+
+
+def test_get_global_contact_positions_shape():
+    pg = _make_probegroup_full(3)
+    pos = pg.get_global_contact_positions()
+    assert pos.shape == (pg.get_contact_count(), pg.ndim)
+
+
+def test_get_global_contact_positions_matches_per_probe():
+    pg = _make_probegroup_full(3)
+    pos = pg.get_global_contact_positions()
+    offset = 0
+    for probe in pg.probes:
+        n = probe.get_contact_count()
+        np.testing.assert_array_equal(pos[offset : offset + n], probe.contact_positions)
+        offset += n
+
+
+def test_get_global_contact_positions_single_probe():
+    pg = _make_probegroup_full(1)
+    pos = pg.get_global_contact_positions()
+    np.testing.assert_array_equal(pos, pg.probes[0].contact_positions)
+
+
+def test_get_global_contact_positions_3d():
+    pg = ProbeGroup()
+    for i in range(2):
+        probe = generate_dummy_probe().to_3d()
+        probe.move([i * 100, i * 80, i * 30])
+        pg.add_probe(probe)
+    pos = pg.get_global_contact_positions()
+    assert pos.shape[1] == 3
+    assert pos.shape[0] == pg.get_contact_count()
+
+
+def test_get_global_contact_positions_reflects_move():
+    """Positions should reflect probe movement."""
+    pg = ProbeGroup()
+    probe = generate_dummy_probe()
+    original_pos = probe.contact_positions.copy()
+    probe.move([50, 60])
+    pg.add_probe(probe)
+    pos = pg.get_global_contact_positions()
+    np.testing.assert_array_equal(pos, original_pos + np.array([50, 60]))
+
+
 if __name__ == "__main__":
     test_probegroup()
     # ~ test_probegroup_3d()