%WAVETEST Example Matlab script for WAVELET, using NINO3 SST dataset
%
% See "http://paos.colorado.edu/research/wavelets/"
% Written January 1998 by C. Torrence
%
% Modified Oct 1999, changed Global Wavelet Spectrum (GWS) to be sideways,
%   changed all "log" to "log2", changed logarithmic axis on GWS to
%   a normal axis.

% SST data is three month seasonal averages
%load 'sst_nino3.dat'   % input SST time series
%sst = sst_nino3;

%addpath('/home/nsteiger/MATLAB/wavelet_analysis/')

load distinct_colors_20.mat

%=====================================
% LOAD ENSEMBLE MEGADROUGHT ESTIMATES

% Load both south american and north america drought indices
load('chile_megaD_ens.mat','svmd','ens_drght','Xa_di','Xa_di_ens','analyrs','dsens')

xa_di_sa=Xa_di;
xa_di_ens_sa=Xa_di_ens;
ds_sa=svmd;
mdagr_sa=ens_drght;
sa_yrs=analyrs;

load('nasw_megaD_ens.mat','svmd','ens_drght','Xa_di','Xa_di_ens','analyrs','dsens')
xa_di_na=Xa_di;
xa_di_ens_na=Xa_di_ens;
ds_na=svmd;
mdagr_na=ens_drght;
na_yrs=analyrs;


%-------------------
% WAVELET COHERENCE
%-------------------

addpath('/home/nsteiger/Dropbox/code/wavelet-coherence/')
addpath('/home/nsteiger/Dropbox/code/wavelet-coherence/faq/')

%% Wavelet coherence (WTC)
% The WTC finds regions in time frequency space where the two
% time series co-vary (but does not necessarily have high power).

% format of the time series...
dn(:,1)=na_yrs;
dn(:,2)=xa_di_na(na_yrs);
ds(:,1)=sa_yrs;
ds(:,2)=xa_di_sa(sa_yrs);



% PLOT TIME SERIES ---------------------------------------------------------------
% position = [left bottom width height]
subplot('position',[0.1 0.6 0.65 0.35])
plot(dn(:,1),dn(:,2),'color',clrs.blue,'linewidth',1.5)
hold on
plot(ds(:,1),ds(:,2),'color',clrs.purple,'linewidth',1.5)
legend('NASW','Chile','location','northwest')
set(gca,'XLim',[na_yrs(1) na_yrs(end)],'YLim',[-6 6])
xlabel('Year','fontsize',16)
ylabel('PDSI','fontsize',16)
set(gca,'FontSize',16)
%title('NINO3.4 Sea Surface Temperature (seasonal)')
hold off


% Compute the wavelet coherence
[Rsq,period,scale,coi,wtcsig,t,dt,Wxy,aad,aas,aah]=wtc(dn,ds);


