import numpy as np
import matplotlib.pyplot as plt
from matplotlib import mlab
from scipy import signal

tau, n = 10, 2**10  # 10 second signal with 1024 points
T = tau/n  # sampling interval

t = np.arange(n)*T
x = 4*np.sin(40*np.pi*t)  # 20 Hz sine
w = signal.hanning(n, False)

f = np.fft.fftfreq(n, T)  # frequency slots
df = f[1] - f[0]
X = np.fft.fft(x*w)/w.sum()  # Windowed FFT scaled to amplitude

f_Pxx, Pxx = signal.welch(x, fs=1/T, window='hanning', nperseg=n, noverlap=0,
                          detrend=False, return_onesided=False,
                          scaling='spectrum')
P_psd, f_psd = mlab.psd(x, NFFT=n, Fs=1/T,  window=w,
                        sides='twosided', scale_by_freq=False)

# shift for plotting:
f, X, f_Pxx, Pxx = (np.fft.fftshift(zz) for zz in (f, X, f_Pxx, Pxx))

fg0, axx0 = plt.subplots(4, 1, sharex=True, num=1, figsize=(7, 7), clear=True)

axx0[0].set_title(r"FFT Magnitude Spectrum of $x(t)=4\sin(40\pi t)$")
axx0[0].plot(f, np.abs(X), 'C0.-')

axx0[1].set_title("Scipy's Welch Algorithm")
axx0[1].plot(f_Pxx, np.sqrt(Pxx), "C1.-")

axx0[2].set_title("psd()")
axx0[2].plot(f_psd, np.sqrt(P_psd), "C2.-")

axx0[3].set_title("magnitude_spectrum()")
axx0[3].magnitude_spectrum(x, Fs=1/T, sides='twosided', color="C3", marker='.',
                           scale='linear')
# Correct result:
# axx0[3].magnitude_spectrum(x, Fs=1/T, sides='twosided', color="C3", marker='.',
#                           scale='linear', window=w/w.sum())

axx0[-1].set_xlabel(r"Frequency $f$")
axx0[0].set_xlim(19, 21)
for ax in axx0:
    ax.grid(True)
    ax.set_ylabel(r"$|X(f)|$")
fg0.tight_layout()
plt.show()