Merge 0a49e45 into 59529e2

spirograph/mpl/__init__.py

@@ -0,0 +1,5 @@
+from ._partitioning import (
+    graph_fraction_of_total_variance,
+    graph_fractional_uncertainty,
+    graph_variance,
+)

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