add plot_forest draft

For everything to work, the following features were added to:
* support for "__variable__" in aesthetics
* default values for common aesthetics
* support for dict of datatree input -> multiple models aligned within the same plot

Author: OriolAbril

docs/source/api/plots.rst

@@ -16,5 +16,6 @@ A complementary introduction and guide to ``plot_...`` functions is available at
    :toctree: generated/
 
    plot_dist
+   plot_forest
    plot_trace
    plot_trace_dist

src/arviz_plots/backend/__init__.py

@@ -12,6 +12,7 @@
 to any type of the plotting backend or even custom objects, but all instances
 of the same placeholder must use the same type (whatever that is).
 """
+import numpy as np
 
 error = NotImplementedError(
     "The `arviz_plots.backend` module itself is for reference only. "
@@ -20,6 +21,22 @@
 )
 
 
+# generation of default values for aesthetics
+def get_default_aes(aes_key, n, kwargs):
+    """Generate `n` default values for a given aesthetics keyword."""
+    if aes_key not in kwargs:
+        if aes_key in {"x", "y"}:
+            return np.arange(n)
+        if aes_key == "alpha":
+            return np.linspace(0.2, 0.7, n)
+        return [None] * n
+    aes_vals = kwargs[aes_key]
+    n_aes_vals = len(aes_vals)
+    if n_aes_vals >= n:
+        return aes_vals[:n]
+    return np.tile(aes_vals, (n // n_aes_vals) + 1)[:n]
+
+
 # object creation and i/o
 def show(chart):
     """Show this :term:`chart`.

src/arviz_plots/backend/bokeh/__init__.py

@@ -7,6 +7,7 @@
 from bokeh.plotting import figure
 from bokeh.plotting import show as _show
 
