From faa495ecdb75f8d0722e98ba523b0e0141d0fd73 Mon Sep 17 00:00:00 2001
From: Ruben Verweij <verweij@physics.leidenuniv.nl>
Date: Mon, 27 May 2019 17:41:09 +0200
Subject: [PATCH] Use _register_get_frame functions

---
 nd2reader/parser.py |  5 +++++
 nd2reader/reader.py | 41 ++++++++++++++++++++++++-----------------
 2 files changed, 29 insertions(+), 17 deletions(-)

diff --git a/nd2reader/parser.py b/nd2reader/parser.py
index 4f17992..de7e7c3 100644
--- a/nd2reader/parser.py
+++ b/nd2reader/parser.py
@@ -92,6 +92,11 @@ class Parser(object):
             Frame: the requested image
 
         """
+        frame_number = 0 if frame_number is None else frame_number
+        field_of_view = 0 if field_of_view is None else field_of_view
+        channel = 0 if channel is None else channel
+        z_level = 0 if z_level is None else z_level
+
         image_group_number = self._calculate_image_group_number(frame_number, field_of_view, z_level)
         try:
             timestamp, raw_image_data = self._get_raw_image_data(image_group_number, channel,
diff --git a/nd2reader/reader.py b/nd2reader/reader.py
index 52078d0..fe263d7 100644
--- a/nd2reader/reader.py
+++ b/nd2reader/reader.py
@@ -52,24 +52,22 @@ class ND2Reader(FramesSequenceND):
         except KeyError:
             return 0
 
-    def get_frame_2D(self, c=0, t=0, z=0, x=0, y=0, v=0):
-        """Gets a given frame using the parser
+    def get_frame_vczyx(self, v=None, c=None, t=None, z=None, x=None, y=None):
+        x = self.metadata["width"] if x <= 0 else x
+        y = self.metadata["height"] if y <= 0 else y
 
-        Args:
-            x: The x-index (pims expects this)
-            y: The y-index (pims expects this)
-            c: The color channel number
-            t: The frame number
-            z: The z stack number
-            v: The field of view index
+        result = []
+        for v in self._get_possible_coords('v', v):
+            for c in self._get_possible_coords('c', c):
+                for z in self._get_possible_coords('z', z):
+                    result.append(self._parser.get_image_by_attributes(t, v, c, z, y, x))
 
-        Returns:
-            numpy.ndarray: The requested frame
+        return np.squeeze(np.array(result, dtype=self._dtype))
 
-        """
-        x = self.metadata["width"] if x <= 0 else x
-        y = self.metadata["height"] if y <= 0 else y
-        return self._parser.get_image_by_attributes(t, v, c, z, y, x)
+    def _get_possible_coords(self, dim, default):
+        if dim in self.sizes:
+            return [default] if default is not None else range(self.sizes[dim])
+        return [None]
 
     @property
     def parser(self):
@@ -158,6 +156,17 @@ class ND2Reader(FramesSequenceND):
         # provide the default
         self.iter_axes = self._guess_default_iter_axis()
 
+        self._register_get_frame(self.get_frame_vczyx, 'vczyx')
+        self._register_get_frame(self.get_frame_vczyx, 'vzyx')
+        self._register_get_frame(self.get_frame_vczyx, 'vcyx')
+        self._register_get_frame(self.get_frame_vczyx, 'vyx')
+
+        self._register_get_frame(self.get_frame_vczyx, 'czyx')
+        self._register_get_frame(self.get_frame_vczyx, 'cyx')
+
+        self._register_get_frame(self.get_frame_vczyx, 'zyx')
+        self._register_get_frame(self.get_frame_vczyx, 'yx')
+
     def _init_axis_if_exists(self, axis, size, min_size=1):
         if size >= min_size:
             self._init_axis(axis, size)
@@ -196,5 +205,3 @@ class ND2Reader(FramesSequenceND):
         self._timesteps = np.array(list(self._parser._raw_metadata.acquisition_times), dtype=np.float) * 1000.0
 
         return self._timesteps
-
-