Make plotting interfaces consistent (return figs)

This commit is contained in:
Ben Letham 2017-04-17 12:45:44 -07:00
parent 813966fdd7
commit 26ca2f7af7

View file

@ -852,13 +852,14 @@ class Prophet(object):
return pd.DataFrame({'ds': dates}) return pd.DataFrame({'ds': dates})
def plot(self, fcst, uncertainty=True, plot_cap=True, xlabel='ds', def plot(self, fcst, ax=None, uncertainty=True, plot_cap=True, xlabel='ds',
ylabel='y'): ylabel='y'):
"""Plot the Prophet forecast. """Plot the Prophet forecast.
Parameters Parameters
---------- ----------
fcst: pd.DataFrame output of self.predict. fcst: pd.DataFrame output of self.predict.
ax: Optional matplotlib axes on which to plot.
uncertainty: Optional boolean to plot uncertainty intervals. uncertainty: Optional boolean to plot uncertainty intervals.
plot_cap: Optional boolean indicating if the capacity should be shown plot_cap: Optional boolean indicating if the capacity should be shown
in the figure, if available. in the figure, if available.
@ -867,10 +868,13 @@ class Prophet(object):
Returns Returns
------- -------
a matplotlib figure. A matplotlib figure.
""" """
if ax is None:
fig = plt.figure(facecolor='w', figsize=(10, 6)) fig = plt.figure(facecolor='w', figsize=(10, 6))
ax = fig.add_subplot(111) ax = fig.add_subplot(111)
else:
fig = ax.get_figure()
ax.plot(self.history['ds'].values, self.history['y'], 'k.') ax.plot(self.history['ds'].values, self.history['y'], 'k.')
ax.plot(fcst['ds'].values, fcst['yhat'], ls='-', c='#0072B2') ax.plot(fcst['ds'].values, fcst['yhat'], ls='-', c='#0072B2')
if 'cap' in fcst and plot_cap: if 'cap' in fcst and plot_cap:
@ -892,13 +896,6 @@ class Prophet(object):
Will plot whichever are available of: trend, holidays, weekly Will plot whichever are available of: trend, holidays, weekly
seasonality, and yearly seasonality. seasonality, and yearly seasonality.
This method returns a list of matplotlib artists. To show the plot,
>>> m.plot_components(fcst)
>>> from matplotlib import pyplot as plt
>>> plt.show()
To save the figure, replace plt.show with plt.savefig.
Parameters Parameters
---------- ----------
fcst: pd.DataFrame output of self.predict. fcst: pd.DataFrame output of self.predict.
@ -914,7 +911,7 @@ class Prophet(object):
Returns Returns
------- -------
A list of matplotlib artists. A matplotlib figure.
""" """
# Identify components to be plotted # Identify components to be plotted
components = [('trend', True), components = [('trend', True),
@ -927,23 +924,21 @@ class Prophet(object):
fig, axes = plt.subplots(npanel, 1, facecolor='w', fig, axes = plt.subplots(npanel, 1, facecolor='w',
figsize=(9, 3 * npanel)) figsize=(9, 3 * npanel))
artists = []
for ax, plot in zip(axes, components): for ax, plot in zip(axes, components):
if plot == 'trend': if plot == 'trend':
artists += self.plot_trend( self.plot_trend(
fcst, ax=ax, uncertainty=uncertainty, plot_cap=plot_cap) fcst, ax=ax, uncertainty=uncertainty, plot_cap=plot_cap)
elif plot == 'holidays': elif plot == 'holidays':
artists += self.plot_holidays(fcst, ax=ax, self.plot_holidays(fcst, ax=ax, uncertainty=uncertainty)
uncertainty=uncertainty)
elif plot == 'weekly': elif plot == 'weekly':
artists += self.plot_weekly(ax=ax, uncertainty=uncertainty, self.plot_weekly(
weekly_start=weekly_start) ax=ax, uncertainty=uncertainty, weekly_start=weekly_start)
elif plot == 'yearly': elif plot == 'yearly':
artists += self.plot_yearly(ax=ax, uncertainty=uncertainty, self.plot_yearly(
yearly_start=yearly_start) ax=ax, uncertainty=uncertainty, yearly_start=yearly_start)
fig.tight_layout() fig.tight_layout()
return artists return fig
def plot_trend(self, fcst, ax=None, uncertainty=True, plot_cap=True): def plot_trend(self, fcst, ax=None, uncertainty=True, plot_cap=True):
"""Plot the trend component of the forecast. """Plot the trend component of the forecast.