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 onset_WLH(precip, axis=1, kmax=12, threshold=5.0, onset_min=20):
"""Return monsoon onset index computed by Wang & LinHo 2002 method.
Smoothes multi-dimensional pentad precipitation data and computes
onset indices at each point.
Parameters
----------
precip : ndarray
Pentad precipitation data with pentad as the first or second
dimension. Maximum 4D: [year, pentad, lat, lon].
axis : {0, 1}, optional
Axis corresponding to pentad dimension.
kmax : int, optional
Maximum Fourier harmonic for smoothing the input data.
threshold : float, optional
Threshold for onset/withdrawal criteria. Same units as precip.
onset_min: int, optional
Minimum pentad index allowed for onset.
Returns
-------
output : dict
Dict with the following fields:
precip_sm : ndarray, smoothed precip data
Rsq : R-squared for smoothed data
onset : ndarray, pentad index of onset
retreat : ndarray, pentad index of retreat
peak : ndarray, pentad index of peak rainfall
smoothing_kmax, threshold : values used in computation
Pentads are indexed 0-72.
Reference
---------
Wang, B., & LinHo. (2002). Rainy Season of the Asian-Pacific
Summer Monsoon. Journal of Climate, 15(4), 386-398.
"""
nmax = 4
ndim = precip.ndim
if ndim > nmax:
raise ValueError('Too many dimensions in precip. Max %dD' % nmax)
# Smooth with truncated Fourier series
precip_sm, Rsq = atm.fourier_smooth(precip, kmax, axis=axis)
# Add singleton dimension for looping
if axis == 0:
precip_sm = np.expand_dims(precip_sm, 0)
elif axis != 1:
raise ValueError('Invalid axis %d. Must be 0 or 1.' % axis)
while precip_sm.ndim < nmax:
precip_sm = np.expand_dims(precip_sm, -1)
# Calculate indices for each year and grid point
dims = precip_sm.shape
dims_out = list(dims)
dims_out.pop(1)
onset = np.nan * np.ones(dims_out)
retreat = np.nan * np.ones(dims_out)
peak = np.nan * np.ones(dims_out)
for y in range(dims[0]):
for i in range(dims[2]):
for j in range(dims[3]):
inds = onset_WLH_1D(precip_sm[y,:,i,j], threshold, onset_min)
onset[y,i,j] = inds[0]
retreat[y,i,j] = inds[1]
peak[y,i,j] = inds[2]
# Pack everything into a dict
output = {}
output['precip_sm'] = precip_sm
output['onset'] = onset
output['retreat'] = retreat
output['peak'] = peak
# Collapse any extra dimensions that were added
if axis == 0:
for key in output:
output[key] = atm.collapse(output[key], 0)
while onset.ndim > ndim:
for key in output:
if key != 'Rsq':
output[key] = atm.collapse(output[key], -1)
# Some more data to output
output['Rsq'] = Rsq
output['smoothing_kmax'] = kmax
output['threshold'] = threshold
return output
elif name == 'MERRA_MFC':
pcp = mfcbar
days = mfcbar.day.values
pentad = False
precip_jan = 0.0 # Use zero for now
elif name == 'MERRA_PRECIP':
pcp = precipbar
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)
early, late = {}, {}
for nm in index1:
early[nm], late[nm] = extreme_years(index1[nm])
# ----------------------------------------------------------------------
# Fourier harmonics
kmax_list = np.arange(2, 21, 1)
nms = [pcp_nm, 'U850', 'V850']
days = np.arange(-138, 227)
ts1 = ts.sel(dayrel=days)
ts_sm = {kmax : xray.Dataset() for kmax in kmax_list}
Rsq = {kmax : {} for kmax in kmax_list}
for kmax in kmax_list:
for nm in nms:
vals, Rsq[kmax][nm] = atm.fourier_smooth(ts1[nm], kmax)
print kmax, nm, Rsq[kmax][nm]
ts_sm[kmax][nm] = xray.DataArray(vals, coords=ts1[nm].coords)
# Find days where smoothed values are closest to actual timeseries
# values at days 0, 15
def closest_day(nm, ts1, ts_sm, d0, buf=20):
val0 = ts1[nm].sel(dayrel=d0).values
sm = atm.subset(ts_sm[nm], {'dayrel' : (d0 - buf, d0 + buf)})
i0 = int(np.argmin(abs(sm - val0)))
day0 = int(sm['dayrel'][i0])
return day0
# Annual + semi-annual harmonics
xticks = np.arange(-120, 230, 30)
sz = 10
titlestr = 'CMAP %d-%dE, %d-%dN ' % (lon1, lon2, lat1, lat2)
precip = precipdat.read_cmap(cmap_file)
precipbar = atm.mean_over_geobox(precip, lat1, lat2, lon1, lon2)
nyears, npentad = precipbar.shape
years = precipbar.year.values
nyears = len(years)
pentads = precipbar.pentad
onset = {}
# Threshold for onset criteria
threshold = 5.0
# Smooth with truncated Fourier series
kmax = 12
kann = 4
pcp_sm, Rsq = atm.fourier_smooth(precipbar, kmax)
pcp_ann, Rsq_ann = atm.fourier_smooth(precipbar, kann)
label_sm, label_ann = [], []
for y in range(nyears):
label_sm.append('kmax=%d, $R^2$=%.2f' % (kmax, Rsq[y]))
label_ann.append('kmax=%d, $R^2$=%.2f' % (kann, Rsq_ann[y]))
i_onset, i_retreat, i_peak = get_onset(years, pcp_sm)
plot_years(years, precipbar, pcp_sm, label_sm, i_onset, i_peak, i_retreat,
titlestr, pcp_ann, label_ann)
onset['fourier'] = i_onset
# Smooth with rolling mean
nroll = 3
pcp_sm = np.zeros(precipbar.shape)