wip slice with contact_annoations

Author: samuelgarcia

probeinterface/probe.py

@@ -820,6 +820,13 @@ def get_slice(self, selection):
             if isinstance(v, np.ndarray):
                 assert v.shape[0] == n
                 d[k] = v[selection].copy()
+            
+            if k == 'contact_annotations':
+                print(type(v))
+                for kk, vv in v.items():
+                    print(type(vv))
+                    print(vv.shape, selection.shape, selection.dtype, n)
+                    d[k][kk] = vv[selection].copy()
 
         sliced_probe = Probe.from_dict(d)
 

tests/test_probe.py

@@ -109,10 +109,13 @@ def test_probe():
     other_3d = Probe.from_dataframe(df)
     assert other_3d.ndim == 3
 
+    print('ici', probe.contact_annotations['impedance'].shape)
     # slice handling
     selection = np.arange(0,18,2)
     # print(selection.dtype.kind)
     sliced_probe = probe.get_slice(selection)
+    print('lai', probe.contact_annotations['impedance'].shape)
+    print('yep', sliced_probe.contact_annotations['impedance'].shape)
 
     #~ from probeinterface.plotting import plot_probe_group, plot_probe
     #~ import matplotlib.pyplot as plt
@@ -148,3 +151,4 @@ def test_set_shanks():
 
     test_set_shanks()
 
+