Add roundtrip tests

Author: alejoe91

src/pynwb/ndx_probeinterface/io.py

@@ -69,38 +69,60 @@ def to_probeinterface(ndx_probeinterface_probe) -> Probe:
     polygon = ndx_probeinterface_probe.planar_contour
 
     positions = []
-    contact_ids = []
     shapes = []
-    shape_params = []
-    shank_ids = []
-    plane_axes = []
-    channel_indices = []
+
+    contact_ids = None
+    shape_params = None
+    shank_ids = None
+    plane_axes = None
+    device_channel_indices = None
+
+    possible_shape_keys = ["radius", "width", "height"]
     for shank in ndx_probeinterface_probe.shanks.values():
         positions.append(shank.contact_table["contact_position"][:])
-        contact_ids.append(shank.contact_table["contact_id"][:])
         shapes.append(shank.contact_table["contact_shape"][:])
-        channel_indices.append(shank.contact_table["device_channel_index_pi"][:])
-        plane_axes.append(shank.contact_table["contact_plane_axes"][:])
-        shank_ids.append([int(shank.shank_id)] * len(shank.contact_table))
-        # WARNING: currently assumes that all the contacts have the same shape
-        shape_word = [shape for shape in shape_words if shape in shank.contact_table[:].columns][0]
-        shape_params.append([{shape_word: val} for val in shank.contact_table[shape_word][:]])
+        if "contact_id" in shank.contact_table.colnames:
+            if contact_ids is None:
+                contact_ids = []
+            contact_ids.append(shank.contact_table["contact_id"][:])
+        if "device_channel_index_pi" in shank.contact_table.colnames:
+            if device_channel_indices is None:
+                device_channel_indices = []
+            device_channel_indices.append(shank.contact_table["device_channel_index_pi"][:])
+        if "contact_plane_axes" in shank.contact_table.colnames:
+            if plane_axes is None:
+                plane_axes = []
+            plane_axes.append(shank.contact_table["contact_plane_axes"][:])
+        if shank_ids is None:
+            shank_ids = []
+        shank_ids.append([str(shank.shank_id)] * len(shank.contact_table))
+        for possible_shape_key in possible_shape_keys:
+            if possible_shape_key in shank.contact_table.colnames:
+                if shape_params is None:
+                    shape_params = []
+                shape_params.append([{possible_shape_key: val} for val in shank.contact_table[possible_shape_key][:]])
 
     positions = [item for sublist in positions for item in sublist]
-    contact_ids = [item for sublist in contact_ids for item in sublist]
     shapes = [item for sublist in shapes for item in sublist]
-    plane_axes = [item for sublist in plane_axes for item in sublist]
-    shank_ids = [item for sublist in shank_ids for item in sublist]
-    channel_indices = [item for sublist in channel_indices for item in sublist]
-    shape_params = [item for sublist in shape_params for item in sublist]
+
+    if contact_ids is not None:
+        contact_ids = [item for sublist in contact_ids for item in sublist]
+    if plane_axes is not None:
+        plane_axes = [item for sublist in plane_axes for item in sublist]
+    if shape_params is not None:
+        shape_params = [item for sublist in shape_params for item in sublist]
+    if shank_ids is not None:
+        shank_ids = [item for sublist in shank_ids for item in sublist]
+    if device_channel_indices is not None:
+        device_channel_indices = [item for sublist in channel_indices for item in sublist]
 
     probeinterface_probe = Probe(ndim=ndim, si_units=unit)
     probeinterface_probe.set_contacts(
         positions=positions, shapes=shapes, shape_params=shape_params, plane_axes=plane_axes, shank_ids=shank_ids
     )
     probeinterface_probe.set_contact_ids(contact_ids=contact_ids)
-    probeinterface_probe.set_device_channel_indices(channel_indices=channel_indices)
-
+    if device_channel_indices is not None:
+        probeinterface_probe.set_device_channel_indices(channel_indices=device_channel_indices)
     probeinterface_probe.set_planar_contour(polygon)
 
     return probeinterface_probe

src/pynwb/tests/test_probe.py

@@ -174,7 +174,7 @@ def tearDown(self):
         for path in [self.path0, self.path1, self.path2]:
             remove_test_file(path)
 
-    def test_roundtrip_from_probe_single_shank(self):
+    def test_roundtrip_nwb_from_probe_single_shank(self):
         device = Probe.from_probe(self.probe0)
         self.nwbfile0.add_device(device)
 
@@ -184,9 +184,9 @@ def test_roundtrip_from_probe_single_shank(self):
         with NWBHDF5IO(self.path0, mode="r", load_namespaces=True) as io:
             read_nwbfile = io.read()
             devices = read_nwbfile.devices
-            self.assertContainerEqual(device, read_nwbfile.devices[device.name])
+            self.assertContainerEqual(device, read_nwbfile.devices[device.name])            
 
-    def test_roundtrip_from_probe_multi_shank(self):
+    def test_roundtrip_nwb_from_probe_multi_shank(self):
         device = Probe.from_probe(self.probe1)
         self.nwbfile1.add_device(device)
 
@@ -198,7 +198,7 @@ def test_roundtrip_from_probe_multi_shank(self):
             devices = read_nwbfile.devices
             self.assertContainerEqual(device, read_nwbfile.devices[device.name])
 
-    def test_roundtrip_from_probegroup(self):
+    def test_roundtrip_nwb_from_probegroup(self):
         devices = Probe.from_probegroup(self.probegroup)
         for device in devices:
             self.nwbfile2.add_device(device)
@@ -211,6 +211,16 @@ def test_roundtrip_from_probegroup(self):
             for device in devices:
                 self.assertContainerEqual(device, read_nwbfile.devices[device.name])
 
+    def test_roundtrip_pi_from_probe_single_shank(self):
+        probe_arr = self.probe0.to_numpy()
+        device = Probe.from_probe(self.probe0)
+        np.testing.assert_array_equal(probe_arr, device.to_probeinterface().to_numpy())    
+
+    def test_roundtrip_pi_from_probe_multi_shank(self):
+        probe_arr = self.probe1.to_numpy()
+        device = Probe.from_probe(self.probe1)
+        np.testing.assert_array_equal(probe_arr, device.to_probeinterface().to_numpy())
+
 
 if __name__ == "__main__":
     # test = TestProbeConstructors()
@@ -223,6 +233,8 @@ def test_roundtrip_from_probegroup(self):
     test1.setUp()
     # test.test_constructor_from_probe_single_shank()
     # test.test_constructor_from_probe_multi_shank()
-    test1.test_roundtrip_from_probe_single_shank()
-    test1.test_roundtrip_from_probe_multi_shank()
-    test1.test_roundtrip_from_probegroup()
+    test1.test_roundtrip_pi_from_probe_single_shank()
+    test1.test_roundtrip_pi_from_probe_multi_shank()
+    
+    # test1.test_roundtrip_from_probe_multi_shank()
+    # test1.test_roundtrip_from_probegroup()