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