def v_components(ubudget, scale=None, eqbuf=5.0):
"""Return mean, eddy-driven, etc components of v for streamfunction.
"""
comp_dict = {'MMC' : 'ADV_AVG', 'PGF' : 'PGF_ST', 'EDDY_ST' : 'EMFC_ST',
'EDDY_TR' : 'EMFC_TR', 'EDDY_CRS' : 'ADV_CRS'}
if scale is not None:
ubudget = ubudget * scale
latname = atm.get_coord(ubudget, 'lat', 'name')
lat = ubudget[latname]
f = atm.coriolis(lat)
f[abs(lat) < eqbuf] = np.nan
v = xray.Dataset()
v['TOT'] = ubudget['COR'] / f
for nm in sorted(comp_dict):
v[nm] = - ubudget[comp_dict[nm]] / f
v['EDDY'] = v['EDDY_CRS'] + v['EDDY_TR'] + v['EDDY_ST']
v['RESID'] = v['TOT'] - v['MMC'] - v['PGF'] - v['EDDY']
return v
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
