% Script to calculate figures 3, 4 and 5 in Peltola et al. (2021): "The 
% high frequency response correction of eddy covariance fluxes. Part 1: 
% an experimental approach and its interdependence with the time-lag 
% estimation" submitted to Atmospheric Measurement Techniques.
%
% O. Peltola (09/2020)
% olli.peltola@fmi.fi
%
% You should be able to run the code by just modifying filein and
% figfilepath to correspond to locations on your computer.
%
% The accompanying file (ECdata.csv) contains 2 hours of EC data measured
% at the Siikaneva site. The data have been already despiked and
% coordinates have been already rotated so that u component is directed
% towards mean wind and mean(w)=mean(v)=0. See more details in the 
% submitted manuscript.
%
% O. Peltola (09/2020)
% olli.peltola@fmi.fi

close all
clear all
fclose('all');
clc

% setting figure style
set(0,'defaultaxesfontsize',12,'defaultaxeslinewidth',1.5,...
   'defaultlinelinewidth',1,'defaultpatchlinewidth',1.5,...
   'defaulttextfontsize',12,'defaultlinemarkersize',3);

%% file locations
% file containing the 10 Hz EC data
filein = 'C:\ECdata.csv';
% folder for the saved figures
figfilepath = 'C:\figs\';

filein = 'C:\LocalData\OPPELTOL\documents\codes\lag_and_LPF\Peltola_et_al_2020_AMT\ECdata.csv';
figfilepath = 'C:\LocalData\OPPELTOL\documents\codes\lag_and_LPF\Peltola_et_al_2020_AMT\';
%% initialising settings


% select which figures to plot (plotswitch=1 => figures 3 & 4, plotswitch=0 => figure 5)
plotswitch = 0;

% response times
if plotswitch
    tauvec = [0.1 0.3 0.7];
else
    tauvec = [0.05:0.02:1];
end

% scalar column
% sccol = 5;% CO2
% sccol = 6;% H2O
sccol = 4;% T

% temporal resolution of data
dt = 0.1;

% amount of bins in which the spectra/cospectra will be averaged
nspecbins = 82;
nspecbeg = 8;

% save figures?
save_figs = true;


% colors for the lines in fig 3
clrs = lines(size(tauvec,2));


fmax = [5e-2];
%% loading in the data

% starting to load
fid = fopen(filein,'r');
tmp = fscanf(fid,'%f,%f,%f,%f,%f,%f\n');
fclose(fid);
dat = reshape(tmp,6,[]);
dat = dat';

u = dat(:,1);
v = dat(:,2);
w = dat(:,3);
c = dat(:,sccol);

% amount of time steps
N = size(c,1);


%% starting to process

% unattenuated covariance
wc = mean((c-mean(c)).*(w-mean(w)));

% negative flux => multiplying c with -1 to get figures correct
if wc<0
    c = -c;
end

% storing the original time series
corig = c;
worig = w;

