def pgradient(var, lat1, lat2, lon1, lon2, plev):
"""Return d/dp of a lat-lon variable."""
pwidth = 100
p1, p2 = plev - pwidth, plev + pwidth
var = atm.subset(var, {'lat' : (lat1, lat2), 'lon' : (lon1, lon2),
'plev' : (p1, p2)}, copy=False)
latlonstr = latlon_filestr(lat1, lat2, lon1, lon2)
attrs = var.attrs
pname = atm.get_coord(var, 'plev', 'name')
pdim = atm.get_coord(var, 'plev', 'dim')
pres = var[pname]
pres = atm.pres_convert(pres, pres.attrs['units'], 'Pa')
dvar_dp = atm.gradient(var, pres, axis=pdim)
dvar_dp = atm.subset(dvar_dp, {pname : (plev, plev)}, copy=False,
squeeze=True)
varnm = 'D%sDP' % var.name
name = '%s%d' % (varnm, plev)
dvar_dp.name = name
attrs['long_name'] = 'd/dp of ' + var.attrs['long_name']
attrs['standard_name'] = 'd/dp of ' + var.attrs['standard_name']
attrs['units'] = ('(%s)/Pa' % attrs['units'])
attrs[pname] = plev
attrs['filestr'] = '%s_%s' % (name, latlonstr)
attrs['varnm'] = varnm
dvar_dp.attrs = attrs
return dvar_dp
def advection(uflow, vflow, omegaflow, u, dudp):
"""Return x, y and p components of advective terms in momentum budget.
"""
a = atm.constants.radius_earth
latlon = latlon_data(u)
latdim, londim = latlon.attrs['latdim'], latlon.attrs['londim']
latrad, coslat = latlon['LATRAD'], latlon['COSLAT']
if londim is not None:
lonrad = latlon['LONRAD']
ds = xray.Dataset()
if londim is not None:
ds['X'] = atm.gradient(u, lonrad, londim) * uflow / (a*coslat)
else:
ds['X'] = 0.0 * u
ds['Y'] = atm.gradient(u*coslat, latrad, latdim) * vflow / (a*coslat)
ds['P'] = omegaflow * dudp
return ds
def fluxdiv(u, v, omega, dudp, domegadp):
"""Return x, y and p components of EMFD terms in momentum budget.
"""
a = atm.constants.radius_earth
latlon = latlon_data(u)
latdim, londim = latlon.attrs['latdim'], latlon.attrs['londim']
latrad, coslat = latlon['LATRAD'], latlon['COSLAT']
coslat = latlon['COSLAT']
coslat_sq = coslat ** 2
if londim is not None:
lonrad = latlon['LONRAD']
ds = xray.Dataset()
if londim is not None:
ds['X'] = atm.gradient(u * u, lonrad, londim) / (a*coslat)
else:
ds['X'] = 0.0 * u
ds['Y'] = atm.gradient(u * v * coslat_sq, latrad, latdim) / (a*coslat_sq)
ds['P'] = omega * dudp + u * domegadp
return ds
def calc_dp(var, plev):
"""Extract subset of pressure levels and calculate d/dp."""
plevs = atm.get_coord(var, 'plev')
pname = atm.get_coord(var, 'plev', 'name')
pdim = atm.get_coord(var, 'plev', 'dim')
ind = (list(plevs)).index(plev)
i1 = max(0, ind - 1)
i2 = min(len(plevs) - 1, ind + 1) + 1
psub = plevs[i1:i2]
varsub = var.sel(**{pname : psub})
pres = atm.pres_convert(psub, 'hPa', 'Pa')
atm.disptime()
print('Computing d/dp for pressure level %d' % plev)
dvar = atm.gradient(varsub, pres, axis=pdim)
dvar = dvar.sel(**{pname : plev})
dvar.name = 'D%sDP' % var.name
atm.disptime()
return dvar
vq_int = get_var(datadir, 'VFLXQV', subset, year)
mfc = atm.moisture_flux_conv(uq_int, vq_int, already_int=True)
mfcbar = mfc.mean(dim='YDim').mean(dim='XDim')
# Test atm.gradient
a = atm.constants.radius_earth.values
latdim, londim = 1, 2
lat = atm.get_coord(uq_int, 'lat')
latrad = np.radians(lat)
latrad[abs(lat) > 89] = np.nan
coslat = xray.DataArray(np.cos(latrad), coords={'YDim': lat})
lon = atm.get_coord(uq_int, 'lon')
lonrad = np.radians(lon)
mfc_x = atm.gradient(uq_int, lonrad, londim) / (a * coslat)
mfc_y = atm.gradient(vq_int * coslat, latrad, latdim) / (a * coslat)
mfc_test = mfc_x + mfc_y
mfc_test = -atm.precip_convert(mfc_test, 'kg/m2/s', 'mm/day')
mfc_test_bar = mfc_test.mean(dim='YDim').mean(dim='XDim')
diff = mfc_test - mfc
print(diff.max())
print(diff.min())
plt.plot(mfcbar)
plt.plot(mfc_test_bar)
print(mfc_test_bar - mfcbar)
# ----------------------------------------------------------------------
# Vertical gradient du/dp
def calc_ubudget(datafiles, ndays, lon1, lon2, plev=200):
"""Calculate momentum budget for daily data in one year.
Keys of datafiles dict must be: U, V, DUDP, H, OMEGA, DOMEGADP, DUDTANA
"""
# Read data
data = xray.Dataset()
for nm in datafiles:
print('Reading ' + datafiles[nm])
with xray.open_dataset(datafiles[nm]) as ds:
if nm in ds.data_vars:
var = ds[nm]
else:
var = ds[nm + '%d' % plev]
if 'Day' in var.dims:
var = var.rename({'Day' : 'day'})
data[nm] = atm.squeeze(var)
data['PHI'] = atm.constants.g.values * data['H']
# Put zeros in for any missing variables (e.g. du/dp)
for nm in ['OMEGA', 'DUDP', 'DOMEGADP', 'DUDTANA']:
if nm not in data.data_vars:
data[nm] = 0.0 * data['U']
# Eddy decomposition
taxis = 0
for nm in data.data_vars:
print('Eddy decomposition for ' + nm)
comp = eddy_decomp(data[nm], ndays, lon1, lon2, taxis)
for compnm in comp:
data[compnm] = comp[compnm]
# Momentum budget calcs
# du/dt = sum of terms in ubudget
ubudget = xray.Dataset()
readme = 'Momentum budget: ACCEL = sum of all other data variables'
ubudget.attrs['readme'] = readme
ubudget.attrs['ndays'] = ndays
ubudget.attrs['lon1'] = lon1
ubudget.attrs['lon2'] = lon2
# Advective terms
keypairs = [ ('AVG', 'AVG'), ('AVG', 'ST'), ('ST', 'AVG')]
print('Computing advective terms')
for pair in keypairs:
print(pair)
ukey, flowkey = pair
u = data['U_' + ukey]
dudp = data['DUDP_' + ukey]
uflow = data['U_' + flowkey]
vflow = data['V_' + flowkey]
omegaflow = data['OMEGA_' + flowkey]
adv = advection(uflow, vflow, omegaflow, u, dudp)
for nm in adv.data_vars:
key = 'ADV_%s_%s_%s' % (ukey, flowkey, nm)
ubudget[key] = - adv[nm]
long_name = 'Advection of %s momentum by %s' % (ukey, flowkey)
ubudget[key].attrs['long_name'] = long_name
# EMFD terms
keys = ['TR', 'ST']
print('Computing EMFD terms')
for key in keys:
print(key)
u = data['U_' + key]
v = data['V_' + key]
omega = data['OMEGA_' + key]
dudp = data['DUDP_' + key]
domegadp = data['DOMEGADP_' + key]
emfd = fluxdiv(u, v, omega, dudp, domegadp)
for nm in emfd.data_vars:
ubudget['EMFC_%s_%s' % (key, nm)] = - emfd[nm]
# Coriolis terms
latlon = latlon_data(data['V_ST'])
lat = latlon['LAT']
f = atm.coriolis(lat)
ubudget['COR_AVG'] = data['V_AVG'] * f
ubudget['COR_ST'] = data['V_ST'] * f
# Pressure gradient terms
a = atm.constants.radius_earth.values
coslat = latlon['COSLAT']
lonrad = latlon['LONRAD']
londim = atm.get_coord(data['PHI_ST'], 'lon', 'dim')
ubudget['PGF_ST'] = - atm.gradient(data['PHI_ST'], lonrad, londim) / (a*coslat)
# Analysis increment for dU/dt
ubudget['ANA'] = data['DUDTANA']
# Time mean
print('Computing rolling time mean')
for nm in ubudget.data_vars:
ubudget[nm] = atm.rolling_mean(ubudget[nm], ndays, axis=taxis, center=True)
# Acceleration
nseconds = 60 * 60 * 24 * ndays
delta_u = np.nan * data['U']
u = data['U'].values
delta_u.values[ndays//2:-ndays//2] = (u[ndays:] - u[:-ndays]) / nseconds
ubudget['ACCEL'] = delta_u
return ubudget, data
lat1, lat2 = 10, 30
subset_dict = {'lat' : (lat1, lat2), 'lon' : (lon1, lon2)}
# ----------------------------------------------------------------------
# Read data
ts = xray.Dataset()
for nm in files:
var = atm.combine_daily_years(nm, files[nm], years, yearname='year',
subset_dict=subset_dict)
var = atm.mean_over_geobox(var, lat1, lat2, lon1, lon2)
units = var.attrs.get('units')
if units in ['kg/m2/s', 'kg m-2 s-1']:
var = atm.precip_convert(var, units, 'mm/day')
elif units in ['kg/m2', 'kg m-2']:
var.attrs['units'] = 'mm'
long_name = var.attrs.get('long_name')
var.attrs = {'units' : units, 'long_name' : long_name}
ts[nm] = utils.wrapyear_all(var, daymin, daymax)
# dW/dt
dwdt = atm.gradient(ts['TQV'], ts['day'].values, axis=-1)
dwdt.attrs['units'] = 'mm/day'
dwdt.attrs['long_name'] = 'd/dt of ' + ts['TQV'].attrs.get('long_name')
ts['DWDT'] = dwdt
ts = ts.rename({'DQVDT_ANA' : 'DWDT_ANA'}).drop('TQV')
for nm, val in zip(['lat1', 'lat2', 'lon1', 'lon2'], [lat1, lat2, lon1, lon2]):
ts.attrs[nm] = val
print('Saving to ' + savefile)
ts.to_netcdf(savefile)
vq_int = get_var(datadir, 'VFLXQV', subset, year)
mfc = atm.moisture_flux_conv(uq_int, vq_int, already_int=True)
mfcbar = mfc.mean(dim='YDim').mean(dim='XDim')
# Test atm.gradient
a = atm.constants.radius_earth.values
latdim, londim = 1, 2
lat = atm.get_coord(uq_int, 'lat')
latrad = np.radians(lat)
latrad[abs(lat) > 89] = np.nan
coslat = xray.DataArray(np.cos(latrad), coords={'YDim' : lat})
lon = atm.get_coord(uq_int, 'lon')
lonrad = np.radians(lon)
mfc_x = atm.gradient(uq_int, lonrad, londim) / (a*coslat)
mfc_y = atm.gradient(vq_int * coslat, latrad, latdim) / (a*coslat)
mfc_test = mfc_x + mfc_y
mfc_test = - atm.precip_convert(mfc_test, 'kg/m2/s', 'mm/day')
mfc_test_bar = mfc_test.mean(dim='YDim').mean(dim='XDim')
diff = mfc_test - mfc
print(diff.max())
print(diff.min())
plt.plot(mfcbar)
plt.plot(mfc_test_bar)
print(mfc_test_bar - mfcbar)
# ----------------------------------------------------------------------
# Vertical gradient du/dp
