stepresponse.py

import sys
from scipy.io import wavfile
import numpy as np
import scipy.fft as fft
import matplotlib.pyplot as plt
import matplotlib.ticker


def plotStepResp(file_path):
    '''
    Requires a mono wav
    '''
    samplerate, data = wavfile.read(file_path)

    data = np.longdouble(data) * np.hanning(data.shape[-1])

    fft_data = fft.fftshift(fft.fft(data))
    freq = fft.fftshift(fft.fftfreq(data.shape[-1], d=1/samplerate))

    fig, axes = plt.subplots(2, 1, sharex=True, tight_layout=True)
    ax = axes[0]

    # generate xticks
    # XXX we don't actually use this array, just the first
    #     and last value.
    xticks = []
    cur_tick = samplerate/2
    while cur_tick > 10:
        xticks.append(cur_tick)
        cur_tick /= 10
    xticks.sort()
    print(xticks)

    ax.set_xscale('log')
    ax.set_xlim(xticks[0], xticks[-1])
    ax.get_xaxis().set_major_locator(
        matplotlib.ticker.LogLocator(subs=[float(xticks[0])]))
    ax.get_xaxis().set_minor_locator(matplotlib.ticker.LogLocator(
        subs=[float(xticks[0]*2*i) for i in range(1, 5)]))
    ax.get_xaxis().set_major_formatter(matplotlib.ticker.ScalarFormatter())
    ax.get_xaxis().set_minor_formatter(matplotlib.ticker.ScalarFormatter())

    half_freq = freq[freq > 0]
    half_fft_data_abs = np.abs(fft_data[freq > 0])

    plot_data = 20*np.log10(half_fft_data_abs/half_fft_data_abs.max())

    # create fit line arbitrarily based in the middle of the plot
    middle_freq = np.sqrt(xticks[-1] * xticks[0])  # "middle" logarithmically
    middle_index = np.argmin(np.abs(half_freq - middle_freq))
    middle_value = plot_data[middle_index]
    print(middle_freq, middle_index, middle_value)

    # ideal step response is 20dB/dec
    fitline = middle_value - 20*np.log10(half_freq/middle_freq)
    diffline = plot_data - fitline

    ymin = 3*((diffline.min() // 3))
    ymax = 3*((diffline.max() // 3) + 1)
    axes[1].set_ylim(ymin, ymax)

    ax.yaxis.set_major_locator(matplotlib.ticker.MultipleLocator(10))
    ax.yaxis.set_minor_locator(matplotlib.ticker.MultipleLocator(2))
    axes[1].yaxis.set_major_locator(matplotlib.ticker.MultipleLocator(3))
    axes[1].yaxis.set_minor_locator(matplotlib.ticker.MultipleLocator(0.5))

    ax.plot(half_freq, plot_data, 'b-')
    ax.plot(half_freq, fitline, 'm:', linewidth=3)

    axes[1].plot(half_freq, diffline, 'b-')
    axes[1].plot(half_freq, +1*np.ones(half_freq.shape[-1]), 'm:')
    axes[1].plot(half_freq, -1*np.ones(half_freq.shape[-1]), 'm:')
    axes[1].plot(half_freq, +3*np.ones(half_freq.shape[-1]), 'm--')
    axes[1].plot(half_freq, -3*np.ones(half_freq.shape[-1]), 'm--')


    ax.grid(which='minor', linestyle='--', b=True)
    ax.grid(which='major', linestyle='-', b=True, alpha=0.9)
    axes[1].grid(which='minor', linestyle='--', b=True)
    axes[1].grid(which='major', linestyle='-', b=True, alpha=0.9)

    plt.show()


if __name__ == "__main__":
    plotStepResp(sys.argv[1])