dpss_fft_test.py

#!/usr/bin/env python2
# -*- coding: utf-8 -*-
"""
Created on Tue Sep 19 16:54:05 2017

@author: larsoner
"""
from scipy.signal.windows import dpss
from scipy import fftpack
import numpy as np
import matplotlib.pyplot as plt

Ms = np.arange(1, 41)
factors = (50, 20, 10, 5, 2.0001)
energy = np.empty((3, len(Ms), len(factors)))
for mi, M in enumerate(Ms):
    for fi, factor in enumerate(factors):
        NW = M / float(factor)
        # Corrected
        win = dpss(M, NW)
        energy[0, mi, fi] = np.sum(win ** 2) / np.sqrt(M)
        # Uncorrected
        win = dpss(M, NW, norm=np.inf)
        energy[1, mi, fi] = np.sum(win ** 2) / np.sqrt(M)
        # FFT correction
        if M % 2 == 0:
            s = fftpack.fft(win)
            k = np.concatenate((np.arange(M/2), np.arange(-M/2, 0)))
            shift = -(M/2 - 0.5)
            s_sh = s * np.exp(-1j * 2 * np.pi * shift / M * k)
            win /= s_sh.mean().real
        energy[2, mi, fi] = np.sum(win ** 2) / np.sqrt(M)
fig, ax = plt.subplots(1, figsize=(4, 8))
hs = ax.plot(Ms, energy[0], '-o', markersize=4,
             markeredgecolor='none')
leg = [hs[-1]]
for hi, hh in enumerate(hs):
    h = ax.plot(Ms, energy[1, :, hi], '-o', markersize=4,
                color=hh.get_color(), markeredgecolor='none',
                alpha=0.33)
    if hi == len(hs) - 1:
        leg.insert(0, h[0])
    h = ax.plot(Ms, energy[2, :, hi], '-o', markersize=4,
                color=hh.get_color(), markeredgecolor='none',
                alpha=0.66)
    if hi == len(hs) - 1:
        leg.append(h[0])
ax.set(xlabel='M (samples)', ylabel=r'Power / $\sqrt{M}$')
ax.legend(leg, ['Uncorrected', r'Corrected: $\frac{M^2}{M^2+NW}$',
                'Corrected (FFT)'])
fig.tight_layout()