def read_yt_zm_pl_daily(var, level, dir_in, dim): """ read daily zonal-mean yt data on a given pressure level across multiple years """ import numpy as np import xarray as xr from ds21grl.misc import month_to_dayofyear, leap_year_test data = np.zeros((dim.years.size, dim.nday, dim.lat.size)) var = translate_var_CESM_to_erai(var) for iyear in range(0, dim.years.size): for imonth in range(0, dim.months.size): print(dim.years[iyear], dim.months[imonth]) timestamp = str(dim.years[iyear]) + '-' + format( dim.months[imonth], "02") filename = dir_in + 'pressure_level/' + var + str( level) + '/' + var + '_' + timestamp + '.nc' ds = xr.open_dataset(filename).sel(method='nearest').mean( dim='lon') data_monthly = ds[var].values day = ds['day'].values ds.close() # remove leap year day if (leap_year_test(dim.years[iyear]) == 1) & (dim.months[imonth] == 2): data_monthly = data_monthly[:28, :] day = day[:28] daysofyear = np.arange(month_to_dayofyear(imonth - 1), month_to_dayofyear(imonth)) data[iyear, daysofyear, :] = data_monthly[:, :] return data
def calc_ykt_daily_VQ_eddy_erai(dir_in, dim): """ Calculates zonally and vertically integrated moisture transport as a function of zonal wavenumber for all eddies (transient + stw). This is a wrapper for the main transport calculation function """ import numpy as np import xarray as xr from ds21grl import const from ds21grl.misc import tic, toc, leap_year_test, month_to_dayofyear # Initialize VQ = np.zeros( (dim.years.size, dim.nday, dim.lat.size, dim.wavenumber.size)) for iyear in range(0, dim.years.size): for imonth in range(0, dim.months.size): tic() # read data for each month timestamp = str(dim.years[iyear]) + '-' + format( dim.months[imonth], "02") filename1 = dir_in + '/model_level/v/v_' + timestamp + '.nc' filename2 = dir_in + '/model_level/q/q_' + timestamp + '.nc' filename3 = dir_in + '/sfc/lnsp/lnsp_' + timestamp + '.nc' ds1 = xr.open_dataset(filename1) ds2 = xr.open_dataset(filename2) ds3 = xr.open_dataset(filename3) day = ds1.day.values dim.P0 = ds1.P0.values dim.hyai = ds1.hyai.values / dim.P0 # set units same as model dim.hybi = ds1.hybi.values dim.lev = ds1.lev.values dim.ilev = ds1.ilev.values V = ds1.v.values Q = ds2.q.values PS = np.exp(ds3.lnsp.values) ds1.close() ds2.close() ds3.close() # remove leap year day if (leap_year_test(dim.years[iyear]) == 1) & (dim.months[imonth] == 2): V = V[:28, :, :, :] Q = Q[:28, :, :, :] PS = PS[:28, :, :] day = day[:28] # define energy + mass flux terms Q = const.Lv * Q dp = compute_dp_ml(PS, dim) dp = np.swapaxes(dp, 0, 1) # shift dims of dp same as V = V * dp / const.go # calc transport VQ_month = calc_ykt_transport(V, Q, dim) # dump into long time series array and convert to PW daysofyear = np.arange(month_to_dayofyear(imonth - 1), month_to_dayofyear(imonth)) VQ[iyear, daysofyear, :, :] = VQ_month[:, :, :] / 10**15 toc() return VQ
def read_yt_zm_sfc_daily(var, dir_in, dim): """ read daily zonal-mean yt surface data across multiple years in monthly format """ import numpy as np import xarray as xr from fxn.misc import month_to_dayofyear, leap_year_test data = np.zeros((dim.years.size, dim.nday, dim.lat.size)) var = translate_var_CESM_to_erai(var) for iyear in range(0, dim.years.size): for imonth in range(0, dim.months.size): timestamp = str(dim.years[iyear]) + '-' + format( dim.months[imonth], "02") print(dim.years[iyear], dim.months[imonth]) if var == 'CFS': filename1 = dir_in + 'sfc/ssr/ssr_' + timestamp + '.nc' filename2 = dir_in + 'sfc/ssrs/ssrs_' + timestamp + '.nc' filename3 = dir_in + 'sfc/str/str_' + timestamp + '.nc' filename4 = dir_in + 'sfc/strc/strc_' + timestamp + '.nc' ds1 = xr.open_dataset(filename1).mean(dim='longitude') ds2 = xr.open_dataset(filename2).mean(dim='longitude') ds3 = xr.open_dataset(filename3).mean(dim='longitude') ds4 = xr.open_dataset(filename4).mean(dim='longitude') day = ds1.day.values data_monthly = ds1['ssr'].values - ds2['ssrc'].values + ds3[ 'str'].values - ds4['strc'].values ds1.close() ds2.close() ds3.close() ds4.close() elif var == 'SWCFS': filename1 = dir_in + 'sfc/ssr/ssr_' + timestamp + '.nc' filename2 = dir_in + 'sfc/ssrs/ssrs_' + timestamp + '.nc' ds1 = xr.open_dataset(filename1).mean(dim='longitude') ds2 = xr.open_dataset(filename2).mean(dim='longitude') day = ds1.day.values data_monthly = ds1['ssr'].values - ds2['ssrc'].values ds1.close() ds2.close() elif var == 'LWCFS': filename1 = dir_in + 'sfc/str/str_' + timestamp + '.nc' filename2 = dir_in + 'sfc/strc/strc_' + timestamp + '.nc' ds1 = xr.open_dataset(filename1).mean(dim='longitude') ds2 = xr.open_dataset(filename2).mean(dim='longitude') day = ds1.day.values data_monthly = ds1['str'].values - ds2['strc'].values ds1.close() ds2.close() else: filename1 = dir_in + 'sfc/' + var + '/' + var + '_' + timestamp + '.nc' ds1 = xr.open_dataset(filename1).mean(dim='longitude') day = ds1.day.values data_monthly = ds1[var].values ds1.close() # remove leap year day if (leap_year_test(dim.years[iyear]) == 1) & (dim.months[imonth] == 2): data_monthly = data_monthly[:28, :] day = day[:28] daysofyear = np.arange(month_to_dayofyear(imonth - 1), month_to_dayofyear(imonth)) data[iyear, daysofyear, :] = data_monthly[:, :] return data