support rel_z in nwb rec extractor
@@ -142,11 +142,16 @@ class NwbRecordingExtractor(BaseRecording):
properties = dict()
for es_ind, (channel_id, electrode_table_index) in enumerate(zip(channel_ids, self._es.electrodes.data)):
if 'rel_x' in self._nwbfile.electrodes:
+ ndim = 2 # assume 2 dimensions
+ if 'rel_z' in self._nwbfile.electrodes: ndim = 3 # if we have rel_z, it is 3 dimensions
+
if 'location' not in properties:
- properties['location'] = np.zeros((self.get_num_channels(), 2), dtype=float)
+ properties['location'] = np.zeros((self.get_num_channels(), ndim), dtype=float)
properties['location'][es_ind, 0] = self._nwbfile.electrodes['rel_x'][electrode_table_index]
if 'rel_y' in self._nwbfile.electrodes:
properties['location'][es_ind, 1] = self._nwbfile.electrodes['rel_y'][electrode_table_index]
+ if 'rel_z' in self._nwbfile.electrodes:
+ properties['location'][es_ind, 2] = self._nwbfile.electrodes['rel_z'][electrode_table_index]
for col in self._nwbfile.electrodes.colnames:
if isinstance(self._nwbfile.electrodes[col][electrode_table_index], ElectrodeGroup):
@@ -167,7 +172,7 @@ class NwbRecordingExtractor(BaseRecording):
if 'offset' not in properties:
properties['offset'] = np.zeros(self.get_num_channels(), dtype=type(offset))
properties['offset'][es_ind] = offset
- elif col in ['x', 'y', 'z', 'rel_x', 'rel_y']:
+ elif col in ['x', 'y', 'z', 'rel_x', 'rel_y', 'rel_z']:
continue
else:
val = self._nwbfile.electrodes[col][electrode_table_index]
