%compute AHT (especially at the equator) in the TRACE experiment
clear all; close all;
load('F:\trace\ZLH.mat') %Latent heat flux
load('F:\trace\ZSH.mat') %Sensible heat flux
load('F:\trace\ZFLNS.mat') %net longwave flux at the surface
load('F:\trace\ZFSNS.mat') %net shortwave flux at the surface
load('F:\trace\ZOLR.mat') %OLR - top of model
load('F:\trace\ZFSNT.mat') %net shortwave at TOA
load('F:\trace\TS.mat') %surface temperature

%make the surface heat flux -- to the atmosphere
SHF= -ZFSNS + ZLH + ZSH + ZFLNS;
TOA = ZFSNT-ZOLR;
NY=length(YEARS);
NN=length(lat);
x=sind(lat); %sin of latitude as a corrdinate for global integrlas (or avverages)
ae=6.36*10^6; %Earth radius
norm=trapz(x, ones(size(x))); %normalization for spatial weighted global averages

mlat=[-90; (lat(2:end)+lat(1:end-1))/2;90]; mNN=length(mlat); %the heat transport will be calculated on this grid 
dlat=diff(mlat);
dA=(2*pi*ae*cosd(lat)).*(ae*dlat*pi/180);

for ny=1:NY; %for ny=1:NY;
    ny;
GTOA(ny)=trapz(x, squeeze(TOA(ny,:)))/norm;
GSHF(ny)=trapz(x, squeeze(SHF(ny,:)))/norm;
GASR(ny)=trapz(x, squeeze(ZFSNT(ny,:)))/norm;
GLW(ny)=trapz(x, squeeze(ZOLR(ny,:)))/norm;

AHT(ny,1)=0; OHT(ny,1)=0;

LW_HT(ny,1)=0; SW_HT(ny,1)=0;
for nn=2:mNN;
AHT(ny,nn) = AHT(ny,nn-1) + dA(nn-1)*(TOA(ny,nn-1) - GTOA(ny) + SHF(ny,nn-1)-GSHF(ny));
OHT(ny,nn) = OHT(ny,nn-1) + dA(nn-1)*(-SHF(ny,nn-1)+GSHF(ny));
LW_HT(ny,nn) = LW_HT(ny,nn-1) + dA(nn-1)*(-ZOLR(ny,nn-1)+GLW(ny));
SW_HT(ny,nn) = SW_HT(ny,nn-1) + dA(nn-1)*(ZFSNT(ny,nn-1)-GASR(ny));
end
end


do=1:22000;
load('F:\trace\zonal_PC.mat', 'YEARS', 'lat', 'ZPC');
load('F:\trace\zonal_PL.mat', 'YEARS', 'lat', 'ZPL');
load('F:\trace\TS.mat', 'YEARS', 'lat', 'ZTS', 'ZSST', 'ZLTS', 'grad_SST');
grad_SST=double(grad_SST);
YEARS=-22000:39;
P=double(ZPC+ZPL);
clear ZPC ZPL
NY=length(YEARS);
flat=-20:0.05:20;
gray=[0.5 0.5 0.5];
lred=[0.6 0 0];
for ny=1:NY;
    FL=squeeze(P(ny,:));
    fP=interp1(lat,FL,flat);
    cdf=cumsum(fP)/sum(fP);
    PCENT(ny)=flat(find2(cdf,0.5));
end

LW=0.7; LWL=2;
FSz=22; FN='Arial';

figure(1); clf;
plot(YEARS(do), PCENT(do), '-','Color', gray, 'LineWidth', LW);
%low pass filter
cut_per=200;
low_PCENT = low_filter(PCENT, 1, cut_per);
hold on; plot(YEARS(do), low_PCENT(do), 'k', 'LineWidth', LWL)
SF_SST=2;
hold on; plot(YEARS(do), SF_SST*grad_SST(do), 'Color',lred,'LineWidth',LW);
low_SST = low_filter(grad_SST, 1, cut_per);
hold on; plot(YEARS(do), SF_SST*low_SST(do), 'r','LineWidth',LWL);

xlabel('Thousand Years','FontSize', FSz, 'FontName', FN)
ylabel('ITCZ location (P_C_E_N_T in ^o)', 'FontSize', FSz, 'FontName', FN);
xticl=-20:5:5;
xtic=xticl*1000;
set(gca, 'FontSize', FSz, 'FontName',FN, 'XTick', xtic, 'XTickLabel', xticl, 'LineWidth', 2.4)
xmn=-22040; xmx=2000;
ymn=-4; ymx=1;
title('CCSM trace timeseries of ITCZ position', 'FontSize', FSz, 'FontName', FN);
axis([xmn xmx ymn ymx]); box on;

nni=find(mlat ==0);
AHT_EQ=AHT(:,nni);
PW=10^-15;
SF_AHT=3;
off_AHT=-1.5;
lblue = [135 206 250]/255
ytic=-4:.5:1;
yticl={ '-4  ';'';'-3  ';'';'-2  ';'';'-1  ';'';'0  ';'';'1  '};
figure(1); hold on; plot(YEARS, -PW*AHT_EQ*3+off_AHT, 'Color',lblue, 'LineWidth', LW)
low_AHT = low_filter(AHT_EQ, 1, cut_per);
figure(1); hold on; plot(YEARS, -PW*low_AHT*3+off_AHT, 'Color','b', 'LineWidth', LWL)
set(gca, 'YTick', ytic, 'YTickLabel', yticl)

%and scatter plots (scatter clouds)
samp=100; %sampling frequency of the low pass filtered ts
st=40;
MS=8; LWM=2; LW=2.3;LWA=2.5;
color_dhoe3_fine;
cmap=cmapf;
[xEdges,yEdges, C, numY, numX]=scattercloud2(PW*AHT_EQ',PCENT', 60,1,'k+',cmap);%,1,1,'k+',cmap)[xEdges_PI,yEdges_PI, C_PI, numY_PI, numX_PI] = scattercloud2(X,Y,20,1,'k+',cmap);

gray=[0.4 0.4 0.4];
hold on;plot(PW*low_AHT(st:samp:end), low_PCENT(st:samp:end), '+','Color', gray, 'MarkerSize', MS, 'LineWidth', LWM) 
[a b ci R] = lreg3(PW*low_AHT(st:samp:end), low_PCENT(st:samp:end));
xmn=-0.5; xmx=1; ymn=-3.5; ymx=1;
axis([xmn xmx ymn ymx]);
hold on; plot([xmn xmx], a+b*[xmn xmx], '--','Color', gray, 'LineWidth', LW);
hold on; plot([xmn xmx], 0*[ymn ymx], 'k-', 'LineWidth', LW);
hold on; plot(0*[xmn xmx], [ymn ymx], 'k-', 'LineWidth', LW);
xlabel('AHT_EQ (PW)', 'FontSize', FSz, 'FontName', FN)
ylabel('ITCZ location (^o)', 'FontSize', FSz, 'FontName', FN);
xtic=[-2:1/3:2];
ytic=-5:.5:2;
yticl={'-5  ';'';'-4  ';'';'-3  ';'';'-2  ';'';'-1  ';'';'0  ';'';'1'};

set(gca, 'FontSize', FSz, 'FontName', FN, 'LineWidth', 2.5, 'XTick',xtic ,'XTickLabel', xtic, 'YTick', ytic, 'YTickLabel', yticl ); box on;

[xEdges,yEdges, C, numY, numX]=scattercloud2(grad_SST',PCENT', 60,1,'k+',cmap);%,1,1,'k+',cmap)[xEdges_PI,yEdges_PI, C_PI, numY_PI, numX_PI] = scattercloud2(X,Y,20,1,'k+',cmap);


hold on;plot(low_SST(st:samp:end), low_PCENT(st:samp:end), '+','Color', gray, 'MarkerSize', MS, 'LineWidth', LWM) 
[a b ci R] = lreg3(low_SST(st:samp:end), low_PCENT(st:samp:end));
xmn=-1.0; xmx=0.5; ymn=-3.5; ymx=1;
axis([xmn xmx ymn ymx]);
hold on; plot([xmn xmx], a+b*[xmn xmx], '--','Color', gray, 'LineWidth', LW);
hold on; plot([xmn xmx], 0*[ymn ymx], 'k-', 'LineWidth', LW);
hold on; plot(0*[xmn xmx], [ymn ymx], 'k-', 'LineWidth', LW);
xlabel('SST contrast (K)', 'FontSize', FSz, 'FontName', FN)
ylabel('ITCZ location (^o)', 'FontSize', FSz, 'FontName', FN);
xtic=[-2:0.5:2];
ytic=-5:.5:2;
yticl={'-5  ';'';'-4  ';'';'-3  ';'';'-2  ';'';'-1  ';'';'0  ';'';'1'};

set(gca, 'FontSize', FSz, 'FontName', FN, 'LineWidth', 2.5, 'XTick',xtic ,'XTickLabel', xtic, 'YTick', ytic, 'YTickLabel', yticl ); box on;

%divide the AHT_EQ into ocean and non ocean components (radiative)
LW=0.2;
OHT_EQ=OHT(:,nni);
RAD_EQ=AHT_EQ+OHT_EQ;
SW_EQ=SW_HT(:,nni);
LW_EQ=LW_HT(:,nni);

figure(10); clf
low_OHT=low_filter(OHT_EQ, 1, cut_per);
low_RAD=low_filter(RAD_EQ, 1, cut_per);
low_LW=low_filter(LW_EQ, 1, cut_per);
low_SW=low_filter(SW_EQ, 1, cut_per);

LWL=2.7;
lgreen=[.7 1 .7];
dgreen=[34 139 34]/255;

lorange = [255 165 0]/255;
dorange=[255 100 0]/255;
purple=[160 32 240]/255;
figure(10); hold on; plot(YEARS, PW*AHT_EQ, 'Color',lblue, 'LineWidth', LW)
figure(10); hold on; plot(YEARS, -PW*OHT_EQ, 'Color',lgreen, 'LineWidth', LW)
figure(10); hold on; plot(YEARS, PW*RAD_EQ, 'Color',lorange, 'LineWidth', LW)

figure(10); hold on; plot(YEARS, PW*low_AHT, 'Color','b', 'LineWidth', LWL)
figure(10); hold on; plot(YEARS, -PW*low_OHT, 'Color',dgreen, 'LineWidth', LWL)
figure(10); hold on; plot(YEARS, PW*low_RAD, 'Color',dorange, 'LineWidth', LWL)

%add in the time means
LWM=1.7;
figure(10); hold on; plot(YEARS, PW*mean(low_AHT)*ones(size(YEARS)), 'b--', 'LineWidth', LWM)
figure(10); hold on; plot(YEARS, -PW*mean(low_OHT)*ones(size(YEARS)), '--','Color',dgreen, 'LineWidth', LWM)
figure(10); hold on; plot(YEARS, PW*mean(low_RAD)*ones(size(YEARS)),'--', 'Color',dorange, 'LineWidth', LWM)
%figure(10); hold on; plot(YEARS, PW*low_LW, 'Color',purple, 'LineWidth', LWL)
xlabel('Years X 1000', 'FontSize', FSz, 'FontName', FN);
xticl=-20:5:0;
xtic=xticl*1000;
ytic=-2:0.5:2;
yticl={'-2  ';'';'-1  ';'';'0  ';'';'-1  ';'';'-2'};
ylabel('AHT_E_Q / -OHT_E_Q / and {RAD_T_O_A} (PW)', 'FontSize', FSz, 'FontName', FN);
set(gca, 'FontSize', FSz, 'FontName', FN, 'LineWidth', 2.5, 'XTick',xtic, 'XTickLabel', xticl, 'YTick', ytic,'YTickLabel', yticl); box on;
ymn=-2; ymx=2; xmn=-22.2*1000; xmx=100;
axis([xmn xmx ymn ymx]);


[xEdges,yEdges, C, numY, numX]=scattercloud2(PW*AHT_EQ',-PW*OHT_EQ', 60,1,'k+',cmap);%,1,1,'k+',cmap)[xEdges_PI,yEdges_PI, C_PI, numY_PI, numX_PI] = scattercloud2(X,Y,20,1,'k+',cmap);

LW=2.8;
hold on;plot(PW*low_AHT(st:samp:end), -PW*low_OHT(st:samp:end), '+','Color', gray, 'MarkerSize', MS, 'LineWidth', LWM) 
[a b ci R] = lreg3(PW*low_AHT(st:samp:end), -PW*low_OHT(st:samp:end));
xmn=-0.5; xmx=1; ymn=-1.5; ymx=0;
axis([xmn xmx ymn ymx]);
hold on; plot([xmn xmx], a+b*[xmn xmx], '--','Color', gray, 'LineWidth', LW);
hold on; plot([xmn xmx], 0*[ymn ymx], 'k-', 'LineWidth', LW);
hold on; plot(0*[xmn xmx], [ymn ymx], 'k-', 'LineWidth', LW);
xlabel('AHT_E_Q (PW)', 'FontSize', FSz, 'FontName', FN)
ylabel('-OHT_E_Q (PW)', 'FontSize', FSz, 'FontName', FN);
xtic=[-0.5:0.25:1.5];

ytic=-1.5:.25:0;
yticl={'-1.5  ';'';'-1.0  ';'';'-0.5  ';'';'0  '};%;'';'-1  ';'';'0  ';'';'1'};

set(gca, 'FontSize', FSz, 'FontName', FN, 'LineWidth', 2.5, 'XTick',xtic ,'XTickLabel', xtic, 'YTick', ytic, 'YTickLabel', yticl ); box on;


[xEdges,yEdges, C, numY, numX]=scattercloud2(PW*AHT_EQ',PW*RAD_EQ', 60,1,'k+',cmap);%,1,1,'k+',cmap)[xEdges_PI,yEdges_PI, C_PI, numY_PI, numX_PI] = scattercloud2(X,Y,20,1,'k+',cmap);

hold on;plot(PW*low_AHT(st:samp:end), PW*low_RAD(st:samp:end), '+','Color', gray, 'MarkerSize', MS, 'LineWidth', LWM) 
[a b ci R] = lreg3(PW*low_AHT(st:samp:end), PW*low_RAD(st:samp:end));
xmn=-0.5; xmx=1; 
ymn=0; ymx=1.5;
axis([xmn xmx ymn ymx]);
hold on; plot([xmn xmx], a+b*[xmn xmx], '--','Color', gray, 'LineWidth', LW);
hold on; plot([xmn xmx], 0*[ymn ymx], 'k-', 'LineWidth', LW);
hold on; plot(0*[xmn xmx], [ymn ymx], 'k-', 'LineWidth', LW);
xlabel('AHT_E_Q (PW)', 'FontSize', FSz, 'FontName', FN)
ylabel('<RAD_T_O_A> (PW)', 'FontSize', FSz, 'FontName', FN);
xtic=[-0.5:0.25:1.5];

ytic=0:.25:1.5;
yticl={'0  ';'';'0.5  ';'';'1.0  ';'';'1.5  '};%;'';'-1  ';'';'0  ';'';'1'};

set(gca, 'FontSize', FSz, 'FontName', FN, 'LineWidth', 2)

%power spectra and phase/frequency relationships
LWA=2.5;
SF_AHT=2.65*10^-15;
SF_SST=2.1;
 [fs, pow_SST] = myfft_s(SF_SST*grad_SST, 1,1);
 [fs, pow_AHT] = myfft_s(SF_AHT*AHT_EQ, 1,1);
 [fs, pow_PCENT] = myfft_s(PCENT, 1,1);
 [fs, pow_PCENT_AHT] = myfft_s(PCENT.*AHT_EQ', 1,1);
 [fs, pow_PCENT_SST] = myfft_s(PCENT.*grad_SST, 1,1);
 %smooth the high frequency part
 
 smooth_per=1100; %period to start smoothing after
 
 sm=20; %smooth number;
 do_no_smooth=find(1./fs >smooth_per);
 do_smooth=find(1./fs <=smooth_per);
 FS = [fs(do_no_smooth) mean(c2g(fs(do_smooth), length(do_smooth)/sm,sm)')]; %the frequency for the smoothed and unsmoothed regions
 POW_SST = [pow_SST(do_no_smooth) mean(c2g(pow_SST(do_smooth), length(do_smooth)/sm,sm)')]; %the frequency for the smoothed and unsmoothed regions
 POW_AHT = [pow_AHT(do_no_smooth) mean(c2g(pow_AHT(do_smooth), length(do_smooth)/sm,sm)')]; %the frequency for the smoothed and unsmoothed regions
 POW_PCENT = [pow_PCENT(do_no_smooth) mean(c2g(pow_PCENT(do_smooth), length(do_smooth)/sm,sm)')]; %the frequency for the smoothed and unsmoothed regions
 POW_PCENT_SST = [pow_PCENT_SST(do_no_smooth) mean(c2g(pow_PCENT_SST(do_smooth), length(do_smooth)/sm,sm)')]; %the frequency for the smoothed and unsmoothed regions
 POW_PCENT_AHT = [pow_PCENT_AHT(do_no_smooth) mean(c2g(pow_PCENT_AHT(do_smooth), length(do_smooth)/sm,sm)')]; %the frequency for the smoothed and unsmoothed regions
 
 
 figure(20); clf;
%  hold on; plot(log(FS), FS.*POW_SST/trapz(log(FS), FS.*POW_SST), 'r', 'LineWidth', LW);
%  hold on; plot(log(FS), FS.*POW_AHT/trapz(log(FS), FS.*POW_AHT), 'b', 'LineWidth', LW);
%  hold on; plot(log(FS), FS.*POW_PCENT/trapz(log(FS), FS.*POW_PCENT), 'k', 'LineWidth', LW);

 hold on; plot(log(FS), FS.*POW_SST, 'r', 'LineWidth', LW);
 hold on; plot(log(FS), FS.*POW_AHT, 'b', 'LineWidth', LW);
 hold on; plot(log(FS), FS.*POW_PCENT,  'k', 'LineWidth', LW);
 

 xticl=[22000 10000 3000 1000 200 100 50 10 2];
 xtic=-log(xticl);
 set(gca, 'FontSize', FSz, 'FontName', FN, 'XTickLabel', xticl, 'XTick', xtic, 'LineWidth', LWA, 'YTickLabel', [])
 box on;
 axis([-10 -0.5 0 1.0*10^-5])
 %and the cohenerence
 NT=length(YEARS);
 [coh_SST,F_coh] = mscohere(PCENT,grad_SST,hamming(NT/2),100,NT/2);
 [coh_AHT,F_coh] = mscohere(PCENT,AHT_EQ,hamming(NT/2),100,NT/2);
 f_coh=F_coh/(2*pi);
 %coh=mscohere(AHT_EQ,grad_SST);
 DO_no_smooth=find(1./f_coh >smooth_per);
 DO_smooth=find(1./f_coh <=smooth_per);
 
 COH_FS = [f_coh(DO_no_smooth)' mean(c2g(f_coh(DO_smooth), length(DO_smooth)/sm,sm)')]; %the frequency for the smoothed and unsmoothed regions
 COH_SST = [coh_SST(DO_no_smooth)' mean(c2g(coh_SST(DO_smooth), length(DO_smooth)/sm,sm)')]; %the frequency for the smoothed and unsmoothed regions
 COH_AHT = [coh_AHT(DO_no_smooth)' mean(c2g(coh_AHT(DO_smooth), length(DO_smooth)/sm,sm)')]; %the frequency for the smoothed and unsmoothed regions
 figure(23);clf;
 %plot(log(f_coh(2:end)),coh(2:end),'linewidth',2); title('Magnitude-squared Coherence');
 hold on; plot(log(COH_FS),COH_SST, 'r', 'LineWidth', LW);  
  hold on; plot(log(COH_FS),COH_AHT, 'b', 'LineWidth', LW);  
  
 xticl=[10000 3000 1000 200 100 50 10 2];
 xtic=-log(xticl);
 set(gca, 'FontSize', FSz, 'FontName', FN, 'XTickLabel', xticl, 'XTick', xtic, 'LineWidth', LWA, 'YTickLabel', [])
box on;
axis([-10 -0.5 0 1])

%%% now repeat by breaking down the AHT into radiative and surface flux
%%% components

 [fs, pow_AHT] = myfft_s(AHT_EQ, 1,1);
 [fs, pow_RAD] = myfft_s(RAD_EQ, 1,1);
 [fs, pow_OHT] = myfft_s(OHT_EQ, 1,1);
 
 %smooth over the given ranges
  FS = [fs(do_no_smooth) mean(c2g(fs(do_smooth), length(do_smooth)/sm,sm)')]; %the frequency for the smoothed and unsmoothed regions
  POW_AHT = [pow_AHT(do_no_smooth) mean(c2g(pow_AHT(do_smooth), length(do_smooth)/sm,sm)')]; %the frequency for the smoothed and unsmoothed regions
  POW_RAD = [pow_RAD(do_no_smooth) mean(c2g(pow_RAD(do_smooth), length(do_smooth)/sm,sm)')]; %the frequency for the smoothed and unsmoothed regions
  POW_OHT = [pow_OHT(do_no_smooth) mean(c2g(pow_OHT(do_smooth), length(do_smooth)/sm,sm)')]; %the frequency for the smoothed and unsmoothed regions

  figure(30);clf;
  hold on; plot(log(FS), FS.*POW_AHT, 'b', 'LineWidth', LW);
  hold on; plot(log(FS), FS.*POW_OHT, 'Color',dgreen, 'LineWidth', LW);
  hold on; plot(log(FS), FS.*POW_RAD,  'Color',lorange, 'LineWidth', LW);
  
  xticl=[22000 10000 3000 1000 200 100 50 10 2];
  xtic=-log(xticl);
  set(gca, 'FontSize', FSz, 'FontName', FN, 'XTickLabel', xticl, 'XTick', xtic, 'LineWidth', LWA)
  axis([-10 -0.5 0 10^24])
  box on; 
  %and coherence
 [coh_RAD,F_coh] = mscohere(AHT_EQ,RAD_EQ,hamming(NT/2),100,NT/2);
 [coh_OHT,F_coh] = mscohere(AHT_EQ,OHT_EQ,hamming(NT/2),100,NT/2);
 f_coh=F_coh/(2*pi);
 %coh=mscohere(AHT_EQ,grad_SST);
 DO_no_smooth=find(1./f_coh >smooth_per);
 DO_smooth=find(1./f_coh <=smooth_per);
 
 COH_FS = [f_coh(DO_no_smooth)' mean(c2g(f_coh(DO_smooth), length(DO_smooth)/sm,sm)')]; %the frequency for the smoothed and unsmoothed regions
 COH_RAD = [coh_RAD(DO_no_smooth)' mean(c2g(coh_RAD(DO_smooth), length(DO_smooth)/sm,sm)')]; %the frequency for the smoothed and unsmoothed regions
 COH_OHT = [coh_OHT(DO_no_smooth)' mean(c2g(coh_OHT(DO_smooth), length(DO_smooth)/sm,sm)')]; %the frequency for the smoothed and unsmoothed regions
 figure(23);clf;
 %plot(log(f_coh(2:end)),coh(2:end),'linewidth',2); title('Magnitude-squared Coherence');
 hold on; plot(log(COH_FS),COH_RAD, 'Color',dorange, 'LineWidth', LW);  
  hold on; plot(log(COH_FS),COH_OHT, 'Color',dgreen, 'LineWidth', LW);  
  
 xticl=[10000 3000 1000 200 100 50 10 2];
 xtic=-log(xticl);
 set(gca, 'FontSize', FSz, 'FontName', FN, 'XTickLabel', xticl, 'XTick', xtic, 'LineWidth', LWA, 'YTickLabel', [])
box on;
axis([-10 -0.5 0 1])
  