add partition plot

CHANGES.rst

@@ -16,6 +16,7 @@ New features and enhancements
 * Use small geojson in the notebook. (:pull:`124`).
 * Add the Colours of Figanos page (:issue:`126`, :pull:`127`).
 * Figanos now adheres to PEPs 517/518/621 using the `flit` backend for building and packaging. (:pull:`135`).
+* New function ``fg.partition`` (:pull:`134`).
 
 Bug fixes
 ^^^^^^^^^

docs/notebooks/figanos_docs.ipynb

@@ -1016,6 +1016,147 @@
     "\n",
     "fg.taylordiagram(sims, std_range=(0, 1.3), contours=5, contours_kw={'colors': 'green'}, plot_kw={'reference': {'marker':'*'}})\n"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a0a3afb0",
+   "metadata": {},
+   "source": [
+    "## Partition plots\n",
+    "\n",
+    "Partition plots show the fraction of uncertainty associated with different components.\n",
+    "Xclim has a few different [partition functions](https://xclim.readthedocs.io/en/stable/api.html#uncertainty-partitioning). \n",
+    "\n",
+    "This tutorial is a reproduction of [xclim's documentation](https://xclim.readthedocs.io/en/stable/notebooks/partitioning.html).\n",
+    "\n",
+    "<div class=\"alert alert-info\">\n",
+    "    \n",
+    "Note that you could also use the [xscen library](https://xscen.readthedocs.io/en/latest/index.html) to build and ensemble from a catalog with `xscen.ensembles.get_partition_input`.\n",
+    "    \n",
+    "</div>"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a91d12a2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Fetch data\n",
+    "from __future__ import annotations\n",
+    "\n",
+    "import pandas as pd\n",
+    "\n",
+    "import xclim.ensembles\n",
+    "\n",
+    "# The directory in the Atlas repo where the data is stored\n",
+    "# host = \"https://github.com/IPCC-WG1/Atlas/raw/main/datasets-aggregated-regionally/data/CMIP6/CMIP6_tas_land/\"\n",
+    "host = \"https://raw.githubusercontent.com/IPCC-WG1/Atlas/main/datasets-aggregated-regionally/data/CMIP6/CMIP6_tas_land/\"\n",
+    "\n",
+    "# The file pattern, e.g. CMIP6_ACCESS-CM2_ssp245_r1i1p1f1.csv\n",
+    "pat = \"CMIP6_{model}_{scenario}_{member}.csv\"\n",
+    "\n",
+    "# Here we'll download data only for a very small demo sample of models and scenarios.\n",
+    "\n",
+    "# Download data for a few models and scenarios.\n",
+    "models = [\"ACCESS-CM2\", \"CMCC-CM2-SR5\", \"CanESM5\"]\n",
+    "\n",
+    "# The variant label is not always r1i1p1f1, so we need to match it to the model.\n",
+    "members = [\"r1i1p1f1\", \"r1i1p1f1\", \"r1i1p1f1\"]\n",
+    "\n",
+    "# GHG concentrations and land-use scenarios\n",
+    "scenarios = [\"ssp245\", \"ssp370\", \"ssp585\"]\n",
+    "\n",
+    "# Create the input ensemble.\n",
+    "data = []\n",
+    "for model, member in zip(models, members):\n",
+    "    for scenario in scenarios:\n",
+    "        url = host + pat.format(model=model, scenario=scenario, member=member)\n",
+    "\n",
+    "        # Fetch data using pandas\n",
+    "        df = pd.read_csv(url, index_col=0, comment=\"#\", parse_dates=True)[\"world\"]\n",
+    "        # Convert to a DataArray, complete with coordinates.\n",
+    "        da = (\n",
+    "            xr.DataArray(df)\n",
+    "            .expand_dims(model=[model], scenario=[scenario])\n",
+    "            .rename(date=\"time\")\n",
+    "        )\n",
+    "        data.append(da)\n",
+    "        \n",
+    "# Combine DataArrays from the different models and scenarios into one.\n",
+    "ens_mon = xr.combine_by_coords(data)[\"world\"]\n",
+    "\n",
+    "# Then resample the monthly time series at the annual frequency\n",
+    "ens = ens_mon.resample(time=\"Y\").mean()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d91c12cf",
+   "metadata": {},
+   "source": [
+    "Compute uncertainties with xclim.\n",
+    "Make sure to activate the `fraction` argument before plotting with figanos."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "645c3bb3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# compute uncertainty\n",
+    "mean, uncertainties = xclim.ensembles.hawkins_sutton(ens, baseline=(\"2016\", \"2030\"))\n",
+    "#frac= xc.ensembles.fractional_uncertainty(uncertainties)\n",
+    "\n",
+    "#FIXME: xc.ensembles.fractional_uncertainty has not been released yet. Until until it is released, here it is.\n",
+    "def fractional_uncertainty(u: xr.DataArray):\n",
+    "    \"\"\"\n",
+    "    Return the fractional uncertainty.\n",
+    "\n",
+    "    Parameters\n",
+    "    ----------\n",
+    "    u: xr.DataArray\n",
+    "      Array with uncertainty components along the `uncertainty` dimension.\n",
+    "\n",
+    "    Returns\n",
+    "    -------\n",
+    "    xr.DataArray\n",
+    "      Fractional, or relative uncertainty with respect to the total uncertainty.\n",
+    "    \"\"\"\n",
+    "    uncertainty = u / u.sel(uncertainty=\"total\") * 100\n",
+    "    uncertainty.attrs.update(u.attrs)\n",
+    "    uncertainty.attrs[\"long_name\"] = \"Fraction of total variance\"\n",
+    "    uncertainty.attrs[\"units\"] = \"%\"\n",
+    "    return uncertainty\n",
+    "\n",
+    "frac= fractional_uncertainty(uncertainties)\n",
+    "\n",
+    " "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "32db8ef5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# plot\n",
+    "fg.partition(\n",
+    "    frac,\n",
+    "    start_year='2016', # change the x-axis\n",
+    "    show_num=True, # put the number of element of each uncertainty source in the legend\n",
+    "    fill_kw={\n",
+    "    \"variability\": {'color':\"#DC551A\"},\n",
+    "    \"model\": {'color':\"#2B2B8B\"},\n",
+    "    \"scenario\": {'color':\"#275620\"},\n",
+    "    \"total\": {'color':\"k\"},},\n",
+    "    line_kw={'lw':3}\n",
+    ")"
+   ]
   }
  ],
  "metadata": {
@@ -1030,7 +1171,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.12"
+   "version": "3.10.6"
   }
  },
  "nbformat": 4,

