|
|
import numpy as npfrom scipy.linalg import lu
import xarray as xr
class InvalidDimException(Exception): "Raised when the program can not identify (index of images, x, y) axes." def __init__(self, dims): if len(dims)>3: self.message = 'The input data must have two or three axes: (index of images(alternative), x, y)' else: self.message = 'Can not identify (index of images(alternative), x, y) from ' + str(dims) super().__init__(self.message)
class DataSizeException(Exception): "Raised when the shape of the data is not correct." def __init__(self): self.message = 'The input data size does not match.' super().__init__(self.message)
class FringeRemoval(): """
FRINGEREMOVAL - Fringe removal and noise reduction from absorption images. Creates an optimal reference image for each absorption image in a set as a linear combination of reference images, with coefficients chosen to minimize the least-squares residuals between each absorption image and the optimal reference image. The coefficients are obtained by solving a linear set of equations using matrix inverse by LU decomposition.
Application of the algorithm is described in C. F. Ockeloen et al, Improved detection of small atom numbers through image processing, arXiv:1007.2136 (2010).
Adapted from a MATLAB script copy provided by Guoxian Su. Original Authors: Shannon Whitlock, Caspar Ockeloen Reference: C. F. Ockeloen, A. F. Tauschinsky, R. J. C. Spreeuw, and S. Whitlock, Improved detection of small atom numbers through image processing, arXiv:1007.2136 May 2009; """
def __init__(self) -> None: """Initialize the class
"""
self.nimgsR = 0 # The number of the reference images self.xdim = 0 # The shape of x axis self.ydim = 0 # The shape of y axis self._mask = None # The mask array to choose the region of interest for fringes removal self._center = None # Set the mask array by center and span self._span = None self.reshape=True # If it is necessary to reshape the data from (index of images(alternative), x, y) to (y, x, index of images(alternative)) self.P = None self.L = None self.U = None @property def center(self): """The getter of the center of region of insterest (ROI)
:return: The center of region of insterest (ROI) :rtype: tuple """
return self._center @center.setter def center(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._mask = None self._center = value
@property def span(self): """The getter of the span of region of insterest (ROI)
:return: The span of region of insterest (ROI) :rtype: tuple """
return self._span
@span.setter def span(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._mask = None self._span = value def reshape_data(self, data): """The function is to reshape the data to the correct shape.
In order to minimize the calculation time, the data has to have a shape of (y, x, index of images(alternative)). However, usually the input data has a shape of (index of images(alternative), x, y). It can also convert the xarray DataArray and Dataset to numpy array.
:param data: The input data. :type data: xarray, numpy array or list :raises InvalidDimException: Raised when the program can not identify (index of images, x, y) axes. :raises InvalidDimException: Raised when the shape of the data is not correct. :return: The data with correct shape :rtype: xarray, numpy array or list """
if data is None: return data
if isinstance(data, type(xr.DataArray())): dims = data.dims if len(dims)>3: raise InvalidDimException(dims) xAxis = None yAxis = None if len(dims) == 2: imageAxis = '' else: imageAxis = None for dim in dims: if (dim == 'x') or ('_x' in dim): xAxis = dim elif (dim == 'y') or ('_y' in dim): yAxis = dim else: imageAxis = dim if (xAxis is None) or (yAxis is None) or (imageAxis is None): raise InvalidDimException(dims) if len(dims) == 2: data = data.transpose(yAxis, xAxis) else: data = data.transpose(yAxis, xAxis, imageAxis) data = data.to_numpy() else: data = np.array(data) if len(data.shape) == 3: data = np.swapaxes(data, 0, 2) # data = np.swapaxes(data, 0, 1) elif len(data.shape) == 2: data = np.swapaxes(data, 0, 1) return data def _reshape_absorption_images(self, data): if data is None: return data
if isinstance(data, type(xr.DataArray())): dims = data.dims if len(dims)>3: raise InvalidDimException(dims) xAxis = None yAxis = None if len(dims) == 2: imageAxis = '' else: imageAxis = None for dim in dims: if (dim == 'x') or ('_x' in dim): xAxis = dim elif (dim == 'y') or ('_y' in dim): yAxis = dim else: imageAxis = dim if (xAxis is None) or (yAxis is None) or (imageAxis is None): raise InvalidDimException(dims) if len(dims) == 2: data = data.transpose(yAxis, xAxis) else: data = data.transpose(yAxis, xAxis, imageAxis) self.nimgs = len(data[imageAxis]) data = data.stack(axis=[yAxis, xAxis]) else: data = np.array(data) if len(data.shape) == 3: data = np.swapaxes(data, 0, 2) # data = np.swapaxes(data, 0, 1) elif len(data.shape) == 2: data = np.swapaxes(data, 0, 1) return data @property def referenceImages(self): res = self._referenceImages.reshape(self.ydim, self.xdim, self.nimgsR) res = np.swapaxes(res, 0, 2) return res @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())): value = value.to_numpy() self.nimgsR = value.shape[2] self.xdim = value.shape[1] self.ydim = value.shape[0] self._referenceImages = (value.reshape(self.xdim * self.ydim, self.nimgsR).astype(np.float32))
def add_reference_images(self, data): """Add a new reference images
:param data: The new reference image. :type data: xarray, numpy array or list :raises DataSizeException: Raised when the shape of the data is not correct. """
if self.reshape: data = self.reshape_data(data) elif isinstance(data, type(xr.DataArray())): data = data.to_numpy() if not ((data.shape[0]==self.ydim) and (data.shape[1]==self.xdim)): raise DataSizeException
data = data.reshape(self.xdim * self.ydim) self._referenceImages = np.append(self._referenceImages, data, axis=1) def _remove_first_reference_images(self): """Remove the first reference images
"""
self._referenceImages = np.delete(self._referenceImages, 0, axis=1) def update_reference_images(self, data): """Update the reference images set by removing the first one and adding a new one at the end.
:param data: The new reference image. :type data: xarray, numpy array or list """
self._remove_first_reference_images() self.add_reference_images(data)
self.decompose_referenceImages() @property def mask(self): return self._mask @mask.setter def mask(self, value): if self.reshape: value = self.reshape_data(value) elif isinstance(value, type(xr.DataArray())): value = value.to_numpy() if not ((value.shape[0]==self.ydim) and (value.shape[1]==self.xdim)): raise DataSizeException self._mask = value self._center = None self._span = None self.k = np.where(self._mask.flatten() == 1)[0] def _auto_mask(self): mask = np.ones((self.ydim, self.xdim), dtype=np.uint8) if not self._center is None: x_start = int(self._center[0] - self._span[0] / 2) x_end = int(self._center[0] + self._span[0] / 2) 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 return mask def decompose_referenceImages(self): if self._mask is None: self.mask = self._auto_mask() self.P, self.L, self.U = lu(self._referenceImages[self.k, :].T @ self._referenceImages[self.k, :], permute_l = False, p_indices = True) def _fringe_removal(self, absorptionImages): b = self.temp @ absorptionImages[self.k] c = np.linalg.solve(self.U, np.linalg.solve(self.L[self.P], b)) optrefimages = (self._referenceImages @ c) return optrefimages def fringe_removal(self, absorptionImages, referenceImages=None, mask=None, reshape=None, dask='parallelized'): """
This function will generate a 'fake' background images, which can help to remove the fringes. Important: Please substract the drak images from the both of images with atoms and without atoms before using this function!!!
:param absorptionImages: A set of images with atoms in absorption imaging :type absorptionImages: xarray, numpy array or list :param referenceImages: A set of images without atoms in absorption imaging, defaults to None :type referenceImages: xarray, numpy array or list, optional :param mask: An array to choose the region of interest for fringes removal, defaults to None, defaults to None :type mask: numpy array, optional :param reshape: If it needs to reshape the data, defaults to None :type reshape: bool, optional :param dask: Please refer to xarray.apply_ufunc() :type dask: {"forbidden", "allowed", "parallelized"}, optional :return: The 'fake' background to help removing the fringes :rtype: xarray array """
if not reshape is None: self.reshape = reshape if not referenceImages is None: self.referenceImages = referenceImages if not mask is None: self.mask = mask if self.P is None: self.decompose_referenceImages() if self.reshape: absorptionImages = self._reshape_absorption_images(absorptionImages) self.temp = self._referenceImages[self.k, :].T optrefimages = xr.apply_ufunc(self._fringe_removal, absorptionImages, input_core_dims=[['axis']], output_core_dims=[['axis']], dask=dask, vectorize=True, output_dtypes=float) return optrefimages.unstack()
|
