Source code for ddd_pylib.image_manip.fast_fourier_transform_operator

import numpy as np
import xarray as xr
from ddd_pylib._base import ImageOperator





class FastFourierTransformOperator(ImageOperator):
    norm_mode = set({"backward", "ortho", "forward"})
    def __init__(self):
        super().__init__()
        self.axis = "X"
        self.norm = None
        self.inverse = False



    def getPrefix(self):
        return "FFT - "




    def setAxis(self, axis):
        if axis not in self.getMainImage().dims:
            raise ValueError("Invalid axis : axis doesn't exist in this image")
        self.axis = axis




    def setNorm(self, norm):
        if norm not in self.norm_mode:
            raise ValueError("Norm mode invalid")
        self.norm_mode = norm




    def switchInverse(self):
        self.inverse = not self.inverse


    def _fourierTransform(self, image):
        index_axis = image.dims.index(self.axis)
        if self.inverse:
            res = np.fft.ifft(a=image.values, axis=index_axis, norm=self.norm)
        else:
            res = np.fft.fft(a=image.values, axis=index_axis, norm=self.norm)
        new_da = image.copy()
        new_da.data = res.real
        return new_da

    def _applyToFrame(self, frameData, t, nT):
        image = frameData[0]
        return self._fourierTransform(image)