def test_shift_cftime_index():
"""Tests that ``CFTimeIndex`` is shifted by the appropriate amount."""
idx = xr.cftime_range("1990", "2000", freq="YS")
da = xr.DataArray(np.random.rand(len(idx)), dims="time", coords=[idx])
expected = idx.shift(3, "YS")
res = shift_cftime_index(da, "time", 3, "YS")
assert (expected == res).all()
def test_shift_cftime_index():
"""Tests that ``CFTimeIndex`` is shifted by the appropriate amount."""
idx = xr.cftime_range('1990', '2000', freq='YS')
da = xr.DataArray(np.random.rand(len(idx)), dims='time', coords=[idx])
expected = idx.shift(3, 'YS')
res = shift_cftime_index(da, 'time', 3, 'YS')
assert (expected == res).all()
def plot_ensemble_perfect_model(
pm, variable=None, ax=None, show_members=False, cmap="tab10"
):
"""Plot datasets from PerfectModelEnsemble.
Args:
pm (PerfectModelEnsemble): PerfectModelEnsemble.
variable (str or None): `variable` to plot. Defaults to the first in data_vars.
ax (plt.axes): Axis to use in plotting. By default, creates a new axis.
show_members (bool): whether to display all members individually.
Defaults to False.
cmap (str): Name of matplotlib-recognized colorbar. Defaults to 'tab10'.
Returns:
ax: plt.axes
"""
_check_only_climpred_dims(pm)
if variable is None:
variable = list(pm.get_initialized().data_vars)[0]
initialized = pm.get_initialized()[variable]
uninitialized = pm.get_uninitialized()
if isinstance(uninitialized, xr.Dataset):
uninitialized = uninitialized[variable]
uninitialized_present = True
else:
uninitialized_present = False
control = pm.get_control()
if isinstance(control, xr.Dataset):
control = control[variable]
calendar = infer_calendar_name(initialized.init)
lead_freq = get_lead_cftime_shift_args(initialized.lead.attrs["units"], 1)[1]
control_color = "gray"
if ax is None:
_, ax = plt.subplots(figsize=(10, 4))
cmap = mpl.cm.get_cmap(cmap, initialized.init.size)
for ii, i in enumerate(initialized.init.values):
dsi = initialized.sel(init=i).rename({"lead": "time"})
if uninitialized_present:
dsu = uninitialized.sel(init=i).rename({"lead": "time"})
# convert lead time into cftime
start_str = i.strftime()[:10]
if initialized.lead.min() == 0:
dsi["time"] = xr.cftime_range(
start=start_str,
freq=lead_freq,
periods=dsi.time.size,
calendar=calendar,
)
elif initialized.lead.min() == 1:
dsi["time"] = xr.cftime_range(
start=start_str,
freq=lead_freq,
periods=dsi.time.size,
calendar=calendar,
)
dsi["time"] = shift_cftime_index(dsi.time, "time", 1, lead_freq)
if uninitialized_present:
dsu["time"] = dsi["time"]
if not show_members:
dsi = dsi.mean("member")
if uninitialized_present:
dsu = dsu.mean("member")
member_alpha = 1
lw = 2
labelstr = "ensemble mean"
else:
member_alpha = 0.5
lw = 1
labelstr = "members"
# plot ensemble mean, first white then color to highlight ensemble mean
if uninitialized_present:
dsu.mean("member").plot(ax=ax, color="white", lw=3, zorder=8, alpha=0.6)
dsu.mean("member").plot(
ax=ax, color=control_color, lw=2, zorder=9, alpha=0.6
)
# plot ensemble mean, first white then color to highlight ensemble mean
dsi.mean("member").plot(ax=ax, color="white", lw=3, zorder=10)
dsi.mean("member").plot(ax=ax, color=cmap(ii), lw=2, zorder=11)
dsi.plot(
ax=ax,
hue="member",
color=cmap(ii),
alpha=member_alpha,
lw=lw,
label=labelstr,
)
if uninitialized_present:
dsu.plot(
ax=ax,
hue="member",
color=control_color,
alpha=member_alpha / 2,
lw=lw,
label="uninitialized " + labelstr,
)
if isinstance(control, xr.DataArray):
control.plot(ax=ax, color=control_color, label="control")
# show only one item per label in legend
handles, labels = ax.get_legend_handles_labels()
by_label = OrderedDict(zip(labels, handles))
ax.legend(by_label.values(), by_label.keys())
ax.set_title(" ")
return ax
def plot_lead_timeseries_hindcast(
he, variable=None, ax=None, show_members=False, cmap="jet"
):
"""Plot datasets from HindcastEnsemble.
Args:
he (HindcastEnsemble): HindcastEnsemble.
variable (str or None): `variable` to plot. Defaults to the first in data_vars.
ax (plt.axes): Axis to use in plotting. By default, creates a new axis.
show_members (bool): whether to display all members individually.
Defaults to False.
cmap (str): Name of matplotlib-recognized colorbar. Defaults to 'jet'.
Returns:
ax: plt.axes
"""
_check_only_climpred_dims(he)
if variable is None:
variable = list(he.get_initialized().data_vars)[0]
hind = he.get_initialized()[variable]
lead_freq = get_lead_cftime_shift_args(hind.lead.attrs["units"], 1)[1]
hist = he.get_uninitialized()
if isinstance(hist, xr.Dataset):
hist = hist[variable]
obs = he._datasets["observations"]
cmap = mpl.cm.get_cmap(cmap, hind.lead.size)
if ax is None:
_, ax = plt.subplots(figsize=(10, 4))
if isinstance(hist, xr.DataArray):
if "member" in hist.dims and not show_members:
hist = hist.mean("member")
member_alpha = 1
lw = 2
else:
member_alpha = 0.4
lw = 1
hist.plot(
ax=ax,
lw=lw,
hue="member",
color="gray",
alpha=member_alpha,
label="uninitialized",
zorder=hind.lead.size + 1,
)
for i, lead in enumerate(hind.lead.values):
h = hind.sel(lead=lead).rename({"init": "time"})
if not show_members and "member" in h.dims:
h = h.mean("member")
lead_alpha = 1
else:
lead_alpha = 0.5
h["time"] = shift_cftime_index(h.time, "time", int(lead), lead_freq)
h.plot(
ax=ax,
hue="member",
color=cmap(i),
label=f"initialized: lead={lead}",
alpha=lead_alpha,
zorder=hind.lead.size - i,
)
if len(obs) > 0:
if isinstance(obs, xr.Dataset):
obs = obs[variable]
obs.plot(
ax=ax,
color="k",
lw=3,
ls="-",
label="observations",
zorder=hind.lead.size + 2,
)
# show only one item per label in legend
handles, labels = ax.get_legend_handles_labels()
by_label = OrderedDict(zip(labels, handles))
ax.legend(
by_label.values(), by_label.keys(), loc="center left", bbox_to_anchor=(1, 0.5)
)
ax.set_title("")
return ax