lgsout = [];
lgsout2 = [];
CFout = [];
CFestout = [];
tauestout = [];
for i=1:size(tauvec,2)
    
    tau = tauvec(i);
    disp(['Using ' num2str(tau) ' s'])
    
    % low pass filtering
    cfilt = lowpass_first_order(corig,dt,tau);

    m  = 2^(nextpow2(N)-1);
    c = corig(N-m-20:N-1);
    w = worig(N-m-20:N-1);
    cfilt = cfilt(N-m-20:N-1);
    

    % fluctuations around the mean
    cfiltp = cfilt-mean(cfilt);
    wp = w-mean(w);
    cp = c-mean(c);
    
    % unattenuated covariance
    wc = mean(wp.*cp);
    %% cross covariance
    [cc,lags] = xcorr(cp,wp,'unbiased');
    [ccf,lags] = xcorr(cfiltp,wp,'unbiased');
    ccf = ccf(abs(lags)<500);
    cc = cc(abs(lags)<500);
    lags = lags(abs(lags)<500);
    f = find(abs(ccf)==max(abs(ccf)));

    % time lag between attenuated scalar and w
    lgs = lags(f);
    lgsout = [lgsout;lgs];
    
    %% Fourier transform, power spectra, cospectra
    
    
    
    % cospectra and power spectra related to unattenuated scalar
    res = calculate_spectra(wp(1+lgs:end),cp(1+lgs:end),m,1/dt);    
    Co1 = res{1}(:,3);
    Po1 = res{1}(:,2);
    freq1 = res{1}(:,1);
    
    % cospectra and power spectra related to attenuated scalar, no lag
    res = calculate_spectra(wp(1+lgs:end),cfiltp(1+lgs:end),m,1/dt);    
    Cofilt1 = res{1}(:,3);
    Pofilt1 = res{1}(:,2);
    
    % cospectra and power spectra related to attenuated scalar, lag
    res = calculate_spectra(wp,cfiltp,m,1/dt,lgs);    
    Cofilt12 = res{1}(:,3);
    Pofilt12 = res{1}(:,2);
    
    % frequency interval
    df = freq1(2)-freq1(1);
    
    %% bin-averaging to logarithmically evenly spaced frequency bins
    % cospectra
    [Co,freq] = binavespec(freq1,Co1,nspecbins,nspecbeg);
    [Cofilt,freq] = binavespec(freq1,Cofilt1,nspecbins,nspecbeg);
    [Cofilt2,freq] = binavespec(freq1,Cofilt12,nspecbins,nspecbeg);

    % power spectra
    [Po,freq] = binavespec(freq1,Po1,nspecbins,nspecbeg);
    [Pofilt,freq] = binavespec(freq1,Pofilt1,nspecbins,nspecbeg);
    [Pofilt2,freq] = binavespec(freq1,Pofilt12,nspecbins,nspecbeg);

    % frequency interval for bin-averaged spectra
    df2 = diff(freq);
    df2 = [df2(end);df2];
    %% reference values for correction factors
    % covariances calculated from the integral of cospectra
    covsspe = [sum(df.*Co1) sum(df.*Cofilt1) sum(df.*Cofilt12)];
    
    % correction factors calculated from the integrals
    CF = 1./[sum(df.*Cofilt1)./sum(df.*Co1) sum(df.*Cofilt12)./sum(df.*Co1) sum(df.*Pofilt1)./sum(df.*Po1)];

    CFout = [CFout;CF];

    %% fitting transfer functions
    % unattenuated frequency range
    fb = find(freq<fmax);
    
    % empirical transfer function from cospectra, where lag was accounted
    % for
    TRF = Cofilt2./Co.*(trapz(freq(fb),Co(fb))./trapz(freq(fb),Cofilt2(fb)));
    % finding realistic values of TRF
    j = find(TRF<1.2 & TRF>-0.5 & freq>1e-2);

    %estimating the response time and LPF induced lag from a fit
    s = fitoptions('Method','NonlinearLeastSquares',...
                   'Lower',[0 0 0.9],...
                   'Upper',[7 5 1.1],...
                   'Startpoint',[0.18 0.2 1],...
                   'Weights',abs(df2(j).*Co(j)));
    % fit type = Fn*Hp*H           
    fitt = fittype('c*(cos(-2*pi*x*b)-2*pi*x*a*sin(-2*pi*x*b))./(1+(2*pi*x*a).^2)','options',s);
    if size(j,2)>=0.5*size(TRF,2)
        [co,gof] = fit(freq(j),TRF(j),fitt);
        
        % LPF induced lag
        lgsout2 = [lgsout2;co.b];
        % response time
        Tc = co.a;
    else
        % not enough data => NaN
        Tc = NaN;
        lgsout2 = [lgsout2;NaN];
    end
    

    
    % empirical transfer function from cospectra, where lag was accounted
    % for
    TRF2 = Cofilt2./Co.*(trapz(freq(fb),Co(fb))./trapz(freq(fb),Cofilt2(fb)));
    % finding realistic values of TRF
    j = find(TRF2<1.2 & TRF2>-0.5 & freq>1e-2);

    %estimating the response time from a fit
    s = fitoptions('Method','NonlinearLeastSquares',...
                   'Lower',[0 0.9],...
                   'Upper',[7 1.1],...
                   'Startpoint',[0.18 1],...
                   'Weights',abs(df2(j).*Co(j)));
    % fit type = Fn*H           
    fitt = fittype('c./(1+(2*pi*x*a).^2)','options',s);
    if size(j,2)>=0.5*size(TRF2,2)
        [co2,gof] = fit(freq(j),TRF2(j),fitt);
        % response time
        Tcf = co2.a;
    else
        % not enough data => NaN
        Tcf = NaN;
    end     


    % empirical transfer function from power spectra
    TRF3 = Pofilt./Po.*(trapz(freq(fb),Po(fb))./trapz(freq(fb),Pofilt(fb)));
    % finding realistic values of TRF
    j = find(TRF3<1.2 & TRF3>-0.5 & freq>1e-2);

    %estimating the response time from a fit
    s = fitoptions('Method','NonlinearLeastSquares',...
                   'Lower',[0 0.9],...
                   'Upper',[7 1.1],...
                   'Startpoint',[0.18 1],...
                   'Weights',abs(df2(j).*Co(j)));
    % fit type = Fn*H           
    fitt = fittype('c./(1+(2*pi*x*a).^2)','options',s);
    if size(j,2)>=0.5*size(TRF3,2)
        [co3,gof] = fit(freq(j),TRF3(j),fitt);
        % response time
        Tcf2 = co3.a;
    else
        % not enough data => NaN
        Tcf2 = NaN;
    end    
    
    % estimated response times
    tauestout = [tauestout;[tau Tc Tcf Tcf2]];
    
    
    % setting normalisation factor to 1 for correction factor calculations
    warning off
    co.c=1;co2.c=1;co3.c=1;
    warning on
    
    %% correction factors estimated with transfer functions
    % these are in the same order as in the paper
    % Method 1) Sabbatini et al. (2018)
    % Method 2) Aubinet et al. (2000)
    % Method 3) Fratini et al. (2012)
    % Method 4) This study
    % correction factors estimated with the transfer functions
    CFest = sum(df.*Co1)./[sum(df.*Co1.*co3(freq1)) sum(df.*Co1.*co2(freq1)) ...
        sum(df.*Co1.*(co3(freq1)).^0.5) sum(df.*Co1.*co(freq1))];
    CFestout = [CFestout;CFest];

    %% plotting
    
    % transfer functions
    figure(999)
    h1 = semilogx(freq,TRF,'ko');
    hold on
    h2 = semilogx(freq,TRF3,'k^');
    h3 = semilogx(freq,co3(freq),'r-');
    h4 = semilogx(freq,co2(freq),'b-');
    h6 = semilogx(freq,(co3(freq)).^0.5,'r--');
    h5 = semilogx(freq,co(freq),'k-');
    xlim([1e-2 max(freq)])
    ylim([-0.5 1.2])                
    legend([h1 h2 h3 h4 h5 h6],{'H_{e,cs}' 'H_{e,ps}' sprintf('H fitted to H_{e,ps}\n\\tau=%.2f s',co3.a) ...
            sprintf('H fitted to H_{e,cs}\n\\tau=%.2f s',co2.a) ...
            sprintf('HH_{p} fitted to H_{e,cs}\n\\tau=%.2f s & \\Deltat=%.2f s',co.a,co.b) 'H^{0.5}, \tau from H_{e,ps}'},'location','SW','box','off')
    title(['\tau=' num2str(tau) ' s'])
    hold off
    xlabel('Frequency (Hz)')
    ylabel('Transfer function (-)')
    
    
    % cospectra
    figure(998)
    h1 = semilogx(freq,freq.*Co./wc,'k-');
    hold on
    h2 = semilogx(freq,freq.*Cofilt2./wc,'ko');
    h3 = semilogx(freq,freq.*Co.*co3(freq)./wc,'r-');
    h4 = semilogx(freq,freq.*Co.*co2(freq)./wc,'b-');
    h5 = semilogx(freq,freq.*Co.*(co3(freq)).^0.5./wc,'r--');
    h6 = loglog(freq,freq.*Co.*co(freq)./wc,'k-');
    xlim([1e-2 max(freq)])
    ylim([-0.03 0.4])                
    legend([h1 h2 h3 h4 h5 h6],{'fCo/cov' 'fCo/cov, T attenuated' 'HfCo/cov, \tau from ps' ...
            'HfCo/cov, \tau from cs' 'H^{0.5}fCo/cov, \tau from ps' ...
            'HH_pfCo/cov, \tau & t_{lpf} from cs'},'location','NE','box','off')
    title(['\tau=' num2str(tau) ' s'])
    hold off
    xlabel('Frequency (Hz)')
    ylabel('fCo/cov (-)')
    
    % relative bias in ogive
    figure(997)
    semilogx(freq1(end:-1:1),(cumsum(df.*Cofilt12(end:-1:1))-cumsum(df.*Co1(end:-1:1).*co(freq1(end:-1:1))))./sum(df.*Cofilt12(end:-1:1)),'k-')
    hold on
    semilogx(freq1(end:-1:1),(cumsum(df.*Cofilt12(end:-1:1))-cumsum(df.*Co1(end:-1:1).*co2(freq1(end:-1:1))))./sum(df.*Cofilt12(end:-1:1)),'b-')
    semilogx(freq1(end:-1:1),(cumsum(df.*Cofilt12(end:-1:1))-cumsum(df.*Co1(end:-1:1).*co3(freq1(end:-1:1))))./sum(df.*Cofilt12(end:-1:1)),'r-')
    semilogx(freq1(end:-1:1),(cumsum(df.*Cofilt12(end:-1:1))-cumsum(df.*Co1(end:-1:1).*co3(freq1(end:-1:1)).^0.5))./sum(df.*Cofilt12(end:-1:1)),'r--')
    semilogx([freq1(1) freq1(end)],[0 0],'k--')
    hold off
    ylabel('Relative bias in ogive (-)')
    title(['\tau=' num2str(tau) ' s'])
    xlabel('Frequency (Hz)')
    xlim([1e-3 max(freq1)])
    ylim([-0.05 0.1])
    hold off
    
    if plotswitch
        % plotting figure 3
        xlm = [-3 15];

        % cross covariances
        figure(3)       
        if i==1
            hndls = [];
            txtcc = {'unattenuated'};
            h = plot(lags*dt,cc./(max(cc)),'k-');
            hndls = [hndls;h];
            hold on
            plot([0 0],[0 2],'--','color',[0.6 0.6 0.6])
            plot(lags(cc==max(cc))*dt,cc(cc==max(cc))./(max(cc)),'ko','markerfacecolor','w','markersize',5)
        end
        h = plot(lags*dt,ccf./(max(cc)),'-','color',clrs(i,:));
        plot(xlm*dt,1./CFest(1).*[1 1],'-.','color',clrs(i,:))
        plot(xlm*dt,1./CFest(4).*[1 1],'--','color',clrs(i,:))
        plot(lags(ccf==max(ccf))*dt,ccf(ccf==max(ccf))./(max(cc)),'o','color',clrs(i,:),'markerfacecolor','w','markersize',5)
        hndls = [hndls;h];
        txtcc = [txtcc;['\tau=' num2str(tau) ' s']];
        xlim(xlm*dt)
        set(gca,'xtick',[-0.3:0.3:1.5])
        ylim([0.2 1.05])
        ylim([min(ccf(lags>xlm(1) & lags<xlm(2))./(max(cc))) 1.02])
        xlabel('time lag (s)')
        ylabel('Normalised cross-covariance (-)')
        if i==size(tauvec,2)
            legend(hndls,txtcc,'location','ne')
        end
        
        % calculating the transfer functions for the figure
        H = 1./(1+(2*pi*freq.*tau).^2);
        sqrtH = H.^0.5;
        HpH = (cos(-2*pi*freq*lgs*dt)-2*pi*freq*tau.*sin(-2*pi*freq*lgs*dt)).*H;
        % plotting figure 4
        Conorm = Co./trapz(freq(fb),Co(fb));
        figure(4)
        if i==1
            % creating axes
            marg_h = [0.17 0.1];
            marg_w = [0.1 0.1];
            gap = [0.01 0.01];
            Nh = 1;Nw = size(tauvec,2);
            axh = (1-sum(marg_h)-(Nh-1)*gap(1))/Nh; 
            axw = (1-sum(marg_w)-(Nw-1)*gap(2))/Nw;
            py = 1-marg_h(2)-axh; 
            px = marg_w(1);
            ha = zeros(size(tauvec,2),1);
            for iax=1:size(tauvec,2)
                ha(iax) = axes('Units','normalized', ...
                'Position',[px py axw axh], ...
                'XTickLabel','', ...
                'YTickLabel','');
                px = px+axw+gap(2);
            end
        end
        ylms = [[-0.2 1.2];[-0.1333 0.8]];
        p1 = polyfit(ylms(2,:)',ylms(1,:)',1);
        p2 = polyfit(ylms(1,:)',ylms(2,:)',1);
        axes(ha(i))
        ax = gca;
        yyaxis(ax,'right')
        semilogx(freq,TRF,'ko','markerfacecolor','w')
        hold on
        semilogx(freq,TRF3,'k^','markerfacecolor','w')
        semilogx(freq,H,'r-')
        semilogx(freq,sqrtH,'b-')
        semilogx(freq,HpH,'k-')
        semilogx(freq([1 size(freq,1)]),[0 0],'k--')
        yyaxis(ax,'left')
        semilogx(freq,freq.*Conorm,'-','color',[0.9 0.9 0.9],'linewidth',3);
        ax.YAxis(1).Color = ax.XAxis.Color;
        ax.YAxis(2).Color = ax.XAxis.Color;
        yyaxis(ax,'right')
        ylim(ylms(1,:))
        xlim([1e-2 5])
        set(gca,'XtickLabelMode','auto')
        set(gca,'ytick',polyval(p1,[0:0.1:0.8]))
        if i==size(tauvec,2)
            set(gca,'yticklabel',[0:0.1:0.8])
            ylabel('f*Co^{w,T}/cov(w,T) (-)')
        else
            set(gca,'yticklabel',[])
        end
        xlabel('Frequency (Hz)')
        yyaxis(ax,'left')
        ylim(ylms(2,:))
        xlim([1e-2 5])
        set(gca,'ytick',polyval(p2,[-0.2:0.2:1.2]))
        if i==1
            set(gca,'yticklabel',[-0.2:0.2:1.2])
            ylabel('Transfer function (-)')
        else
            set(gca,'yticklabel',[])
        end
        title(['\tau=' num2str(tau) ' s, t_{lpf}=' num2str(lgsout(end)*dt) ' s'],'fontweight','normal')
        if i==1
            ylabel('Transfer function')
            legend({'f*Co^{w,T}/cov(w,T)' 'H_{e,cs}' 'H_{e,ps}' 'H' 'H^{0.5}' 'H*H_p'},'location','sw','fontsize',8)
        end
        xlabel('Frequency (Hz)')
    end

