def eddy_decomp(var, nt, lon1, lon2, taxis=0):
"""Decompose variable into mean and eddy fields."""
lonname = atm.get_coord(var, 'lon', 'name')
tstr = 'Time mean (%d-%s rolling)' % (nt, var.dims[taxis])
lonstr = atm.latlon_labels([lon1, lon2], 'lon', deg_symbol=False)
lonstr = 'zonal mean (' + '-'.join(lonstr) + ')'
name, attrs, coords, dims = atm.meta(var)
varbar = atm.rolling_mean(var, nt, axis=taxis, center=True)
varbarzon = atm.subset(varbar, {lonname : (lon1, lon2)})
varbarzon = varbarzon.mean(dim=lonname)
varbarzon.attrs = attrs
comp = xray.Dataset()
comp[name + '_AVG'] = varbarzon
comp[name + '_AVG'].attrs['component'] = tstr + ', ' + lonstr
comp[name + '_ST'] = varbar - varbarzon
comp[name + '_ST'].attrs = attrs
comp[name + '_ST'].attrs['component'] = 'Stationary eddy'
comp[name + '_TR'] = var - varbar
comp[name + '_TR'].attrs = attrs
comp[name + '_TR'].attrs['component'] = 'Transient eddy'
return comp
def get_mfc_box(mfcfiles, precipfiles, evapfiles, years, nroll, lat1, lat2,
lon1, lon2):
"""Return daily tseries MFC, precip and evap averaged over lat-lon box.
"""
subset_dict = {'lat' : (lat1, lat2), 'lon' : (lon1, lon2)}
databox = {}
if mfcfiles is not None:
mfc = atm.combine_daily_years('MFC', mfcfiles, years, yearname='year',
subset_dict=subset_dict)
databox['MFC'] = mfc
if precipfiles is not None:
pcp = atm.combine_daily_years('PRECTOT', precipfiles, years, yearname='year',
subset_dict=subset_dict)
databox['PCP'] = pcp
if evapfiles is not None:
evap = atm.combine_daily_years('EVAP', evapfiles, years, yearname='year',
subset_dict=subset_dict)
databox['EVAP'] = evap
nms = databox.keys()
for nm in nms:
var = databox[nm]
var = atm.precip_convert(var, var.attrs['units'], 'mm/day')
var = atm.mean_over_geobox(var, lat1, lat2, lon1, lon2)
databox[nm + '_UNSM'] = var
databox[nm + '_ACC'] = np.cumsum(var, axis=1)
if nroll is None:
databox[nm] = var
else:
databox[nm] = atm.rolling_mean(var, nroll, axis=-1, center=True)
tseries = xray.Dataset(databox)
return tseries
def annotate_theta_e(days, latmax, ax=None, nroll=7):
if ax is None:
ax = plt.gca()
if nroll is not None:
latmax = atm.rolling_mean(latmax, nroll, center=True)
latmax_0 = latmax.sel(dayrel=0)
ax.plot(days, latmax, 'k', linewidth=2, label='Latitude of Max')
ax.legend(loc='lower right', fontsize=10)
s = atm.latlon_labels(latmax_0, latlon='lat', fmt='%.1f')
ax.annotate(s, xy=(0, latmax_0), xycoords='data',
xytext=(-50, 50), textcoords='offset points',
arrowprops=dict(arrowstyle="->"))
def get_data(varnm, datafiles, regdays, seasons, lon1, lon2, nroll=None):
var, onset, retreat = utils.load_dailyrel(datafiles[varnm])
if nroll is not None:
var = atm.rolling_mean(var, nroll, axis=1, center=True)
# Seasonal averages and daily lat-lon data
data = xray.Dataset()
for season in seasons:
key = varnm + '_' + season
data[key] = ssn_average(var, onset, retreat, season)
# Daily data on regdays
data[varnm + '_DAILY'] = var.sel(dayrel=regdays)
# Sector mean data
var_sector = atm.dim_mean(var, 'lon', lon1, lon2)
alldata = {'data_latlon' : data, 'var_sector' : var_sector,
'onset' : onset, 'retreat' : retreat}
return alldata
days_ssn = atm.season_days('JJAS')
subset_dict = {'lat' : (lat1, lat2), 'lon' : (lon1, lon2)}
ts = xray.Dataset()
ssn = xray.Dataset()
varnms = {'PCP' : 'PRECTOT', 'MFC' : 'MFC', 'EVAP' : 'EVAP', 'W' : 'TQV',
'ANA' : 'DQVDT_ANA'}
for nm in files:
var = atm.combine_daily_years(varnms[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')
var_sm = atm.rolling_mean(var, nroll, axis=-1, center=True)
ts[nm] = var_sm
if nm == 'W':
# Difference in precipitable water from beginning to end of season
var_ssn = var_sm.sel(day=days_ssn)
ssn['d' + nm] = (var_ssn[:,-1] - var_ssn[:, 0]) / len(days_ssn)
# dW/dt
dvar = np.nan * np.ones(var.shape)
for y, year in enumerate(years):
dvar[y] = np.gradient(var[y])
ts['d%s/dt' % nm] = xray.DataArray(dvar, coords=var.coords)
else:
# Seasonal mean and daily timeseries
ssn[nm] = var.sel(day=days_ssn).mean(dim='day')
# Net precip and residuals
var = ds['VFLXQV'].load()
data[nm] = var * atm.constants.Lv.values
else:
data[nm] = ds[nm].load()
# Scale units and rename variables
data = data * scale
nms = data.data_vars.keys()
for nm in nms:
data = data.rename({nm : nm.replace('FLX', '')})
# Take subset and smooth with rolling mean
daydim = atm.get_coord(data['VMSE'], 'dayrel', 'dim')
for nm in data.data_vars:
data[nm] = atm.rolling_mean(data[nm], nroll, axis=daydim, center=True)
# Average over equatorial region
data_eq = atm.dim_mean(data, 'lat', eqlat1, eqlat2)
# Cross-equatorial flues integrated over sectors
a = atm.constants.radius_earth.values
eq_int = xray.Dataset()
eq_int.attrs['units'] = sector_units
lonranges = [(40, 60), (40, 100), (lon1, lon2)]
eq_int.attrs['lonranges'] = ['%dE-%dE' % lonrange for lonrange in lonranges]
for lonrange in lonranges:
lon1, lon2 = lonrange
dist = a * np.radians(lon2 - lon1)
for nm in data_eq.data_vars:
key = nm + '_%dE-%dE' % (lon1, lon2)
days = precipbar.day.values
pentad = False
precip_jan = 0.0 # Use zero for now
years = pcp.year.values
key = 'WLH_%s_kmax%d' % (name, kmax)
print(key)
pcp_sm, Rsq = atm.fourier_smooth(pcp, kmax)
index[key] = get_onset_WLH(years, days, pcp_sm, threshold, key, pentad,
precip_jan)
# Smooth with rolling mean
key = 'WLH_%s_nroll%d' % (name, nroll[name])
print(key)
pcp_sm = atm.rolling_mean(pcp, nroll[name], axis=-1, center=True)
index[key] = get_onset_WLH(years, days, pcp_sm.values, threshold, key, pentad,
precip_jan)
# Unsmoothed pentad timeserires
key = 'WLH_%s_unsmth' % name
print(key)
index[key] = get_onset_WLH(years, days, pcp, threshold, key, pentad,
precip_jan)
# ----------------------------------------------------------------------
# OCI index (Wang et al 2009) and SJKE index (Boos and Emmanuel 2009)
ds = atm.combine_daily_years(['U', 'V'], ocifiles, years)
ds = ds.rename({'Year' : 'year', 'Day' : 'day'})
u850 = atm.squeeze(ds['U'])
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
# Read data and calculate indices
# Precipitation
precip = precipdat.read_cmap(pcpfile, yearmin=min(years), yearmax=max(years))
pcp_box = atm.mean_over_geobox(precip, lat1, lat2, lon1, lon2)
# -- Interpolate to daily resolution
days = np.arange(1, 367)
pcp_i = np.nan * np.ones((len(years), len(days)))
for y, year in enumerate(years):
pcp_i[y] = np.interp(days, pcp_box['day'], pcp_box[y])
coords = {'day' : days, 'year' : years}
pcp = xray.DataArray(pcp_i, dims=['year', 'day'], coords=coords)
# Monsoon onset, retreat indices
index = utils.get_onset_indices(onset_nm, indfiles, years)
mfc = atm.rolling_mean(index['ts_daily'], nroll, center=True)
onset = index['onset']
ssn_length=index['length'].mean(dim='year')
data = {}
data['MFC'] = utils.daily_rel2onset(mfc, onset, npre, npost)
data[pcp_nm] = utils.daily_rel2onset(pcp, onset, npre, npost)
data['MFC_ACC'] = utils.daily_rel2onset(index['tseries'], onset, npre, npost)
for nm in varnms:
print('Loading ' + relfiles[nm])
with xray.open_dataset(relfiles[nm]) as ds:
if nm == 'PSI':
data[nm] = atm.streamfunction(ds['V'])
psimid = atm.subset(data[nm], {'plev' : (pmid, pmid)},
squeeze=True)
index_all[nm] = ds.load()
index = index_all[onset_nm]
index['length'] = index['retreat'] - index['onset']
onset_all = pd.DataFrame()
for nm in index_all:
onset_all[nm] = index_all[nm]['onset'].to_series()
# Onset/retreat at grid points
print('Loading ' + ptsfile)
with xray.open_dataset(ptsfile) as index_pts:
index_pts.load()
for nm in index_pts.data_vars:
if pts_xroll is not None:
index_pts[nm] = atm.rolling_mean(index_pts[nm], pts_xroll, axis=-1,
center=True)
if pts_yroll is not None:
index_pts[nm] = atm.rolling_mean(index_pts[nm], pts_yroll, axis=-2,
center=True)
# Regression of gridpoint indices onto large-scale index
print('Regression of gridpoint indices onto large-scale index')
pts_reg, pts_mask = {}, {}
for nm in index_pts.data_vars:
ind = index[nm].sel(year=index_pts['year'])
pts_reg[nm] = atm.regress_field(index_pts[nm], ind, axis=0)
pts_reg[nm]['pts_mask'] = (pts_reg[nm]['p'] >= 0.05)
# Mask out grid points where CHP index is ill-defined
def applymask(ds, mask_in):
for nm in ds.data_vars:
