SavGol test

gistfile1.py

# -*- coding: utf-8 -*-
"""
Examine if a 2x window length is appropriate for savgol filtering.
"""

from scipy.signal import savgol_filter
from scipy import fftpack
import numpy as np
import matplotlib.pyplot as plt
plt.ion()

fs = 1000.
n = 10000
orders = np.arange(2, 9)
window_freqs = np.array([5, 10, 25, 50])
window_durs = 2. / window_freqs  # use 2x as long a window, maybe?
n_runs = 100
spectra = 0
orig_spectra = 0
rng = np.random.RandomState(0)
for ri in range(n_runs):
    x = rng.randn(n)
    run_spectra = []
    for order in orders:
        specs = []
        for window_dur in window_durs:
            window_len = (int(np.round(window_dur * fs)) // 2) * 2 + 1
            specs.append(fftpack.fft(savgol_filter(x, window_len, order)))
        run_spectra.append(specs)
    spectra += np.abs(np.array(run_spectra)) ** 2
    orig_spectra += np.abs(fftpack.fft(x)) ** 2
freqs = fftpack.fftfreq(n, 1. / fs)
mask = freqs > 0.5
freqs = freqs[mask]
spectra /= n_runs
orig_spectra /= n_runs
spectra = 20 * np.log10(spectra[..., mask])
orig_spectra = 20 * np.log10(orig_spectra[..., mask])

xlim = [freqs[0], 4. * window_freqs.max()]
fig, axs = plt.subplots(len(orders), len(window_durs), squeeze=False)
for oi in range(len(orders)):
    for wi in range(len(window_durs)):
        freq = window_freqs[wi]
        ax = axs[oi, wi]
        ax.plot(freqs, orig_spectra, 'y')
        ax.plot([freqs.min(), freqs.max()],
                [np.nanmedian(orig_spectra) - 20] * 2, 'k--')
        ax.plot(freqs, spectra[oi, wi], 'k')
        ax.plot([freq, freq], ax.get_ylim(), 'k:')
        ax.set_xlim(xlim)
        ax.set_xscale('log')
        if wi == 0:
            ax.set_ylabel(orders[oi])
        else:
            ax.set_yticks([])
        if oi == len(orders) - 1:
            ax.set_xlabel('%s Hz' % freq)
        else:
            ax.set_xticks([])