def add_labels(grp, labels, pos, fontsize, fontweight='bold'):
"""Add labels to subplots in a figure
Parameters
----------
grp : atm.FigGroup object
Instance of atm.FigGroup class for the subplots to annotate
labels : list
Labels to annotate, e.g. ['a', 'b', 'c']
pos : list of tuples
Positions for each label
fontsize, fontweight
Text parameters for labels
"""
try:
n = len(pos[0])
except TypeError:
pos = [pos] * len(labels)
row, col = 0, 0
for i in range(len(labels)):
grp.subplot(row, col)
atm.text(labels[i], pos[i], fontsize=fontsize, fontweight=fontweight)
col += 1
if col == grp.ncol:
col = 0
row += 1
return None
def plot_hist(ind, binwidth=5, incl_daystr=True, ax=None, pos=(0.05, 0.7),
kw={'alpha' : 0.3, 'color' : 'k'}):
"""Plot histogram of onset days.
"""
if ax is None:
ax = plt.gca()
def daystr(day):
day = round(day)
mm, dd = atm.jday_to_mmdd(day)
mon = atm.month_str(mm)
return '%.0f (%s-%.0f)' % (day, mon, dd)
if isinstance(ind, pd.Series) or isinstance(ind, xray.DataArray):
ind = ind.values
b1 = np.floor(np.nanmin(ind) / binwidth) * binwidth
b2 = np.ceil(np.nanmax(ind) / binwidth) * binwidth
bin_edges = np.arange(b1, b2 + 1, binwidth)
n, bins, _ = ax.hist(ind, bin_edges, **kw)
ax.set_xlabel('Day of Year')
ax.set_ylabel('Num of Occurrences')
if incl_daystr:
dmean = daystr(np.nanmean(ind))
dmin = daystr(np.nanmin(ind))
dmax = daystr(np.nanmax(ind))
else:
dmean = '%.0f' % np.nanmean(ind)
dmin = '%.0f' % np.nanmin(ind)
dmax = '%.0f' % np.nanmax(ind)
s = 'Mean %s\n' % dmean + 'Std %.0f\n' % np.nanstd(ind)
s = s + 'Min %s\n' % dmin + 'Max %s' % dmax
x0, y0 = pos
atm.text(s, (x0, y0), ax=ax, horizontalalignment='left')
def add_labels(grp, labels, pos, fontsize, fontweight='bold'):
# Expand pos to list for each subplot, if needed
try:
n = len(pos[0])
except TypeError:
pos = [pos] * (grp.nrow * grp.ncol)
i = 0
for row in range(grp.nrow):
for col in range(grp.ncol):
grp.subplot(row, col)
atm.text(labels[i], pos[i], fontsize=fontsize,
fontweight=fontweight)
i += 1
def plotyear(chp, y, xlims, ylims):
days = chp["day"].values
d_onset = chp["onset"][y]
d_retreat = chp["retreat"][y]
plt.plot(days, chp["tseries"][y])
plt.plot(days, chp["tseries_fit"][y])
plt.plot([d_onset, d_onset], ylims, "k")
plt.plot([d_retreat, d_retreat], ylims, "k")
plt.xlim(xlims[0], xlims[1])
plt.ylim(ylims[0], ylims[1])
plt.title(years[y])
atm.text("Onset: %d\nRetreat: %d" % (d_onset, d_retreat), (0.03, 0.83))
plt.grid()
def plotyear(chp, y, xlims, ylims):
days = chp['day'].values
d_onset = chp['onset'][y]
d_retreat = chp['retreat'][y]
plt.plot(days, chp['tseries'][y])
plt.plot(days, chp['tseries_fit_onset'][y])
plt.plot(days, chp['tseries_fit_retreat'][y])
plt.plot([d_onset, d_onset], ylims, 'k')
plt.plot([d_retreat, d_retreat], ylims, 'k')
plt.xlim(xlims[0], xlims[1])
plt.ylim(ylims[0], ylims[1])
plt.title(years[y])
atm.text('Onset: %d\nRetreat: %d' % (d_onset, d_retreat), (0.03, 0.83))
plt.grid()
def plot_maps(var, days, grp, cmin=0, cmax=20, cint=1, axlims=(5, 35, 60, 100),
cmap='PuBuGn', res='c', extend='max', cticks=None,
daypos=(0.05, 0.85)):
"""Lat-lon maps of precip on selected days."""
clev = np.arange(cmin, cmax + cint/2.0, cint)
if cticks is None:
cticks = np.arange(cmin, clev.max() + 1, 2)
lat1, lat2, lon1, lon2 = axlims
for day in days:
grp.next()
pcp = var.sel(dayrel=day)
m = atm.init_latlon(lat1, lat2, lon1, lon2, resolution=res)
m = atm.contourf_latlon(pcp, m=m, clev=clev, axlims=axlims, cmap=cmap,
colorbar=False, extend=extend)
atm.text(day, daypos, fontsize=12, fontweight='bold')
# plt.colorbar(ax=grp.axes.ravel().tolist(), orientation='vertical',
# shrink=0.8, ticks=cticks)
atm.colorbar_multiplot(orientation='vertical', shrink=0.8, ticks=cticks)
fix_axes(axlims)
def plot_index_years(index, nrow=3, ncol=4,
fig_kw={'figsize' : (11, 7), 'sharex' : True,
'sharey' : True},
gridspec_kw={'left' : 0.1, 'right' : 0.95, 'wspace' : 0.05,
'hspace' : 0.1},
incl_fit=False, suptitle='', xlabel='Day', ylabel='Index',
xlims=None, ylims=None, xticks=np.arange(0, 401, 100),
grid=True):
"""Plot daily timeseries of monsoon onset/retreat index each year.
"""
years = atm.get_coord(index, 'year')
days = atm.get_coord(index, 'day')
grp = atm.FigGroup(nrow, ncol, fig_kw=fig_kw, gridspec_kw=gridspec_kw,
suptitle=suptitle)
for year in years:
grp.next()
ind = atm.subset(index, {'year' : (year, year)}, squeeze=True)
ts = ind['tseries']
d0_list = [ind['onset'], ind['retreat']]
plt.plot(days, ts, 'k')
for d0 in d0_list:
plt.axvline(d0, color='k')
if incl_fit and 'tseries_fit_onset' in ind:
plt.plot(days, ind['tseries_fit_onset'], 'r')
if incl_fit and 'tseries_fit_retreat' in ind:
plt.plot(days, ind['tseries_fit_retreat'], 'b')
atm.text(year, (0.05, 0.9))
atm.ax_lims_ticks(xlims=xlims, ylims=ylims, xticks=xticks)
plt.grid(grid)
if grp.row == grp.nrow - 1:
plt.xlabel(xlabel)
if grp.col == 0:
plt.ylabel(ylabel)
return grp
ts = index[onset_nm]["tseries"].sel(year=year)
ts_onset = index[onset_nm]["tseries_fit_onset"].sel(year=year)
ts_retreat = index[onset_nm]["tseries_fit_retreat"].sel(year=year)
days = index[onset_nm]["tseries"]["day"]
figsize = (5, 3.5)
gs_kw = {"left": 0.15, "bottom": 0.15}
plt.subplots(1, 1, figsize=figsize, gridspec_kw=gs_kw)
plt.plot(days, ts, "k")
plt.plot(days, ts_onset, "r")
plt.plot(days, ts_retreat, "b")
plt.grid()
plt.xlim(0, 366)
plt.xlabel("Day")
plt.ylabel("MFC_ACC (mm)")
atm.text(year, (0.05, 0.9))
# ----------------------------------------------------------------------
# Summarize onset, retreat, indices in dataframes
onset = pd.DataFrame()
for nm in index:
onset[nm] = index[nm]["onset"].to_series()
corr = onset.corr()
retreat = index[onset_nm]["retreat"].to_series()
length = retreat - onset[onset_nm]
# Labels for onset indices correlations
labels = {}
for nm in onset.columns:
tseries.plot(ax=ax, grid=True, legend=False)
ax.set_title(title, loc='left', fontsize=11)
# Plot climatology and a few individual years
plotyears = [None, years[0], years[1], years[2]]
plt.figure(figsize=(12, 9))
suptitle = version.upper() + ' MFC Budget (%s) - Daily Tseries' % latlonstr
plt.suptitle(suptitle)
nrow, ncol = 2, 2
for y, year in enumerate(plotyears):
ax = plt.subplot(nrow, ncol, y + 1)
plot_tseries(ts, year, ax=ax)
if y == 0:
ax.legend(loc='upper left', fontsize=9)
s = 'RESID = MFC-P+E-dW/dt-ANA'
atm.text(s, (0.03, 0.5), fontsize=9)
# ----------------------------------------------------------------------
# Plot daily timeseries of W
W = ts['W']
d1, d2 = days_ssn[0], days_ssn[-1]
W1 = W.sel(day=d1)
W2 = W.sel(day=d2)
plotyears = years[:4]
clrs = ['b', 'r', 'g', 'k']
plt.figure()
for y, year in enumerate(plotyears):
plt.plot(W['day'], W[y], clrs[y], label=year)
plt.plot([d1, d2], [W1[y], W2[y]], clrs[y], marker='.')
plt.grid()
def plot_tseries_together(data, onset=None, years=None, suptitle='',
figsize=(14,10), legendsize=10,
legendloc='lower right', nrow=3, ncol=4,
yearnm='year', daynm='day', standardize=True,
label_attr=None, data_style=None, onset_style=None,
show_days=False):
"""Plot multiple daily timeseries together each year.
Parameters
----------
data : xray.Dataset
Dataset of timeseries variables to plot together.
onset : ndarray or dict of ndarrays, optional
Array of onset day for each year, or dict of onset arrays (e.g.
to compare onset days from different methods).
years : ndarray, optional
Subset of years to include. If omitted, all years are included.
suptitle : str, optional
Supertitle for plot.
figsize : 2-tuple, optional
Size of each figure.
legendsize : int, optional
Font size for legend
legendloc : str, optional
Legend location
nrow, ncol : int, optional
Number of rows, columns in each figure.
yearnm, daynm : str, optional
Name of year and day dimensions in data.
standardize : bool, optional
If True, standardize each timeseries by dividing by its
standard deviation.
label_attr : str, optional
Attribute of each data variable to use for labels. If omitted,
then the variable name is used.
data_style, onset_style : list or dict, optional
Matlab-style strings for each data variable or onset index.
show_days : bool, optional
If True, annotate each subplot with a textbox showing the
onset days.
"""
if years is None:
# All years
years = data[yearnm].values
data = atm.subset(data, {yearnm : (years, None)})
if label_attr is not None:
labels = {nm : data[nm].attrs[label_attr] for nm in data.data_vars}
if onset is not None:
if isinstance(onset, dict):
if onset_style is None:
onset_style = {key : 'k' for key in onset.keys()}
else:
onset = {'onset' : onset}
if onset_style is None:
onset_style = {'onset' : 'k'}
textpos = {key : (0.05, 0.9 - 0.1*i) for i, key in enumerate(onset)}
# Plot each year
for y, year in enumerate(years):
df = atm.subset(data, {yearnm : (year, None)}).to_dataframe()
df.drop(yearnm, axis=1, inplace=True)
if label_attr is not None:
df.rename(columns=labels, inplace=True)
if standardize:
for key in df.columns:
df[key] = (df[key] - np.nanmean(df[key])) / np.nanstd(df[key])
ylabel = 'Standardized Timeseries'
else:
ylabel = 'Timeseries'
if y % (nrow * ncol) == 0:
fig, axes = plt.subplots(nrow, ncol, figsize=figsize, sharex=True)
plt.subplots_adjust(left=0.08, right=0.95, wspace=0.2, hspace=0.2)
plt.suptitle(suptitle)
yplot = 1
else:
yplot += 1
i, j = atm.subplot_index(nrow, ncol, yplot)
ax = axes[i-1, j-1]
df.plot(ax=ax, style=data_style)
ax.grid()
if yplot == 1:
ax.legend(fontsize=legendsize, loc=legendloc)
else:
ax.legend_.remove()
if onset is not None:
for key in onset:
d0 = onset[key][y]
ax.plot([d0, d0], ax.get_ylim(), onset_style[key])
if show_days:
atm.text(d0, textpos[key], ax=ax, color=onset_style[key])
if j == 1:
ax.set_ylabel(ylabel)
if i == nrow:
ax.set_xlabel('Day')
else:
ax.set_xlabel('')
ax.set_title(year)
if key.startswith('CHP'):
for nm, clr in zip(['onset', 'retreat'], ['r', 'g']):
pred = index[keylong]['tseries_fit_' + nm][y]
pred = atm.subset(pred, {'day' : (df.index, None)})
ax.plot(df.index, pred, clr)
if std_ts:
ax.set_ylim(ylim1, ylim2)
else:
plt.autoscale(tight=True)
ax.plot([d0, d0], ax.get_ylim(), style)
txt = '%d' % d0
if d1 is not None:
ax.plot([d1, d1], ax.get_ylim(), style)
txt = txt + ', %d' % d1
ax.grid()
atm.text(txt, (0.05, 0.9), ax=ax, color=style)
if i == 0:
ax.set_title(year)
if i == nrow - 1:
ax.set_xlabel('Day')
else:
ax.set_xlabel('')
if y % ncol == 0:
ax.set_ylabel(key)
iplot += ncol
saveclose('tseries_stacked_', isave, exts)
# Correlations between daily timeseries
style = {'MFC' : 'k', pcp_nm : 'k--'}
fig_kw = {'figsize' : (14, 10)}
gridspec_kw = {'left' : 0.04, 'right' : 0.94, 'bottom' : 0.06, 'top' : 0.95,
'hspace' : 0.07, 'wspace' : 0.07}
suptitle = 'Daily and Cumulative Precip/MFC for %s Onset' % onset_nm
nrow, ncol = 3, 4
grp = atm.FigGroup(nrow, ncol, fig_kw=fig_kw, gridspec_kw=gridspec_kw,
suptitle=suptitle)
for y, year in enumerate(years):
grp.next()
if grp.row == 0 and grp.col == ncol - 1:
legend = True
else:
legend = False
data1 = xray.Dataset({'MFC' : data['MFC'][y], pcp_nm : data[pcp_nm][y]})
data2 = data['MFC_ACC'][y]
axes = lineplots(data1, data2, style, xlims, xticks, ylims, yticks,
y2_lims=y2_lims, legend=legend,
length=index['length'][y])
title = '%d\n%d, %d' % (year, onset[y], index['retreat'][y])
atm.text(title, (0.03, 0.85), fontsize=10)
if grp.row < nrow - 1:
for ax in axes:
clear_labels('x', ax)
for i, ax in enumerate(axes):
if i == 0 and grp.col > 0:
clear_labels('y', ax)
if i > 0 and grp.col < ncol - 1:
clear_labels('y', ax)
yrly_index(onset_all, legend=True)
else:
for i in [0, 1]:
grp.next()
plt.axis('off')
ax = plt.subplot(2, 1, 1)
df = index[['length', 'retreat', 'onset']].to_dataframe()
plt.boxplot(df.values, vert=False, labels=['Length', 'Retreat', 'Onset'],
whis='range')
plt.xlabel('Day of Year | Number of Days')
plt.xlim(120, 320)
plt.xticks(np.arange(120, 321, 20))
pos = ax.get_position()
pos2 = [pos.x0, pos.y0 + 0.05, pos.width, pos.height]
ax.set_position(pos2)
atm.text('a', (-0.16, 1.01), fontsize=labelsize, fontweight='bold')
# Plot daily tseries
legend = True
if ind_nm == 'onset':
dlist = [15]
else:
dlist = None
daily_tseries(tseries, index, pcp_nm, npre, npost, legend, grp, ind_nm=ind_nm,
dlist=dlist)
# Add a-d labels
if ind_nm == 'onset':
labels = ['a', 'b', 'c', 'd']
add_labels(grp, labels, labelpos, labelsize)
else: