plot_rank: add ref_line, bar, vlines and marker_vlines kwargs

CHANGELOG.md

@@ -5,6 +5,7 @@
 * Added `to_dataframe` method to InferenceData ([1395](https://github.com/arviz-devs/arviz/pull/1395))
 * Added `__getitem__` magic to InferenceData ([1395](https://github.com/arviz-devs/arviz/pull/1395))
 * Added group argument to summary ([1408](https://github.com/arviz-devs/arviz/pull/1408))
+* Add `ref_line`, `bar`, `vlines` and `marker_vlines` kwargs to `plot_rank` ([1419](https://github.com/arviz-devs/arviz/pull/1419))
 * Add observed argument to (un)plot observed data in `plot_ppc` ([1422](https://github.com/arviz-devs/arviz/pull/1422))
 
 ### Maintenance and fixes

arviz/plots/backends/bokeh/rankplot.py

@@ -23,10 +23,31 @@ def plot_rank(
     colors,
     ref_line,
     labels,
+    ref_line_kwargs,
+    bar_kwargs,
+    vlines_kwargs,
+    marker_vlines_kwargs,
     backend_kwargs,
     show,
 ):
     """Bokeh rank plot."""
+    if ref_line_kwargs is None:
+        ref_line_kwargs = {}
+    ref_line_kwargs.setdefault("line_dash", "dashed")
+    ref_line_kwargs.setdefault("line_color", "black")
+
+    if bar_kwargs is None:
+        bar_kwargs = {}
+    bar_kwargs.setdefault("line_color", "white")
+
+    if vlines_kwargs is None:
+        vlines_kwargs = {}
+    vlines_kwargs.setdefault("line_width", 2)
+    vlines_kwargs.setdefault("line_dash", "solid")
+
+    if marker_vlines_kwargs is None:
+        marker_vlines_kwargs = {}
+
     if backend_kwargs is None:
         backend_kwargs = {}
 
@@ -62,6 +83,7 @@ def plot_rank(
         gap = 1
         width = bin_ary[1] - bin_ary[0]
 
+        bar_kwargs.setdefault("width", width)
         # Center the bins
         bin_ary = (bin_ary[1:] + bin_ary[:-1]) / 2
 
@@ -74,34 +96,30 @@ def plot_rank(
                     x=bin_ary,
                     top=y_ticks[-1] + counts,
                     bottom=y_ticks[-1],
-                    width=width,
                     fill_color=colors[idx],
-                    line_color="white",
+                    **bar_kwargs,
                 )
                 if ref_line:
-                    hline = Span(
-                        location=y_ticks[-1] + counts.mean(), line_dash="dashed", line_color="black"
-                    )
+                    hline = Span(location=y_ticks[-1] + counts.mean(), **ref_line_kwargs)
                     ax.add_layout(hline)
             if labels:
                 ax.yaxis.axis_label = "Chain"
         elif kind == "vlines":
             ymin = np.full(len(all_counts), all_counts.mean())
             for idx, counts in enumerate(all_counts):
-                ax.circle(bin_ary, counts, fill_color=colors[idx], line_color=colors[idx])
-
-                x_locations = [(bin, bin) for bin in bin_ary]
-                y_locations = [(ymin[idx], counts_) for counts_ in counts]
-                ax.multi_line(
-                    x_locations,
-                    y_locations,
-                    line_dash="solid",
+                ax.circle(
+                    bin_ary,
+                    counts,
+                    fill_color=colors[idx],
                     line_color=colors[idx],
-                    line_width=3,
+                    **marker_vlines_kwargs,
                 )
+                x_locations = [(bin, bin) for bin in bin_ary]
+                y_locations = [(ymin[idx], counts_) for counts_ in counts]
+                ax.multi_line(x_locations, y_locations, line_color=colors[idx], **vlines_kwargs)
 
             if ref_line:
-                hline = Span(location=all_counts.mean(), line_dash="dashed", line_color="black")
+                hline = Span(location=all_counts.mean(), **ref_line_kwargs)
                 ax.add_layout(hline)
 
         if labels:

arviz/plots/backends/matplotlib/rankplot.py

@@ -20,10 +20,32 @@ def plot_rank(
     colors,
     ref_line,
     labels,
+    ref_line_kwargs,
+    bar_kwargs,
+    vlines_kwargs,
+    marker_vlines_kwargs,
     backend_kwargs,
     show,
 ):
     """Matplotlib rankplot.."""
+    if ref_line_kwargs is None:
+        ref_line_kwargs = {}
+    ref_line_kwargs.setdefault("linestyle", "--")
+    ref_line_kwargs.setdefault("color", "k")
+
+    if bar_kwargs is None:
+        bar_kwargs = {}
+    bar_kwargs.setdefault("align", "center")
+
+    if vlines_kwargs is None:
+        vlines_kwargs = {}
+    vlines_kwargs.setdefault("lw", 2)
+
+    if marker_vlines_kwargs is None:
+        marker_vlines_kwargs = {}
+        marker_vlines_kwargs.setdefault("marker", "o")
+        marker_vlines_kwargs.setdefault("lw", 0)
+
     if backend_kwargs is None:
         backend_kwargs = {}
 
@@ -52,6 +74,8 @@ def plot_rank(
         gap = all_counts.max() * 1.05
         width = bin_ary[1] - bin_ary[0]
 
+        bar_kwargs.setdefault("width", width)
+        bar_kwargs.setdefault("edgecolor", ax.get_facecolor())
         # Center the bins
         bin_ary = (bin_ary[1:] + bin_ary[:-1]) / 2
 
@@ -63,23 +87,22 @@ def plot_rank(
                     bin_ary,
                     counts,
                     bottom=y_ticks[-1],
-                    width=width,
-                    align="center",
                     color=colors[idx],
-                    edgecolor=ax.get_facecolor(),
+                    **bar_kwargs,
                 )
                 if ref_line:
-                    ax.axhline(y=y_ticks[-1] + counts.mean(), linestyle="--", color="k")
+                    ax.axhline(y=y_ticks[-1] + counts.mean(), **ref_line_kwargs)
             if labels:
                 ax.set_ylabel("Chain", fontsize=ax_labelsize)
         elif kind == "vlines":
             ymin = all_counts.mean()
+
             for idx, counts in enumerate(all_counts):
-                ax.plot(bin_ary, counts, "o", color=colors[idx])
-                ax.vlines(bin_ary, ymin, counts, lw=2, colors=colors[idx])
+                ax.plot(bin_ary, counts, color=colors[idx], **marker_vlines_kwargs)
+                ax.vlines(bin_ary, ymin, counts, colors=colors[idx], **vlines_kwargs)
             ax.set_ylim(0, all_counts.mean() * 2)
             if ref_line:
-                ax.axhline(y=all_counts.mean(), linestyle="--", color="k")
+                ax.axhline(y=ymin, **ref_line_kwargs)
 
         if labels:
             ax.set_xlabel("Rank (all chains)", fontsize=ax_labelsize)

arviz/plots/rankplot.py

@@ -24,6 +24,10 @@ def plot_rank(
     figsize=None,
     ax=None,
     backend=None,
+    ref_line_kwargs=None,
+    bar_kwargs=None,
+    vlines_kwargs=None,
+    marker_vlines_kwargs=None,
     backend_kwargs=None,
     show=None,
 ):
@@ -80,6 +84,18 @@ def plot_rank(
         its own array of plot areas (and return it).
     backend: str, optional
         Select plotting backend {"matplotlib","bokeh"}. Default "matplotlib".
+    ref_line_kwargs : dict, optional
+        Reference line keyword arguments, passed to :meth:`mpl:matplotlib.axes.Axes.axhline` or
+        :meth:`bokeh:bokeh.model.Span`.
+    bar_kwargs : dict, optional
+        Bars keyword arguments, passed to :meth:`mpl:matplotlib.axes.Axes.bar` or
+        :meth:`bokeh:bokeh.plotting.figure.Figure.vbar`.
+    vlines_kwargs : dict, optional
+        Vlines keyword arguments, passed to :meth:`mpl:matplotlib.axes.Axes.vlines` or
+        :meth:`bokeh:bokeh.plotting.figure.Figure.multi_line`.
+    marker_vlines_kwargs : dict, optional
+        Marker for the vlines keyword arguments, passed to :meth:`mpl:matplotlib.axes.Axes.plot` or
+        :meth:`bokeh:bokeh.plotting.figure.Figure.circle`.
     backend_kwargs: bool, optional
         These are kwargs specific to the backend being used. For additional documentation
         check the plotting method of the backend.
@@ -121,6 +137,13 @@ def plot_rank(
         >>> az.plot_rank(centered_data, var_names="mu", kind='vlines', ax=ax[0])
         >>> az.plot_rank(noncentered_data, var_names="mu", kind='vlines', ax=ax[1])
 
+    Change the aesthetics using kwargs
+
+    .. plot::
+        :context: close-figs
+
+        >>> az.plot_rank(noncentered_data, var_names="mu", kind="vlines",
+        >>>              vlines_kwargs={'lw':0}, marker_vlines_kwargs={'lw':3});
     """
     if transform is not None:
         data = transform(data)
@@ -161,6 +184,10 @@ def plot_rank(
         colors=colors,
         ref_line=ref_line,
         labels=labels,
+        ref_line_kwargs=ref_line_kwargs,
+        bar_kwargs=bar_kwargs,
+        vlines_kwargs=vlines_kwargs,
+        marker_vlines_kwargs=marker_vlines_kwargs,
         backend_kwargs=backend_kwargs,
         show=show,
     )

arviz/tests/base_tests/test_plots_bokeh.py

@@ -1030,8 +1030,19 @@ def test_plot_posterior_point_estimates(models, point_estimate):
         {"var_names": "mu"},
         {"var_names": ("mu", "tau"), "coords": {"theta_dim_0": [0, 1]}},
         {"var_names": "mu", "ref_line": True},
+        {
+            "var_names": "mu",
+            "ref_line_kwargs": {"line_width": 2, "line_color": "red"},
+            "bar_kwargs": {"width": 50},
+        },
         {"var_names": "mu", "ref_line": False},
         {"var_names": "mu", "kind": "vlines"},
+        {
+            "var_names": "mu",
+            "kind": "vlines",
+            "vlines_kwargs": {"line_width": 0},
+            "marker_vlines_kwargs": {"radius": 20},
+        },
     ],
 )
 def test_plot_rank(models, kwargs):

arviz/tests/base_tests/test_plots_matplotlib.py

@@ -848,8 +848,19 @@ def test_plot_autocorr_var_names(models, var_names):
         {"var_names": "mu"},
         {"var_names": ("mu", "tau"), "coords": {"theta_dim_0": [0, 1]}},
         {"var_names": "mu", "ref_line": True},
+        {
+            "var_names": "mu",
+            "ref_line_kwargs": {"lw": 2, "color": "C2"},
+            "bar_kwargs": {"width": 0.7},
+        },
         {"var_names": "mu", "ref_line": False},
         {"var_names": "mu", "kind": "vlines"},
+        {
+            "var_names": "mu",
+            "kind": "vlines",
+            "vlines_kwargs": {"lw": 0},
+            "marker_vlines_kwargs": {"lw": 3},
+        },
     ],
 )
 def test_plot_rank(models, kwargs):