From 4841c09eb78ca478e6bc7d0d143d585013a21643 Mon Sep 17 00:00:00 2001
From: Ruben Verweij <verweij@physics.leidenuniv.nl>
Date: Mon, 6 Mar 2017 13:38:49 +0100
Subject: [PATCH] Correctly parse experiment data for NDAqcuisition and add
 get_timesteps function to reader

---
 nd2reader/raw_metadata.py | 18 ++++++++++--------
 nd2reader/reader.py       | 18 ++++++++++++++++++
 2 files changed, 28 insertions(+), 8 deletions(-)

diff --git a/nd2reader/raw_metadata.py b/nd2reader/raw_metadata.py
index 1ce26ad..e317c8a 100644
--- a/nd2reader/raw_metadata.py
+++ b/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')]
diff --git a/nd2reader/reader.py b/nd2reader/reader.py
index cf99aed..0952054 100644
--- a/nd2reader/reader.py
+++ b/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']]