dpss_pop_test.py

from scipy.signal.windows import dpss
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)
        # Other correction (pop middle value)
        if M % 2 == 0:
            win = dpss(M + 1, NW)
            win = np.concatenate([win[:M//2], win[-M//2:]])
        energy[2, mi, fi] = np.sum(win ** 2) / np.sqrt(M)
fig, ax = plt.subplots(1)
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 (pop)'])
fig.tight_layout()