
%==========================================================================
% THIS IS THE DRIVING SCRIPT FOR PDSI CALCULATIONS
% Nathan Steiger, Apr 2016
%==========================================================================

addpath('/home/nsteiger/MATLAB/annual-da/')

tic

% MODEL TYPE
%mdl='echam5';
%mdl='gfdlhist';
%mdl='cesm_lme';enstyp='VOLC_GRA'; ens='005';disp(['CESM ensemble = ' ens])
mdl='cesm_lme'; ens='009';disp(['CESM ensemble = ' ens])
%mdl='gfdlesm2m_pi';
%mdl='gfdlesm2g_pi';
%mdl='ccsm4_lm';

disp(['GCM = ' mdl])

% AWC OPTION
awc_o='c';% constant
%awc_o='v';% variable

% LOAD THE MODEL DATA

disp('Loading data...')

if strcmp(mdl,'echam5')
    varpth='/d3/nsteiger/echam5wiso/';
    
    syr=1871;eyr=2011;lyrs=length(syr:eyr);
    
    load('/d1/nsteiger/climate-data/echam5/e5_latlon.mat')
    
    % land area
    load([varpth 'sftlf_fx_ECHAM5-wiso.mat'])

    disp('Hacking PDSI so that it is computed over ocean too...')
    sftlf=100*ones(size(sftlf));
    
    % radiation (W/m^2)
    trads=reshape(permute(flip(ncread([varpth 'trads_mon_1871_2011.nc'],'trads'),2),[2 1 3]),length(lat)*length(lon),12*lyrs)';
    srads=reshape(permute(flip(ncread([varpth 'srads_mon_1871_2011.nc'],'srads'),2),[2 1 3]),length(lat)*length(lon),12*lyrs)';
    
    % t_emperature (C):
    t_c=reshape(permute(flip(ncread([varpth 't2m_mon_1871_2011.nc'],'temp2'),2),[2 1 3]),length(lat)*length(lon),12*lyrs)'-273.15;
    %t_f = t_c*33.8; % C to  F
    
    % Surface pressure (Pa):
    ps=reshape(permute(flip(ncread([varpth 'ps_mon_1871_2011.nc'],'aps'),2),[2 1 3]),length(lat)*length(lon),12*lyrs)';
    
    % Vapor pressure (Pa):
    huss=reshape(permute(flip(squeeze(ncread([varpth 'qsurf_mon_1871_2011.nc'],'q')),2),[2 1 3]),length(lat)*length(lon),12*lyrs)';
    
    % Surface wind (m/s) (CMIP5 output is at 10m)
    u10=reshape(permute(flip(ncread([varpth 'wind10_mon_1871_2011.nc'],'wind10'),2),[2 1 3]),length(lat)*length(lon),12*lyrs)';
    
    % precipitation, will be converted from kg*m^-2*s^-1 to (mm/day):
    pr=reshape(permute(flip(ncread([varpth 'precip_mon_1871_2011.nc'],'precip'),2),[2 1 3]),length(lat)*length(lon),12*lyrs)';
    
