Script for calculation of SSW composites (split of displacement)

composite_construction.py

def open_date_file(file_path):
    df = pd.read_csv(file_path, index_col=0, parse_dates=True)
    df['BeginDate'] = df.BeginDate.apply(lambda t: pd.to_datetime(t, format='%Y-%m-%d'))
    return df


var = sys.argv[1]
path = ''
max_lag = 30
type_ls = ['split', 'displ']
files_path = path+var+'/'+var+'_6hrPlev_CMAM_CMAM30-SD_r1i1p1_????????00-*.nc'
print('opening '+files_path)

with xr.open_mfdataset(files_path) as ds:
    print(ds.chunks)
    print('anomalies calculation')
    climatology = ds.groupby('time.month').mean('time')
    anomalies = ds.groupby('time.month') - climatology
    anomalies = anomalies.chunk({'time': 1460})
 
    print('composites construction')
    for di, ssw_type in enumerate(type_ls):
        print(ssw_type)
        df_dates = open_date_file('ssw_dates_'+ssw_type+'.csv')
        xa_ls = []
        for il, lag in enumerate(xrange(-max_lag,max_lag+1)):
            dates = df_dates.BeginDate+pd.Timedelta(str(lag)+' days')
            comp = anomalies[var].sel(time = dates.values).mean('time')
            xa_ls.append(comp)
            
        print('concatenation')
        xa_comp = xr.concat(xa_ls, dim = 'lag')
        xa_ls = []
        xa_comp['lag'] = range(-max_lag, max_lag+1)
        ds_comp = xa_comp.to_dataset(name = 'zg')
       
        print('saving')
        ds_comp.to_netcdf(cesta+'composites/'+var+'_anomalies_comp_'+type_ls[di]+'.nc')
        

print('done')