analyseScript/Analyser/ImagingAnalyser.py

import numpy as np
import xarray as xr
import copy


class ImageAnalyser():

    def __init__(self) -> None:
        self._image_name = {
            'atoms': 'atoms',
            'background': 'background',
            'dark': 'dark',
            'OD':'OD',
        }
        self._center = None
        self._span = None
        self._fraction = None
    
    @property
    def image_name(self):
        return self._image_name
    
    @image_name.setter
    def image_name(self, value):
        self._image_name.update(value)

    @property
    def center(self):
        return self._center
    
    @center.setter
    def center(self, value):
        self._center = value

    @property
    def span(self):
        return self._span

    @span.setter
    def span(self, value):
        self._span = value

    @property
    def fraction(self):
        return self._fraction

    @fraction.setter
    def fraction(self, value):
        self._fraction = value
    
    def get_offset_from_corner(self, dataArray, x_fraction=None, y_fraction=None, fraction=None, xAxisName='x', yAxisName='y'):
        
        if fraction is None:
            if x_fraction is None:
                x_fraction = self._fraction[0]
        
            if y_fraction is None:
                y_fraction = self._fraction[1]
        else:
            x_fraction = fraction[0]
            y_fraction = fraction[1]

        x_number = dataArray[xAxisName].shape[0]
        y_number = dataArray[yAxisName].shape[0]

        mean = dataArray.isel(x=slice(0, int(x_number * x_fraction)), y=slice(0 , int(y_number * y_fraction))).mean(dim=[xAxisName, yAxisName])
        mean += dataArray.isel(x=slice(0, int(x_number * x_fraction)), y=slice(int(y_number - y_number * y_fraction) , int(y_number))).mean(dim=[xAxisName, yAxisName])
        mean += dataArray.isel(x=slice(int(x_number - x_number * x_fraction) , int(x_number)), y=slice(0 , int(y_number * y_fraction))).mean(dim=[xAxisName, yAxisName])
        mean += dataArray.isel(x=slice(int(x_number - x_number * x_fraction) , int(x_number)), y=slice(int(y_number - y_number * y_fraction) , int(y_number))).mean(dim=[xAxisName, yAxisName])
    
        return mean / 4
    
    def substract_offset(self, dataArray, **kwargs):
        res = dataArray - self.get_offset_from_corner(dataArray, **kwargs)
        res.attrs = copy.copy(dataArray.attrs)
        return res
    
    def crop_image(self, dataSet, center=None, span=None):

        if center is None:
            center = self._center
        if span is None:
            span = self._span

        x_start = int(center[0] - span[0] / 2)
        x_end = int(center[0] + span[0] / 2)
        y_end = int(center[1] + span[1] / 2)
        y_start = int(center[1] - span[1] / 2)
        
        dataSet.attrs['x_start'] = x_start
        dataSet.attrs['x_end'] = x_end
        dataSet.attrs['y_end'] = y_end
        dataSet.attrs['y_start'] = y_start
        dataSet.attrs['x_center'] = center[0]
        dataSet.attrs['y_center'] = center[1]
        dataSet.attrs['x_span'] = span[0]
        dataSet.attrs['y_span'] = span[1]
        
        if isinstance(dataSet, type(xr.Dataset())):
            for key in list(dataSet.data_vars):
                dataSet[key].attrs['x_start'] = x_start
                dataSet[key].attrs['x_end'] = x_end
                dataSet[key].attrs['y_end'] = y_end
                dataSet[key].attrs['y_start'] = y_start
                dataSet[key].attrs['x_center'] = center[0]
                dataSet[key].attrs['y_center'] = center[1]
                dataSet[key].attrs['x_span'] = span[0]
                dataSet[key].attrs['y_span'] = span[1]

        return dataSet.isel(x=slice(x_start, x_end), y=slice(y_start, y_end))
    
    def get_OD(self, imageAtom, imageBackground, imageDrak):

        numerator = np.atleast_1d(imageBackground - imageDrak)
        denominator = np.atleast_1d(imageAtom - imageDrak)
        
        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):
        return dataSet.sum(dim=['x', 'y'], **kwargs)

    def get_absorption_images(self, dataSet, dask='allowed', keep_attrs=True, **kwargs):

        kwargs.update(
            {
                'dask': dask,
                'keep_attrs': keep_attrs,
            }
        )

        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)

        return dataSet

    def remove_background(self, dataSet, dask='allowed', keep_attrs=True, **kwargs):

        kwargs.update(
            {
                'dask': dask,
                'keep_attrs': keep_attrs,
            }
        )

        xr.apply_ufunc(self.get_OD, dataSet[self._image_name['atoms']], dataSet[self._image_name['background']], dataSet[self._image_name['dark']], **kwargs)