elseif strcmp(mdl,'gfdlhist')
    
    varpth='/d1/nsteiger/climate-data/gfdl-esm2m/';
    
    syr=1861;eyr=2005;lyrs=length(syr:eyr);
    
    lat=ncread([varpth 'tas_Amon_GFDL-ESM2M_historical_r1i1p1_186101-200512.nc'],'lat');
    lon=ncread([varpth 'tas_Amon_GFDL-ESM2M_historical_r1i1p1_186101-200512.nc'],'lon');
    
    % land area
    sftlf=permute(ncread([varpth 'sftlf_fx_GFDL-ESM2M_historical_r0i0p0.nc'],'sftlf'),[2 1]);
    
    % radiation (W/m^2)
    rsds=reshape(permute(ncread([varpth 'rsds_Amon_GFDL-ESM2M_historical_r1i1p1_186101-200512.nc'],'rsds'),[2 1 3]),length(lat)*length(lon),12*lyrs)';
    rsus=reshape(permute(ncread([varpth 'rsus_Amon_GFDL-ESM2M_historical_r1i1p1_186101-200512.nc'],'rsus'),[2 1 3]),length(lat)*length(lon),12*lyrs)';
    rlds=reshape(permute(ncread([varpth 'rlds_Amon_GFDL-ESM2M_historical_r1i1p1_186101-200512.nc'],'rlds'),[2 1 3]),length(lat)*length(lon),12*lyrs)';
    rlus=reshape(permute(ncread([varpth 'rlus_Amon_GFDL-ESM2M_historical_r1i1p1_186101-200512.nc'],'rlus'),[2 1 3]),length(lat)*length(lon),12*lyrs)';
    
    % t_emperature (C):
    t_c=reshape(permute(ncread([varpth 'tas_Amon_GFDL-ESM2M_historical_r1i1p1_186101-200512.nc'],'tas'),[2 1 3]),length(lat)*length(lon),12*lyrs)'-273.15;
    %t_f = t_c*33.8; % C to  F
    
    % Surface pressure (Pa):
    ps=reshape(permute(ncread([varpth 'ps_Amon_GFDL-ESM2M_historical_r1i1p1_186101-200512.nc'],'ps'),[2 1 3]),length(lat)*length(lon),12*lyrs)';
    
    % Vapor pressure (Pa):
    huss=reshape(permute(ncread([varpth 'huss_Amon_GFDL-ESM2M_historical_r1i1p1_186101-200512.nc'],'huss'),[2 1 3]),length(lat)*length(lon),12*lyrs)';
    
    % Surface wind (m/s) (CMIP5 output is at 10m)
    u10=reshape(permute(ncread([varpth 'sfcWind_Amon_GFDL-ESM2M_historical_r1i1p1_186101-200512.nc'],'sfcWind'),[2 1 3]),length(lat)*length(lon),12*lyrs)';
    
    % precipitation, will be converted from kg*m^-2*s^-1 to (mm/day):
    pr=reshape(permute(ncread([varpth 'pr_Amon_GFDL-ESM2M_historical_r1i1p1_186101-200512.nc'],'pr'),[2 1 3]),length(lat)*length(lon),12*lyrs)';

elseif strcmp(mdl,'gfdlesm2g_pi')
    
    varpth='/d1/nsteiger/climate-data/gfdl-esm2g-pi/';
    
    syr=1;eyr=500;lyrs=length(syr:eyr);
    
    lat=ncread([varpth 'tas_Amon_GFDL-ESM2G_piControl_r1i1p1_000101-050012.nc'],'lat');
    lon=ncread([varpth 'tas_Amon_GFDL-ESM2G_piControl_r1i1p1_000101-050012.nc'],'lon');
    
    % land area
    %sftlf=permute(ncread([varpth 'sftlf_fx_GFDL-ESM2G_piControl_r0i0p0.nc'],'sftlf'),[2 1]);
    sftlf=permute(ncread('/d1/nsteiger/climate-data/gfdl-esm2m-pi/sftlf_fx_GFDL-ESM2M_piControl_r0i0p0.nc','sftlf'),[2 1]); % using "M" so they have the same points...
    
    % radiation (W/m^2)
    rsds=reshape(permute(ncread([varpth 'rsds_Amon_GFDL-ESM2G_piControl_r1i1p1_000101-050012.nc'],'rsds'),[2 1 3]),length(lat)*length(lon),12*lyrs)';
    rsus=reshape(permute(ncread([varpth 'rsus_Amon_GFDL-ESM2G_piControl_r1i1p1_000101-050012.nc'],'rsus'),[2 1 3]),length(lat)*length(lon),12*lyrs)';
    rlds=reshape(permute(ncread([varpth 'rlds_Amon_GFDL-ESM2G_piControl_r1i1p1_000101-050012.nc'],'rlds'),[2 1 3]),length(lat)*length(lon),12*lyrs)';
    rlus=reshape(permute(ncread([varpth 'rlus_Amon_GFDL-ESM2G_piControl_r1i1p1_000101-050012.nc'],'rlus'),[2 1 3]),length(lat)*length(lon),12*lyrs)';
    
    % t_emperature (C):
    t_c=reshape(permute(ncread([varpth 'tas_Amon_GFDL-ESM2G_piControl_r1i1p1_000101-050012.nc'],'tas'),[2 1 3]),length(lat)*length(lon),12*lyrs)'-273.15;
    %t_f = t_c*33.8; % C to  F
    
    % Surface pressure (Pa):
    ps=reshape(permute(ncread([varpth 'ps_Amon_GFDL-ESM2G_piControl_r1i1p1_000101-050012.nc'],'ps'),[2 1 3]),length(lat)*length(lon),12*lyrs)';
    
    % Vapor pressure (Pa):
    huss=reshape(permute(ncread([varpth 'huss_Amon_GFDL-ESM2G_piControl_r1i1p1_000101-050012.nc'],'huss'),[2 1 3]),length(lat)*length(lon),12*lyrs)';
    
    % Surface wind (m/s) (CMIP5 output is at 10m)
    u10=reshape(permute(ncread([varpth 'sfcWind_Amon_GFDL-ESM2G_piControl_r1i1p1_000101-050012.nc'],'sfcWind'),[2 1 3]),length(lat)*length(lon),12*lyrs)';
    
    % precipitation, will be converted from kg*m^-2*s^-1 to (mm/day):
    pr=reshape(permute(ncread([varpth 'pr_Amon_GFDL-ESM2G_piControl_r1i1p1_000101-050012.nc'],'pr'),[2 1 3]),length(lat)*length(lon),12*lyrs)';

