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