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)})
