Source code for mojo.analysis.windowedFFT

"""
windowedFFT - Perform a fast fourier transformation
===================================================

Created on 11.01.2013

author: well_jn

Perform a fast fourier transformation of a real function.
You have the choice of multiple __window functions.
For no windowing use the default values of the constructor.
Choices of __window functions:
HAMMING,
HANNING,
BARTLETT,
BLACKMAN,
FLAT

Example::

    from analysis.windowedFFT import WindowedFFT

    time = np.arange(0, 1, dt)
    myFFT = WindowedFFT(func_timeDomain = np.sin(2 * np.pi * freq * time), d_time = dt, window_type = 'HAMMING', window_size = 50)
    freq = myFFT.getFrequencies()
    func_freq = myFFT.getSpectrum()

Create a time array and with this a sine function.
Create windowedFFT object with this sine and the choice
for the windowing. You get the frequencies for the positive range of the FFT and
the single single-sided amplitude spectrum of the sine.

requires
--------
numpy
"""


import numpy as np


class WindowedFFT:
    """
    Class to setup a FFT with windowing to get the single-sided amplitude spectrum of a real function
    """

    def __init__(self, func_timeDomain, d_time, window_type='FLAT', window_size=0, window_overlap=0):
        '''
        Constructor of windowedFFTObject object.
        '''
        self.__func_timeDomain = func_timeDomain
        self.__d_time = d_time

        self.__window_type = None
        self.__window_size = 0
        self.__window_overlap = 0
        self.__window = None

        self.__fft_range = 0
        self.__positive_range = 0

        self.setWindowType(window_type)
        self.setWindowSize(window_size)
        self.setWindowOverlap(window_overlap)

    def __setupSingleSidedFFT(self):
        """
        Check if the necessary information were set by the user and
        calculate some additional stuff.
        """
        self.__fft_range = self.__window_size // 2 + 1
        self.__positive_range = self.__fft_range // 2

    def __getWindow(self):
        """
        Depending on the windowing string call the numpy function to get
        an array with values of the __window.
        """
        if self.__window_type == 'HAMMING':

            self.__window = np.hamming(self.__window_size)

        elif self.__window_type == 'HANNING':

            self.__window = np.hanning(self.__window_size)

        elif self.__window_type == 'BARTLETT':

            self.__window = np.bartlett(self.__window_size)

        elif self.__window_type == 'BLACKMAN':

            self.__window = np.blackman(self.__window_size)

        elif self.__window_type == 'FLAT':

            self.__window = np.kaiser(self.__window_size, 0)

        else:

            raise ValueError("invalid __window type '%s'!" % self.__window_type)

        self.__window /= np.sum(self.__window) / self.__window_size

    def setWindowType(self, window_type):
        """
        Set which type of windowing you want to use.
        Possible choices are:
        - HAMMING
        - HANNING
        - BARTLETT
        - BLACKMAN
        - FLAT
        """
        self.__window_type = window_type.upper()

    def setWindowSize(self, window_size: int):
        """
        Set the length of the used window.
        """
        self.__window_size = window_size

        if self.__window_size > self.__func_timeDomain.size:

            raise ValueError('Window size is larger than the length of the data.')

        elif self.__window_size < 0:

            raise ValueError('Window size must be a positive integer.')

        elif self.__window_size == 0:

            self.__window_size = self.__func_timeDomain.size

        self.__setupSingleSidedFFT()

    def setWindowOverlap(self, window_overlap: int):
        """
        Set the overlapping of the window.
        """
        if window_overlap < 0:

            raise ValueError('Window overlap must be a positive integer.')

        elif window_overlap >= self.__window_size:

            raise ValueError('Window overlap is larger than the window size.')

        self.__window_overlap = window_overlap

    def getSingleSidedFrequencies(self):
        """
        Get corresponding frequencies for the resulting spectrum.
        """
        freq = np.fft.fftfreq(self.__fft_range, 2 * self.__d_time)
        freq = freq[list(range(self.__positive_range))]

        return freq

    def getSingleSidedSpectrum(self):
        """
        Compute a single-sided amplitude spectrum of a real __func_timeDomain(t) with windowing.
        """

        self.__getWindow()

        idx = 0
        func_freq = np.zeros(self.__fft_range, dtype=complex)

        while idx + self.__window_size <= self.__func_timeDomain.size:

            func_timeDomain_window = self.__func_timeDomain[idx:idx + self.__window_size]

            func_timeDomain_window = self.__window * func_timeDomain_window

            func_freq_window = np.fft.rfft(func_timeDomain_window) / self.__window_size

            func_freq += func_freq_window

            idx += (self.__window_size - self.__window_overlap)

        func_freq = 2 * func_freq[list(range(self.__positive_range))]

        return func_freq