end

if plotswitch==0
    % plotting figure 5
    figure(5)
    subplot(1,3,1)
    h1 = plot(tauvec,CFout(:,2),'k-');
    hold on
%     h2 = plot(tauvec,1./AFout(:,1),'--','color',[0.5 0.5 0.5]);
    h3 = plot(tauvec,CFestout(:,1),'ko');
    h4 = plot(tauvec,CFestout(:,2),'bo');
    h5 = plot(tauvec,CFestout(:,3),'k*');
    h6 = plot(tauvec,CFestout(:,4),'ro');
    ylim([1 max(CFestout(:))])
    xlim([0 max(tauvec)])
    xlabel('Response time, \tau (s)')
    ylabel('Correction factor (-)')
    legend([h1 h3 h4 h5 h6],{'reference' 'Method 1' 'Method 2' 'Method 3' 'Method 4'},'location','NW')   
    text(max(tauvec)-(max(tauvec))*0.12,1+(max(CFestout(:))-1)*0.1,'a)')
    
    subplot(1,3,2)
    h1 = plot([0 max(tauvec)],[0 max(tauvec)],'k--');
    hold on
    h2 = plot(tauvec,tauestout(:,2),'ro');
    h3 = plot(tauvec,tauestout(:,3),'bo');
    h4 = plot(tauvec,tauestout(:,4),'ko');
    xlim([0 max(tauvec)])
    ylim([0 max(tauvec)])
    xlabel('Response time, \tau (s)')
    ylabel('Estimated response time (s)')
