From be1be6aefb2a2376604faad403fc3bad3639f380 Mon Sep 17 00:00:00 2001
From: Ruben Verweij <ruben@lighthacking.nl>
Date: Wed, 16 Oct 2019 14:41:52 +0200
Subject: [PATCH] Update get_frame_vczyx function based on feedback in issue
 #24

---
 nd2reader/reader.py | 59 +++++++++++++++++++++++++++++++++++++--------
 1 file changed, 49 insertions(+), 10 deletions(-)

diff --git a/nd2reader/reader.py b/nd2reader/reader.py
index 1d4b3c4..45aefe6 100644
--- a/nd2reader/reader.py
+++ b/nd2reader/reader.py
@@ -1,3 +1,4 @@
+from pims import Frame
 from pims.base_frames import FramesSequenceND
 
 from nd2reader.exceptions import EmptyFileError
@@ -58,16 +59,54 @@ class ND2Reader(FramesSequenceND):
         return self.get_frame_vczyx(v=v, c=c, t=t, z=z, x=x, y=y)
 
     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
-
-        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))
-
-        return np.squeeze(np.array(result, dtype=self._dtype))
+        """Retrieve a frame based on the specified coordinates
+        Axes order is set by self.bundle_axes, x and y coordinates are ignored,
+        because we always return Frame objects.
+        """
+        # remove 'x', 'y' from bundle axes and set to width, height
+        bundle_axes = list(self.bundle_axes)
+        try:
+            bundle_axes.remove('x')
+        except ValueError:
+            pass
+        try:
+            bundle_axes.remove('y')
+        except ValueError:
+            pass
+        x = self.metadata["width"]
+        y = self.metadata["height"]
+
+        # make coords dictionary based on function input
+        coords = dict(v=v, c=c, t=t, z=z)
+
+        # Set appropriate values for None and bundle_axes coords
+        for dim in coords:
+            coords[dim] = self._get_possible_coords(dim, coords[dim])
+
+        # Initialize empty array of Frames of right shape
+        if len(bundle_axes) > 0:
+            shape = tuple((len(coords[dim]) for dim in bundle_axes))
+            results = np.empty(shape, dtype=Frame)
+        else:
+            results = np.empty((1,), dtype=Frame)
+
+        # order for the get_image_by_attributes function
+        argument_order = dict(t=0, v=1, c=2, z=3)
+
+        # Now, collect the results in the right order
+        for index, _ in np.ndenumerate(results):
+            current_coords = [0, 0, 0, 0, y, x]
+            for dim in coords:
+                if dim in bundle_axes:
+                    dim_val = coords[dim][index[bundle_axes.index(dim)]]
+                else:
+                    dim_val = coords[dim][0]
+                current_coords[argument_order[dim]] = dim_val
+
+            # Actually get the corresponding Frame
+            results[index] = Frame(self._parser.get_image_by_attributes(*current_coords), metadata=self.metadata)
+
+        return results
 
     def _get_possible_coords(self, dim, default):
         if dim in self.sizes: