Merge f727174 into 67d806c

.zenodo.json

@@ -28,6 +28,11 @@
       "name": "Braun, Marco",
       "affiliation": "Ouranos",
       "orcid": "0000-0001-5061-3217"
+    },
+    {
+      "name": "Dupuis, Éric",
+      "affiliation": "Ouranos, Montréal, Québec, Canada",
+      "orcid": "0000-0001-7976-4596"
     }
   ],
   "keywords": [

AUTHORS.rst

@@ -20,3 +20,4 @@ Contributors
 
 * Gabriel Rondeau-Genesse <rondeau-genesse.gabriel@ouranos.ca> `@RondeauG <https://github.com/RondeauG>`_
 * Marco Braun <Braun.Marco@ouranos.ca> `@vindelico <https://github.com/vindelico>`_
+* Éric Dupuis <dupuis.eric@ouranos.ca> `@coxipi <https://github.com/coxipi>`_

CHANGELOG.rst

@@ -4,7 +4,7 @@ Changelog
 
 0.4.0 (unreleased)
 ------------------
-Contributors to this version: Trevor James Smith (:user:`Zeitsperre`), Marco Braun (:user:`vindelico`), Pascal Bourgault (:user:`aulemahal`), Sarah-Claude Bourdeau-Goulet (:user:`Sarahclaude`)
+Contributors to this version: Trevor James Smith (:user:`Zeitsperre`), Marco Braun (:user:`vindelico`), Pascal Bourgault (:user:`aulemahal`), Sarah-Claude Bourdeau-Goulet (:user:`Sarahclaude`), Éric Dupuis (:user:`coxipi`)
 
 New features and enhancements
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
@@ -13,13 +13,16 @@ New features and enhancements
 * Added style sheet ``transparent.mplstyle`` (:issue:`183`, :pull:`185`)
 * Fix ``NaN`` issues, extreme values in sizes legend and added edgecolors in ``fg.matplotlib.scattermap``  (:pull:`184`).
 * New function ``fg.data`` for fetching package data and defined `matplotlib` style definitions. (:pull:`211`).
+* ``fg.taylordiagram`` can now accept datasets with many dimensions (not only `taylor_params`), provided that they all share the same `ref_std` (e.g. normalized taylor diagrams)  (:pull:`214`).
+* A new optional way to organize points in a `fg.taylordiagram`  with  `colors_key`, `markers_key`  : DataArrays with a common dimension value or a common attrtibute are grouped with the same color/marker (:pull:`214`).
 
 Breaking changes
 ^^^^^^^^^^^^^^^^
 * `figanos` no longer supports Python 3.8. (:pull:`210`).
 * `figanos` now uses a `'src' layout <https://packaging.python.org/en/latest/discussions/src-layout-vs-flat-layout>`_ for the package. (:pull:`210`).
 * `cartopy` has been pinned above v0.23.0 due to a licensing issue. (:pull:`210`).
 * `twine` and `wheel` have been removed from the `dev` requirements. (:pull:`210`).
+* ``fg.taylordiagram`` returns a tuple of `(fig, floating_ax, legend)` instead of only `floating_ax`. (:pull:`214`).
 
 Internal changes
 ^^^^^^^^^^^^^^^^

src/figanos/matplotlib/plot.py

@@ -1827,6 +1827,8 @@ def taylordiagram(
     legend_kw: dict[str, Any] | None = None,
     std_label: str | None = None,
     corr_label: str | None = None,
+    colors_key: str | None = None,
+    markers_key: str | None = None,
 ):
     """Build a Taylor diagram.
 
@@ -1854,10 +1856,16 @@ def taylordiagram(
         Label for the standard deviation (x and y) axes.
     corr_label : str, optional
         Label for the correlation axis.
+    colors_key : str, optional
+        Attribute or dimension of DataArrays used to separate DataArrays into groups with different colors. If present,
+        it overrides the "color" key in `plot_kw`.
+    markers_key : str, optional
+        Attribute or dimension of DataArrays used to separate DataArrays into groups with different markers. If present,
+        it overrides the "marker" key in `plot_kw`.
 
     Returns
     -------
-    matplotlib.axes.Axes
+    (plt.figure, mpl_toolkits.axisartist.floating_axes.FloatingSubplot, plt.legend)
     """
     plot_kw = empty_dict(plot_kw)
     fig_kw = empty_dict(fig_kw)
@@ -1875,8 +1883,7 @@ def taylordiagram(
         data = {"_no_label": data}  # mpl excludes labels starting with "_" from legend
         plot_kw = {"_no_label": empty_dict(plot_kw)}
     elif not plot_kw:
-        plot_kw = {}
-
+        plot_kw = {k: {} for k in data.keys()}
     # check type
     for key, v in data.items():
         if not isinstance(v, xr.DataArray):
@@ -1886,6 +1893,30 @@ def taylordiagram(
         if key == "reference":
             raise ValueError("'reference' is not allowed as a key in data.")
 
+    # If there are other dimensions than 'taylor_param', create a bigger dict with them
+    data_keys = list(data.keys())
+    for data_key in data_keys:
+        da = data[data_key]
+        dims = list(set(da.dims) - {"taylor_param"})
+        if dims != []:
+            da = da.stack(pl_dims=dims)
+            for i, dim_key in enumerate(da.pl_dims.values):
+                if isinstance(dim_key, list) or isinstance(dim_key, tuple):
+                    dim_key = "-".join([str(k) for k in dim_key])
+                da0 = da.isel(pl_dims=i)
+                # if colors_key/markers_key is a dimension, add it as an attribute for later use
+                if markers_key in dims:
+                    da0.attrs[markers_key] = da0[markers_key].values.item()
+                if colors_key in dims:
+                    da0.attrs[colors_key] = da0[colors_key].values.item()
+                new_data_key = (
+                    f"{data_key}-{dim_key}" if data_key != "_no_label" else dim_key
+                )
+                data[new_data_key] = da0
+                plot_kw[new_data_key] = empty_dict(plot_kw[f"{data_key}"])
+            data.pop(data_key)
+            plot_kw.pop(data_key)
+
     # remove negative correlations
     initial_len = len(data)
     removed = [
@@ -1929,10 +1960,9 @@ def taylordiagram(
     # make labels
     if not std_label:
         try:
-            std_label = (
-                get_localized_term("standard deviation")
-                + f" ({list(data.values())[0].units})"
-            )
+            units = list(data.values())[0].units
+            std_label = get_localized_term("standard deviation")
+            std_label = std_label if units == "" else f"{std_label} ({units})"
         except AttributeError:
             std_label = get_localized_term("standard deviation").capitalize()
 
@@ -1949,12 +1979,11 @@ def taylordiagram(
     transform = PolarAxes.PolarTransform()
 
     # Setup the axis, here we map angles in degrees to angles in radius
-    rlocs = np.array([0, 0.2, 0.4, 0.6, 0.8, 1])
-    rlocs_deg = rlocs * 90
-    tlocs = rlocs_deg * np.pi / 180  # convert degrees to radians
+    # Correlation labels
+    rlocs = np.array([0, 0.2, 0.4, 0.6, 0.7, 0.8, 0.9, 0.95, 0.99, 1])
+    tlocs = np.arccos(rlocs)  # Conversion to polar angles
     gl1 = gf.FixedLocator(tlocs)  # Positions
-    tf1 = gf.DictFormatter(dict(zip(tlocs, np.flip(rlocs).astype(str))))
-
+    tf1 = gf.DictFormatter(dict(zip(tlocs, map(str, rlocs))))
     # Standard deviation axis extent
     radius_min = std_range[0] * max(max_std)
     radius_max = std_range[1] * max(max_std)
@@ -2013,7 +2042,8 @@ def taylordiagram(
     # rmse contours from reference standard deviation
     if contours:
         radii, angles = np.meshgrid(
-            np.linspace(radius_min, radius_max), np.linspace(0, np.pi / 2)
+            np.linspace(radius_min, radius_max),
+            np.linspace(0, np.pi / 2),
         )
         # Compute centered RMS difference
         rms = np.sqrt(ref_std**2 + radii**2 - 2 * ref_std * radii * np.cos(angles))
@@ -2023,47 +2053,86 @@ def taylordiagram(
 
         ax.clabel(ct, ct.levels, fontsize=8)
 
-        ct_line = Line2D(
+        # points.append(ct_line)
+        ct_line = ax.plot(
             [0],
             [0],
             ls=contours_kw["linestyles"],
             lw=1,
             c="k" if "colors" not in contours_kw else contours_kw["colors"],
             label="rmse",
         )
-        points.append(ct_line)
+        points.append(ct_line[0])
 
     # get color options
     style_colors = matplotlib.rcParams["axes.prop_cycle"].by_key()["color"]
     if len(data) > len(style_colors):
         style_colors = style_colors * math.ceil(len(data) / len(style_colors))
     cat_colors = Path(__file__).parents[1] / "data/ipcc_colors/categorical_colors.json"
+    # get marker options (only used if `markers_key` is set)
+    style_markers = "oDv^<>p*hH+x|_"
+    if len(data) > len(style_markers):
+        style_markers = style_markers * math.ceil(len(data) / len(style_markers))
+
+    # set colors and markers styles based on discrimnating attributes (if specified)
+    if colors_key or markers_key:
+        if colors_key:
+            # get_scen_color : look for SSP, RCP, CMIP model color
+            colorsd = {
+                k: get_scen_color(k, cat_colors) or style_colors[i]
+                for i, k in enumerate({da.attrs[colors_key] for da in data.values()})
+            }
+        if markers_key:
+            markersd = {
+                k: style_markers[i]
+                for i, k in enumerate({da.attrs[markers_key] for da in data.values()})
+            }
+
+        for key, da in data.items():
+            if colors_key:
+                plot_kw[key]["color"] = colorsd[da.attrs[colors_key]]
+            if markers_key:
+                plot_kw[key]["marker"] = markersd[da.attrs[markers_key]]
 
     # plot scatter
     for (key, da), i in zip(data.items(), range(len(data))):
         # look for SSP, RCP, CMIP model color
-        if get_scen_color(key, cat_colors):
-            plot_kw[key].setdefault("color", get_scen_color(key, cat_colors))
-        else:
-            plot_kw[key].setdefault("color", style_colors[i])
-
-        # convert corr to polar coordinates
-        plot_corr = (1 - da.sel(taylor_param="corr").values) * 90 * np.pi / 180
-
+        if colors_key is None:
+            plot_kw[key].setdefault(
+                "color", get_scen_color(key, cat_colors) or style_colors[i]
+            )
         # set defaults
         plot_kw[key] = {"label": key} | plot_kw[key]
 
+        # legend will be handled later in this case
+        if markers_key or colors_key:
+            plot_kw[key]["label"] = ""
+
         # plot
         pt = ax.scatter(
-            plot_corr, da.sel(taylor_param="sim_std").values, **plot_kw[key]
+            np.arccos(da.sel(taylor_param="corr").values),
+            da.sel(taylor_param="sim_std").values,
+            **plot_kw[key],
         )
         points.append(pt)
 
     # legend
     legend_kw.setdefault("loc", "upper right")
-    fig.legend(points, [pt.get_label() for pt in points], **legend_kw)
-
-    return floating_ax
+    legend = fig.legend(points, [pt.get_label() for pt in points], **legend_kw)
+
+    # plot new legend if markers/colors represent a certain dimension
+    if colors_key or markers_key:
+        handles = list(floating_ax.get_legend_handles_labels()[0])
+        if markers_key:
+            for k, m in markersd.items():
+                handles.append(Line2D([0], [0], color="k", label=k, marker=m, ls=""))
+        if colors_key:
+            for k, c in colorsd.items():
+                handles.append(Line2D([0], [0], color=c, label=k, ls="-"))
+        legend.remove()
+        legend = fig.legend(handles=handles, **legend_kw)
+
+    return fig, floating_ax, legend
 
 
 def hatchmap(