+from .. import get_default_aes as get_agnostic_default_aes
 from .legend import legend
 
 __all__ = [
@@ -31,6 +32,27 @@ class UnsetDefault:
 unset = UnsetDefault()
 
 
+# generation of default values for aesthetics
+def get_default_aes(aes_key, n, kwargs):
+    """Generate `n` *bokeh valid* default values for a given aesthetics keyword."""
+    if aes_key not in kwargs:
+        if "color" in aes_key:
+            # fmt: off
+            vals = [
+                '#3f90da', '#ffa90e', '#bd1f01', '#94a4a2', '#832db6',
+                '#a96b59', '#e76300', '#b9ac70', '#717581', '#92dadd'
+            ]
+            # fmt: on
+        elif aes_key in {"linestyle", "line_dash"}:
+            vals = ["solid", "dashed", "dotted", "dashdot"]
+        elif aes_key == "marker":
+            vals = ["circle", "cross", "triangle", "x", "diamond"]
+        else:
+            return get_agnostic_default_aes(aes_key, n, {})
+        return get_agnostic_default_aes(aes_key, n, {aes_key: vals})
+    return get_agnostic_default_aes(aes_key, n, kwargs)
+
+
 # object creation and i/o
 def show(chart):
     """Show the provided bokeh layout."""

src/arviz_plots/backend/matplotlib/__init__.py

@@ -11,10 +11,12 @@
 from matplotlib.cbook import normalize_kwargs
 from matplotlib.collections import PathCollection
 from matplotlib.lines import Line2D
+from matplotlib.pyplot import rcParams
 from matplotlib.pyplot import show as _show
 from matplotlib.pyplot import subplots
 from matplotlib.text import Text
 
+from .. import get_default_aes as get_agnostic_default_aes
 from .legend import legend
 
 __all__ = [
@@ -39,6 +41,33 @@ class UnsetDefault:
 unset = UnsetDefault()
 
 
+# generation of default values for aesthetics
+def get_default_aes(aes_key, n, kwargs):
+    """Generate `n` *bokeh valid* default values for a given aesthetics keyword."""
+    if aes_key not in kwargs:
+        default_prop_cycle = rcParams["axes.prop_cycle"].by_key()
+        if ("color" in aes_key) or aes_key == "c":
+            # fmt: off
+            vals = [
+                '#3f90da', '#ffa90e', '#bd1f01', '#94a4a2', '#832db6',
+                '#a96b59', '#e76300', '#b9ac70', '#717581', '#92dadd'
+            ]
+            # fmt: on
+            vals = default_prop_cycle.get("color", vals)
+        elif aes_key in {"linestyle", "ls"}:
+            vals = ["-", "--", ":", "-."]
+            vals = default_prop_cycle.get("linestyle", vals)
+        elif aes_key in {"marker", "m"}:
+            vals = ["o", "+", "^", "x", "d"]
+            vals = default_prop_cycle.get("marker", vals)
+        elif aes_key in default_prop_cycle:
+            vals = default_prop_cycle[aes_key]
+        else:
+            return get_agnostic_default_aes(aes_key, n, {})
+        return get_agnostic_default_aes(aes_key, n, {aes_key: vals})
+    return get_agnostic_default_aes(aes_key, n, kwargs)
+
+
 # object creation and i/o
 def show(chart):  # pylint: disable=unused-argument
     """Show all existing matplotlib figures."""

src/arviz_plots/plot_collection.py

@@ -8,6 +8,16 @@
 from datatree import DataTree
 
 
+def concat_model_dict(data):
+    """Merge multiple Datasets into a single one along a new model dimension."""
+    if isinstance(data, dict):
+        ds_list = data.values()
+        if not all(isinstance(ds, xr.Dataset) for ds in ds_list):
+            raise TypeError("Provided data must be a Dataset or dictionary of Datasets")
+        data = xr.concat(ds_list, dim="model").assign_coords(model=list(data))
+    return data
+
+
 def sel_subset(sel, present_dims):
     """Subset a dictionary of dim: coord values.
 
@@ -97,8 +107,15 @@ def _get_aes_dict_from_dt(aes_dt):
     an aes DataTree directly when initializating a PlotCollection object.
     This method is used to generate the more basic dictionary from the DataTree.
     """
+    child_list = list(aes_dt.children.values())
     aes = {}
-    for ds in aes_dt.children.values():
+    aes_in_all_vars = set.intersection(*[set(child.data_vars) for child in child_list])
+    aes = {
+        aes_key: ["__variable__"]
+        for aes_key in aes_in_all_vars
+        if any(child[aes_key].item(0) != child_list[0][aes_key].item(0) for child in child_list)
+    }
+    for ds in child_list:
         for aes_key, values in ds.items():
             if not values.dims:
                 continue
@@ -211,7 +228,7 @@ def data(self):
     @data.setter
     def data(self, value):
         # might want/be possible to make some checks on the data before setting it
-        self._data = value
+        self._data = concat_model_dict(value)
 
     @property
     def aes_set(self):
@@ -301,31 +318,66 @@ def generate_aes_dt(self, aes=None, **kwargs):
         but it will always be possible to set their value manually.
         """
         if aes is None:
-            aes = {}
+            aes = self._aes
+            kwargs = self._kwargs
         self._aes = aes
         self._kwargs = kwargs
-        self._aes_dt = DataTree()
-        for var_name, da in self.data.items():
-            ds = xr.Dataset()
-            for aes_key, dims in aes.items():
-                aes_vals = kwargs.get(aes_key, [None])
-                aes_dims = [dim for dim in dims if dim in da.dims]
-                aes_raw_shape = [da.sizes[dim] for dim in aes_dims]
-                if not aes_raw_shape:
-                    ds[aes_key] = aes_vals[0]
-                    continue
-                n_aes = np.prod(aes_raw_shape)
-                n_aes_vals = len(aes_vals)
-                if n_aes_vals > n_aes:
-                    aes_vals = aes_vals[:n_aes]
-                elif n_aes_vals < n_aes:
-                    aes_vals = np.tile(aes_vals, (n_aes // n_aes_vals) + 1)[:n_aes]
-                ds[aes_key] = xr.DataArray(
-                    np.array(aes_vals).reshape(aes_raw_shape),
-                    dims=aes_dims,
-                    coords={dim: da.coords[dim] for dim in dims if dim in da.coords},
+        if not hasattr(self, "backend"):
+            plot_bknd = import_module(".backend", package="arviz_plots")
+        else:
+            plot_bknd = import_module(f".backend.{self.backend}", package="arviz_plots")
+        get_default_aes = plot_bknd.get_default_aes
+        ds_dict = {var_name: xr.Dataset() for var_name in self.data.data_vars}
+        for aes_key, dims in aes.items():
+            if "__variable__" in dims:
+                total_aes_vals = int(
+                    np.sum(
+                        [
+                            np.prod([size for dim, size in da.sizes.items() if dim in dims])
+                            for da in self.data.values()
+                        ]
+                    )
+                )
+                aes_vals = get_default_aes(aes_key, total_aes_vals, kwargs)
+                aes_cumulative = 0
+                for var_name, da in self.data.items():
+                    ds = ds_dict[var_name]
+                    aes_dims = [dim for dim in dims if dim in da.dims]
+                    aes_raw_shape = [da.sizes[dim] for dim in aes_dims]
+                    if not aes_raw_shape:
+                        ds[aes_key] = np.asarray(aes_vals)[
+                            aes_cumulative : aes_cumulative + 1
+                        ].squeeze()
+                        aes_cumulative += 1
+                        continue
+                    n_aes = np.prod(aes_raw_shape)
+                    ds[aes_key] = xr.DataArray(
+                        np.array(aes_vals[aes_cumulative : aes_cumulative + n_aes]).reshape(
+                            aes_raw_shape
+                        ),
+                        dims=aes_dims,
+                        coords={dim: da.coords[dim] for dim in dims if dim in da.coords},
+                    )
+                    aes_cumulative += n_aes
+            else:
+                total_aes_vals = int(
+                    np.prod([self.data.sizes[dim] for dim in self.data.dims if dim in dims])
                 )
-            DataTree(name=var_name, parent=self._aes_dt, data=ds)
+                aes_vals = get_default_aes(aes_key, total_aes_vals, kwargs)
+                for var_name, da in self.data.items():
+                    ds = ds_dict[var_name]
+                    aes_dims = [dim for dim in dims if dim in da.dims]
+                    aes_raw_shape = [da.sizes[dim] for dim in aes_dims]
+                    if not aes_raw_shape:
+                        ds[aes_key] = aes_vals[0]
+                        continue
+                    n_aes = np.prod(aes_raw_shape)
+                    ds[aes_key] = xr.DataArray(
+                        np.array(aes_vals[:n_aes]).reshape(aes_raw_shape),
+                        dims=aes_dims,
+                        coords={dim: da.coords[dim] for dim in dims if dim in da.coords},
+                    )
+        self._aes_dt = DataTree.from_dict(ds_dict)
 
     @property
     def base_loop_dims(self):
@@ -385,8 +437,7 @@ def wrap(
             plot_grid_kws = {}
         if backend is None:
             backend = rcParams["plot.backend"]
-        if isinstance(data, dict):
-            data = xr.concat(data.values(), dim="model").assign_coords(model=list(data))
+        data = concat_model_dict(data)
 
         n_plots, plots_per_var = _process_facet_dims(data, cols)
         if n_plots <= col_wrap:
@@ -501,8 +552,7 @@ def grid(
         repeated_dims = [col for col in cols if col in rows]
         if repeated_dims:
             raise ValueError("The same dimension can't be used for both cols and rows.")
-        if isinstance(data, dict):
-            data = xr.concat(data.values(), dim="model").assign_coords(model=list(data))
+        data = concat_model_dict(data)
 
         n_cols, cols_per_var = _process_facet_dims(data, cols)
         n_rows, rows_per_var = _process_facet_dims(data, rows)

src/arviz_plots/plots/__init__.py

@@ -1,7 +1,8 @@
 """Batteries-included ArviZ plots."""
 
 from .distplot import plot_dist
+from .forestplot import plot_forest
 from .tracedistplot import plot_trace_dist
 from .traceplot import plot_trace
 
-__all__ = ["plot_dist", "plot_trace", "plot_trace_dist"]
+__all__ = ["plot_dist", "plot_forest", "plot_trace", "plot_trace_dist"]

src/arviz_plots/plots/distplot.py

@@ -3,10 +3,9 @@
 import xarray as xr
 from arviz_base import rcParams
 from arviz_base.labels import BaseLabeller
-from arviz_base.utils import _var_names
 
 from arviz_plots.plot_collection import PlotCollection
-from arviz_plots.plots.utils import filter_aes
+from arviz_plots.plots.utils import filter_aes, process_group_variables_coords
 from arviz_plots.visuals import (
     ecdf_line,
     labelled_title,
@@ -23,6 +22,7 @@ def plot_dist(
     var_names=None,
     filter_vars=None,
     group="posterior",
+    coords=None,
     sample_dims=None,
     kind=None,
     point_estimate=None,
@@ -46,8 +46,9 @@ def plot_dist(
 
     Parameters
     ----------
-    dt : DataTree
-        Input data
+    dt : DataTree or dict of {str : DataTree}
+        Input data. In case of dictionary input, the keys are taken to be model names.
+        In such cases, a dimension "model" is generated and can be used to map to aesthetics.
     var_names: str or list of str, optional
         One or more variables to be plotted.
         Prefix the variables by ~ when you want to exclude them from the plot.
@@ -57,6 +58,7 @@ def plot_dist(
         If “regex”, interpret var_names as regular expressions on the real variables names.
     group : str, default "posterior"
         Group to be plotted.
+    coords : dict, optional
     sample_dims : iterable, optional
         Dimensions to reduce unless mapped to an aesthetic.
         Defaults to ``rcParams["data.sample_dims"]``
@@ -106,6 +108,20 @@ def plot_dist(
     Returns
     -------
     PlotCollection
+
+    Examples
+    --------
+    Default plot_dist for a single model:
+
+    .. plot::
+        :context: close-figs
+
+        >>> from arviz_plots import plot_dist
+        >>> from arviz_base import load_arviz_data
+        >>> centered = load_arviz_data('centered_eight')
+        >>> non_centered = load_arviz_data('non_centered_eight')
+        >>> pc = plot_dist(centered)
+
     """
     if ci_kind not in ["hdi", "eti", None]:
         raise ValueError("ci_kind must be either 'hdi' or 'eti'")
@@ -130,27 +146,37 @@ def plot_dist(
     if stats_kwargs is None:
         stats_kwargs = {}
 
-    distribution = dt[group].ds
-    var_names = _var_names(var_names, distribution, filter_vars)
-
-    if var_names is not None:
-        distribution = dt[group].ds[var_names]
+    distribution = process_group_variables_coords(
+        dt, group=group, var_names=var_names, filter_vars=filter_vars, coords=coords
+    )
 
     if plot_collection is None:
         if backend is None:
             backend = rcParams["plot.backend"]
         pc_kwargs.setdefault("col_wrap", 5)
         pc_kwargs.setdefault(
-            "cols", ["__variable__"] + [dim for dim in distribution.dims if dim not in sample_dims]
+            "cols",
+            ["__variable__"]
+            + [dim for dim in distribution.dims if (dim not in sample_dims) and (dim != "model")],
         )
+        if "model" in distribution:
+            pc_kwargs["aes"] = pc_kwargs.get("aes", {}).copy()
+            pc_kwargs["aes"].setdefault("color", ["model"])
         plot_collection = PlotCollection.wrap(
             distribution,
             backend=backend,
             **pc_kwargs,
         )
 
     if aes_map is None:
-        aes_map = {kind: plot_collection.aes_set}
+        if "model" in distribution:
+            aes_map = {
+                kind: plot_collection.aes_set,
+                "credible_interval": ["color"],
+                "point_estimate": ["color"],
+            }
+        else:
+            aes_map = {kind: plot_collection.aes_set}
     if "point_estimate" in aes_map and "point_estimate_text" not in aes_map:
         aes_map["point_estimate_text"] = aes_map["point_estimate"]
     if labeller is None: