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

w_1 = 40  # frequency of the 1st component of the signal (Hz)
w_2 = 60  # frequency of the 2nd component of the signal (Hz)

a = 0.5  # magnitude of the 1st component of the signal
b = 1.0  # magnitude of the 2nd component of the signal

t = np.array([i for i in range(1, 301)]) / 1000  # time samples (s)
fs = 1 / (t[1] - t[0])  # sampling frequency (Hz)

x = a * np.cos(2 * np.pi * w_1 * t) + b * np.sin(2 * np.pi * w_2 * t)  # considered signal
n = .2 * np.random.randn(len(t))  # white Gaussian noise
y = x + n

windows = [None, 'hamming']

for window in windows:
    if window is None:
        label = 'rect'
    else:
        label = window
    f, Pxx_den = signal.periodogram(y, fs=fs, scaling='spectrum', nfft=2048, window=window)
    plt.semilogy(f[1:200], Pxx_den[1:200], label=label)

plt.ylabel('Spectrum')
plt.xlabel('Frequencies (Hz)')
plt.title('Periodogram')
plt.grid(True)
plt.legend()
plt.savefig("Periodogram windows.svg")