Added marker and line options for forestplot (#2220)

* Added marker and line options for forestplot

* Added missing argument to plot

* Replaced explicit plot options with plot_kwargs

* Pushed plot_kwargs into _plot_tree

Author: fonnesbeck

pymc3/plots/forestplot.py

@@ -80,7 +80,7 @@ def _make_rhat_plot(trace, ax, title, labels, varnames, include_transformed):
     return ax
 
 
-def _plot_tree(ax, y, ntiles, show_quartiles):
+def _plot_tree(ax, y, ntiles, show_quartiles, **plot_kwargs):
     """Helper to plot errorbars for the forestplot.
 
     Parameters
@@ -101,22 +101,32 @@ def _plot_tree(ax, y, ntiles, show_quartiles):
     """
     if show_quartiles:
         # Plot median
-        ax.plot(ntiles[2], y, 'bo', markersize=4)
+        ax.plot(ntiles[2], y, color=plot_kwargs.get('color', 'blue'), 
+                        marker=plot_kwargs.get('marker', 'o'), 
+                        markersize=plot_kwargs.get('markersize', 4))
         # Plot quartile interval
-        ax.errorbar(x=(ntiles[1], ntiles[3]), y=(y, y), linewidth=2, color='b')
+        ax.errorbar(x=(ntiles[1], ntiles[3]), y=(y, y), 
+                        linewidth=plot_kwargs.get('linewidth', 2), 
+                        color=plot_kwargs.get('color', 'blue'))
 
     else:
         # Plot median
-        ax.plot(ntiles[1], y, 'bo', markersize=4)
+        ax.plot(ntiles[1], y, marker=plot_kwargs.get('marker', 'o'), 
+                            color=plot_kwargs.get('color', 'blue'),
+                            markersize=plot_kwargs.get('markersize', 4))
 
     # Plot outer interval
-    ax.errorbar(x=(ntiles[0], ntiles[-1]), y=(y, y), linewidth=1, color='b')
+    ax.errorbar(x=(ntiles[0], ntiles[-1]), y=(y, y), 
+                linewidth=int(plot_kwargs.get('linewidth', 2)/2), 
+                color=plot_kwargs.get('color', 'blue'))
+                
     return ax
 
 
 def forestplot(trace_obj, varnames=None, transform=identity_transform, alpha=0.05, quartiles=True,
                rhat=True, main=None, xtitle=None, xlim=None, ylabels=None,
-               chain_spacing=0.05, vline=0, gs=None, plot_transformed=False):
+               chain_spacing=0.05, vline=0, gs=None, plot_transformed=False,
+               **plot_kwargs):
     """
     Forest plot (model summary plot).
 
@@ -160,6 +170,9 @@ def forestplot(trace_obj, varnames=None, transform=identity_transform, alpha=0.0
     plot_transformed : bool
         Flag for plotting automatically transformed variables in addition to
         original variables (defaults to False).
+    plot_kwargs : dict
+        Optional arguments for plot elements. Currently accepts 'fontsize',
+        'linewidth', 'color', 'marker', and 'markersize'.
 
     Returns
     -------
@@ -242,10 +255,12 @@ def forestplot(trace_obj, varnames=None, transform=identity_transform, alpha=0.0
             if k > 1:
                 for q in np.transpose(quants).squeeze():
                     # Multiple y values
-                    interval_plot = _plot_tree(interval_plot, y, q, quartiles)
+                    interval_plot = _plot_tree(interval_plot, y, q, quartiles,
+                                    **plot_kwargs)
                     y -= 1
             else:
-                interval_plot = _plot_tree(interval_plot, y, quants, quartiles)
+                interval_plot = _plot_tree(interval_plot, y, quants, quartiles,
+                                **plot_kwargs)
 
             # Increment index
             var += k
@@ -264,7 +279,7 @@ def forestplot(trace_obj, varnames=None, transform=identity_transform, alpha=0.0
 
     # Add variable labels
     interval_plot.set_yticks([-l for l in range(len(labels))])
-    interval_plot.set_yticklabels(labels)
+    interval_plot.set_yticklabels(labels, fontsize=plot_kwargs.get('fontsize', None))
 
     # Add title
     plot_title = ""
@@ -273,7 +288,7 @@ def forestplot(trace_obj, varnames=None, transform=identity_transform, alpha=0.0
     elif main:
         plot_title = main
     if plot_title:
-        interval_plot.set_title(plot_title)
+        interval_plot.set_title(plot_title, fontsize=plot_kwargs.get('fontsize', None))
 
     # Add x-axis label
     if xtitle is not None:
@@ -293,7 +308,7 @@ def forestplot(trace_obj, varnames=None, transform=identity_transform, alpha=0.0
             spine.set_color('none')  # don't draw spine
 
     # Reference line
-    interval_plot.axvline(vline, color='k', linestyle='--')
+    interval_plot.axvline(vline, color='k', linestyle=':')
 
     # Genenerate Gelman-Rubin plot
     if plot_rhat: