%TRANSHEATFLUX.M
%Computes transient heat flux statistics from NCEP/NCAR derived reanalysis data
%Requires file vt.mon.mean.3_10.nc
%Brian Tang, 4/30/09
%slightly modified by Malte Jansen 4/28/10

clear all;
close all;

%GLOBAL VARIABLES AND CONSTANTS (DO NOT ALTER!)
global nctbx_options
nctbx_options.theAutoscale = 1;
nctbx_options.theAutoNaN = 1;
earthrad = 6.37E6; %m
cp = 1005.7; %J/(kg K)
g = 9.81; %m/s^2

%USER DEFINED OPTIONS
%1: TRANSIENT HEAT FLUX AT A SINGLE LEVEL IN K m/s
%2: VERTICALLY AVERAGED TRANSIENT HEAT FLUX IN K m/s
%3: ZONALLY AVERAGED TRANSIENT HEAT FLUX IN K m/s
%4: ZONALLY AND VERTICALLY INTEGRATED TRANSIENT (INTERNAL) ENERGY TRANSPORT IN PW
opt = 1;
month = [1 1]; %Beginning and End months for averaging period

%READ IN DATA
%ncfloat('time', 'level', 'lat', 'lon'); %% 1009152 elements
try
    load('vt.mon.mean.3_10.mat');
catch
    filein = 'vt.mon.mean.3_10.nc';
    fvt = netcdf(filein,'nowrite');
    vt = fvt{'vt'}(:,:,:,:);
    p = fvt{'level'}(:);
    lat = fvt{'lat'}(:);
    lon = fvt{'lon'}(:);
end

%AVERAGE OVER SPECIFIED MONTHS
vt = squeeze(mean(vt(month(1):month(2),:,:,:),1));

%ARRAY FOR THICKNESS OF INDIVIDUAL LAYERS USED FOR AVERAGING AND INTEGRATION
phalf = (p(1:end-1)+p(2:end))/2;
dp = [1012-phalf(1); phalf(1:end-1)-phalf(2:end); phalf(end)-50];

cint = 2; %contour interval;
%PLOT DATA
if (opt==1)
   disp('TRANSIENT HEAT FLUX AT A SINGLE LEVEL');
   disp(['Available levels to plot are: ', num2str(p')]);
   plotlev = input('Select level at which to plot: ');
   while(isempty(find(p==plotlev)))
      disp('Unavailable level chosen.');
      plotlev = input('Select level at which to plot: ');
   end
   
   vtplev = squeeze(vt(find(p==plotlev),:,:));
   [c,h] = contourf(lon,lat,vtplev,[cint*floor(min(min(vtplev))/cint):cint:cint*ceil(max(max(vtplev))/cint)]);
   hold on; 
   clabel(c,h,'fontsize',13);
   world([lon(1) lon(end) lat(end) lat(1)]);
   grid on;
   xlabel('Longitude','fontsize',14);
   ylabel('Latitude','fontsize',14);
   title(['Transient Heat Flux [K m/s] @ ',num2str(plotlev),'mb, Months: ',num2str(month(1)),'-',num2str(month(2))],'fontsize',14);
   set(gca,'fontsize',14,'layer','top');
   hold off;
elseif (opt==2)
   disp('VERTICALLY AVERAGED TRANSIENT HEAT FLUX');

   vtpave = squeeze(sum(vt.*repmat(dp,[1 length(lat) length(lon)]),1)./sum(repmat(dp,[1 length(lat) length(lon)]),1));
   [c,h] = contourf(lon,lat,vtpave,[cint*floor(min(min(vtpave))/cint):cint:cint*ceil(max(max(vtpave))/cint)]);
   hold on;
   clabel(c,h,'fontsize',13);
   world([lon(1) lon(end) lat(end) lat(1)]);
   grid on;
   xlabel('Longitude','fontsize',14);
   ylabel('Latitude','fontsize',14);
   title(['Vertically Averaged Transient Heat Flux [K m/s], Months: ',num2str(month(1)),'-',num2str(month(2))],'fontsize',14);
   set(gca,'fontsize',14,'layer','top');
   hold off; 
elseif (opt==3)
   disp('ZONALLY AVERAGED TRANSIENT HEAT FLUX');
   %***********************************************
   %Compute zonally averaged transient heat flux in K m/s.
   %Variables you will need: vt('level', 'lat', 'lon') [K m/s] - transient heat flux
   
   %vtzave = ;
   
   %***********************************************
   [c,h] = contourf(lat,p,vtzave,[cint*floor(min(min(vtzave))/cint):cint:cint*ceil(max(max(vtzave))/cint)]);
   clabel(c,h,'fontsize',13);
   grid on;
   xlabel('Latitude','fontsize',14);
   ylabel('Pressure','fontsize',14);
   set(gca,'ydir','reverse');
   title(['Zonally Averaged Transient Heat Flux [K m/s], Months: ',num2str(month(1)),'-',num2str(month(2))],'fontsize',14);
   set(gca,'fontsize',14,'layer','top');
elseif (opt==4)
   disp('ZONALLY AND VERTICALLY INTEGRATED TRANSIENT HEAT (ENERGY) FLUX');
   %***********************************************
   %Compute vertically and zonally integrated transient (internal) energy flux in PW.
   %Variables you will need: 	vt('level', 'lat', 'lon') [K m/s] - transient heat flux
   %				cp [J/(kg K) - specific heat of dry air
   %				g [m/s^2] - gravity
   %				earthrad [m] - Earth's radius
   %				lat [deg] - latitude (here in degrees!)
   %				dp [hPa] - thickness of layers integrating over (convert to Pa!!)
   
   %vtzave = ;
   %H = ;
   
   %***********************************************
   plot(H./1E15,lat,'linewidth',2);
   grid on;
   xlabel('Transient Energy Flux [PW]','fontsize',14);
   ylabel('Latitude','fontsize',14);
   title(['Transient Energy Flux [PW], Months: ',num2str(month(1)),'-',num2str(month(2))],'fontsize',14);
   set(gca,'fontsize',14,'layer','top');
end
