def test_probegroup(): probegroup = ProbeGroup() nchan = 0 for i in range(3): probe = generate_dummy_probe() probe.move([i * 100, i * 80]) n = probe.get_contact_count() probe.set_device_channel_indices(np.arange(n)[::-1] + nchan) shank_ids = np.ones(n) shank_ids[:n // 2] *= i * 2 shank_ids[n // 2:] *= i * 2 + 1 probe.set_shank_ids(shank_ids) probegroup.add_probe(probe) nchan += n indices = probegroup.get_global_device_channel_indices() ids = probegroup.get_global_contact_ids() df = probegroup.to_dataframe() #~ print(df['global_contact_ids']) arr = probegroup.to_numpy(complete=False) other = ProbeGroup.from_numpy(arr) arr = probegroup.to_numpy(complete=True) other = ProbeGroup.from_numpy(arr) d = probegroup.to_dict() other = ProbeGroup.from_dict(d) #~ from probeinterface.plotting import plot_probe_group, plot_probe #~ import matplotlib.pyplot as plt #~ plot_probe_group(probegroup) #~ plot_probe_group(other) #~ plt.show() # checking automatic generation of ids with new dummy probes probegroup.probes = [] for i in range(3): probegroup.add_probe(generate_dummy_probe()) probegroup.auto_generate_contact_ids() probegroup.auto_generate_probe_ids() for p in probegroup.probes: assert p.contact_ids is not None assert 'probe_id' in p.annotations
def set_probes(self, probe_or_probegroup, group_mode='by_probe', in_place=False): """ Attach a Probe to a recording. For this Probe.device_channel_indices is used to link contacts to recording channels. If some contacts of the Probe are not connected (device_channel_indices=-1) then the recording is "sliced" and only connected channel are kept. The probe order is not kept. Channel ids are re-ordered to match the channel_ids of the recording. Parameters ---------- probe_or_probegroup: Probe, list of Probe, or ProbeGroup The probe(s) to be attached to the recording group_mode: str 'by_probe' or 'by_shank'. Adds grouping property to the recording based on the probes ('by_probe') or shanks ('by_shanks') in_place: bool False by default. Useful internally when extractor do self.set_probegroup(probe) Returns ------- sub_recording: BaseRecording A view of the recording (ChannelSliceRecording or clone or itself) """ from spikeinterface import ChannelSliceRecording assert group_mode in ( 'by_probe', 'by_shank'), "'group_mode' can be 'by_probe' or 'by_shank'" # handle several input possibilities if isinstance(probe_or_probegroup, Probe): probegroup = ProbeGroup() probegroup.add_probe(probe_or_probegroup) elif isinstance(probe_or_probegroup, ProbeGroup): probegroup = probe_or_probegroup elif isinstance(probe_or_probegroup, list): assert all([isinstance(e, Probe) for e in probe_or_probegroup]) probegroup = ProbeGroup() for probe in probe_or_probegroup: probegroup.add_probe(probe) else: raise ValueError('must give Probe or ProbeGroup or list of Probe') # handle not connected channels assert all(probe.device_channel_indices is not None for probe in probegroup.probes), \ 'Probe must have device_channel_indices' # this is a vector with complex fileds (dataframe like) that handle all contact attr arr = probegroup.to_numpy(complete=True) # keep only connected contact ( != -1) keep = arr['device_channel_indices'] >= 0 if np.any(~keep): warn( 'The given probes have unconnected contacts: they are removed') arr = arr[keep] inds = arr['device_channel_indices'] order = np.argsort(inds) inds = inds[order] # check if np.max(inds) >= self.get_num_channels(): raise ValueError( 'The given Probe have "device_channel_indices" that do not match channel count' ) new_channel_ids = self.get_channel_ids()[inds] arr = arr[order] # create recording : channel slice or clone or self if in_place: if not np.array_equal(new_channel_ids, self.get_channel_ids()): raise Exception( 'set_proce(inplace=True) must have all channel indices') sub_recording = self else: if np.array_equal(new_channel_ids, self.get_channel_ids()): sub_recording = self.clone() else: sub_recording = ChannelSliceRecording(self, new_channel_ids) # create a vector that handle all conatcts in property sub_recording.set_property('contact_vector', arr, ids=None) # planar_contour is saved in annotations for probe_index, probe in enumerate(probegroup.probes): contour = probe.probe_planar_contour if contour is not None: sub_recording.set_annotation( f'probe_{probe_index}_planar_contour', contour, overwrite=True) # duplicate positions to "locations" property ndim = probegroup.ndim locations = np.zeros((arr.size, ndim), dtype='float64') for i, dim in enumerate(['x', 'y', 'z'][:ndim]): locations[:, i] = arr[dim] sub_recording.set_property('location', locations, ids=None) # handle groups groups = np.zeros(arr.size, dtype='int64') if group_mode == 'by_probe': for group, probe_index in enumerate(np.unique(arr['probe_index'])): mask = arr['probe_index'] == probe_index groups[mask] = group 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(arr[['probe_index', 'shank_ids']])): mask = (arr['probe_index'] == a['probe_index']) & ( arr['shank_ids'] == a['shank_ids']) groups[mask] = group sub_recording.set_property('group', groups, ids=None) return sub_recording