elseif strcmp(mdl,'gfdlesm2m_pi')
    
    varpth='/d1/nsteiger/climate-data/gfdl-esm2m-pi/';
    
    syr=1;eyr=500;lyrs=length(syr:eyr);
    
    lat=ncread([varpth 'tas_Amon_GFDL-ESM2M_piControl_r1i1p1_000101-050012.nc'],'lat');
    lon=ncread([varpth 'tas_Amon_GFDL-ESM2M_piControl_r1i1p1_000101-050012.nc'],'lon');
    
    % land area
    sftlf=permute(ncread([varpth 'sftlf_fx_GFDL-ESM2M_piControl_r0i0p0.nc'],'sftlf'),[2 1]);
    
    % radiation (W/m^2)
    rsds=reshape(permute(ncread([varpth 'rsds_Amon_GFDL-ESM2M_piControl_r1i1p1_000101-050012.nc'],'rsds'),[2 1 3]),length(lat)*length(lon),12*lyrs)';
    rsus=reshape(permute(ncread([varpth 'rsus_Amon_GFDL-ESM2M_piControl_r1i1p1_000101-050012.nc'],'rsus'),[2 1 3]),length(lat)*length(lon),12*lyrs)';
    rlds=reshape(permute(ncread([varpth 'rlds_Amon_GFDL-ESM2M_piControl_r1i1p1_000101-050012.nc'],'rlds'),[2 1 3]),length(lat)*length(lon),12*lyrs)';
    rlus=reshape(permute(ncread([varpth 'rlus_Amon_GFDL-ESM2M_piControl_r1i1p1_000101-050012.nc'],'rlus'),[2 1 3]),length(lat)*length(lon),12*lyrs)';
    
    % t_emperature (C):
    t_c=reshape(permute(ncread([varpth 'tas_Amon_GFDL-ESM2M_piControl_r1i1p1_000101-050012.nc'],'tas'),[2 1 3]),length(lat)*length(lon),12*lyrs)'-273.15;
    %t_f = t_c*33.8; % C to  F
    
    % Surface pressure (Pa):
    ps=reshape(permute(ncread([varpth 'ps_Amon_GFDL-ESM2M_piControl_r1i1p1_000101-050012.nc'],'ps'),[2 1 3]),length(lat)*length(lon),12*lyrs)';
    
    % Vapor pressure (Pa):
    huss=reshape(permute(ncread([varpth 'huss_Amon_GFDL-ESM2M_piControl_r1i1p1_000101-050012.nc'],'huss'),[2 1 3]),length(lat)*length(lon),12*lyrs)';
    
    % Surface wind (m/s) (CMIP5 output is at 10m)
    u10=reshape(permute(ncread([varpth 'sfcWind_Amon_GFDL-ESM2M_piControl_r1i1p1_000101-050012.nc'],'sfcWind'),[2 1 3]),length(lat)*length(lon),12*lyrs)';
    
    % precipitation, will be converted from kg*m^-2*s^-1 to (mm/day):
    pr=reshape(permute(ncread([varpth 'pr_Amon_GFDL-ESM2M_piControl_r1i1p1_000101-050012.nc'],'pr'),[2 1 3]),length(lat)*length(lon),12*lyrs)';