%     legend([h1 h2 h3 h4],{'1:1 line' 'HH_{p} fit to H_{e,cs}' 'H fit to H_{e,cs}' 'H fit to H_{e,ps}'},'location','NW')
    legend([h1 h4 h3 h2],{'1:1 line' 'Methods 1 & 3' 'Method 2' 'Method 4'},'location','NW')
    text(max(tauvec)-(max(tauvec)-min(tauvec))*0.12,(max(tauvec))*0.1,'b)')
    
    subplot(1,3,3)
    plot(tauvec,(lgsout)./10,'ko')
    hold on
    plot(tauvec,lgsout2,'ro')
    xlim([min(tauvec) max(tauvec)])
    ylim([0 max([(lgsout(:))./10;lgsout2(:)])])
    ylim([0 0.6])
    xlabel('Response time, \tau (s)')
    ylabel('Low-pass filtering induced time lag, t_{lpf} (s)')
    legend({sprintf('Cross-covariance\nmaximisation') 'Method 4'},'location','NW')
    text(max(tauvec)-(max(tauvec)-min(tauvec))*0.12,(0.6-0)*0.1,'c)')


    
end
%% saving figs
if save_figs

    if plotswitch
        figure(4)
        set(gcf,'units','centimeters')
%         set(gcf,'paperposition',[0 0 12 29.7/8]*3)% [left bottom width height]
        set(gcf,'paperposition',[0 0 21 29.7/4])% [left bottom width height]        
    %         set(gcf,'paperposition',[0 0 21 29.7/3])% [left bottom width height]
        print(gcf,'-dpng','-r300',[figfilepath 'Figure_4_Peltola_et_al_2021_AMT']);
        print(gcf,'-depsc',[figfilepath 'Figure_4_Peltola_et_al_2021_AMT']);
        saveas(gcf,[figfilepath 'Figure_4_Peltola_et_al_2021_AMT.fig']);
        
        figure(3)
        set(gcf,'units','centimeters')
        set(gcf,'paperposition',[0 0 21/2 29.7/3])% [left bottom width height]
    %         set(gcf,'paperposition',[0 0 21 29.7/3])% [left bottom width height]
        print(gcf,'-dpng','-r300',[figfilepath 'Figure_3_Peltola_et_al_2021_AMT']);
        print(gcf,'-depsc',[figfilepath 'Figure_3_Peltola_et_al_2021_AMT']);
        saveas(gcf,[figfilepath 'Figure_3_Peltola_et_al_2021_AMT.fig']);
    else
        figure(5)
        set(gcf,'units','centimeters')
        set(gcf,'paperposition',[0 0 12 29.7/8]*3)% [left bottom width height]
    %         set(gcf,'paperposition',[0 0 21 29.7/3])% [left bottom width height]
        print(gcf,'-dpng','-r300',[figfilepath 'Figure_5_Peltola_et_al_2021_AMT']);
        print(gcf,'-depsc',[figfilepath 'Figure_5_Peltola_et_al_2021_AMT']);
        saveas(gcf,[figfilepath 'Figure_5_Peltola_et_al_2021_AMT.fig']);
    end
    
end
