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
Пример #2
0
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
Пример #3
0
    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)
Пример #4
0
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)