elseif strcmp(mdl,'cesm_lme')
    
    varpth=['/d6/haibo/lme/cam5/BLMTRC5CN/ALLFORCEINGS/' ens];
    ensnm=ens;
    %varpth=['/d6/haibo/lme/cam5/BLMTRC5CN/' enstyp '/' ens];
    %ensnm=[enstyp '.' ens];

    %syr=850;eyr=1849;lyrs=length(syr:eyr);
    syr=850;eyr=2005;lyrs=length(syr:eyr);
    
    lat=ncread([varpth '/TREFHT/b.e11.BLMTRC5CN.f19_g16.' ensnm '.cam.h0.TREFHT.085001-184912.nc'],'lat');
    lon=ncread([varpth '/TREFHT/b.e11.BLMTRC5CN.f19_g16.' ensnm '.cam.h0.TREFHT.085001-184912.nc'],'lon');
    
    % land area
    sftlf=100*permute(ncread([varpth '/QSOIL/b.e11.BLMTRC5CN.f19_g16.' ensnm '.clm2.h0.QSOIL.085001-184912.nc'],'landmask'),[2 1]);
    
    % radiation (W/m^2)
    FSNS1=ncread([varpth '/FSNS/b.e11.BLMTRC5CN.f19_g16.' ensnm '.cam.h0.FSNS.085001-184912.nc'],'FSNS');
    FSNS2=ncread([varpth '/FSNS/b.e11.BLMTRC5CN.f19_g16.' ensnm '.cam.h0.FSNS.185001-200512.nc'],'FSNS');
    FSNS=reshape(permute(cat(3,FSNS1,FSNS2),[2 1 3]),length(lat)*length(lon),12*lyrs)';
    FLNS1=ncread([varpth '/FLNS/b.e11.BLMTRC5CN.f19_g16.' ensnm '.cam.h0.FLNS.085001-184912.nc'],'FLNS');
    FLNS2=ncread([varpth '/FLNS/b.e11.BLMTRC5CN.f19_g16.' ensnm '.cam.h0.FLNS.185001-200512.nc'],'FLNS');
    FLNS=reshape(permute(cat(3,FLNS1,FLNS2),[2 1 3]),length(lat)*length(lon),12*lyrs)';
    % t_emperature (C):
    Xv_mon1=ncread([varpth '/TREFHT/b.e11.BLMTRC5CN.f19_g16.' ensnm '.cam.h0.TREFHT.085001-184912.nc'],'TREFHT')-273.15;
    Xv_mon2=ncread([varpth '/TREFHT/b.e11.BLMTRC5CN.f19_g16.' ensnm '.cam.h0.TREFHT.185001-200512.nc'],'TREFHT')-273.15;
    t_c=reshape(permute(cat(3,Xv_mon1,Xv_mon2),[2 1 3]),length(lat)*length(lon),12*lyrs)';
    
    %t_c=reshape(permute(ncread([varpth '/TREFHT/b.e11.BLMTRC5CN.f19_g16.' ensnm '.cam.h0.TREFHT.085001-184912.nc'],'TREFHT'),[2 1 3]),length(lat)*length(lon),12*lyrs)'-273.15;
    %t_f = t_c*33.8; % C to  F
   
 
    % Surface pressure (Pa):
    ps1=ncread([varpth '/PS/b.e11.BLMTRC5CN.f19_g16.' ensnm '.cam.h0.PS.085001-184912.nc'],'PS');
    ps2=ncread([varpth '/PS/b.e11.BLMTRC5CN.f19_g16.' ensnm '.cam.h0.PS.185001-200512.nc'],'PS');
    ps=reshape(permute(cat(3,ps1,ps2),[2 1 3]),length(lat)*length(lon),12*lyrs)';
    
    % Specific humidity (kg/kg):
    huss1=ncread([varpth '/QBOT/b.e11.BLMTRC5CN.f19_g16.' ensnm '.clm2.h0.QBOT.085001-184912.nc'],'QBOT');
    huss2=ncread([varpth '/QBOT/b.e11.BLMTRC5CN.f19_g16.' ensnm '.clm2.h0.QBOT.185001-200512.nc'],'QBOT');
    huss=reshape(permute(cat(3,huss1,huss2),[2 1 3]),length(lat)*length(lon),12*lyrs)';
    disp('QBOT is only over land...')
    
    % Surface wind (m/s) (CMIP5 output is at 10m)
    u101=ncread([varpth '/WIND/b.e11.BLMTRC5CN.f19_g16.' ensnm '.clm2.h0.WIND.085001-184912.nc'],'WIND');
    u102=ncread([varpth '/WIND/b.e11.BLMTRC5CN.f19_g16.' ensnm '.clm2.h0.WIND.185001-200512.nc'],'WIND');
    u10=reshape(permute(cat(3,u101,u102),[2 1 3]),length(lat)*length(lon),12*lyrs)';
    disp('Wind is only over land...')
    
    % precipitation, will be converted from kg*m^-2*s^-1 to (mm/day): NOTE: kg/m^2/s - the same as mm/s
    %pr=reshape(permute(ncread([varpth '/RAIN/b.e11.BLMTRC5CN.f19_g16.' ensnm '.clm2.h0.RAIN.085001-184912.nc'],'RAIN'),[2 1 3]),length(lat)*length(lon),12*lyrs)';
    % USE FULL PRECIPITATION VARIABLE AND BIAS CORRECTION 
    Xv_mon01=ncread(['/d1/nsteiger/climate-data/cesm-lme/pr/b.e11.BLMTRC5CN.f19_g16.' ens '.cam.h0.PRECC.085001-184912.nc'],'PRECC');
    Xv_mon02=ncread(['/d1/nsteiger/climate-data/cesm-lme/pr/b.e11.BLMTRC5CN.f19_g16.' ens '.cam.h0.PRECL.085001-184912.nc'],'PRECL');
    Xv_mon001=ncread(['/d1/nsteiger/climate-data/cesm-lme/pr/b.e11.BLMTRC5CN.f19_g16.' ens '.cam.h0.PRECC.185001-200512.nc'],'PRECC');
    Xv_mon002=ncread(['/d1/nsteiger/climate-data/cesm-lme/pr/b.e11.BLMTRC5CN.f19_g16.' ens '.cam.h0.PRECL.185001-200512.nc'],'PRECL');
    Xv_mon1=cat(3,Xv_mon01,Xv_mon001);
    Xv_mon2=cat(3,Xv_mon02,Xv_mon002);
    Xv_mon=Xv_mon1+Xv_mon2; % precip = convective + large scale
    Xv_mon=permute(Xv_mon,[2 1 3]);

    pr=reshape(Xv_mon,length(lat)*length(lon),12*lyrs)';
    % Convert pr to kg/(m^2*s) = mm/s from mm/day
    pr=pr./86400;

    %pr=reshape(permute(ncread([varpth '/RAIN/b.e11.BLMTRC5CN.f19_g16.' ensnm '.clm2.h0.RAIN.085001-184912.nc'],'RAIN'),[2 1 3]),length(lat)*length(lon),12*lyrs)';
    %disp('Precip is only over land...')
    
