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

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


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

@author: cnheider
"""

__all__ = ["FastFourierTransformSpectrogramPlot"]

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 FastFourierTransformSpectrogramPlot(MplDrawer): """ TODO: CENTER Align fft maybe, to mimick librosa stft Short Time Fourier Transform (STFT), with step size of 1 and window lenght of n_fft, and no window function ( TODO: Hanning Smoothing)"""
[docs] def __init__( self, n_fft: int = 64, sampling_rate=int(1.0 / 0.0005), buffer_size_sec: float = 1.0, title: str = "", vertical: bool = True, reverse: bool = False, figure_size=(9, 9), cmap="viridis", render: bool = True, **kwargs ): """ :param n_fft: :type n_fft: :param sampling_rate: :type sampling_rate: :param buffer_size_sec: :type buffer_size_sec: :param title: :type title: :param vertical: :type vertical: :param reverse: :type reverse: :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.n_positive_fft = (self.n_fft + 1) // 2 self.sampling_rate = sampling_rate if buffer_size_sec is not None: self.buffer_size_sec = buffer_size_sec else: self.buffer_size_sec = self.n_fft / sampling_rate * 2 # self.window_function = window_function(self.n_fft) self.buffer_array_size = int(sampling_rate * self.buffer_size_sec) assert self.buffer_array_size >= self.n_fft time_s = numpy.linspace( 0, self.buffer_size_sec, self.buffer_array_size, endpoint=False ) raw_array = numpy.zeros(self.buffer_array_size) zeroes_img = numpy.zeros((self.n_positive_fft, self.buffer_array_size - n_fft)) self.zeroes_padding = numpy.zeros((self.n_positive_fft, n_fft)) gs = GridSpec(3, 2, width_ratios=[100, 2]) ( self.raw_ax, _, self.angle_ax, self.angle_cbar_ax, self.spec_ax, self.spec_cbar_ax, ) = [pyplot.subplot(gs[i]) for i in range(6)] (self.raw_line2d,) = self.raw_ax.plot(time_s, raw_array) self.raw_ax.set_xlim([time_s[0], time_s[-1]]) self.raw_ax.set_ylabel("Signal [Magnitude]") max_freq = numpy.max(numpy.fft.fftfreq(self.n_fft, 1 / sampling_rate)) self.dft_angle_img = self.angle_ax.imshow( zeroes_img, vmin=-math.pi, vmax=math.pi, interpolation="hanning", aspect="auto", extent=[time_s[0], time_s[-1], max_freq, 0], cmap=cmap, ) self.angle_ax.set_ylabel("Phase [Hz]") _ = self.fig.colorbar(self.dft_angle_img, cax=self.angle_cbar_ax) self.angle_cbar_ax.set_ylabel("Angle (Radians)", rotation=90) self.dft_mag_img = self.spec_ax.imshow( zeroes_img, vmin=0, vmax=1, interpolation="hanning", aspect="auto", extent=[time_s[0], time_s[-1], max_freq, 0], cmap=cmap, ) self.spec_ax.set_ylabel("Frequency [Hz]") self.spec_ax.set_xlabel("Time [Sec]") _ = self.fig.colorbar(self.dft_mag_img, cax=self.spec_cbar_ax) self.spec_cbar_ax.set_ylabel("Magnitude (dB)", rotation=90) self.vertical = vertical self.reverse = reverse pyplot.xlim(time_s[0], time_s[-1]) 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:""" y_data = self.raw_line2d.get_ydata() if not self.reverse: y_data = numpy.hstack((y_data[1:], signal_sample)) y_data_view = y_data[-self.n_fft :] else: y_data = numpy.hstack((signal_sample, y_data[:-1])) y_data_view = y_data[: self.n_fft] self.raw_line2d.set_ydata(y_data) cur_lim = self.raw_ax.get_ylim() self.raw_ax.set_ylim( [min(cur_lim[0], signal_sample), max(cur_lim[1], signal_sample)] ) # if self.window_function is not None: # y_data_view *= self.window_function f_coef = numpy.fft.fft(y_data_view, n=self.n_fft)[ : self.n_positive_fft ].reshape( -1, 1 ) # Only select the positive frequencies phase_data = self.dft_angle_img.get_array() new_phase = numpy.angle(f_coef) # new_phase = f_coef.imag if not self.reverse: phase_data = numpy.concatenate( (phase_data[:, 1 : -self.n_fft], new_phase, self.zeroes_padding), axis=-1, ) else: phase_data = numpy.concatenate( (self.zeroes_padding, new_phase, phase_data[:, self.n_fft : -1]), axis=-1, ) self.dft_angle_img.set_array(phase_data) magnitude_data = self.dft_mag_img.get_array() new_mag = 10 * numpy.log10( ( numpy.abs(f_coef) # f_coef.real ** 2 ) + FLOAT_EPS ) # db if not self.reverse: magnitude_data = numpy.concatenate( (magnitude_data[:, 1 : -self.n_fft], new_mag, self.zeroes_padding), axis=-1, ) else: magnitude_data = numpy.concatenate( (self.zeroes_padding, new_mag, magnitude_data[:, self.n_fft : -1]), axis=-1, ) self.dft_mag_img.set_clim( vmin=numpy.min(magnitude_data), vmax=numpy.max(magnitude_data) ) self.dft_mag_img.set_array(magnitude_data) # self.spec_cbar_ax.set_ylabel("Magnitude (Linear)", rotation=90) self.spec_cbar_ax.set_ylabel("Magnitude (dB)", rotation=90)
if __name__ == "__main__": def a() -> None: """ :rtype: None """ duration_sec = 4 mul = 1000 sampling_Hz = 44.1 sampling_rate = int(sampling_Hz * mul) # Hz delta = 1 / sampling_rate n_fft = 128 # 1024 s = FastFourierTransformSpectrogramPlot( n_fft=n_fft, sampling_rate=sampling_rate, buffer_size_sec=delta * n_fft * 4 ) 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()