Code for dsp.stackexchange.com question

fft_plots.py

#! /usr/bin/env python3

import sys

import numpy as np
import matplotlib
import scipy
from matplotlib import pyplot as plt

import math

def dB20(array):
    with np.errstate(divide='ignore'):
        return 20 * np.log10(array)

if True:
    plt.figure(figsize=(10, 20))
    num_plots           = 6
    plot_nr             = 0

    nr_input_bits       = 16 
    sample_rate_in      = 44100
    amplitude           = 1
    freq_hz             = 1000.621

    #============================================================
    # Trivial FFT: 1 second, 44.1k sample points, no window, window
    #============================================================

    N           = int(sample_rate_in)
    x           = np.linspace(0., 1.0, N, endpoint=False)

    y           = np.sin(2 * np.pi * freq_hz * x)
    quant       = 2**(nr_input_bits-1)

    y           = np.round(quant * y) / quant
    w           = np.blackman(len(y)); corr = 2.8

    H_no_win    = np.fft.fft(y) / len(y)
    H_win       = corr * np.fft.fft(y * w) / len(y)

    freqs       = np.fft.fftfreq(len(y))
    x_mask      = freqs >= 0

    plt.figure(figsize=(10, 6))
    num_plots           = 2
    plot_nr             = 0

    plot_nr += 1
    plt.subplot(num_plots,1,plot_nr)
    plt.gca().set_xlim([0.0, sample_rate_in/2])
    plt.grid(True)
    plt.plot(freqs[x_mask] * sample_rate_in, dB20(np.abs(H_win))[x_mask])
    plt.plot(freqs[x_mask] * sample_rate_in, dB20(np.abs(H_no_win))[x_mask])
    plt.title("44.1kHz, 1 second, without and with window")

    plot_nr += 1
    plt.subplot(num_plots,1,plot_nr)
    plt.gca().set_xlim([0.75 * freq_hz, 1.25*freq_hz])
    plt.grid(True)
    plt.plot(freqs[x_mask] * sample_rate_in, dB20(np.abs(H_win))[x_mask])
    plt.plot(freqs[x_mask] * sample_rate_in, dB20(np.abs(H_no_win))[x_mask])
    plt.title("44.1kHz, 1 second, without and with window, zoomed around peak")

    plt.tight_layout()
    plt.savefig("fig1.svg")
    plt.savefig("fig1.png")

    #============================================================
    # 100 FFTs of 44.1k sample points each. Average each bin.
    #============================================================

    num_ffts    = 100
    N           = int(num_ffts * sample_rate_in)
    x           = np.linspace(0., num_ffts, N, endpoint=False)
    y           = np.sin(2 * np.pi * freq_hz * x)
    quant       = 2**(nr_input_bits-1)
    y           = np.round(quant * y) / quant

    H_avg       = np.zeros(sample_rate_in)
    for chunk in np.array_split(y, num_ffts):
        w           = np.blackman(len(chunk)); corr = 2.8
        H_chunk     = corr * np.fft.fft(chunk * w) / len(chunk)
        H_avg       = np.add(H_avg, np.abs(H_chunk))

    H_avg = H_avg / num_ffts

    plt.figure(figsize=(10, 8))
    num_plots           = 2
    plot_nr             = 0

    plot_nr += 1
    plt.subplot(num_plots,1,plot_nr)
    plt.gca().set_xlim([0.0, sample_rate_in/2])
    plt.grid(True)
    plt.plot(freqs[x_mask] * sample_rate_in, dB20(np.abs(H_win))[x_mask])
    plt.plot(freqs[x_mask] * sample_rate_in, dB20(np.abs(H_avg))[x_mask])
    plt.title("44.1kHz, 100 seconds, 100 FFTs of 1s averaged")

    plot_nr += 1
    plt.subplot(num_plots,1,plot_nr)
    plt.gca().set_xlim([0.75 * freq_hz, 1.25*freq_hz])
    plt.grid(True)
    plt.plot(freqs[x_mask] * sample_rate_in, dB20(np.abs(H_win))[x_mask])
    plt.plot(freqs[x_mask] * sample_rate_in, dB20(np.abs(H_avg))[x_mask])
    plt.title("44.1kHz, 100 seconds, 100 FFTs of 1s averaged, zoomed in around peak")

    plt.tight_layout()
    plt.savefig("fig2.svg")
    plt.savefig("fig2.png")

    #============================================================
    # 44.1k * 100 sample points
    #============================================================

    length_s    = 100
    N           = int(length_s * sample_rate_in)
    x           = np.linspace(0., length_s, N, endpoint=False)
    y           = np.sin(2 * np.pi * freq_hz * x)
    quant       = 2**(nr_input_bits-1)
    y           = np.round(quant * y) / quant
    w           = np.blackman(len(y)); corr = 2.8

    H_large     = corr * np.fft.fft(y * w) / len(y)

    freqs_large = np.fft.fftfreq(len(y))
    x_mask_large= freqs_large >= 0

    plt.figure(figsize=(10, 8))
    num_plots           = 2
    plot_nr             = 0

    plot_nr += 1
    plt.subplot(num_plots,1,plot_nr)
    plt.gca().set_xlim([0.0, sample_rate_in/2])
    plt.grid(True)
    plt.plot(freqs_large[x_mask_large] * sample_rate_in, dB20(np.abs(H_large))[x_mask_large])
    plt.title("44.1kHz, 100 seconds, 1 FFT")

    plot_nr += 1
    plt.subplot(num_plots,1,plot_nr)
    plt.gca().set_xlim([0.75 * freq_hz, 1.25*freq_hz])
    plt.grid(True)
    plt.plot(freqs_large[x_mask_large] * sample_rate_in, dB20(np.abs(H_large))[x_mask_large])
    plt.title("44.1kHz, 100 seconds, 1 FFT, zoomed around peak")

    plt.tight_layout()
    plt.savefig("fig3.svg")
    plt.savefig("fig3.png")

    #============================================================
    # Merged bins
    #============================================================

    H_merged    = np.zeros(len(H_large) // 100)

    for bin_nr, bin_val in enumerate(H_large):
        H_merged[bin_nr//100] += np.abs(bin_val)

    plt.figure(figsize=(10, 8))
    num_plots           = 2
    plot_nr             = 0

    plot_nr += 1
    plt.subplot(num_plots,1,plot_nr)
    plt.gca().set_xlim([0.0, sample_rate_in/2])
    plt.grid(True)
    plt.plot(freqs[x_mask] * sample_rate_in, dB20(np.abs(H_merged))[x_mask])
    plt.title("44.1kHz, 100 seconds, 1 FFT, bins merged")

    plot_nr += 1
    plt.subplot(num_plots,1,plot_nr)
    plt.gca().set_xlim([0.75 * freq_hz, 1.25*freq_hz])
    plt.grid(True)
    plt.plot(freqs[x_mask] * sample_rate_in, dB20(np.abs(H_merged))[x_mask])
    plt.title("44.1kHz, 100 seconds, 1 FFT, bins_merged, zoomed around peak")

    plt.tight_layout()
    plt.savefig("fig4.svg")
    plt.savefig("fig4.png")

    #============================================================