figanos/matplotlib/__init__.py

@@ -5,6 +5,7 @@
     gridmap,
     hatchmap,
     heatmap,
+    partition,
     scattermap,
     stripes,
     taylordiagram,

figanos/matplotlib/plot.py

@@ -2037,3 +2037,137 @@ def hatchmap(
         set_plot_attrs(use_attrs, dattrs, ax, wrap_kw={"max_line_len": 60})
 
     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
+
+
+def partition(
+    data: xr.DataArray | xr.Dataset,
+    ax: matplotlib.axes.Axes | None = None,
+    start_year: str | None = None,
+    show_num: bool = True,
+    fill_kw: dict[str, Any] | None = None,
+    line_kw: dict[str, Any] | None = None,
+    fig_kw: dict[str, Any] | None = None,
+    legend_kw: dict[str, Any] | None = None,
+) -> matplotlib.axes.Axes:
+    """Figure of the partition of total uncertainty by components.
+
+    Uncertainty fractions can be computed with xclim
+    (https://xclim.readthedocs.io/en/stable/api.html#uncertainty-partitioning).
+    Make sure the use `fraction=True` in the xclim function call.
+
+
+    Parameters
+    ----------
+    data: xr.DataArray or xr.Dataset
+      Variance over time of the different components of uncertainty.
+      Output of a `xclim.ensembles._partitioning` function.
+    ax : matplotlib axis, optional
+        Matplotlib axis on which to plot
+    start_year: str
+      If None, the x-axis will be the time in year.
+      If str, the x-axis will show the number of year since start_year.
+    show_num: bool
+        If True, show the number of elements for each uncertainty components in parenthesis in the legend.
+        `data` should have attributes named after the components with a list its the elements.
+    fill_kw: dict
+        Keyword arguments passed to `ax.fill_between`.
+        It is possible to pass a dictionary of keywords for each component (uncertainty coordinates).
+    line_kw: dict
+        Keyword arguments passed to `ax.plot` for the lines in between the components.
+        The default is {color="k", lw=2}.
+    fig_kw: dict
+        Keyword arguments passed to `plt.subplots`.
+    legend_kw: dict
+        Keyword arguments passed to `ax.legend`.
+
+    Returns
+    -------
+    mpl.axes.Axes
+    """
+    if isinstance(data, xr.Dataset):
+        if len(data.data_vars) > 1:
+            warnings.warn(
+                "data is xr.Dataset; only the first variable will be used in plot"
+            )
+        data = data[list(data.keys())[0]].squeeze()
+
+    if data.attrs["units"] != "%":
+        raise ValueError(
+            "The units are not %. Use `fraction=True` in the xclim function call."
+        )
+
+    fill_kw = empty_dict(fill_kw)
+    line_kw = empty_dict(line_kw)
+    fig_kw = empty_dict(fig_kw)
+    legend_kw = empty_dict(legend_kw)
+
+    # select data to plot
+    if isinstance(data, xr.DataArray):
+        data = data.squeeze()
+    elif isinstance(data, xr.Dataset):  # in case, it was saved to disk before plotting.
+        if len(data.data_vars) > 1:
+            warnings.warn(
+                "data is xr.Dataset; only the first variable will be used in plot"
+            )
+        data = data[list(data.keys())[0]].squeeze()
+    else:
+        raise TypeError("`data` must contain a xr.DataArray or xr.Dataset")
+
+    if ax is None:
+        fig, ax = plt.subplots(**fig_kw)
+
+    # Select data from reference year onward
+    if start_year:
+        data = data.sel(time=slice(start_year, None))
+
+        # Lead time coordinate
+        time = _add_lead_time_coord(data, start_year)
+        ax.set_xlabel(f"Lead time (years from {start_year})")
+    else:
+        time = data.time.dt.year
+
+    # fill_kw that are direct (not with uncertainty as key)
+    fk_direct = {k: v for k, v in fill_kw.items() if (k not in data.uncertainty.values)}
+
+    # Draw areas
+    past_y = 0
+    black_lines = []
+    for u in data.uncertainty.values:
+        if u not in ["total", "variability"]:
+            present_y = past_y + data.sel(uncertainty=u)
+            num = len(data.attrs.get(u, []))  # compatible with pre PR PR #1529
+            label = f"{u} ({num})" if show_num and num else u
+            ax.fill_between(
+                time, past_y, present_y, label=label, **fill_kw.get(u, fk_direct)
+            )
+            black_lines.append(present_y)
+            past_y = present_y
+    ax.fill_between(
+        time, past_y, 100, label="variability", **fill_kw.get("variability", fk_direct)
+    )
+
+    # Draw black lines
+    line_kw.setdefault("color", "k")
+    line_kw.setdefault("lw", 2)
+    ax.plot(time, np.array(black_lines).T, **line_kw)
+
+    ax.xaxis.set_major_locator(matplotlib.ticker.MultipleLocator(20))
+    ax.xaxis.set_minor_locator(matplotlib.ticker.AutoMinorLocator(n=5))
+
+    ax.yaxis.set_major_locator(matplotlib.ticker.MultipleLocator(10))
+    ax.yaxis.set_minor_locator(matplotlib.ticker.AutoMinorLocator(n=2))
+
+    ax.set_ylabel(f"{data.attrs['long_name']} ({data.attrs['units']})")  #
+
+    ax.set_ylim(0, 100)
+    ax.legend(**legend_kw)
+
+    return ax