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: