add tool function

Analyser/ImagingAnalyser.py

@@ -72,7 +72,7 @@ class ImageAnalyser():
     def substract_offset(self, dataArray, **kwargs):
         return dataArray - self.get_offset_from_corner(dataArray, **kwargs)
     
-    def crop_image(self, dataset, center=None, span=None):
+    def crop_image(self, dataSet, center=None, span=None):
 
         if center is None:
             center = self._center
@@ -84,7 +84,7 @@ class ImageAnalyser():
         y_end = int(center[1] + span[1] / 2)
         y_start = int(center[1] - span[1] / 2)
 
-        return dataset.isel(x=slice(x_start, x_end), y=slice(y_start, y_end))
+        return dataSet.isel(x=slice(x_start, x_end), y=slice(y_start, y_end))
     
     def get_OD(self, imageAtom, imageBackground, imageDrak):
 
@@ -104,7 +104,7 @@ class ImageAnalyser():
     def get_Ncount(self, imageOD):
         return np.sum(imageOD)
 
-    def get_absorption_images(self, dataset, dask='allowed', keep_attrs=True, **kwargs):
+    def get_absorption_images(self, dataSet, dask='allowed', keep_attrs=True, **kwargs):
 
         kwargs.update(
             {
@@ -113,15 +113,15 @@ class ImageAnalyser():
             }
         )
 
-        dataset = dataset.assign(
+        dataSet = dataSet.assign(
             {
-                self._image_name['OD']: xr.apply_ufunc(self.get_OD, dataset[self._image_name['atoms']], dataset[self._image_name['background']], dataset[self._image_name['dark']], **kwargs)
+                self._image_name['OD']: xr.apply_ufunc(self.get_OD, dataSet[self._image_name['atoms']], dataSet[self._image_name['background']], dataSet[self._image_name['dark']], **kwargs)
             }
         )
 
-        return dataset
+        return dataSet
 
-    def remove_background(self, dataset, dask='allowed', keep_attrs=True, **kwargs):
+    def remove_background(self, dataSet, dask='allowed', keep_attrs=True, **kwargs):
 
         kwargs.update(
             {
@@ -130,7 +130,7 @@ class ImageAnalyser():
             }
         )
 
-        xr.apply_ufunc(self.get_OD, dataset[self._image_name['atoms']], dataset[self._image_name['background']], dataset[self._image_name['dark']], **kwargs)
+        xr.apply_ufunc(self.get_OD, dataSet[self._image_name['atoms']], dataSet[self._image_name['background']], dataSet[self._image_name['dark']], **kwargs)
 
 
     

ToolFunction/ToolFunction.py

@@ -1,18 +1,32 @@
-def get_mask():
-    pass
+import numpy as np
+import glob
+from datetime import date
 
 
-def remove_bad_shots():
-    pass
+def get_mask(dataArray):
+    return np.ones(dataArray.shape, dtype=bool)
 
 
-def auto_rechunk():
-    pass
+def remove_bad_shots(dataArray, **kwargs):
+    dataArray.loc[dict(kwargs)] = np.nan
 
 
-def get_h5_file_path(folderpath, maxFileNum):
-    pass
+def auto_rechunk(dataSet):
+    kwargs = {
+        key: "auto"
+        for key in dataSet.dims.keys()
+    }
+    return dataSet.chunk(**kwargs)
 
 
-def get_folder_path():
-    pass
+def get_h5_file_path(folderpath, maxFileNum=None, filename='*.h5',):
+    filepath = np.sort(glob.glob(folderpath + filename))
+    if maxFileNum is None:
+        return filepath
+    else:
+        return filepath[:maxFileNum]
+
+
+def get_date():
+    today = date.today()
+    return today.strftime("%y/%m/%d")