diff --git a/spirograph/mpl/__init__.py b/spirograph/mpl/__init__.py new file mode 100644 index 00000000..39ca2020 --- /dev/null +++ b/spirograph/mpl/__init__.py @@ -0,0 +1,5 @@ +from ._partitioning import ( + graph_fraction_of_total_variance, + graph_fractional_uncertainty, + graph_variance, +) diff --git a/spirograph/mpl/_partitioning.py b/spirograph/mpl/_partitioning.py new file mode 100644 index 00000000..86277f69 --- /dev/null +++ b/spirograph/mpl/_partitioning.py @@ -0,0 +1,161 @@ +""" +# Uncertainty partitioning + +Create figures similar to those from Hawkins & Sutton (2009). All these functions expect DataArrays with dimensions +`time` and `uncertainty`, as generated by `xclim.ensembles.partitioning.hawkins_sutton`. +""" +import matplotlib as mpl +import numpy as np +from cycler import cycler +from matplotlib import pyplot as plt + +colors = { + "variability": "#DC551A", + "model": "#2B2B8B", + "scenario": "#275620", + "total": "k", +} +names = {"variability": "Internal variability"} + + +def graph_variance(variance, ref="2000", ax=None): + """Figure showing the variance over time. + + Parameters + ---------- + variance: xr.DataArray + Variance over time of the different sources of uncertainty. + ref: str + Reference year. Defines year 0. + ax: mpl.axes.Axes + Axes in which to draw. + + Returns + ------- + mpl.axes.Axes + """ + if ax is None: + fig, ax = plt.subplots(1, 1) + + da = variance.copy() + da = da.sel(time=slice(ref, None)) + + keys = da.uncertainty.values + + cmap = mpl.colors.ListedColormap(list([colors[k] for k in keys])) + labels = [names.get(k, k.capitalize()) for k in keys] + + # Lead time coordinate + add_lead_time_coord(da, ref) + + # The list of colors to use + ax.set_prop_cycle(cycler(color=cmap.colors)) + lines = da.plot(x="Lead time", hue="uncertainty", ax=ax) + + ax.set_ylabel("Variance") + plt.legend(lines, labels) + + return ax + + +def graph_fractional_uncertainty(mm, variance, ref="2000", ax=None): + """Figure from Hawkins and Sutton (2009) showing fractional uncertainty. + + Parameters + ---------- + variance: xr.DataArray + Variance over time of the different sources of uncertainty. + mm: xr.DataArray + Model mean. + ref: str + Reference year. Defines year 0. + ax: mpl.axes.Axes + Axes in which to draw. + + Returns + ------- + mpl.axes.Axes + """ + if ax is None: + fig, ax = plt.subplots(1, 1) + + # Compute fractional uncertainty + da = 1.654 * np.sqrt(variance) / mm + + da = da.sel(time=slice(ref, None)) + + keys = da.uncertainty.values + cmap = mpl.colors.ListedColormap(list([colors[k] for k in keys])) + labels = [names.get(k, k.capitalize()) for k in keys] + + # Lead time coordinate + add_lead_time_coord(da, ref) + + # The list of colors to use + ax.set_prop_cycle(cycler(color=cmap.colors)) + lines = da.plot(x="Lead time", hue="uncertainty", ax=ax) + + ax.set_ylabel("Fractional uncertainty") + plt.legend(lines, labels) + + return ax + + +def graph_fraction_of_total_variance(variance, ref="2000", ax=None): + """Figure from Hawkins and Sutton (2009) showing fraction of total variance. + + Parameters + ---------- + variance: xr.DataArray + Variance over time of the different sources of uncertainty. + ref: str + Reference year. Defines year 0. + ax: mpl.axes.Axes + Axes in which to draw. + + Returns + ------- + mpl.axes.Axes + """ + if ax is None: + fig, ax = plt.subplots(1, 1) + + # Compute fraction + da = variance / variance.sel(uncertainty="total") * 100 + + # Select data from reference year onward + da = da.sel(time=slice(ref, None)) + + # Lead time coordinate + lead_time = add_lead_time_coord(da, ref) + + # Draw areas + y1 = da.sel(uncertainty="model") + y2 = da.sel(uncertainty="scenario") + y1 + ax.fill_between(lead_time, 0, y1, color=colors["model"]) + ax.fill_between(lead_time, y1, y2, color=colors["scenario"]) + ax.fill_between(lead_time, y2, 100, color=colors["variability"]) + + # Draw black lines + ax.plot(lead_time, np.array([y1, y2]).T, color="k", lw=2) + + ax.xaxis.set_major_locator(mpl.ticker.MultipleLocator(20)) + ax.xaxis.set_minor_locator(mpl.ticker.AutoMinorLocator(n=5)) + + ax.yaxis.set_major_locator(mpl.ticker.MultipleLocator(10)) + ax.yaxis.set_minor_locator(mpl.ticker.AutoMinorLocator(n=2)) + + ax.set_xlabel(f"Lead time [years from {ref}]") + ax.set_ylabel("Fraction of total variance [%]") + + ax.set_ylim(0, 100) + + return ax + + +def add_lead_time_coord(da, ref): + """Add a lead time coordinate to the data. Modifies da in-place.""" + lead_time = da.time.dt.year - int(ref) + da["Lead time"] = lead_time + da["Lead time"].attrs["units"] = f"years from {ref}" + return lead_time