[neuralynx] load only spikes/events within time range of segments and adjust tests

Author: sprenger

neo/rawio/neuralynxrawio/neuralynxrawio.py

@@ -597,7 +597,13 @@ def _spike_count(self, block_index, seg_index, unit_index):
         data = self._spike_memmap[chan_uid]
         ts = data['timestamp']
 
-        ts0, ts1 = self._timestamp_limits[seg_index]
+        ts0 = self.segment_t_start(block_index, seg_index)
+        ts1 = self.segment_t_stop(block_index, seg_index)
+
+        # rescale to integer sampling of time
+        ts0 = int((ts0 + self.global_t_start) * 1e6)
+        ts1 = int((ts1 + self.global_t_start) * 1e6)
+
 
         # only count spikes inside the timestamp limits, inclusive, and for the specified unit
         keep = (ts >= ts0) & (ts <= ts1) & (unit_id == data['unit_id'])
@@ -612,11 +618,15 @@ def _get_spike_timestamps(self, block_index, seg_index, unit_index, t_start, t_s
         data = self._spike_memmap[chan_uid]
         ts = data['timestamp']
 
-        ts0, ts1 = self._timestamp_limits[seg_index]
-        if t_start is not None:
-            ts0 = int((t_start + self.global_t_start) * 1e6)
-        if t_start is not None:
-            ts1 = int((t_stop + self.global_t_start) * 1e6)
+        ts0, ts1 = t_start, t_stop
+        if ts0 is None:
+            ts0 = self.segment_t_start(block_index, seg_index)
+        if ts1 is None:
+            ts1 = self.segment_t_stop(block_index, seg_index)
+
+        # rescale to integer sampling of time
+        ts0 = int((ts0 + self.global_t_start) * 1e6)
+        ts1 = int((ts1 + self.global_t_start) * 1e6)
 
         keep = (ts >= ts0) & (ts <= ts1) & (unit_id == data['unit_id'])
         timestamps = ts[keep]
@@ -628,17 +638,20 @@ def _rescale_spike_timestamp(self, spike_timestamps, dtype):
         spike_times -= self.global_t_start
         return spike_times
 
-    def _get_spike_raw_waveforms(self, block_index, seg_index, unit_index,
-                                 t_start, t_stop):
+    def _get_spike_raw_waveforms(self, block_index, seg_index, unit_index, t_start, t_stop):
         chan_uid, unit_id = self.internal_unit_ids[unit_index]
         data = self._spike_memmap[chan_uid]
         ts = data['timestamp']
 
-        ts0, ts1 = self._timestamp_limits[seg_index]
-        if t_start is not None:
-            ts0 = int((t_start + self.global_t_start) * 1e6)
-        if t_start is not None:
-            ts1 = int((t_stop + self.global_t_start) * 1e6)
+        ts0, ts1 = t_start, t_stop
+        if ts0 is None:
+            ts0 = self.segment_t_start(block_index, seg_index)
+        if ts1 is None:
+            ts1 = self.segment_t_stop(block_index, seg_index)
+
+        # rescale to integer sampling of time
+        ts0 = int((ts0 + self.global_t_start) * 1e6)
+        ts1 = int((ts1 + self.global_t_start) * 1e6)
 
         keep = (ts >= ts0) & (ts <= ts1) & (unit_id == data['unit_id'])
 
@@ -656,7 +669,14 @@ def _event_count(self, block_index, seg_index, event_channel_index):
         event_id, ttl_input = self.internal_event_ids[event_channel_index]
         chan_id = self.header['event_channels'][event_channel_index]['id']
         data = self._nev_memmap[chan_id]
-        ts0, ts1 = self._timestamp_limits[seg_index]
+
+        ts0 = self.segment_t_start(block_index, seg_index)
+        ts1 = self.segment_t_stop(block_index, seg_index)
+
+        # rescale to integer sampling of time
+        ts0 = int((ts0 + self.global_t_start) * 1e6)
+        ts1 = int((ts1 + self.global_t_start) * 1e6)
+
         ts = data['timestamp']
         keep = (ts >= ts0) & (ts <= ts1) & (data['event_id'] == event_id) & \
                (data['ttl_input'] == ttl_input)
@@ -667,12 +687,16 @@ def _get_event_timestamps(self, block_index, seg_index, event_channel_index, t_s
         event_id, ttl_input = self.internal_event_ids[event_channel_index]
         chan_id = self.header['event_channels'][event_channel_index]['id']
         data = self._nev_memmap[chan_id]
-        ts0, ts1 = self._timestamp_limits[seg_index]
 
-        if t_start is not None:
-            ts0 = int((t_start + self.global_t_start) * 1e6)
-        if t_start is not None:
-            ts1 = int((t_stop + self.global_t_start) * 1e6)
+        ts0, ts1 = t_start, t_stop
+        if ts0 is None:
+            ts0 = self.segment_t_start(block_index, seg_index)
+        if ts1 is None:
+            ts1 = self.segment_t_stop(block_index, seg_index)
+
+        # rescale to integer sampling of time
+        ts0 = int((ts0 + self.global_t_start) * 1e6)
+        ts1 = int((ts1 + self.global_t_start) * 1e6)
 
         ts = data['timestamp']
         keep = (ts >= ts0) & (ts <= ts1) & (data['event_id'] == event_id) & \

neo/test/iotest/test_neuralynxio.py

@@ -331,12 +331,15 @@ def test_incomplete_block_handling_v632(self):
         self.assertEqual(len(block.segments), n_gaps + 1)
         # self.assertEqual(len(block.channel_indexes[0].analogsignals), n_gaps + 1)
 
-        for t, gt in zip(nio._ncs_seg_timestamp_limits.t_start, [8408.806811, 8427.832053,
-                                                                 8487.768561]):
-            self.assertEqual(np.round(t, 4), np.round(gt, 4))
-        for t, gt in zip(nio._ncs_seg_timestamp_limits.t_stop, [8427.831990, 8487.768498,
-                                                                8515.816549]):
-            self.assertEqual(np.round(t, 4), np.round(gt, 4))
+        expected_segment_starts = [8124.582909, 8427.832053, 8487.768561]
+        expected_segment_stops = [8427.831990, 8487.768498, 10794.133994]
+        for seg_idx in range(n_gaps+1):
+
+            t_start = nio.segment_t_start(0, seg_idx) + nio.global_t_start
+            t_stop = nio.segment_t_stop(0, seg_idx) + nio.global_t_start
+
+            self.assertEqual(np.round(t_start, 4), np.round(expected_segment_starts[seg_idx], 4))
+            self.assertEqual(np.round(t_stop, 4), np.round(expected_segment_stops[seg_idx], 4))
 
 
 class TestGaps(CommonNeuralynxIOTest, unittest.TestCase):