clear all, close all, clc
% H. Garnier
% Adapted from
% S. Brunton and N. Kutz
% Data-Driven Science and Engineering: Machine Learning, Dynamical Systems, and Control"
% Cambridge, 2022 

%% Create a simple signal with two frequencies
Ts = 1e-3;
fs = 1/Ts;
t = 0:Ts:2;
x0 = sin(2*pi*50*t) + sin(2*pi*120*t); % Sum of 2 noise-free sines
x = x0 + 2.5*randn(size(t));  %  Add a large amount of white Gaussian noise

%% Compute the Fast Fourier Transform FFT
N = length(t);
X = fft(x,N);       % Compute the fast Fourier transform
PSD_X = X.*conj(X)/N; % Power spectrum (power per freq)
X0 = fft(x0,N);       % Compute the fast Fourier transform
PSD_X0 = X0.*conj(X0)/N; % Power spectrum (power per freq)

freq = fs/N*(0:N); % Create x-axis of frequencies in Hz
L = 1:floor(N/2);   % Only plot the first half of freqs

%% Use the PSD to filter out noise
indices = PSD_X>100;  % Find all freqs with large power
PSD_Xdenoised = PSD_X.*indices;  % Zero out all others
X_f = indices.*X;  % Zero out small Fourier coeffs. in X
xfilt = ifft(X_f); % Inverse FFT for filtered time signal

%% PLOTS
subplot(2,2,1)
plot(t,x0,'k','LineWidth',1.2), hold on
plot(t,x,'r','LineWidth',1.5)
axis([0 .1 -10 10])
grid
legend('Noise-free','Noisy')
xlabel('Time (sec)')

subplot(2,2,3)
plot(t,x0,'k','LineWidth',1.5), hold on
plot(t,xfilt,'-b','LineWidth',1.2)
axis([0 .1 -5 5])
grid
legend('Noise-free','Denoised')
xlabel('Time (sec)')

subplot(2,2,2)
plot(freq(L),PSD_X0(L),'k','LineWidth',1.5), hold on
plot(freq(L),PSD_X(L),'-r','LineWidth',1.2)
grid
legend('Noise-free','Noisy')
xlabel('Frequency (Hz)')

subplot(2,2,4)
plot(freq(L),PSD_X0(L),'k','LineWidth',1.5), hold on
plot(freq(L),PSD_Xdenoised(L),'-b','LineWidth',1.2)
grid
legend('Noise-free','Denoised')
xlabel('Frequency (Hz)')

