better support for aesthetics in labels elements

Author: OriolAbril

src/arviz_plots/plots/forestplot.py

@@ -353,7 +353,6 @@ def plot_forest(
         )
 
     # add labels and shading first, so forest plot is rendered on top
-    _, lab_aes, lab_ignore = filter_aes(plot_collection, aes_map, "labels", sample_dims)
     cumulative_label = []
     x = 0
     for label in labellable_dims:
@@ -366,13 +365,6 @@ def plot_forest(
         shade_extend = 0.5 if add_factor == 0 else 0.3
         if label not in labels:
             continue
-        lab_kwargs = plot_kwargs.get("labels", {}).copy()
-        if "color" not in lab_aes:
-            lab_kwargs.setdefault("color", "black")
-        if x == 0:
-            lab_kwargs.setdefault("horizontal_align", "left")
-        if x == len(labels) - 1:
-            lab_kwargs.setdefault("horizontal_align", "right")
         if label == "__variable__":
             y_max = y_ds.max() + shade_extend
             y_min = y_ds.min() - shade_extend
@@ -413,6 +405,23 @@ def plot_forest(
                 ignore_aes=shade_ignore,
                 **shade_kwargs,
             )
+        _, lab_aes, lab_ignore = filter_aes(plot_collection, aes_map, "labels", sample_dims)
+        extra_ignore_aes = []
+        for aes_key in lab_aes:
+            if aes_key == "overlay":
+                continue
+            aes_ds = plot_collection.get_aes_as_dataset(aes_key)
+            if set(aes_ds.dims).difference(cumulative_label):
+                extra_ignore_aes.append(aes_key)
+        lab_aes = set(lab_aes).difference(extra_ignore_aes)
+        lab_ignore = set(lab_ignore).union(extra_ignore_aes)
+        lab_kwargs = plot_kwargs.get("labels", {}).copy()
+        if "color" not in lab_aes:
+            lab_kwargs.setdefault("color", "black")
+        if x == 0:
+            lab_kwargs.setdefault("horizontal_align", "left")
+        if x == len(labels) - 1:
+            lab_kwargs.setdefault("horizontal_align", "right")
         plot_collection.map(
             annotate_label,
             f"{label.strip('_')}_label",

tests/test_plots.py

@@ -2,7 +2,7 @@
 """Test batteries-included plots."""
 import numpy as np
 import pytest
-from arviz_base import from_dict
+from arviz_base import from_dict, load_arviz_data
 
 from arviz_plots import plot_dist, plot_forest, plot_trace, visuals
 
@@ -120,13 +120,15 @@ def test_plot_forest_extendable(self, datatree, backend):
         assert pc.viz["plot"].sizes["column"] == 3
         assert all("ess" in child.data_vars for child in pc.viz.children.values())
 
-    def test_plot_forest_color_shading(self, datatree2, backend):
-        pc = plot_forest(
-            datatree2,
-            pc_kwargs={"aes": {"color": ["__variable__"]}},
-            aes_map={"labels": ["color"]},
-            shade_label="hierarchy",
-            backend=backend,
-        )
-        assert "plot" in pc.viz.data_vars
-        assert all("shade" in child.data_vars for child in pc.viz.children.values())
+    def test_plot_forest_aes_labels_shading(self, backend):
+        post = load_arviz_data("rugby_field").posterior.ds.sel(draw=slice(None, 100))
+        for pseudo_dim in ("__variable__", "field", "team"):
+            pc = plot_forest(
+                post,
+                pc_kwargs={"aes": {"color": [pseudo_dim]}},
+                aes_map={"labels": ["color"]},
+                shade_label=pseudo_dim,
+                backend=backend,
+            )
+            assert "plot" in pc.viz.data_vars
+            assert all("shade" in child.data_vars for child in pc.viz.children.values())