Skip to content

Commit

Permalink
add option to plot balance inline in jupyter notebooks
Browse files Browse the repository at this point in the history
  • Loading branch information
kellieotto committed May 24, 2018
1 parent ce5ab30 commit 1512a0b
Showing 1 changed file with 7 additions and 2 deletions.
9 changes: 7 additions & 2 deletions pscore_match/match.py
Original file line number Diff line number Diff line change
Expand Up @@ -213,7 +213,7 @@ def _match_info(self):
self.matches['dropped'] = np.setdiff1d(list(range(self.nobs)),
np.append(self.matches['treated'], self.matches['control']))

def plot_balance(self, covariates, test=['t', 'rank'], filename='balance-plot',
def plot_balance(self, covariates, test=['t', 'rank'], notebook=False, filename='balance-plot',
**kwargs):
"""
Plot the p-values for covariate balance before and after matching
Expand All @@ -227,6 +227,8 @@ def plot_balance(self, covariates, test=['t', 'rank'], filename='balance-plot',
test : array-like or str
Statistical test to compare treatment and control covariate distributions.
Options are 't' for a two sample t-test or 'rank' for Wilcoxon rank sum test
notebook : bool
Optional, plot inline in a Jupyter notebook? Default False
filename : str
Optional, name of file to save plot in. Default 'balance-plot'
kwargs : dict
Expand Down Expand Up @@ -356,7 +358,10 @@ def plot_balance(self, covariates, test=['t', 'rank'], filename='balance-plot',
hovermode='closest'
)
fig = go.Figure(data=data, layout=layout)
plotly.offline.plot(fig, filename=filename, show_link=False, **kwargs)
if notebook:
plotly.offline.iplot(fig, filename=filename, show_link=False, **kwargs)
else:
plotly.offline.plot(fig, filename=filename, show_link=False, **kwargs)


################################################################################
Expand Down

0 comments on commit 1512a0b

Please sign in to comment.