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()