Correctly parse experiment data for NDAqcuisition and add get_timesteps function to reader

nd2reader/raw_metadata.py

@@ -279,14 +279,14 @@ class RawMetadata(object):
         :param loop_data:
         :return:
         """
-        if six.b('uiPeriodCount') not in loop_data or loop_data[six.b('uiPeriodCount')] == 0:
-            return []
+        loops = [loop_data]
+        if six.b('uiPeriodCount') in loop_data and loop_data[six.b('uiPeriodCount')] > 0:
+            # special ND experiment
+            if six.b('pPeriod') not in loop_data:
+                return []
 
-        if six.b('pPeriod') not in loop_data:
-            return []
-
-        # take the first dictionary element, it contains all loop data
-        loops = loop_data[six.b('pPeriod')][list(loop_data[six.b('pPeriod')].keys())[0]]
+            # take the first dictionary element, it contains all loop data
+            loops = loop_data[six.b('pPeriod')][list(loop_data[six.b('pPeriod')].keys())[0]]
 
         # take into account the absolute time in ms
         time_offset = 0
@@ -298,7 +298,9 @@ class RawMetadata(object):
             duration = loop[six.b('dDuration')]
 
             # uiLoopType == 6 is a stimulation loop
-            is_stimulation = loop[six.b('uiLoopType')] == 6
+            is_stimulation = False
+            if six.b('uiLoopType') in loop:
+                is_stimulation = loop[six.b('uiLoopType')] == 6
 
             # sampling interval in ms
             interval = loop[six.b('dAvgPeriodDiff')]

nd2reader/reader.py

@@ -99,3 +99,21 @@ class ND2Reader(FramesSequenceND):
 
         # provide the default
         self.iter_axes = 't'
+
+    def get_timesteps(self):
+        """
+        Get the timesteps of the experiment
+        :return:
+        """
+        timesteps = np.array([])
+        current_time = 0.0
+        for loop in self.metadata['experiment']['loops']:
+            if loop['stimulation']:
+                continue
+
+            timesteps = np.concatenate(
+                (timesteps, np.arange(current_time, current_time + loop['duration'], loop['sampling_interval'])))
+            current_time += loop['duration']
+
+        # if experiment did not finish, number of timesteps is wrong. Take correct amount of leading timesteps.
+        return timesteps[:self.metadata['num_frames']]