Repository: arviz-devs/arviz

Issue:
Add interactive legend to bokeh forestplot
## Tell us about it
plot_density with bokeh backend has a super cool legend that allows to hide lines by clicking on the corresponding element in the legend. plot_forest could do the same.

## Thoughts on implementation
I guess searching for the PR that added the legend for plot_density will be a good starting point. Will try to extend and add more info at some point, but if someone has time to browse the history and track down the relevant pieces of code it would be great



Discussion:

Resolving patch:
diff --git a/arviz/plots/backends/bokeh/forestplot.py b/arviz/plots/backends/bokeh/forestplot.py
--- a/arviz/plots/backends/bokeh/forestplot.py
+++ b/arviz/plots/backends/bokeh/forestplot.py
@@ -7,7 +7,7 @@
 import matplotlib.pyplot as plt
 import numpy as np
 from bokeh.models import Band, ColumnDataSource, DataRange1d
-from bokeh.models.annotations import Title
+from bokeh.models.annotations import Title, Legend
 from bokeh.models.tickers import FixedTicker
 
 from ....sel_utils import xarray_var_iter
@@ -126,6 +126,8 @@ def plot_forest(
 
     axes = np.atleast_2d(axes)
 
+    plotted = defaultdict(list)
+
     if kind == "forestplot":
         plot_handler.forestplot(
             hdi_prob,
@@ -134,6 +136,7 @@ def plot_forest(
             markersize,
             axes[0, 0],
             rope,
+            plotted,
         )
     elif kind == "ridgeplot":
         plot_handler.ridgeplot(
@@ -146,6 +149,7 @@ def plot_forest(
             ridgeplot_truncate,
             ridgeplot_quantiles,
             axes[0, 0],
+            plotted,
         )
     else:
         raise TypeError(
@@ -155,11 +159,17 @@ def plot_forest(
 
     idx = 1
     if ess:
-        plot_handler.plot_neff(axes[0, idx], markersize)
+        plotted_ess = defaultdict(list)
+        plot_handler.plot_neff(axes[0, idx], markersize, plotted_ess)
+        if legend:
+            plot_handler.legend(axes[0, idx], plotted_ess)
         idx += 1
 
     if r_hat:
-        plot_handler.plot_rhat(axes[0, idx], markersize)
+        plotted_r_hat = defaultdict(list)
+        plot_handler.plot_rhat(axes[0, idx], markersize, plotted_r_hat)
+        if legend:
+            plot_handler.legend(axes[0, idx], plotted_r_hat)
         idx += 1
 
     for i, ax_ in enumerate(axes.ravel()):
@@ -177,6 +187,7 @@ def plot_forest(
         ax_.y_range = DataRange1d(bounds=backend_config["bounds_y_range"], min_interval=2)
 
     labels, ticks = plot_handler.labels_and_ticks()
+    ticks = [int(tick) if (tick).is_integer() else tick for tick in ticks]
 
     axes[0, 0].yaxis.ticker = FixedTicker(ticks=ticks)
     axes[0, 0].yaxis.major_label_overrides = dict(zip(map(str, ticks), map(str, labels)))
@@ -193,6 +204,8 @@ def plot_forest(
     ].group_offset
     axes[0, 0].y_range._property_values["end"] = y_max  # pylint: disable=protected-access
 
+    if legend:
+        plot_handler.legend(axes[0, 0], plotted)
     show_layout(axes, show)
 
     return axes
@@ -269,30 +282,44 @@ def labels_and_ticks(self):
         def label_idxs():
             labels, idxs = [], []
             for plotter in val:
-                sub_labels, sub_idxs, _, _ = plotter.labels_ticks_and_vals()
+                sub_labels, sub_idxs, _, _, _ = plotter.labels_ticks_and_vals()
                 labels.append(sub_labels)
                 idxs.append(sub_idxs)
             return np.concatenate(labels), np.concatenate(idxs)
 
         return label_idxs()
 
-    def display_multiple_ropes(self, rope, ax, y, linewidth, var_name, selection):
+    def legend(self, ax, plotted):
+        """Add interactive legend with colorcoded model info."""
+        legend_it = []
+        for (model_name, glyphs) in plotted.items():
+            legend_it.append((model_name, glyphs))
+
+        legend = Legend(items=legend_it)
+        ax.add_layout(legend, "right")
+        ax.legend.click_policy = "hide"
+
+    def display_multiple_ropes(
+        self, rope, ax, y, linewidth, var_name, selection, plotted, model_name
+    ):
         """Display ROPE when more than one interval is provided."""
         for sel in rope.get(var_name, []):
             # pylint: disable=line-too-long
             if all(k in selection and selection[k] == v for k, v in sel.items() if k != "rope"):
                 vals = sel["rope"]
-                ax.line(
-                    vals,
-                    (y + 0.05, y + 0.05),
-                    line_width=linewidth * 2,
-                    color=[
-                        color
-                        for _, color in zip(
-                            range(3), cycle(plt.rcParams["axes.prop_cycle"].by_key()["color"])
-                        )
-                    ][2],
-                    line_alpha=0.7,
+                plotted[model_name].append(
+                    ax.line(
+                        vals,
+                        (y + 0.05, y + 0.05),
+                        line_width=linewidth * 2,
+                        color=[
+                            color
+                            for _, color in zip(
+                                range(3), cycle(plt.rcParams["axes.prop_cycle"].by_key()["color"])
+                            )
+                        ][2],
+                        line_alpha=0.7,
+                    )
                 )
                 return ax
 
@@ -307,6 +334,7 @@ def ridgeplot(
         ridgeplot_truncate,
         ridgeplot_quantiles,
         ax,
+        plotted,
     ):
         """Draw ridgeplot for each plotter.
 
@@ -332,11 +360,13 @@ def ridgeplot(
             Defaults to None.
         ax : Axes
             Axes to draw on
+        plotted : dict
+            Contains glyphs for each model
         """
         if alpha is None:
             alpha = 1.0
         for plotter in list(self.plotters.values())[::-1]:
-            for x, y_min, y_max, hdi_, y_q, color in plotter.ridgeplot(
+            for x, y_min, y_max, hdi_, y_q, color, model_name in plotter.ridgeplot(
                 hdi_prob, mult, ridgeplot_kind
             ):
                 if alpha == 0:
@@ -353,14 +383,16 @@ def ridgeplot(
                         facecolor = color
                         alpha = [alpha if ci else 0 for ci in ((x >= hdi_[0]) & (x <= hdi_[1]))]
                     y_min = np.ones_like(x) * y_min
-                    ax.vbar(
-                        x=x,
-                        top=y_max - y_min,
-                        bottom=y_min,
-                        width=0.9,
-                        line_color=border,
-                        color=facecolor,
-                        fill_alpha=alpha,
+                    plotted[model_name].append(
+                        ax.vbar(
+                            x=x,
+                            top=y_max - y_min,
+                            bottom=y_min,
+                            width=0.9,
+                            line_color=border,
+                            color=facecolor,
+                            fill_alpha=alpha,
+                        )
                     )
                 else:
                     tr_x = x[(x >= hdi_[0]) & (x <= hdi_[1])]
@@ -375,41 +407,52 @@ def ridgeplot(
                         line_width=0,
                     )
                     patch.level = "overlay"
+                    plotted[model_name].append(patch)
                     if ridgeplot_truncate:
-                        ax.line(
-                            x, y_max, line_dash="solid", line_width=linewidth, line_color=border
+                        plotted[model_name].append(
+                            ax.line(
+                                x, y_max, line_dash="solid", line_width=linewidth, line_color=border
+                            )
                         )
-                        ax.line(
-                            x, y_min, line_dash="solid", line_width=linewidth, line_color=border
+                        plotted[model_name].append(
+                            ax.line(
+                                x, y_min, line_dash="solid", line_width=linewidth, line_color=border
+                            )
                         )
                     else:
-                        ax.line(
-                            tr_x,
-                            tr_y_max,
-                            line_dash="solid",
-                            line_width=linewidth,
-                            line_color=border,
+                        plotted[model_name].append(
+                            ax.line(
+                                tr_x,
+                                tr_y_max,
+                                line_dash="solid",
+                                line_width=linewidth,
+                                line_color=border,
+                            )
                         )
-                        ax.line(
-                            tr_x,
-                            tr_y_min,
-                            line_dash="solid",
-                            line_width=linewidth,
-                            line_color=border,
+                        plotted[model_name].append(
+                            ax.line(
+                                tr_x,
+                                tr_y_min,
+                                line_dash="solid",
+                                line_width=linewidth,
+                                line_color=border,
+                            )
                         )
                 if ridgeplot_quantiles is not None:
                     quantiles = [x[np.sum(y_q < quant)] for quant in ridgeplot_quantiles]
-                    ax.diamond(
-                        quantiles,
-                        np.ones_like(quantiles) * y_min[0],
-                        line_color="black",
-                        fill_color="black",
-                        size=markersize,
+                    plotted[model_name].append(
+                        ax.diamond(
+                            quantiles,
+                            np.ones_like(quantiles) * y_min[0],
+                            line_color="black",
+                            fill_color="black",
+                            size=markersize,
+                        )
                     )
 
         return ax
 
-    def forestplot(self, hdi_prob, quartiles, linewidth, markersize, ax, rope):
+    def forestplot(self, hdi_prob, quartiles, linewidth, markersize, ax, rope, plotted):
         """Draw forestplot for each plotter.
 
         Parameters
@@ -424,6 +467,8 @@ def forestplot(self, hdi_prob, quartiles, linewidth, markersize, ax, rope):
             Size of marker in center of forestplot line
         ax : Axes
             Axes to draw on
+        plotted : dict
+            Contains glyphs for each model
         """
         if rope is None or isinstance(rope, dict):
             pass
@@ -465,23 +510,32 @@ def forestplot(self, hdi_prob, quartiles, linewidth, markersize, ax, rope):
             qlist = [endpoint, 50, 100 - endpoint]
 
         for plotter in self.plotters.values():
-            for y, selection, values, color in plotter.treeplot(qlist, hdi_prob):
+            for y, model_name, selection, values, color in plotter.treeplot(qlist, hdi_prob):
                 if isinstance(rope, dict):
-                    self.display_multiple_ropes(rope, ax, y, linewidth, plotter.var_name, selection)
+                    self.display_multiple_ropes(
+                        rope, ax, y, linewidth, plotter.var_name, selection, plotted, model_name
+                    )
 
                 mid = len(values) // 2
                 param_iter = zip(
                     np.linspace(2 * linewidth, linewidth, mid, endpoint=True)[-1::-1], range(mid)
                 )
                 for width, j in param_iter:
-                    ax.line(
-                        [values[j], values[-(j + 1)]], [y, y], line_width=width, line_color=color
+                    plotted[model_name].append(
+                        ax.line(
+                            [values[j], values[-(j + 1)]],
+                            [y, y],
+                            line_width=width,
+                            line_color=color,
+                        )
+                    )
+                plotted[model_name].append(
+                    ax.circle(
+                        x=values[mid],
+                        y=y,
+                        size=markersize * 0.75,
+                        fill_color=color,
                     )
-                ax.circle(
-                    x=values[mid],
-                    y=y,
-                    size=markersize * 0.75,
-                    fill_color=color,
                 )
         _title = Title()
         _title.text = "{:.1%} HDI".format(hdi_prob)
@@ -489,18 +543,20 @@ def forestplot(self, hdi_prob, quartiles, linewidth, markersize, ax, rope):
 
         return ax
 
-    def plot_neff(self, ax, markersize):
+    def plot_neff(self, ax, markersize, plotted):
         """Draw effective n for each plotter."""
         max_ess = 0
         for plotter in self.plotters.values():
-            for y, ess, color in plotter.ess():
+            for y, ess, color, model_name in plotter.ess():
                 if ess is not None:
-                    ax.circle(
-                        x=ess,
-                        y=y,
-                        fill_color=color,
-                        size=markersize,
-                        line_color="black",
+                    plotted[model_name].append(
+                        ax.circle(
+                            x=ess,
+                            y=y,
+                            fill_color=color,
+                            size=markersize,
+                            line_color="black",
+                        )
                     )
                 if ess > max_ess:
                     max_ess = ess
@@ -515,12 +571,16 @@ def plot_neff(self, ax, markersize):
 
         return ax
 
-    def plot_rhat(self, ax, markersize):
+    def plot_rhat(self, ax, markersize, plotted):
         """Draw r-hat for each plotter."""
         for plotter in self.plotters.values():
-            for y, r_hat, color in plotter.r_hat():
+            for y, r_hat, color, model_name in plotter.r_hat():
                 if r_hat is not None:
-                    ax.circle(x=r_hat, y=y, fill_color=color, size=markersize, line_color="black")
+                    plotted[model_name].append(
+                        ax.circle(
+                            x=r_hat, y=y, fill_color=color, size=markersize, line_color="black"
+                        )
+                    )
         ax.x_range._property_values["start"] = 0.9  # pylint: disable=protected-access
         ax.x_range._property_values["end"] = 2.1  # pylint: disable=protected-access
 
@@ -586,10 +646,8 @@ def iterator(self):
                 datum_list = list(datum_iter)
                 for _, selection, isel, values in datum_list:
                     selection_list.append(selection)
-                    if not selection:
+                    if not selection or not len(selection_list) % len(datum_list):
                         var_name = self.var_name
-                    elif len(selection_list) == len(datum_list):
-                        var_name = self.var_name + ":"
                     else:
                         var_name = ""
                     label = self.labeller.make_label_flat(var_name, selection, isel)
@@ -602,14 +660,11 @@ def iterator(self):
         y = self.y_start
         for idx, (label, model_data) in enumerate(label_dict.items()):
             for model_name, value_list in model_data.items():
-                if model_name:
-                    row_label = "{}: {}".format(model_name, label)
-                else:
-                    row_label = label
+                row_label = self.labeller.make_model_label(model_name, label)
                 for values in value_list:
-                    yield y, row_label, label, selection_list[idx], values, self.model_color[
-                        model_name
-                    ]
+                    yield y, row_label, model_name, label, selection_list[
+                        idx
+                    ], values, self.model_color[model_name]
                     y += self.chain_offset
                 y += self.var_offset
             y += self.group_offset
@@ -617,30 +672,32 @@ def iterator(self):
     def labels_ticks_and_vals(self):
         """Get labels, ticks, values, and colors for the variable."""
         y_ticks = defaultdict(list)
-        for y, label, _, _, vals, color in self.iterator():
-            y_ticks[label].append((y, vals, color))
-        labels, ticks, vals, colors = [], [], [], []
+        for y, label, model_name, _, _, vals, color in self.iterator():
+            y_ticks[label].append((y, vals, color, model_name))
+        labels, ticks, vals, colors, model_names = [], [], [], [], []
         for label, data in y_ticks.items():
             labels.append(label)
             ticks.append(np.mean([j[0] for j in data]))
             vals.append(np.vstack([j[1] for j in data]))
+            model_names.append(data[0][3])
             colors.append(data[0][2])  # the colors are all the same
-        return labels, ticks, vals, colors
+        return labels, ticks, vals, colors, model_names
 
     def treeplot(self, qlist, hdi_prob):
         """Get data for each treeplot for the variable."""
-        for y, _, _, selection, values, color in self.iterator():
+        for y, _, model_name, _, selection, values, color in self.iterator():
             ntiles = np.percentile(values.flatten(), qlist)
             ntiles[0], ntiles[-1] = hdi(values.flatten(), hdi_prob, multimodal=False)
-            yield y, selection, ntiles, color
+            yield y, model_name, selection, ntiles, color
 
     def ridgeplot(self, hdi_prob, mult, ridgeplot_kind):
         """Get data for each ridgeplot for the variable."""
-        xvals, hdi_vals, yvals, pdfs, pdfs_q, colors = [], [], [], [], [], []
+        xvals, hdi_vals, yvals, pdfs, pdfs_q, colors, model_names = [], [], [], [], [], [], []
 
-        for y, *_, values, color in self.iterator():
+        for y, _, model_name, *_, values, color in self.iterator():
             yvals.append(y)
             colors.append(color)
+            model_names.append(model_name)
             values = values.flatten()
             values = values[np.isfinite(values)]
 
@@ -669,23 +726,25 @@ def ridgeplot(self, hdi_prob, mult, ridgeplot_kind):
             hdi_vals.append(hdi_)
 
         scaling = max(np.max(j) for j in pdfs)
-        for y, x, hdi_val, pdf, pdf_q, color in zip(yvals, xvals, hdi_vals, pdfs, pdfs_q, colors):
-            yield x, y, mult * pdf / scaling + y, hdi_val, pdf_q, color
+        for y, x, hdi_val, pdf, pdf_q, color, model_name in zip(
+            yvals, xvals, hdi_vals, pdfs, pdfs_q, colors, model_names
+        ):
+            yield x, y, mult * pdf / scaling + y, hdi_val, pdf_q, color, model_name
 
     def ess(self):
         """Get effective n data for the variable."""
-        _, y_vals, values, colors = self.labels_ticks_and_vals()
-        for y, value, color in zip(y_vals, values, colors):
-            yield y, _ess(value), color
+        _, y_vals, values, colors, model_names = self.labels_ticks_and_vals()
+        for y, value, color, model_name in zip(y_vals, values, colors, model_names):
+            yield y, _ess(value), color, model_name
 
     def r_hat(self):
         """Get rhat data for the variable."""
-        _, y_vals, values, colors = self.labels_ticks_and_vals()
-        for y, value, color in zip(y_vals, values, colors):
+        _, y_vals, values, colors, model_names = self.labels_ticks_and_vals()
+        for y, value, color, model_name in zip(y_vals, values, colors, model_names):
             if value.ndim != 2 or value.shape[0] < 2:
-                yield y, None, color
+                yield y, None, color, model_name
             else:
-                yield y, _rhat(value), color
+                yield y, _rhat(value), color, model_name
 
     def y_max(self):
         """Get max y value for the variable."""


Test patch: