def get_onset_WLH(years, days, pcp_sm, threshold, titlestr, pentad=True,
pcp_jan=None):
nyears = len(years)
i_onset = np.zeros(nyears)
i_retreat = np.zeros(nyears)
i_peak = np.zeros(nyears)
for y, year in enumerate(years):
vals = indices.onset_WLH_1D(pcp_sm[y], threshold, precip_jan=pcp_jan)
i_onset[y], i_retreat[y], i_peak[y] = vals
# Convert from pentads to day of year
if pentad:
d_onset = [int(atm.pentad_to_jday(i, pmin=0)) for i in i_onset]
d_retreat = [int(atm.pentad_to_jday(i, pmin=0)) for i in i_retreat]
else:
d_onset = [np.nan if np.isnan(i) else days[int(i)] for i in i_onset]
d_retreat = [np.nan if np.isnan(i) else days[int(i)] for i in i_retreat]
# Pack into Dataset
index = xray.Dataset()
days = xray.DataArray(days, {'day' : days})
years = xray.DataArray(years, {'year' : years})
index['tseries'] = xray.DataArray(pcp_sm, dims=['year', 'day'],
coords={'year' : years, 'day': days})
index['onset'] = xray.DataArray(d_onset, coords={'year' : years})
index['retreat'] = xray.DataArray(d_retreat, coords={'year' : years})
index.attrs['title'] = titlestr
return index
def plot_all_WLH(wlh, y=0, axlims=(0,50,50,180), cmap='jet', clines=True):
onset = wlh['onset']
Rsq = wlh['Rsq']
lat = wlh['lat']
lon = wlh['lon']
kmax = wlh['smoothing_kmax']
years=wlh['years']
if wlh['climatology']:
titlestr = 'CMAP %d-%d Climatology' % (years.min(), years.max())
else:
onset = onset[y]
Rsq = Rsq[y]
titlestr = 'CMAP %d' % years[y]
# Note: add 1 to pentad indices to index from 1-73 for comparison
# with Wang & LinHo
onset = onset + 1
# Calculate onset dates from pentads
nlat, nlon = onset.shape
onset_date = np.nan * np.ones((nlat, nlon))
for i in range(nlat):
for j in range(nlon):
jday = atm.pentad_to_jday(onset[i, j], pmin=1)
mon, day = atm.jday_to_mmdd(jday)
onset_date[i, j] = 100*mon + day
# -- Smooth with cubic spline
# lat_i = np.arange(-89.5, 90, 0.5)
# lon_i = np.arange(0, 360, 0.5)
lat_i = np.arange(-90, 90, 4.)
lon_i = np.arange(0, 360, 4.)
onset_date = atm.interp_latlon(onset_date, lat_i, lon_i, lat, lon, order=3)
clev_date = [501, 511, 521, 601, 611, 621, 701, 711]
clev_label = {}
for d in clev_date:
clev_label[d] = '%02d-%02d' % (d//100, d%100)
# Plot maps
manual = False
plt.figure(figsize=(14,10))
cmin, cmax = 20, 55
plt.subplot(211)
_, pc = atm.pcolor_latlon(onset, lat, lon, cmap=cmap, axlims=axlims)
pc.set_clim(cmin, cmax)
if clines:
_, cs = atm.contour_latlon(onset_date, lat_i, lon_i, clev=clev_date,
axlims=axlims)
plt.clabel(cs, clev_date, fmt=clev_label, manual=manual)
plt.title(titlestr + ' Onset Pentad & Day')
else:
plt.title(titlestr + ' Onset Pentad')
plt.subplot(212)
clev = np.arange(0, 1.1, 0.025)
_, pc = atm.pcolor_latlon(Rsq, lat, lon, cmap=cmap, axlims=axlims)
pc.set_clim(0., 1.)
plt.title('$R^2$ for kmax = %d' % kmax)
def read_cmap(datafile, yearmin=None, yearmax=None, pentad_day=3):
"""Read CMAP pentad data for selected years.
Parameters
----------
datafile : str
File path for CMAP data in NetCDF format.
yearmin, yearmax : ints, optional
Min and max years (inclusive) to extract.
pentad_day : int, optional
Which day of pentad to use for day coordinate.
Returns
-------
precip : xray.DataArray
CMAP pentad data for selected years, reshaped to
[year, pentad, lat, lon].
"""
NPYR = 73 # Number of pentads per year
YRMIN = 1979 # First year of CMAP data
if yearmin is None:
yearmin = YRMIN
# Read data
with xray.open_dataset(datafile) as ds:
precip = ds['precip'].load()
# Discard incomplete year at end
ny = precip.shape[0] // NPYR
precip = precip[:ny*NPYR]
# Split into individual years
years = np.arange(YRMIN, YRMIN + ny)
pentads = np.arange(1, NPYR + 1)
precip = atm.split_timedim(precip, NPYR, time0_name='year',
time0_vals=years, time1_name='pentad',
time1_vals=pentads)
# Extract selected years
if yearmax is None:
years = np.arange(yearmin, YRMIN + ny)
else:
years = np.arange(yearmin, yearmax + 1)
precip = precip.sel(year=years)
precip.coords['day'] = atm.pentad_to_jday(precip['pentad'], pmin=1,
day=pentad_day)
precip = precip.swap_dims({'pentad' : 'day'})
return precip
coords={'year' : years, 'day': days})
index['onset'] = xray.DataArray(d_onset, coords={'year' : years})
index['retreat'] = xray.DataArray(d_retreat, coords={'year' : years})
index.attrs['title'] = titlestr
return index
# Threshold and smoothing parameters
threshold = 5.0
kmax = 12
nroll = {'CMAP' : 3, 'MERRA_MFC' : 7, 'MERRA_PRECIP' : 7}
# Read CMAP pentad precip
cmap = precipdat.read_cmap(cmapfile)
cmapbar = atm.mean_over_geobox(cmap, lat1, lat2, lon1, lon2)
cmapdays = [atm.pentad_to_jday(p, pmin=1) for p in cmap.pentad.values]
# MERRA moisture flux convergence
mfc = atm.combine_daily_years('MFC', mfcfiles, years, yearname='year')
mfcbar = atm.mean_over_geobox(mfc, lat1, lat2, lon1, lon2)
# MERRA precip
subset_dict = {'lon' : (lon1, lon2), 'lat' : (lat1, lat2)}
precip = atm.combine_daily_years('PRECTOT', precipfiles, years, yearname='year',
subset_dict=subset_dict)
precip = atm.precip_convert(precip, precip.attrs['units'], 'mm/day')
precipbar = atm.mean_over_geobox(precip, lat1, lat2, lon1, lon2)
# Compute indices for each dataset
for name in ['CMAP', 'MERRA_MFC', 'MERRA_PRECIP']:
print('****' + name + '******')
