helpful plots python

helpful_plots.py

"""
Helpful plots for signal processing.

Created on Fri Jan 08 20:19:58 2016
"""
import matplotlib.pyplot as plt
import numpy as np
from numpy.fft import rfft, rfftfreq
from scipy.signal import get_window


def waveform(sig, t=None):
    """
    Plot waveform vs time with grids and margins.

    Parameters
    ----------
    sig : array_like
        Signal to plot.
    t : array_like
        Timebase for signal.
    """
    if t is None:
        t = np.arange(len(sig))

    plt.figure('Waveform')
    plt.plot(t, sig)

    plt.title('Signal waveform')
    plt.xlabel('Time [s]')
    plt.ylabel('Amplitude')
    plt.margins(0.01, 0.1)
    plt.grid(True, color='0.7', linestyle='-', which='major')
    plt.grid(True, color='0.9', linestyle='-', which='minor')
    plt.show()


def spectrum(sig, t=None, fs=None, axis='log', N=None, window=9,
             label=None):
    """
    Plot spectrum of signal vs frequency.

    Parameters
    ----------
    sig : array_like
        Signal to plot.
    t : array_like
        Timebase for signal.
    fs : float
        Sampling frequency.
    axis : {'log', 'lin'}
        Horizontal axis type, linear or logarithmic.
    N : int
        Number of FFT points to compute. Only 1/2 this many points are plotted.
    window : str, float, tuple, or array_like
        Window function, handled by `scipy.signal.get_window`. Defaults to
        Kaiser with beta=9.
    label : str
        A label for the spectrum to appear in the legend.
    """
    if N is None:
        N = len(sig)

    if t is not None and fs is not None:
        raise ValueError("Only 't' or 'fs' can be specified, not both")
    elif t is None and fs is None:
        fs = len(sig)
    elif fs is None:
        fs = 1/(t[1] - t[0])

    freqs = rfftfreq(N, 1/fs)

    w = get_window(window, len(sig))

    # RFFT is only half of spectrum. Correct for window attenuation.
    ampl = 2/w.sum() * np.absolute(rfft(w * sig, N))

    plt.figure('Spectrum')
    plt.plot(freqs, 20*np.log10(ampl), label=label)
    if axis == 'log':
        plt.xscale('log')

    if label:
        plt.legend(loc='best')

    plt.title('Signal spectrum')
    plt.xlabel('Frequency [Hz]')
    plt.ylabel('Amplitude [dBFS]')  # AES17 definition
    plt.margins(0.01, 0.1)
    plt.grid(True, color='0.7', linestyle='-', which='major')
    plt.grid(True, color='0.9', linestyle='-', which='minor')
    plt.show()


if __name__ == '__main__':
    from numpy import sin, cos, pi
    t = np.linspace(0, 1, 10000, endpoint=False)
    sig = sin(2*pi*4*t) + 0.25*(cos(2*pi*18*t))  # 0 dBFS and -12 dBFS
    waveform(sig, t)
    spectrum(sig, t, N=40000)