%-- SOLUTIONS to Day 3 Exercises
%-- Peter S. Shawhan

%-- Read in the data and the template
data = load('day2data.txt');
template = load('day2template.txt');
%-- Construct arrays with times of samples
tdata = (0:length(data)-1)/16384;
ttemplate = (0:length(template)-1)/16384;

%-------------------------------
display('%----- Part 1a')
cpubase = cputime;
%-- Create an empty array to store all the correlation results.  If you don't
%-- do this, the loop spends most of its time extending the results array!
corr = zeros(1,length(data)-length(template));
%-- Now loop over time offsets
for ioff=0:length(data)-length(template);
  %-- Grab the section of data to be correlated against
  seldata = data(ioff+1:ioff+length(template));
  %-- There are a few different ways to take the dot product
  %- This version takes the sum of an element-by-element product vector:
%%  corr(ioff+1) = sum( seldata.*template' );
  %- This version uses the matrix product of a row vector with a column vector:
  corr(ioff+1) = seldata*template;
end
filtering_cputime = cputime - cpubase

display('%----- Part 1b')
plot(corr);
title('(1b) Correlation, i.e. output from matched filter in time domain')
%-- Now a clear spike is seen!  Maybe also a smaller one near the right
%--    side of the plot...

%-- Also can estimate the SNR by normalizing to the rms of the filter output
%-- for just noise, grabbed from the first part of the output time series
noise_est = sqrt(mean(corr(1:50000).^2));
SNR = max(corr) / noise_est

input('(1b) A clear signal spike is seen; Press RETURN to continue');

%-------------------------------
% Part 2

display('%----- Part 2a')
rms_orig = sqrt( mean(data.^2))

display('%----- Part 2b')
%-- Generate an array with integers spaced by 16
sel = (1:16:length(data));
dsdata = data(sel);
%-- Another way to do that:
% dsdata = downsample(data,16);

%-- Calculate the rms
rms_simply_decimated = sqrt( mean(dsdata.^2))
%-- This rms is essentially the same as before, because taking a subsample of
%-- the original data points does not change their characteristic spread

display('%----- Part 2c')
[b,a] = butter(6,512/8192,'low');
            %-- The 'low' argument is optional -- makes lowpass by default
filtdata = filter(b,a,data);
%-- Select every 16th sample
dsdata = filtdata(sel);
plot(tdata(sel),dsdata);
title('(2c) Downsampled from original time series by factor of 16')
input('(2c) Plot of properly down-sampled data; Press RETURN to continue');

display('%----- Part 2d')
rms_properly_downsampled = sqrt( mean(dsdata.^2))
%-- The new RMS is somewhat smaller than the original, because we filtered out
%-- the high-frequency parts of the original time series

display('%----- Part 2e')
%-- Note that we give pwelch the NEW sampling rate
pwelch(dsdata,[],[],[],1024);
input('(2e) PSD of downsampled data, fNyquist=512 Hz; Press RETURN to continue');

display('%----- Part 2f')
%-- To down-sample the template, it shouldn't be necessary to filter first
%-- because there shouldn't be any signal power at high frequency.  I will
%-- go ahead and filter before selecting every 16th sample, but I would expect
%-- the results to be the same if I didn't filter.
filttemplate = filter(b,a,template);
tsel = (1:16:length(template));
dstemplate = filttemplate(tsel);

%-- Do the matched filter on the downsampled data and template
cpubase = cputime;
dscorr = zeros(1,length(dsdata)-length(dstemplate));
for ioff=0:length(dsdata)-length(dstemplate);
  seldata = dsdata(ioff+1:ioff+length(dstemplate));
  dscorr(ioff+1) = sum( seldata.*dstemplate' );
end
filtering_cputime = cputime - cpubase

plot(dscorr);

%-- Also can estimate the SNR by normalizing to the rms of the filter output
%-- for just noise, grabbed from the first part of the output time series
noise_est = sqrt(mean(dscorr(1:50000/16).^2));
SNR = max(dscorr) / noise_est

title('(2f) Correlation after downsampling by a factor of 16')
input('(2f) Output from time-domain matched filter after down-sampling; Press RETURN to continue');

%-------------------------------
% Part 3

display('%----- Part 3a')
cpubase = cputime;
fftdata = fft(data);

display('%----- Part 3b')
%-- Extend the template to be the same length as the data time series.
%-- Here is one way to do that:
exttemplate = template;
exttemplate(length(data)) = 0;
%-- Here is another way -- note how we add ROWS to the array, with a semicolon
%% exttemplate = [ template; zeros(length(data)-length(template),1) ];
%-- Now calculate the FFT of the extended template
ffttemplate = fft(exttemplate);

display('%----- Part 3c')
%-- NOTE: the ' operator produces the complex-conjugate transpose!
product = fftdata .* ffttemplate';
%-- (If ffttemplate were already a row vector like fftdata, then we would
%-- have used:      product = fftdata .* conj(ffttemplate);      )
freqcorr = ifft(product);
filtering_cputime = cputime - cpubase

display('%----- Part 3d')
plot(tdata,freqcorr)
%-- Identical to what we got from time-domain matched filtering
title('(3d) Output of matched filtering in freq domain')

%-- Also can estimate the SNR by normalizing to the rms of the filter output
%-- for just noise, grabbed from the first part of the output time series
noise_est = sqrt(mean(freqcorr(1:50000).^2));
SNR = max(freqcorr) / noise_est
vs_nonoptimal_SNR = SNR;

input('(3d) Freq-domain filter output is same as before; Press RETURN to continue');

%-------------------------------
% Part 4

display('%----- Part 4a')
noisespectrum = load('day3asd.txt');
%-- NOTE: That file contains the AMPLITUDE spectral density, so we have
%-- to square it to get the PSD.
noisepsd = noisespectrum.^2;

display('%----- Part 4b')
%-- Combine the data, template, and noise PSD
product = fftdata .* ffttemplate' ./ noisepsd;
optcorr = ifft(product);

display('%----- Part 4c')
plot(tdata,optcorr)
title('(4c) Output of OPTIMAL matched filtering in freq domain')

%-- The optimal correlation is only slightly different from the non-optimal
%-- correlation we got before.  That's because the template spans a limited
%-- frequency range, ~35 to ~250 Hz, and the noise power spectrum does not
%-- vary so much over that frequency range.  So the optimal SNR is only ~14%
%-- better than what we got from non-optimal filtering.

display('%----- Part 4d')
%-- Also can estimate the SNR
noise_est = sqrt(mean(optcorr(1:50000).^2));
optimal_SNR = max(optcorr) / noise_est
%-- Compare to non-optimal SNR that we got in part 3
vs_nonoptimal_SNR

input('(4d) Optimal-filtered SNR is ~14% higher than non-optimal; Press RETURN to continue');

display('%----- Part 4e')

%-- Estimating the PSD from the data is a little subtle.  We need to end up
%-- with a frequency series that has the same resolution as the Fourier-
%-- tranformed data, because we're going to multiply those arrays together
%-- element-by-element.  We can get that by simply using the periodogram of
%-- the data.  However, that isn't a great estimate of the PSD because of the
%-- large variance of the PSD in any given frequency bin.  So a somewhat better
%-- thing to do is to smooth out the periodogram.  We'll try it both ways.

%-- First, use the periodogram as the PSD estimate
periodo = abs(fftdata).^2;
%-- Combine the data, template, and noise PSD
product = fftdata .* ffttemplate' ./ periodo;
pcorr = ifft(product);
noise_est = sqrt(mean(pcorr(1:50000).^2));
periodo_SNR = max(pcorr) / noise_est

%-- Now, smooth the periodogram and use that
speriodo = smooth(periodo)';
%-- Combine the data, template, and noise PSD
product = fftdata .* ffttemplate' ./ speriodo;
%-- Smoothing breaks the symmetry of the frequency-domain representation,
%-- so the inverse Fourier transform has a small imaginary part
spcorr = real(ifft(product));
noise_est = sqrt(mean(spcorr(1:50000).^2));
smoothed_periodo_SNR = max(spcorr) / noise_est

%-- Also plot the analytic PSD and the smoothed periodogram overlaid
fvals=(0:length(speriodo)-1)/(length(speriodo)/16384);
kmx=8000;
semilogy(fvals(2:kmx),speriodo(2:kmx)/(16*16384),fvals(2:kmx),noisepsd(2:kmx));
title('(4e) Periodogram vs. analytic PSD, showing signal contribution from 40-220 Hz');

display('(4e) Plot shows bump from ~40-220 Hz due to signal in the data');