% PLOT WAVELET COHERENCE ---------------------------------------------------------
% position = [left bottom width height]
subplot('position',[0.1 0.1 0.65 0.4])
Yticks = 2.^(fix(log2(min(period))):fix(log2(max(period))));
H=imagesc(t,log2(period),Rsq);%#ok
%[c,H]=safecontourf(t,log2(period),Rsq,[0:.05:1]);
%set(H,'linestyle','none')
set(gca,'clim',[0 1])
%HCB=colorbar;%#ok
set(gca,'YLim',log2([min(period),max(period)]), ...
'YDir','reverse', 'layer','top', ...
'YTick',log2(Yticks(:)), ...
'YTickLabel',num2str(Yticks'), ...
'layer','top')
ylabel('Period')
hold on
%phase plot
aWxy=angle(Wxy);
aaa=aWxy;
aaa(Rsq<.5)=NaN; %remove phase indication where Rsq is low
%[xx,yy]=meshgrid(t(1:5:end),log2(period));
phs_dt=round(length(t)/aad(1)); tidx=max(floor(phs_dt/2),1):phs_dt:length(t);
phs_dp=round(length(period)/aad(2)); pidx=max(floor(phs_dp/2),1):phs_dp:length(period);
phaseplot(t(tidx),log2(period(pidx)),aaa(pidx,tidx),aas,aah);
if ~all(isnan(wtcsig))
   [c,h] = contour(t,log2(period),wtcsig,[1 1],'k');%#ok
   set(h,'linewidth',2)
end
%suptitle([sTitle ' coherence']);
tt=[t([1 1])-dt*.5;t;t([end end])+dt*.5];
hcoi=fill(tt,log2([period([end 1]) coi period([1 end])]),'w');
set(hcoi,'alphadatamapping','direct','facealpha',.5)
hold off
%title(['Wavelet coherence: NASW-Chile (PDSI)'] )
xlabel('Year','fontsize',16)
ylabel('Period','fontsize',16)
set(gca,'FontSize',16)


% PLOT GLOBAL COHERENCE ------------------------------------------------------------

% compute global coherence
[gwss periods ptile] = globalcoher(dn,ds);

% position = [left bottom width height]
subplot('position',[0.77 0.1 0.2 0.4])
plot(gwss,log2(periods),'color',clrs.red,'linewidth',2)
hold on
plot(ptile,log2(periods),'--','color',clrs.grey,'linewidth',2)
hold off
legend('Coh.','AR-null')
xlabel('Coherence','fontsize',16)
%title('Global Wavelet Spectrum')
set(gca,'YLim',log2([min(periods),max(periods)]), ...
	'YDir','reverse', ...
	'YTick',log2(Yticks(:)), ...
	'YTickLabel','')
set(gca,'XLim',[0,1.25*max(gwss)])
set(gca,'FontSize',16)


fig = gcf;fig.PaperUnits = 'inches';
fig.PaperPosition = [0 2 10 7.5]; % [left bottom width height]

%print(['/home/nsteiger/Dropbox/sa-phyda/waveletcoh_nasw_chile_pdsi.png'],'-dpng','-r300');


% GLOBAL COHERENCE BY ITSELF

% position = [left bottom width height]
figure
plot(log2(periods),gwss,'color',clrs.black,'linewidth',2)
hold on
plot(log2(periods),ptile,'--','color',clrs.grey,'linewidth',2)
hold off
legend('Coh(NASW,SASW)','AR-null','location','northwest')
ylabel('Coherence','fontsize',18)
xlabel('Period (Years)','fontsize',18)
%title('Global Wavelet Spectrum')
%set(gca,'XLim',log2([min(periods),max(periods)]), ...
set(gca,'XLim',log2([min(periods),max(periods(periods<180))]), ...
	'XTick',log2(Yticks(:)), ...
        'XTickLabel',num2str(Yticks'))
xtickangle(45)
set(gca,'YLim',[0,1.1*max(gwss)])
set(gca,'FontSize',16)


fig = gcf;fig.PaperUnits = 'inches';
fig.PaperPosition = [0 2 3.7 4.5]; % [left bottom width height]

print(['/home/nsteiger/Dropbox/coh_nasw_chile_pdsi.png'],'-dpng','-r300');




return



%% Cross wavelet transform (XWT)
% The XWT finds regions in time frequency space where
% the time series show high common power.

figure
xwt(dn,ds)
title(['XWT: NASW-Chile (PDSI)'] )


return


% SHOW EACH WAVELET TOGETHER IN A SINGLE PLOT


%------------------------------------------------------ Computation

% NORTH AMERICA

xyrs=na_yrs;
Xd=xa_di_na(na_yrs);

% normalize by standard deviation (not necessary, but makes it easier
% to compare with plot on Interactive Wavelet page, at
% "http://paos.colorado.edu/research/wavelets/plot/"
variance = std(Xd)^2;
Xd = (Xd - mean(Xd))/sqrt(variance) ;


n = length(Xd);
dt = 1 ;
time = [0:length(Xd)-1]*dt + xyrs(1)+1 ;  % construct time array
xlim = [xyrs(1),xyrs(end)];  % plotting range
pad = 1;      % pad the time series with zeroes (recommended)
dj = 0.25;    % this will do 4 sub-octaves per octave
s0 = 2*dt;    % this says start at a scale of 2*dt 
j1 = 7/dj;    % this says do 7 powers-of-two with dj sub-octaves each
atc=autocorr(Xd);
lag1 = atc(2);  % lag-1 autocorrelation for red noise background
mother = 'Morlet';

% Wavelet transform:
[wave,period,scale,coi] = wavelet(Xd,dt,pad,dj,s0,j1,mother);
power = (abs(wave)).^2 ;        % compute wavelet power spectrum

% Significance levels: (variance=1 for the normalized SST)
[signif,fft_theor] = wave_signif(1.0,dt,scale,0,lag1,-1,-1,mother);
sig95 = (signif')*(ones(1,n));  % expand signif --> (J+1)x(N) array
sig95 = power ./ sig95;         % where ratio > 1, power is significant

% Global wavelet spectrum & significance levels:
global_ws = variance*(sum(power')/n);   % time-average over all times
dof = n - scale;  % the -scale corrects for padding at edges
global_signif = wave_signif(variance,dt,scale,1,lag1,-1,dof,mother);

% Save variable specific variables
periodn=period;
powern=power;
sig95n=sig95;
timen=time;
coin=coi;
global_wsn=global_ws;
global_signifn=global_signif;


% SOUTH AMERICA

xyrs=sa_yrs;
Xd=xa_di_sa(sa_yrs);

% normalize by standard deviation (not necessary, but makes it easier
% to compare with plot on Interactive Wavelet page, at
% "http://paos.colorado.edu/research/wavelets/plot/"
variance = std(Xd)^2;
Xd = (Xd - mean(Xd))/sqrt(variance) ;


n = length(Xd);
dt = 1 ;
time = [0:length(Xd)-1]*dt + xyrs(1)+1 ;  % construct time array
xlim = [xyrs(1),xyrs(end)];  % plotting range
pad = 1;      % pad the time series with zeroes (recommended)
dj = 0.25;    % this will do 4 sub-octaves per octave
s0 = 2*dt;    % this says start at a scale of 2*dt 
j1 = 7/dj;    % this says do 7 powers-of-two with dj sub-octaves each
atc=autocorr(Xd);
lag1 = atc(2);  % lag-1 autocorrelation for red noise background
mother = 'Morlet';

% Wavelet transform:
[wave,period,scale,coi] = wavelet(Xd,dt,pad,dj,s0,j1,mother);
power = (abs(wave)).^2 ;        % compute wavelet power spectrum

% Significance levels: (variance=1 for the normalized SST)
[signif,fft_theor] = wave_signif(1.0,dt,scale,0,lag1,-1,-1,mother);
sig95 = (signif')*(ones(1,n));  % expand signif --> (J+1)x(N) array
sig95 = power ./ sig95;         % where ratio > 1, power is significant

% Global wavelet spectrum & significance levels:
global_ws = variance*(sum(power')/n);   % time-average over all times
dof = n - scale;  % the -scale corrects for padding at edges
global_signif = wave_signif(variance,dt,scale,1,lag1,-1,dof,mother);

% Save variable specific variables
periods=period;
powers=power;
sig95s=sig95;
times=time;
cois=coi;
global_wss=global_ws;
global_signifs=global_signif;



%------------------------------------------------------ Plotting

%--- Plot time series

%subplot('position',[0.1 0.75 0.65 0.2])
%subplot('position',[0.1 0.07 0.65 0.2])



%--- Contour plot wavelet power spectrum
% position = [left bottom width height]
subplot('position',[0.1 0.6 0.65 0.3])
levels = [0.0625,0.125,0.25,0.5,1,2,4,8,16] ;
Yticks = 2.^(fix(log2(min(periodn))):fix(log2(max(periodn))));
%contour(time,log2(period),log2(power),log2(levels));  %*** or use 'contourfill'
%contourf(time,log2(period),log2(power),log2(levels));  %*** or use 'contourfill'
imagesc(timen,log2(periodn),log2(powern));  %*** uncomment for 'image' plot
xlabel('Time (year)')
ylabel('Period (years)')
title('NASW PDSI Wavelet Power Spectrum')
set(gca,'XLim',xlim(:))
set(gca,'YLim',log2([min(periodn),max(periodn)]), ...
	'YDir','reverse', ...
	'YTick',log2(Yticks(:)), ...
	'YTickLabel',Yticks)
% 95% significance contour, levels at -99 (fake) and 1 (95% signif)
hold on
contour(timen,log2(periodn),sig95n,[-99,1],'k');
hold on
% cone-of-influence, anything "below" is dubious
plot(timen,log2(coin),'k')
hold off

%--- Plot global wavelet spectrum
subplot('position',[0.77 0.6 0.2 0.3])
plot(global_wsn,log2(periodn))
hold on
plot(global_signifn,log2(periodn),'--')
hold off
xlabel('Power (degC^2)')
title('Global Wavelet Spectrum')
set(gca,'YLim',log2([min(periodn),max(periodn)]), ...
	'YDir','reverse', ...
	'YTick',log2(Yticks(:)), ...
	'YTickLabel','')
set(gca,'XLim',[0,1.25*max(global_wsn)])



%--- Contour plot wavelet power spectrum
% position = [left bottom width height]
subplot('position',[0.1 0.18 0.65 0.3])
levels = [0.0625,0.125,0.25,0.5,1,2,4,8,16] ;
Yticks = 2.^(fix(log2(min(periods))):fix(log2(max(periods))));
%contour(time,log2(period),log2(power),log2(levels));  %*** or use 'contourfill'
%contourf(time,log2(period),log2(power),log2(levels));  %*** or use 'contourfill'
imagesc(times,log2(periods),log2(powers));  %*** uncomment for 'image' plot
xlabel('Time (year)')
ylabel('Period (years)')
title('Chile PDSI Wavelet Power Spectrum')
set(gca,'XLim',xlim(:))
set(gca,'YLim',log2([min(periods),max(periods)]), ...
	'YDir','reverse', ...
	'YTick',log2(Yticks(:)), ...
	'YTickLabel',Yticks)
% 95% significance contour, levels at -99 (fake) and 1 (95% signif)
hold on
contour(times,log2(periods),sig95s,[-99,1],'k');
hold on
% cone-of-influence, anything "below" is dubious
plot(times,log2(cois),'k')
hold off

%--- Plot global wavelet spectrum
subplot('position',[0.77 0.18 0.2 0.3])
plot(global_wss,log2(periods))
hold on
plot(global_signifs,log2(periods),'--')
hold off
xlabel('Power (degC^2)')
title('Global Wavelet Spectrum')
set(gca,'YLim',log2([min(periods),max(periods)]), ...
	'YDir','reverse', ...
	'YTick',log2(Yticks(:)), ...
	'YTickLabel','')
set(gca,'XLim',[0,1.25*max(global_wss)])


fig = gcf;fig.PaperUnits = 'inches';
fig.PaperPosition = [0 2 9 6.75]; % [left bottom width height]

print(['/home/nsteiger/Dropbox/sa-phyda/wavelet_nasw_chile_pdsi_' num2str(xyrs(1)) '_' num2str(xyrs(end)) '.png'],'-dpng','-r300');





return




% COMPUTE EACH WAVELET AND SUBPLOTS SEPARATELY


%------------------------------------------------------ Computation


xyrs=na_yrs;
Xd=xa_di_na(na_yrs);

% normalize by standard deviation (not necessary, but makes it easier
% to compare with plot on Interactive Wavelet page, at
% "http://paos.colorado.edu/research/wavelets/plot/"
variance = std(Xd)^2;
Xd = (Xd - mean(Xd))/sqrt(variance) ;


n = length(Xd);
dt = 1 ;
time = [0:length(Xd)-1]*dt + xyrs(1)+1 ;  % construct time array
xlim = [xyrs(1),xyrs(end)];  % plotting range
pad = 1;      % pad the time series with zeroes (recommended)
dj = 0.25;    % this will do 4 sub-octaves per octave
s0 = 2*dt;    % this says start at a scale of 2*dt 
j1 = 7/dj;    % this says do 7 powers-of-two with dj sub-octaves each
atc=autocorr(Xd);
lag1 = atc(2);  % lag-1 autocorrelation for red noise background
mother = 'Morlet';

% Wavelet transform:
[wave,period,scale,coi] = wavelet(Xd,dt,pad,dj,s0,j1,mother);
power = (abs(wave)).^2 ;        % compute wavelet power spectrum

% Significance levels: (variance=1 for the normalized SST)
[signif,fft_theor] = wave_signif(1.0,dt,scale,0,lag1,-1,-1,mother);
sig95 = (signif')*(ones(1,n));  % expand signif --> (J+1)x(N) array
sig95 = power ./ sig95;         % where ratio > 1, power is significant

% Global wavelet spectrum & significance levels:
global_ws = variance*(sum(power')/n);   % time-average over all times
dof = n - scale;  % the -scale corrects for padding at edges
global_signif = wave_signif(variance,dt,scale,1,lag1,-1,dof,mother);

% Scale-average between El Nino periods of 2--8 years
avg = find((scale >= 2) & (scale < 8));
Cdelta = 0.776;   % this is for the MORLET wavelet
scale_avg = (scale')*(ones(1,n));  % expand scale --> (J+1)x(N) array
scale_avg = power ./ scale_avg;   % [Eqn(24)]
scale_avg = variance*dj*dt/Cdelta*sum(scale_avg(avg,:));   % [Eqn(24)]
scaleavg_signif = wave_signif(variance,dt,scale,2,lag1,-1,[2,7.9],mother);


%------------------------------------------------------ Plotting

%--- Plot time series
subplot('position',[0.1 0.75 0.65 0.2])
plot(time,Xd)
set(gca,'XLim',xlim(:))
xlabel('Time (year)')
ylabel('NASW PDSI')
%title('NINO3.4 Sea Surface Temperature (seasonal)')
hold off

%--- Contour plot wavelet power spectrum
subplot('position',[0.1 0.37 0.65 0.28])
levels = [0.0625,0.125,0.25,0.5,1,2,4,8,16] ;
Yticks = 2.^(fix(log2(min(period))):fix(log2(max(period))));
%contour(time,log2(period),log2(power),log2(levels));  %*** or use 'contourfill'
%contourf(time,log2(period),log2(power),log2(levels));  %*** or use 'contourfill'
imagesc(time,log2(period),log2(power));  %*** uncomment for 'image' plot
xlabel('Time (year)')
ylabel('Period (years)')
title('NASW PDSI Wavelet Power Spectrum')
set(gca,'XLim',xlim(:))
set(gca,'YLim',log2([min(period),max(period)]), ...
	'YDir','reverse', ...
	'YTick',log2(Yticks(:)), ...
	'YTickLabel',Yticks)
% 95% significance contour, levels at -99 (fake) and 1 (95% signif)
hold on
contour(time,log2(period),sig95,[-99,1],'k');
hold on
% cone-of-influence, anything "below" is dubious
plot(time,log2(coi),'k')
hold off

%--- Plot global wavelet spectrum
subplot('position',[0.77 0.37 0.2 0.28])
plot(global_ws,log2(period))
hold on
plot(global_signif,log2(period),'--')
hold off
xlabel('Power (degC^2)')
title('Global Wavelet Spectrum')
set(gca,'YLim',log2([min(period),max(period)]), ...
	'YDir','reverse', ...
	'YTick',log2(Yticks(:)), ...
	'YTickLabel','')
set(gca,'XLim',[0,1.25*max(global_ws)])

%--- Plot 2--8 yr scale-average time series
subplot('position',[0.1 0.07 0.65 0.2])
plot(time,scale_avg)
set(gca,'XLim',xlim(:))
xlabel('Time (year)')
ylabel('Avg variance (degC^2)')
title('2-8 yr Scale-average Time Series')
hold on
plot(xlim,scaleavg_signif+[0,0],'--')
hold off

fig = gcf;fig.PaperUnits = 'inches';
fig.PaperPosition = [0 2 9 6.75]; % [left bottom width height]

%print(['/home/nsteiger/Dropbox/sa-phyda/wavelet_subs_naswpdsi_' num2str(xyrs(1)) '_' num2str(xyrs(end)) '.png'],'-dpng','-r300');



%------------------------------------------------------ Computation


xyrs=sa_yrs;
Xd=xa_di_sa(sa_yrs);

% normalize by standard deviation (not necessary, but makes it easier
% to compare with plot on Interactive Wavelet page, at
% "http://paos.colorado.edu/research/wavelets/plot/"
variance = std(Xd)^2;
Xd = (Xd - mean(Xd))/sqrt(variance) ;


n = length(Xd);
dt = 1 ;
time = [0:length(Xd)-1]*dt + xyrs(1)+1 ;  % construct time array
xlim = [xyrs(1),xyrs(end)];  % plotting range
pad = 1;      % pad the time series with zeroes (recommended)
dj = 0.25;    % this will do 4 sub-octaves per octave
s0 = 2*dt;    % this says start at a scale of 2*dt 
j1 = 7/dj;    % this says do 7 powers-of-two with dj sub-octaves each
atc=autocorr(Xd);
lag1 = atc(2);  % lag-1 autocorrelation for red noise background
mother = 'Morlet';

% Wavelet transform:
[wave,period,scale,coi] = wavelet(Xd,dt,pad,dj,s0,j1,mother);
power = (abs(wave)).^2 ;        % compute wavelet power spectrum

% Significance levels: (variance=1 for the normalized SST)
[signif,fft_theor] = wave_signif(1.0,dt,scale,0,lag1,-1,-1,mother);
sig95 = (signif')*(ones(1,n));  % expand signif --> (J+1)x(N) array
sig95 = power ./ sig95;         % where ratio > 1, power is significant

% Global wavelet spectrum & significance levels:
global_ws = variance*(sum(power')/n);   % time-average over all times
dof = n - scale;  % the -scale corrects for padding at edges
global_signif = wave_signif(variance,dt,scale,1,lag1,-1,dof,mother);

% Scale-average between El Nino periods of 2--8 years
avg = find((scale >= 2) & (scale < 8));
Cdelta = 0.776;   % this is for the MORLET wavelet
scale_avg = (scale')*(ones(1,n));  % expand scale --> (J+1)x(N) array
scale_avg = power ./ scale_avg;   % [Eqn(24)]
scale_avg = variance*dj*dt/Cdelta*sum(scale_avg(avg,:));   % [Eqn(24)]
scaleavg_signif = wave_signif(variance,dt,scale,2,lag1,-1,[2,7.9],mother);


%------------------------------------------------------ Plotting

%--- Plot time series
subplot('position',[0.1 0.75 0.65 0.2])
plot(time,Xd)
set(gca,'XLim',xlim(:))
xlabel('Time (year)')
ylabel('Chile PDSI')
%title('NINO3.4 Sea Surface Temperature (seasonal)')
hold off

%--- Contour plot wavelet power spectrum
subplot('position',[0.1 0.37 0.65 0.28])
levels = [0.0625,0.125,0.25,0.5,1,2,4,8,16] ;
Yticks = 2.^(fix(log2(min(period))):fix(log2(max(period))));
%contour(time,log2(period),log2(power),log2(levels));  %*** or use 'contourfill'
%contourf(time,log2(period),log2(power),log2(levels));  %*** or use 'contourfill'
imagesc(time,log2(period),log2(power));  %*** uncomment for 'image' plot
xlabel('Time (year)')
ylabel('Period (years)')
title('Chile PDSI Wavelet Power Spectrum')
set(gca,'XLim',xlim(:))
set(gca,'YLim',log2([min(period),max(period)]), ...
	'YDir','reverse', ...
	'YTick',log2(Yticks(:)), ...
	'YTickLabel',Yticks)
% 95% significance contour, levels at -99 (fake) and 1 (95% signif)
hold on
contour(time,log2(period),sig95,[-99,1],'k');
hold on
% cone-of-influence, anything "below" is dubious
plot(time,log2(coi),'k')
hold off

%--- Plot global wavelet spectrum
subplot('position',[0.77 0.37 0.2 0.28])
plot(global_ws,log2(period))
hold on
plot(global_signif,log2(period),'--')
hold off
xlabel('Power (degC^2)')
title('Global Wavelet Spectrum')
set(gca,'YLim',log2([min(period),max(period)]), ...
	'YDir','reverse', ...
	'YTick',log2(Yticks(:)), ...
	'YTickLabel','')
set(gca,'XLim',[0,1.25*max(global_ws)])

%--- Plot 2--8 yr scale-average time series
subplot('position',[0.1 0.07 0.65 0.2])
plot(time,scale_avg)
set(gca,'XLim',xlim(:))
xlabel('Time (year)')
ylabel('Avg variance (degC^2)')
title('2-8 yr Scale-average Time Series')
hold on
plot(xlim,scaleavg_signif+[0,0],'--')
hold off

fig = gcf;fig.PaperUnits = 'inches';
fig.PaperPosition = [0 2 9 6.75]; % [left bottom width height]

%print(['/home/nsteiger/Dropbox/sa-phyda/wavelet_subs_chilepdsi_' num2str(xyrs(1)) '_' num2str(xyrs(end)) '.png'],'-dpng','-r300');














return

%--- Contour plot wavelet power spectrum
figure
levels = [0.0625,0.125,0.25,0.5,1,2,4,8,16] ;
Yticks = 2.^(fix(log2(min(period))):fix(log2(max(period))));
%contour(time,log2(period),log2(power),log2(levels));  %*** or use 'contourfill'
%contourf(time,log2(period),log2(power),log2(levels));  %*** or use 'contourfill'
imagesc(time,log2(period),log2(power));  %*** uncomment for 'image' plot
xlabel('Time (year)')
ylabel('Period (years)')
title('NINO3.4 SST Wavelet Power Spectrum')
set(gca,'XLim',xlim(:))
set(gca,'YLim',log2([min(period),max(period)]), ...
	'YDir','reverse', ...
	'YTick',log2(Yticks(:)), ...
	'YTickLabel',Yticks)
% 95% significance contour, levels at -99 (fake) and 1 (95% signif)
hold on
contour(time,log2(period),sig95,[-99,1],'k');
hold on
% cone-of-influence, anything "below" is dubious
plot(time,log2(coi),'k')
hold off

fig = gcf;fig.PaperUnits = 'inches';
fig.PaperPosition = [0 2 7 4]; % [left bottom width height]

%print(['/home/nsteiger/Dropbox/sa-phyda/wavelet_nino34_' num2str(analyrs1(1)) '_' num2str(analyrs1(end)) '.png'],'-dpng','-r300');







return



%===========================================================================
% LOAD ENSO INDICES

% Load full ensembles for the nino and amo variables
efl='cesm_lme010_r12000_p998_state_2ozn_avg_AprMar_prxydtst_5_prxtp_tca_2896_swtchbld100_05-Jan-2018_14:24:35.mat';
% LOAD JJA FOR AMO ANALYSIS
%efl='cesm_lme010_r12000_p998_state_2ozn_avg_JunAug_prxydtst_5_prxtp_tca_2896_swtchbld100_05-Jan-2018_14:11:15.mat';

pth='/d2/nsteiger/output-da/hydroclimate/subrecons_5Jan18/';
%load([pth,efl],'r_o','r_f','state_tp','xmeta','Xa_m','Xa_ens','reconYrs','p_yrs')
load([pth,efl])


% What years to look at?
analyrs1=1000:1925;
[~,~,ir1] = intersect(analyrs1,r_o:r_f); % indices of recon segment
% monthly indices
ovrlprng=rng_intr([r_o,r_f],[analyrs1(1),analyrs1(end)]);
mon_r=r_o:(1/12):(r_f+11/12);
irm1_a=find(abs(ovrlprng(1)-mon_r)<1e-6);
irm1_b=find(abs(ovrlprng(2)-1/12-mon_r)<1e-6)+12;



% k=find(state_tp=='o');
% xa_amo_ens=permute(squeeze(Xa_ens(xmeta{k}.id_X(1),:,:)),[2 1]);
% % Ensemble mean of AMO
% xa_amo_mn=mean(xa_amo_ens,2);
% amo_off=mean(mean(xa_amo_ens(ir2,:),2));


k=find(state_tp=='2');
xa_gmt_ens=permute(squeeze(Xa_ens(xmeta{k}.id_X(1),:,:)),[2 1]);

k=find(state_tp=='n'); % specific nino variable
if strcmp(xmeta{k}.indinfo{3}.indxnm,'Nino 3.4'); i=3; else; error('Wrong Nino specification'); end
xa_nino_ens=reshape(permute(Xa_ens((xmeta{k}.id_X(1)+(i-1)*12):(xmeta{k}.id_X(1)+i*12-1),:,:),[1 3 2]),12*reconYrs,length(p_yrs));
xa_nino_mn=reshape(Xa_m((xmeta{k}.id_X(1)+(i-1)*12):(xmeta{k}.id_X(1)+i*12-1),:),12*reconYrs,1);
% Nino offset with respect to reference period
nino_off=mean(mean(xa_nino_ens(irm1_a:irm1_b,:),2));

% Annual nino from monthly data (leaves year ends not quite correct)
xa_nino_ann=mon2ann(xa_nino_mn,4,3);
% Annual nino from spatial field
%xa_nino_ann=nino(Xt2m,lat,lon,3,'n');

% SEASONAL AVERAGE: MAM, JJA, SON, DJF
x0=xa_nino_mn(irm1_a:irm1_b);
sst=mean(reshape(x0(3:end-1),3,[]))'; % mean starting in March ending in Dec



%------------------------------------------------------ Computation

% normalize by standard deviation (not necessary, but makes it easier
% to compare with plot on Interactive Wavelet page, at
% "http://paos.colorado.edu/research/wavelets/plot/"
variance = std(sst)^2;
sst = (sst - mean(sst))/sqrt(variance) ;

n = length(sst);
dt = 0.25 ;
time = [0:length(sst)-1]*dt + analyrs1(1)+1 ;  % construct time array
xlim = [analyrs1(1),analyrs1(end)];  % plotting range
pad = 1;      % pad the time series with zeroes (recommended)
dj = 0.25;    % this will do 4 sub-octaves per octave
s0 = 2*dt;    % this says start at a scale of 6 months
j1 = 7/dj;    % this says do 7 powers-of-two with dj sub-octaves each
lag1 = 0.72;  % lag-1 autocorrelation for red noise background
mother = 'Morlet';

% Wavelet transform:
[wave,period,scale,coi] = wavelet(sst,dt,pad,dj,s0,j1,mother);
power = (abs(wave)).^2 ;        % compute wavelet power spectrum

% Significance levels: (variance=1 for the normalized SST)
[signif,fft_theor] = wave_signif(1.0,dt,scale,0,lag1,-1,-1,mother);
sig95 = (signif')*(ones(1,n));  % expand signif --> (J+1)x(N) array
sig95 = power ./ sig95;         % where ratio > 1, power is significant

% Global wavelet spectrum & significance levels:
global_ws = variance*(sum(power')/n);   % time-average over all times
dof = n - scale;  % the -scale corrects for padding at edges
global_signif = wave_signif(variance,dt,scale,1,lag1,-1,dof,mother);

% Scale-average between El Nino periods of 2--8 years
avg = find((scale >= 2) & (scale < 8));
Cdelta = 0.776;   % this is for the MORLET wavelet
scale_avg = (scale')*(ones(1,n));  % expand scale --> (J+1)x(N) array
scale_avg = power ./ scale_avg;   % [Eqn(24)]
scale_avg = variance*dj*dt/Cdelta*sum(scale_avg(avg,:));   % [Eqn(24)]
scaleavg_signif = wave_signif(variance,dt,scale,2,lag1,-1,[2,7.9],mother);


%------------------------------------------------------ Plotting

%--- Plot time series
subplot('position',[0.1 0.75 0.65 0.2])
plot(time,sst)
set(gca,'XLim',xlim(:))
xlabel('Time (year)')
ylabel('NINO3.4 SST (degC)')
title('NINO3.4 Sea Surface Temperature (seasonal)')
hold off

%--- Contour plot wavelet power spectrum
subplot('position',[0.1 0.37 0.65 0.28])
levels = [0.0625,0.125,0.25,0.5,1,2,4,8,16] ;
Yticks = 2.^(fix(log2(min(period))):fix(log2(max(period))));
%contour(time,log2(period),log2(power),log2(levels));  %*** or use 'contourfill'
%contourf(time,log2(period),log2(power),log2(levels));  %*** or use 'contourfill'
imagesc(time,log2(period),log2(power));  %*** uncomment for 'image' plot
xlabel('Time (year)')
ylabel('Period (years)')
title('NINO3.4 SST Wavelet Power Spectrum')
set(gca,'XLim',xlim(:))
set(gca,'YLim',log2([min(period),max(period)]), ...
	'YDir','reverse', ...
	'YTick',log2(Yticks(:)), ...
	'YTickLabel',Yticks)
% 95% significance contour, levels at -99 (fake) and 1 (95% signif)
hold on
contour(time,log2(period),sig95,[-99,1],'k');
hold on
% cone-of-influence, anything "below" is dubious
plot(time,log2(coi),'k')
hold off

%--- Plot global wavelet spectrum
subplot('position',[0.77 0.37 0.2 0.28])
plot(global_ws,log2(period))
hold on
plot(global_signif,log2(period),'--')
hold off
xlabel('Power (degC^2)')
title('Global Wavelet Spectrum')
set(gca,'YLim',log2([min(period),max(period)]), ...
	'YDir','reverse', ...
	'YTick',log2(Yticks(:)), ...
	'YTickLabel','')
set(gca,'XLim',[0,1.25*max(global_ws)])

% %--- Plot 2--8 yr scale-average time series
% subplot('position',[0.1 0.07 0.65 0.2])
% plot(time,scale_avg)
% set(gca,'XLim',xlim(:))
% xlabel('Time (year)')
% ylabel('Avg variance (degC^2)')
% title('2-8 yr Scale-average Time Series')
% hold on
% plot(xlim,scaleavg_signif+[0,0],'--')
% hold off

fig = gcf;fig.PaperUnits = 'inches';
fig.PaperPosition = [0 2 9 6.75]; % [left bottom width height]

%print(['/home/nsteiger/Dropbox/sa-phyda/wavelet_subs_nino34_' num2str(analyrs1(1)) '_' num2str(analyrs1(end)) '.png'],'-dpng','-r300');



%--- Contour plot wavelet power spectrum
figure
levels = [0.0625,0.125,0.25,0.5,1,2,4,8,16] ;
Yticks = 2.^(fix(log2(min(period))):fix(log2(max(period))));
%contour(time,log2(period),log2(power),log2(levels));  %*** or use 'contourfill'
%contourf(time,log2(period),log2(power),log2(levels));  %*** or use 'contourfill'
imagesc(time,log2(period),log2(power));  %*** uncomment for 'image' plot
xlabel('Time (year)')
ylabel('Period (years)')
title('NINO3.4 SST Wavelet Power Spectrum')
set(gca,'XLim',xlim(:))
set(gca,'YLim',log2([min(period),max(period)]), ...
	'YDir','reverse', ...
	'YTick',log2(Yticks(:)), ...
	'YTickLabel',Yticks)
% 95% significance contour, levels at -99 (fake) and 1 (95% signif)
hold on
contour(time,log2(period),sig95,[-99,1],'k');
hold on
% cone-of-influence, anything "below" is dubious
plot(time,log2(coi),'k')
hold off

fig = gcf;fig.PaperUnits = 'inches';
fig.PaperPosition = [0 2 7 4]; % [left bottom width height]

%print(['/home/nsteiger/Dropbox/sa-phyda/wavelet_nino34_' num2str(analyrs1(1)) '_' num2str(analyrs1(end)) '.png'],'-dpng','-r300');























