optimize the fringes removal function

Analyser/FringeRemoval.py

@@ -195,6 +195,10 @@ class FringeRemoval():
     
     @referenceImages.setter
     def referenceImages(self, value):
+        if value is None:
+            self._referenceImages = None
+            return 
+        
         if self.reshape:
             value = self.reshape_data(value)
         elif isinstance(value, type(xr.DataArray())):
@@ -257,6 +261,10 @@ class FringeRemoval():
             raise DataSizeException
         
         self._mask = value
+        
+        self._center = None
+        self._span = None
+        
         self.k = np.where(self._mask.flatten() == 1)[0]
     
     def _auto_mask(self):
@@ -267,7 +275,7 @@ class FringeRemoval():
             y_end = int(self._center[1] + self._span[1] / 2)
             y_start = int(self._center[1] - self._span[1] / 2)
             
-        mask[y_start:y_end, x_start:x_end] = 0
+            mask[y_start:y_end, x_start:x_end] = 0
         
         return mask
     

Analyser/ImagingAnalyser.py

@@ -2,6 +2,10 @@ import numpy as np
 import xarray as xr
 import copy
 
+from DataContainer.ReadData import read_hdf5_file
+from Analyser.FringeRemoval import FringeRemoval
+from ToolFunction.ToolFunction import get_scanAxis
+
 
 class ImageAnalyser():
     """A class for operate with and analyse images
@@ -15,11 +19,15 @@ class ImageAnalyser():
             'background': 'background',
             'dark': 'dark',
             'OD':'OD',
+            'optimumBackground':'optimumBackground'
         }
         self._center = None
         self._span = None
         self._fraction = None
 
+        self._fringeRemoval = FringeRemoval()
+        self.fringeRemovalReferenceImages = None
+        
     @property
     def image_name(self):
         """The getter of the names of three standard images for absorption images
@@ -150,7 +158,7 @@ class ImageAnalyser():
         res.attrs = copy.copy(dataArray.attrs)
         return res
     
-    def crop_image(self, dataSet, center=None, span=None):
+    def crop_image(self, dataSet, center=None, span=None, fringeRemoval=False):
         """Crop the image according to the region of interest (ROI).
 
         :param dataSet: The images
@@ -193,6 +201,13 @@ class ImageAnalyser():
                 dataSet[key].attrs['x_span'] = span[0]
                 dataSet[key].attrs['y_span'] = span[1]
         
+        if fringeRemoval:
+            scanAxis = list(get_scanAxis(self.fringeRemovalReferenceImages))
+            if not scanAxis[1] is None:
+                self._fringeRemoval.referenceImages = self.fringeRemovalReferenceImages.isel(x=slice(x_start, x_end), y=slice(y_start, y_end)).stack(_imgIdx=scanAxis)
+            else:
+                self._fringeRemoval.referenceImages = self.fringeRemovalReferenceImages.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):
@@ -221,6 +236,32 @@ class ImageAnalyser():
         else:
             return imageOD
     
+    def get_OD_no_dark(self, imageAtom, imageBackground):
+        """Calculate the OD image for absorption imaging without dark images.
+
+        :param imageAtom: The image with atoms
+        :type imageAtom: numpy array
+        :param imageBackground: The image without atoms
+        :type imageBackground: numpy array
+        :param imageDrak: The image without light
+        :type imageDrak: numpy array
+        :return: The OD images
+        :rtype: numpy array
+        """
+
+        numerator = np.atleast_1d(imageBackground)
+        denominator = np.atleast_1d(imageAtom)
+        
+        numerator[numerator == 0] = 1
+        denominator[denominator == 0] = 1
+        imageOD = np.abs(np.divide(denominator, numerator)) 
+        imageOD= -np.log(imageOD)
+
+        if len(imageOD) == 1:
+            return imageOD[0]
+        else:
+            return imageOD
+
     def get_Ncount(self, dataSet, dim=['x', 'y'], **kwargs):
         """Sum all the value in the image to give the Ncount.
 
