Repository: arviz-devs/arviz

Issue:
Is it possible to have a different rope for every variable dimension
Hello,

I have been looking at the forest_plot documentation and examples and it seems pretty clear to me that you can have a different rope for each variable. But say I have a different rope for each dimension of a variable say mu 0, mu 1... which all correspond to the same:
```python
mu_o = pm.Normal('mu_o', mu=data_fo.mean(), sd=sd_prior, dims='resid')
```
Thank you all very much and keep up the good work!
Sergio



Discussion:
There may be a bug in the `rope` code. There are examples of the usage you ask for in the documentation of `plot_posterior`, and I have been able to use it with:

    rope={"mu": [{"mu_dim": 2, "rope": (-1, 1)}, {"mu_dim": 4, "rope": (0, 3)}]}

![image](https://user-images.githubusercontent.com/23738400/93254624-eb1f2d00-f798-11ea-90e3-ad5af34c600b.png)


but it does not work for `plot_forest` cc @aloctavodia 
Following on the other issue #1381 . I have created this forestplot:
```python
az.plot_forest(
    (posterior.sel(state="o"), posterior.sel(state="fo")),
    var_names="mu",
    combined=True, 
    colors=["C0", "C1"], 
    model_names=["o", "fo"],
    textsize=9,
    #rope=rope
)
```
If I add a rope like the one you suggested:
```python
rope={"mu": [{"mu_dim": 'A42', "rope": (-1, 1)}, {"mu_dim": 'T33', "rope": (0, 3)}]}
```
I get this error:
```python
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
<ipython-input-124-7b15dd7be9eb> in <module>
----> 1 az.plot_forest(
      2     (posterior.sel(state="o"), posterior.sel(state="fo")),
      3     var_names="mu",
      4     combined=True,
      5     colors=["C0", "C1"],

~/data_partition/bin2/anaconda3/envs/nmr_assign_state/lib/python3.8/site-packages/arviz/plots/forestplot.py in plot_forest(data, kind, model_names, var_names, filter_vars, transform, coords, combined, hdi_prob, rope, quartiles, ess, r_hat, colors, textsize, linewidth, markersize, ridgeplot_alpha, ridgeplot_overlap, ridgeplot_kind, ridgeplot_quantiles, figsize, ax, backend, backend_config, backend_kwargs, show, credible_interval)
    226     # TODO: Add backend kwargs
    227     plot = get_plotting_function("plot_forest", "forestplot", backend)
--> 228     axes = plot(**plot_forest_kwargs)
    229     return axes

~/data_partition/bin2/anaconda3/envs/nmr_assign_state/lib/python3.8/site-packages/arviz/plots/backends/matplotlib/forestplot.py in plot_forest(ax, datasets, var_names, model_names, combined, colors, figsize, width_ratios, linewidth, markersize, kind, ncols, hdi_prob, quartiles, rope, ridgeplot_overlap, ridgeplot_alpha, ridgeplot_kind, ridgeplot_quantiles, textsize, ess, r_hat, backend_kwargs, show)
     91     axes = np.atleast_1d(axes)
     92     if kind == "forestplot":
---> 93         plot_handler.forestplot(
     94             hdi_prob, quartiles, xt_labelsize, titlesize, linewidth, markersize, axes[0], rope,
     95         )

~/data_partition/bin2/anaconda3/envs/nmr_assign_state/lib/python3.8/site-packages/arviz/plots/backends/matplotlib/forestplot.py in forestplot(self, hdi_prob, quartiles, xt_labelsize, titlesize, linewidth, markersize, ax, rope)
    325             for y, rope_var, values, color in plotter.treeplot(qlist, hdi_prob):
    326                 if isinstance(rope, dict):
--> 327                     self.display_multiple_ropes(rope, ax, y, linewidth, rope_var)
    328 
    329                 mid = len(values) // 2

~/data_partition/bin2/anaconda3/envs/nmr_assign_state/lib/python3.8/site-packages/arviz/plots/backends/matplotlib/forestplot.py in display_multiple_ropes(self, rope, ax, y, linewidth, rope_var)
    221     def display_multiple_ropes(self, rope, ax, y, linewidth, rope_var):
    222         """Display ROPE when more than one interval is provided."""
--> 223         vals = dict(rope[rope_var][0])["rope"]
    224         ax.plot(
    225             vals,

KeyError: 'mu T112'
```
Really appreciate your help
I am quite sure this is a bug. Given that plot_posterior accepts the rope dict, we should see where they differ and fix the forest plot version so it works and behaves like the one in `plot_posterior`. 

It may also be a good idea once this is done to add a link in the description of `rope` argument in forest plot to plot_posterior to help understand the rope scalar/list/dict possibilities
Yes, the documentation is a bit confusing. 

What would you like me to do? Can you reproduce this unexpected behaviour in a simple test-case?
I think the first step would be creating a minimal example to reproduce the bug. If my intuition is correct, using the same data and rope argument as in [`plot_posterior`](https://arviz-devs.github.io/arviz/generated/arviz.plot_posterior.html#arviz.plot_posterior) docs should serve this goal. This example in particular looks like the one that will show this issue:

```
rope = {'mu': [{'rope': (-2, 2)}], 'theta': [{'school': 'Choate', 'rope': (2, 4)}]}
az.plot_posterior(data, var_names=['mu', 'theta'], rope=rope)
```

The next step should probably be comparing the source codes of `plot_posterior` and `plot_forest`, I don't know why but instead of having the same behaviour when interpreting the rope argument, one works correctly and the other crashes. Hopefully, fixing the differences between the two will solve the issue as `plot_posterior` seems to work correctly.
I am affraid I don't have much knowledge to work on this... If there is anything I can help you with let me know. 

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
@@ -266,22 +266,25 @@ def label_idxs():
 
         return label_idxs()
 
-    def display_multiple_ropes(self, rope, ax, y, linewidth, rope_var):
+    def display_multiple_ropes(self, rope, ax, y, linewidth, var_name, selection):
         """Display ROPE when more than one interval is provided."""
-        vals = dict(rope[rope_var][0])["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"])
+        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,
                 )
-            ][2],
-            line_alpha=0.7,
-        )
-        return ax
+                return ax
 
     def ridgeplot(
         self,
@@ -452,9 +455,9 @@ def forestplot(self, hdi_prob, quartiles, linewidth, markersize, ax, rope):
             qlist = [endpoint, 50, 100 - endpoint]
 
         for plotter in self.plotters.values():
-            for y, rope_var, values, color in plotter.treeplot(qlist, hdi_prob):
+            for y, selection, values, color in plotter.treeplot(qlist, hdi_prob):
                 if isinstance(rope, dict):
-                    self.display_multiple_ropes(rope, ax, y, linewidth, rope_var)
+                    self.display_multiple_ropes(rope, ax, y, linewidth, plotter.var_name, selection)
 
                 mid = len(values) // 2
                 param_iter = zip(
@@ -560,6 +563,7 @@ def iterator(self):
             skip_dims = set()
 
         label_dict = OrderedDict()
+        selection_list = []
         for name, grouped_datum in zip(self.model_names, grouped_data):
             for _, sub_data in grouped_datum:
                 datum_iter = xarray_var_iter(
@@ -569,6 +573,7 @@ def iterator(self):
                     reverse_selections=True,
                 )
                 for _, selection, values in datum_iter:
+                    selection_list.append(selection)
                     label = make_label(self.var_name, selection, position="beside")
                     if label not in label_dict:
                         label_dict[label] = OrderedDict()
@@ -577,14 +582,16 @@ def iterator(self):
                     label_dict[label][name].append(values)
 
         y = self.y_start
-        for label, model_data in label_dict.items():
+        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
                 for values in value_list:
-                    yield y, row_label, label, values, self.model_color[model_name]
+                    yield y, row_label, label, selection_list[idx], values, self.model_color[
+                        model_name
+                    ]
                     y += self.chain_offset
                 y += self.var_offset
             y += self.group_offset
@@ -592,7 +599,7 @@ 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():
+        for y, label, _, _, vals, color in self.iterator():
             y_ticks[label].append((y, vals, color))
         labels, ticks, vals, colors = [], [], [], []
         for label, data in y_ticks.items():
@@ -604,10 +611,10 @@ def labels_ticks_and_vals(self):
 
     def treeplot(self, qlist, hdi_prob):
         """Get data for each treeplot for the variable."""
-        for y, _, label, values, color in self.iterator():
+        for y, _, _, 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, label, ntiles, color
+            yield y, selection, ntiles, color
 
     def ridgeplot(self, hdi_prob, mult, ridgeplot_kind):
         """Get data for each ridgeplot for the variable."""
diff --git a/arviz/plots/backends/matplotlib/forestplot.py b/arviz/plots/backends/matplotlib/forestplot.py
--- a/arviz/plots/backends/matplotlib/forestplot.py
+++ b/arviz/plots/backends/matplotlib/forestplot.py
@@ -230,19 +230,22 @@ def label_idxs():
 
         return label_idxs()
 
-    def display_multiple_ropes(self, rope, ax, y, linewidth, rope_var):
+    def display_multiple_ropes(self, rope, ax, y, linewidth, var_name, selection):
         """Display ROPE when more than one interval is provided."""
-        vals = dict(rope[rope_var][0])["rope"]
-        ax.plot(
-            vals,
-            (y + 0.05, y + 0.05),
-            lw=linewidth * 2,
-            color="C2",
-            solid_capstyle="round",
-            zorder=0,
-            alpha=0.7,
-        )
-        return ax
+        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.plot(
+                    vals,
+                    (y + 0.05, y + 0.05),
+                    lw=linewidth * 2,
+                    color="C2",
+                    solid_capstyle="round",
+                    zorder=0,
+                    alpha=0.7,
+                )
+                return ax
 
     def ridgeplot(
         self,
@@ -376,9 +379,9 @@ def forestplot(
             qlist = [endpoint, 50, 100 - endpoint]
 
         for plotter in self.plotters.values():
-            for y, rope_var, values, color in plotter.treeplot(qlist, hdi_prob):
+            for y, selection, values, color in plotter.treeplot(qlist, hdi_prob):
                 if isinstance(rope, dict):
-                    self.display_multiple_ropes(rope, ax, y, linewidth, rope_var)
+                    self.display_multiple_ropes(rope, ax, y, linewidth, plotter.var_name, selection)
 
                 mid = len(values) // 2
                 param_iter = zip(
@@ -502,6 +505,7 @@ def iterator(self):
             skip_dims = set()
 
         label_dict = OrderedDict()
+        selection_list = []
         for name, grouped_datum in zip(self.model_names, grouped_data):
             for _, sub_data in grouped_datum:
                 datum_iter = xarray_var_iter(
@@ -511,6 +515,7 @@ def iterator(self):
                     reverse_selections=True,
                 )
                 for _, selection, values in datum_iter:
+                    selection_list.append(selection)
                     label = make_label(self.var_name, selection, position="beside")
                     if label not in label_dict:
                         label_dict[label] = OrderedDict()
@@ -519,14 +524,16 @@ def iterator(self):
                     label_dict[label][name].append(values)
 
         y = self.y_start
-        for label, model_data in label_dict.items():
+        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
                 for values in value_list:
-                    yield y, row_label, label, values, self.model_color[model_name]
+                    yield y, row_label, label, selection_list[idx], values, self.model_color[
+                        model_name
+                    ]
                     y += self.chain_offset
                 y += self.var_offset
             y += self.group_offset
@@ -534,7 +541,7 @@ 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():
+        for y, label, _, _, vals, color in self.iterator():
             y_ticks[label].append((y, vals, color))
         labels, ticks, vals, colors = [], [], [], []
         for label, data in y_ticks.items():
@@ -546,10 +553,10 @@ def labels_ticks_and_vals(self):
 
     def treeplot(self, qlist, hdi_prob):
         """Get data for each treeplot for the variable."""
-        for y, _, label, values, color in self.iterator():
+        for y, _, _, 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, label, ntiles, color
+            yield y, selection, ntiles, color
 
     def ridgeplot(self, hdi_prob, mult, ridgeplot_kind):
         """Get data for each ridgeplot for the variable."""
diff --git a/arviz/plots/forestplot.py b/arviz/plots/forestplot.py
--- a/arviz/plots/forestplot.py
+++ b/arviz/plots/forestplot.py
@@ -136,10 +136,23 @@ def plot_forest(
         >>>                            var_names=["^the"],
         >>>                            filter_vars="regex",
         >>>                            combined=True,
-        >>>                            ridgeplot_overlap=3,
         >>>                            figsize=(9, 7))
         >>> axes[0].set_title('Estimated theta for 8 schools model')
 
+    Forestpĺot with ropes
+
+    .. plot::
+        :context: close-figs
+
+        >>> rope = {'theta': [{'school': 'Choate', 'rope': (2, 4)}], 'mu': [{'rope': (-2, 2)}]}
+        >>> axes = az.plot_forest(non_centered_data,
+        >>>                            rope=rope,
+        >>>                            var_names='~tau',
+        >>>                            combined=True,
+        >>>                            figsize=(9, 7))
+        >>> axes[0].set_title('Estimated theta for 8 schools model')
+
+
     Ridgeplot
 
     .. plot::


Test patch:
diff --git a/arviz/tests/base_tests/test_plots_bokeh.py b/arviz/tests/base_tests/test_plots_bokeh.py
--- a/arviz/tests/base_tests/test_plots_bokeh.py
+++ b/arviz/tests/base_tests/test_plots_bokeh.py
@@ -490,7 +490,10 @@ def test_plot_ess_no_divergences(models):
         (
             {
                 "var_names": ["mu", "tau"],
-                "rope": {"mu": [{"rope": (-0.1, 0.1)}], "tau": [{"rope": (0.2, 0.5)}]},
+                "rope": {
+                    "mu": [{"rope": (-0.1, 0.1)}],
+                    "theta": [{"school": "Choate", "rope": (0.2, 0.5)}],
+                },
             },
             1,
         ),
diff --git a/arviz/tests/base_tests/test_plots_matplotlib.py b/arviz/tests/base_tests/test_plots_matplotlib.py
--- a/arviz/tests/base_tests/test_plots_matplotlib.py
+++ b/arviz/tests/base_tests/test_plots_matplotlib.py
@@ -279,8 +279,11 @@ def test_plot_trace_futurewarning(models, prop):
         ({"kind": "ridgeplot", "r_hat": True, "ess": True, "ridgeplot_alpha": 0}, 3),
         (
             {
-                "var_names": ["mu", "tau"],
-                "rope": {"mu": [{"rope": (-0.1, 0.1)}], "tau": [{"rope": (0.2, 0.5)}]},
+                "var_names": ["mu", "theta"],
+                "rope": {
+                    "mu": [{"rope": (-0.1, 0.1)}],
+                    "theta": [{"school": "Choate", "rope": (0.2, 0.5)}],
+                },
             },
             1,
         ),