Source code for draugr.drawers.mpl_drawers.spectral.fast_fourier_transform

#!/usr/bin/env python3
# -*- coding: utf-8 -*-


__author__ = "Christian Heider Nielsen"
__doc__ = """
Created on 27/04/2019

@author: cnheider
"""

__all__ = ["FastFourierTransformPlot"]

import math

import numpy
from matplotlib import pyplot
from matplotlib.gridspec import GridSpec

from draugr.drawers.mpl_drawers.mpldrawer import MplDrawer
from draugr.tqdm_utilities import progress_bar

FLOAT_EPS = numpy.finfo(float).eps


[docs]class FastFourierTransformPlot(MplDrawer): """ Plots last computed fft of data"""
[docs] def __init__( self, n_fft: int = 64, sampling_rate=int(1.0 / 0.0005), title: str = "", figure_size=(9, 9), render: bool = True, **kwargs ): """ :param n_fft: :type n_fft: :param sampling_rate: :type sampling_rate: :param title: :type title: :param placement: :type placement: :param fig_size: :type fig_size: :param render: :type render:""" super().__init__(render=render, figure_size=figure_size, **kwargs) if not render: return self.fig = pyplot.figure(figsize=figure_size) self.n_fft = n_fft self.abs_n_fft = (self.n_fft + 1) // 2 self.sampling_rate = sampling_rate freq_bins = numpy.arange(self.n_fft) raw_array = numpy.zeros(self.n_fft) self.zeroes_padding = numpy.zeros((self.abs_n_fft, n_fft)) gs = GridSpec(3, 1) (self.raw_ax, self.angle_ax, self.mag_ax) = [ pyplot.subplot(gs[i]) for i in range(3) ] freqs = numpy.fft.fftfreq(self.n_fft, 1 / sampling_rate) (self.dft_raw_img,) = self.raw_ax.plot(freq_bins, raw_array) self.raw_ax.set_xlabel("Time (Sec)") self.raw_ax.set_ylabel("Amplitude") (self.dft_angle_img,) = self.angle_ax.plot(freq_bins, raw_array) self.angle_ax.set_xlabel("Phase [Hz]") self.angle_ax.set_ylabel("Angle (Radians)") self.angle_ax.set_ylim([-math.pi, math.pi]) # self.angle_ax.set_xticks(freqs) (self.dft_mag_img,) = self.mag_ax.plot(freq_bins, raw_array) self.mag_ax.set_xlabel("Frequency [Hz]") self.mag_ax.set_xlabel("Magnitude (dB)") # self.mag_ax.set_xticks(freqs) pyplot.title(title) pyplot.tight_layout()
def _draw(self, signal_sample: float, delta: float = 1 / 120) -> None: """ :param signal_sample: :param delta: 1 / 60 for 60fps :return:""" raw_array = self.dft_raw_img.get_ydata() raw_array = numpy.hstack((signal_sample, raw_array[:-1])) self.dft_raw_img.set_ydata(raw_array) cur_lim = self.raw_ax.get_ylim() self.raw_ax.set_ylim( [min(cur_lim[0], signal_sample), max(cur_lim[1], signal_sample)] ) f_coef = numpy.fft.fft(raw_array, n=self.n_fft) self.dft_angle_img.set_ydata(numpy.angle(f_coef)) mag = 10 * numpy.log10((numpy.abs(f_coef) ** 2) + FLOAT_EPS) self.dft_mag_img.set_ydata(mag) self.mag_ax.set_ylim([min(mag), max(mag)])
if __name__ == "__main__": def a() -> None: """ :rtype: None """ duration_sec = 4 mul = 1000 sampling_Hz = 44 sampling_rate = sampling_Hz * mul # Hz delta = 1 / sampling_rate n_fft = 64 s = FastFourierTransformPlot(n_fft=n_fft, sampling_rate=sampling_rate) for t in progress_bar(numpy.arange(0, duration_sec, delta)): ts = 2 * numpy.pi * t s1 = numpy.sin(ts * 1 * sampling_Hz / 2**4 * mul) s2 = numpy.sin(ts * 3 * sampling_Hz / 2**3 * mul + 0.33 * numpy.pi) s3 = numpy.sin(ts * 5 * sampling_Hz / 2**2 * mul + 0.66 * numpy.pi) signal = s1 + s2 + s3 signal /= 3 # signal += (numpy.random.random() - 0.5) * 2 * 1 / 2 # Noise s.draw(signal, delta=delta) a()