def single_WLH(cmap_file, yrmin, yrmax, lat0, lon0, loc_nm, kmax, kann, onset_min=20): # Single grid point and single year/climatology precip = precipdat.read_cmap(cmap_file, yrmin, yrmax) if yrmax > yrmin: precip = precip.mean(axis=0) titlestr = 'CMAP %d-%d' % (yrmin, yrmax) else: precip = precip[0] titlestr = 'CMAP %d' % yrmin latval, ilat0 = atm.find_closest(precip.lat, lat0) lonval, ilon0 = atm.find_closest(precip.lon, lon0) d = 1.25 latstr = atm.latlon_labels(latval+d, 'lat', '%.1f', deg_symbol=False) lonstr = atm.latlon_labels(lonval+d, 'lon', '%.1f', deg_symbol=False) titlestr = '%s %s (%s, %s)' % (titlestr, loc_nm, latstr, lonstr) pcp = precip[:, ilat0, ilon0] pcp_sm, Rsq = atm.fourier_smooth(pcp, kmax) pcp_ann, Rsq_ann = atm.fourier_smooth(pcp, kann) i_onset, i_retreat, i_peak = onset_WLH_1D(pcp_sm, threshold, onset_min) plot_single_WLH(pcp, pcp_sm, pcp_ann, Rsq, Rsq_ann, i_onset, i_retreat, i_peak, kmax, kann, titlestr) return pcp, pcp_sm, pcp_ann, Rsq, Rsq_ann, i_onset, i_retreat, i_peak
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 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="->"))
import atmos as atm import precipdat import merra datadir = atm.homedir() + 'datastore/merra/daily/' years = np.arange(1979, 2015) # Day range to extract (inclusive) # Apr 1 - Oct 1 non-leap year | Mar 31 - Sep 30 leap year daymin = 91 daymax = 274 # lat, lon range to extract lon1, lon2 = 40, 120 lat1, lat2 = -60, 60 lons = atm.latlon_labels([lon1, lon2], 'lon', deg_symbol=False) lats = atm.latlon_labels([lat1, lat2], 'lat', deg_symbol=False) latlon = '%s-%s_%s-%s' % (lons[0], lons[1], lats[0], lats[1]) savefile = datadir + ('merra_precip_%s_days%d-%d_%d-%d.nc' % (latlon, daymin, daymax, years.min(), years.max())) subset_dict = {'day' : (daymin, daymax), 'lat' : (lat1, lat2), 'lon' : (lon1, lon2)} for y, year in enumerate(years): datafile = datadir + 'merra_precip_%d.nc' % year print('Loading ' + datafile) with xray.open_dataset(datafile) as ds: precip1 = atm.subset(ds['PRECTOT'], subset_dict) precip1 = precip1.load() precip1.coords['year'] = year
uv = atm.subset(uv, {'day' : (daymin, daymax)}) uv = atm.squeeze(uv) uv = uv.drop('Height') # Average over box tseries = xray.Dataset() for nm in varnames: var = atm.squeeze(uv[nm]) label = nm + '200_box' tseries[label] = atm.mean_over_geobox(var, lat1, lat2, lon1, lon2) # Ro averaged in 10 degree latitude bins edges = range(-30, 31, 10) binvals = [(edge1, edge2) for edge1, edge2 in zip(edges[:-1], edges[1:])] for i, lats in enumerate(binvals): latlabel = '-'.join(atm.latlon_labels(np.array(lats), 'lat', deg_symbol=False)) key = 'Ro200_' + latlabel tseries[key] = atm.mean_over_geobox(uv['Ro'], lats[0], lats[1], lon1, lon2) # Apply 7-day rolling mean to each tseries nroll = 7 axis = -1 center = True for key in tseries.data_vars.keys(): tseries[key] = atm.rolling_mean(tseries[key], nroll, axis, center) # Add other timeseries for key in index.keys(): tseries[short[key]] = atm.subset(index[key]['tseries'], {'day' : (daymin, daymax)}) tseries['MFC_box'] = tseries['W_MP_n7']
for nm in ['MFC', 'MFC_ACC', pcp_nm]: tseries[nm] = databar[nm] for nm in varnms: print(nm) var = atm.subset(databar[nm], {'dayrel' : (-npre, npost)}) lat = atm.get_coord(var, 'lat') if nm == 'PSI': var = atm.subset(var, {'lat' : (-25, 10)}) latname = atm.get_coord(var, 'lat', 'name') pname = atm.get_coord(var, 'plev', 'name') var_out = var.max(dim=latname).max(dim=pname) tseries['PSIMAX'] = atm.rolling_mean(var_out, nroll, center=True) else: for lat0 in lat_extract: lat0_str = atm.latlon_labels(lat0, 'lat', deg_symbol=False) key = nm + '_' + lat0_str val, ind = atm.find_closest(lat, lat0) var_out = atm.squeeze(var[:, ind]) tseries[key] = atm.rolling_mean(var_out, nroll, center=True) # ---------------------------------------------------------------------- # Functions for plotting fmt_axes = atm.ax_lims_ticks clear_labels = atm.clear_labels to_dataset = atm.to_dataset def contourf_latday(var, clev=None, title='', nc_pref=40, grp=None, xlims=(-120, 200), xticks=np.arange(-120, 201, 30), ylims=(-60, 60), yticks=np.arange(-60, 61, 20),
# Daily timeseries ts = xray.Dataset() for nm in ['GPCP', 'PRECTOT']: ts[nm] = atm.mean_over_geobox(data[nm], lat1, lat2, lon1, lon2) ts['MFC'] = utils.daily_rel2onset(index_all['CHP_MFC']['daily_ts'], index[ind_nm], npre, npost) ts['CMFC'] = utils.daily_rel2onset(index_all['CHP_MFC']['tseries'], index[ind_nm], npre, npost) # Extract variables at specified latitudes for nm, lat0 in lat_extract.iteritems(): var = atm.dim_mean(data[nm], 'lon', lon1, lon2) lat = atm.get_coord(var, 'lat') lat0_str = atm.latlon_labels(lat0, 'lat', deg_symbol=False) # key = nm + '_' + lat0_str key = nm lat_closest, _ = atm.find_closest(lat, lat0) print '%s %.2f %.2f' % (nm, lat0, lat_closest) ts[key] = atm.subset(var, {'lat' : (lat_closest, None)}, squeeze=True) # Compute climatology and smooth with rolling mean if 'year' in ts.dims: ts = ts.mean(dim='year') if nroll is not None: for nm in ts.data_vars: ts[nm] = atm.rolling_mean(ts[nm], nroll, center=True) tseries = atm.subset(ts, {'dayrel' : (-npre, npost)})