From 8be62b0dc80f65f65e88abc913c1ddc38d772a9b Mon Sep 17 00:00:00 2001
From: Maarten <54527765+MaartenBransen@users.noreply.github.com>
Date: Fri, 8 Sep 2023 15:30:05 +0200
Subject: [PATCH] more efficient xy slicing

---
 scm_confocal/visitech.py | 41 +++++++++++++++++++++-------------------
 1 file changed, 22 insertions(+), 19 deletions(-)

diff --git a/scm_confocal/visitech.py b/scm_confocal/visitech.py
index db5b35c..65239ac 100644
--- a/scm_confocal/visitech.py
+++ b/scm_confocal/visitech.py
@@ -816,7 +816,8 @@ def __repr__(self):
         return "<scm_confocal.visitech_series('{}',{},{},{})>".format(
             self.filename,self.zsize,self.zstep,self.backsteps)
     
-    def load_data(self,indices=slice(None,None,None),dtype=np.uint16):
+    def load_data(self,indices=slice(None,None,None),dtype=np.uint16,
+                  xslice=None,yslice=None):
         """
         load images from datafile into 3D numpy array
         
@@ -833,7 +834,14 @@ def load_data(self,indices=slice(None,None,None),dtype=np.uint16):
         if type(indices) == slice:
             indices = range(self.nf)[indices]
 
-        data = np.array(self.datafile[indices])
+        if xslice is None and yslice is None:
+            data = np.array(self.datafile[indices])
+        elif xslice is None:
+            data = np.array([self.datafile[i][yslice] for i in indices])
+        elif yslice is None:
+            data = np.array([self.datafile[i][:,xslice] for i in indices])
+        else:
+            data = np.array([self.datafile[i][yslice,xslice] for i in indices])
 
         if not data.dtype == dtype:
             print('rescaling data to type',dtype)
@@ -957,16 +965,6 @@ def load_stack(self,dim_range={},dtype=np.uint16,remove_backsteps=True,
             elif val==slice(None):
                 dim_range.pop(key)
 
-        #warn for inefficient x and y trimming
-        if 'x-axis' in dim_range or 'y-axis' in dim_range:
-            print("[WARNING] scm_confocal.visitech_faststack.load_stack: "+
-                  "Loading only part of the data along dimensions 'x-axis' "+
-                  "and/or 'y-axis' not implemented. Data will be loaded fully"+
-                  " into memory before discarding values outside of the "+
-                  "slice range specified for the x-axis and/or y-axis. "+
-                  "Other axes for which a range is specified will still "+
-                  "be treated normally, avoiding unneccesary memory use.")
-
         #remove values outside of dim_range from indices
         if 'time' in dim_range:
             indices = indices[dim_range['time']]
@@ -983,18 +981,23 @@ def load_stack(self,dim_range={},dtype=np.uint16,remove_backsteps=True,
         # self.get_timestamps(load_stack_indices=True)
         self._stack_indices = indices
 
+        #add None arguments for x and y range if not slicing those
+        if not 'y-axis' in dim_range:
+            dim_range['y-axis'] = None
+        if not 'x-axis' in dim_range:
+            dim_range['x-axis'] = None
+
         #load and reshape data
-        stack = self.load_data(indices=indices.ravel(),dtype=dtype)
+        stack = self.load_data(
+            indices=indices.ravel(),
+            dtype=dtype,
+            xslice=dim_range['x-axis'],
+            yslice=dim_range['y-axis'],
+        )
         shape = (indices.shape[0],indices.shape[1],stack.shape[1],
                  stack.shape[2])
         stack = stack.reshape(shape)
 
-        #trim x and y axis
-        if 'y-axis' in dim_range:
-            stack = stack[:,:,dim_range['y-axis']]
-        if 'x-axis' in dim_range:
-            stack = stack[:,:,:,dim_range['x-axis']]
-
         return stack
 
     def yield_stack(self,dim_range={},dtype=np.uint16,remove_backsteps=True,