Skip to content

Added marker and line options for forestplot #2220

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
May 26, 2017
Merged

Added marker and line options for forestplot #2220

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
6a1184f
Added marker and line options for forestplot
fonnesbeck May 25, 2017
e53d44e
Added missing argument to plot
fonnesbeck May 25, 2017
86a7538
Replaced explicit plot options with plot_kwargs
fonnesbeck May 25, 2017
3c7e5b8
Pushed plot_kwargs into _plot_tree
fonnesbeck May 25, 2017
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 26 additions & 11 deletions pymc3/plots/forestplot.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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).

Expand Down Expand Up @@ -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
-------
Expand Down Expand Up @@ -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
Expand All @@ -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 = ""
Expand All @@ -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:
Expand All @@ -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:
Expand Down