elseif strcmp(mdl,'ccsm4_lm')
    
    varpth='/d1/nsteiger/climate-data/ccsm4-lm/';
    
    syr=850;eyr=1850;lyrs=length(syr:eyr);
    
    lat=ncread([varpth 'b40.lm850-1850.1deg.001.cam2.h0.TREFHT.085001-185012.nc'],'lat');
    lon=ncread([varpth 'b40.lm850-1850.1deg.001.cam2.h0.TREFHT.085001-185012.nc'],'lon');
    
    % land area
    sftlf=double(100*permute(ncread([varpth 'b40.lm850-1850.1deg.001.clm2.h0.QSOIL.085001-185012.nc'],'landmask'),[2 1]));
    
    % radiation (W/m^2)
    FSNS=reshape(permute(ncread([varpth 'b40.lm850-1850.1deg.001.cam2.h0.FSNS.085001-185012.nc'],'FSNS'),[2 1 3]),length(lat)*length(lon),12*lyrs)';
    FLNS=reshape(permute(ncread([varpth 'b40.lm850-1850.1deg.001.cam2.h0.FLNS.085001-185012.nc'],'FLNS'),[2 1 3]),length(lat)*length(lon),12*lyrs)';
        
    % t_emperature (C):
    t_c=reshape(permute(ncread([varpth 'b40.lm850-1850.1deg.001.cam2.h0.TREFHT.085001-185012.nc'],'TREFHT'),[2 1 3]),length(lat)*length(lon),12*lyrs)'-273.15;
    %t_f = t_c*33.8; % C to  F
    
    % Surface pressure (Pa):
    ps=reshape(permute(ncread([varpth 'b40.lm850-1850.1deg.001.cam2.h0.PS.085001-185012.nc'],'PS'),[2 1 3]),length(lat)*length(lon),12*lyrs)';
    
    % Specific humidity (kg/kg):
    huss=reshape(permute(ncread([varpth 'b40.lm850-1850.1deg.001.clm2.h0.QBOT.085001-185012.nc'],'QBOT'),[2 1 3]),length(lat)*length(lon),12*lyrs)';
    disp('QBOT is only over land...')
    
    % Surface wind (m/s) (CMIP5 output is at 10m)
    u10=reshape(permute(ncread([varpth 'b40.lm850-1850.1deg.001.clm2.h0.WIND.085001-185012.nc'],'WIND'),[2 1 3]),length(lat)*length(lon),12*lyrs)';
    disp('Wind is only over land...')
    
    % precipitation, [m/s] convert to mm/s == kg/m^2/s
    pr=reshape(permute(ncread([varpth 'b40.lm850-1850.1deg.001.cam2.h0.PRECT.085001-185012.nc'],'PRECT'),[2 1 3]),length(lat)*length(lon),12*lyrs)';
    pr=pr*1e3; %convert to mm/s == kg/m^2/s
    
end