@@ -233,7 +274,7 @@ class ImageAnalyser():
         """
         return dataSet.sum(dim=['x', 'y'], **kwargs)
 
-    def get_absorption_images(self, dataSet, dask='allowed', keep_attrs=True, **kwargs):
+    def get_absorption_images(self, dataSet, fringeRemoval=False, dask='allowed', keep_attrs=True, **kwargs):
         """Calculate the OD images for absorption imaging.
 
         :param dataSet: The data from absorption imaging.
@@ -253,11 +294,31 @@ class ImageAnalyser():
             }
         )
 
-        dataSet = dataSet.assign(
+        if fringeRemoval:
-            {
+            
-                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)
+            scanAxis = list(get_scanAxis(dataSet))
-            }
+            if not scanAxis[1] is None:
-        )
+                OptimumRef = self._fringeRemoval.fringe_removal(dataSet[self._image_name['atoms']].stack(_imgIdx=scanAxis))
+            else:
+                OptimumRef = self._fringeRemoval.fringe_removal(dataSet[self._image_name['atoms']])
+            
+            dataSet = dataSet.assign(
+                {
+                    self._image_name['optimumBackground']: OptimumRef
+                }
+            )
+            
+            dataSet = dataSet.assign(
+                {
+                    self._image_name['OD']: xr.apply_ufunc(self.get_OD_no_dark, dataSet[self._image_name['atoms']], dataSet[self._image_name['optimumBackground']], **kwargs)
+                }
+            )
+        else:
+            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)
+                }
+            )
         
         # dataSet[self._image_name['OD']].attrs.update(dataSet.attrs)
 
@@ -283,5 +344,67 @@ class ImageAnalyser():
 
         xr.apply_ufunc(self.get_OD, dataSet[self._image_name['atoms']], dataSet[self._image_name['background']], dataSet[self._image_name['dark']], **kwargs)
     
+    @property
+    def fringeRemovalCenter(self):
+        """The getter of the center of region of insterest (ROI)
 
+        :return: The center of region of insterest (ROI)
+        :rtype: tuple
+        """
+        return self._fringeRemoval.center
     
+    @fringeRemovalCenter.setter
+    def fringeRemovalCenter(self, value):
+        """The setter of the center of region of insterest (ROI)
+
+        :param value: The center of region of insterest (ROI)
+        :type value: tuple
+        """
+        self._fringeRemoval.center = value
+
+    @property
+    def fringeRemovalSpan(self):
+        """The getter of the span of region of insterest (ROI) 
+
+        :return: The span of region of insterest (ROI) 
+        :rtype: tuple
+        """
+        return self._fringeRemoval.span
+
+    @fringeRemovalSpan.setter
+    def fringeRemovalSpan(self, value):
+        """The setter of the span of region of insterest (ROI) 
+
+        :param value: The span of region of insterest (ROI) 
+        :type value: tuple
+        """
+        self._fringeRemoval.span = value
+    
+    # @property
+    # def fringeRemovalReferenceImages(self):
+    #     return self._fringeRemoval.referenceImages
+    
+    # @fringeRemovalReferenceImages.setter
+    # def fringeRemovalReferenceImages(self, value):
+    #     self._fringeRemoval.referenceImages = value
+        
+    def load_fringe_removal_background_from_hdf5(self, img_dir, SequenceName, data, shotNum, group, crop=False, load=False, **kwargs):
+        folderPath = img_dir + SequenceName + "/" + data
+        filePath = folderPath + "/" + shotNum + "/*.h5"
+        
+        dataSet = read_hdf5_file(filePath, group, **kwargs)
+        scanAxis = dataSet.scanAxis
+        
+        dataSet = dataSet[self._image_name['background']]
+        dataSet.attrs['scanAxis'] = scanAxis
+        
+        if crop:
+            dataSet = self.crop_image(dataSet)
+        
+        if load:
+            self.fringeRemovalReferenceImages = dataSet.load()
+        else:
+            self.fringeRemovalReferenceImages = dataSet
+    
+    def load_fringe_removal_background_from_database():
+        pass