From e081db52e13189e0ff35e77028f9b84e427b2186 Mon Sep 17 00:00:00 2001 From: Isaac Laughlin Date: Fri, 3 Mar 2017 11:29:12 -0800 Subject: [PATCH] Refactoring mpl code to address #62, #63 --- python/fbprophet/forecaster.py | 164 ++++++++++++++++++--------------- 1 file changed, 91 insertions(+), 73 deletions(-) diff --git a/python/fbprophet/forecaster.py b/python/fbprophet/forecaster.py index d940bd7..ac43274 100644 --- a/python/fbprophet/forecaster.py +++ b/python/fbprophet/forecaster.py @@ -34,6 +34,8 @@ except ImportError: # fb-block 2 class Prophet(object): + # Plotting default color for R/Python consistency + forecast_color = '#0072B2' def __init__( self, growth='linear', @@ -616,16 +618,15 @@ class Prophet(object): ------- a matplotlib figure. """ - forecast_color = '#0072B2' fig = plt.figure(facecolor='w', figsize=(10, 6)) ax = fig.add_subplot(111) ax.plot(self.history['ds'].values, self.history['y'], 'k.') - ax.plot(fcst['ds'].values, fcst['yhat'], ls='-', c=forecast_color) + ax.plot(fcst['ds'].values, fcst['yhat'], ls='-', c=self.forecast_color) if 'cap' in fcst: ax.plot(fcst['ds'].values, fcst['cap'], ls='--', c='k') if uncertainty: ax.fill_between(fcst['ds'].values, fcst['yhat_lower'], - fcst['yhat_upper'], color=forecast_color, alpha=0.2) + fcst['yhat_upper'], color=self.forecast_color, alpha=0.2) ax.grid(True, which='major', c='gray', ls='-', lw=1, alpha=0.2) ax.set_xlabel(xlabel) ax.set_ylabel(ylabel) @@ -648,87 +649,104 @@ class Prophet(object): a matplotlib figure. """ # Identify components to be plotted - plot_trend = True - plot_holidays = self.holidays is not None - plot_weekly = 'weekly' in fcst - plot_yearly = 'yearly' in fcst + components = [('plot_trend', True), + ('plot_holidays', self.holidays is not None), + ('plot_weekly', 'weekly' in fcst), + ('plot_yearly', 'yearly' in fcst)] + components = [(plot, cond) for plot, cond in components if cond] + npanel = len(components) + + fig, axes = plt.subplots(npanel, 1, facecolor='w', figsize=(9, 3 * npanel)) - npanel = plot_trend + plot_holidays + plot_weekly + plot_yearly - forecast_color = '#0072B2' - fig = plt.figure(facecolor='w', figsize=(9, 3 * npanel)) - panel_num = 1 - ax = fig.add_subplot(npanel, 1, panel_num) - ax.plot(fcst['ds'].values, fcst['trend'], ls='-', c=forecast_color) + artists = [] + for ax, plot in zip(axes, [getattr(self, plot) for plot, _ in components]): + artists += plot(fcst, ax=ax, uncertainty=uncertainty) + + fig.tight_layout() + return artists + + def plot_trend(self, fcst, ax=None, uncertainty=True, **plotargs): + artists = [] + if not ax: + ax = fig.add_subplot(111) + artists += ax.plot(fcst['ds'].values, fcst['trend'], ls='-', c=self.forecast_color) if 'cap' in fcst: - ax.plot(fcst['ds'].values, fcst['cap'], ls='--', c='k') + artists += ax.plot(fcst['ds'].values, fcst['cap'], ls='--', c='k') if uncertainty: - ax.fill_between( + artists += [ax.fill_between( fcst['ds'].values, fcst['trend_lower'], fcst['trend_upper'], - color=forecast_color, alpha=0.2) + color=self.forecast_color, alpha=0.2)] ax.grid(True, which='major', c='gray', ls='-', lw=1, alpha=0.2) ax.xaxis.set_major_locator(MaxNLocator(nbins=7)) ax.set_xlabel('ds') ax.set_ylabel('trend') + return artists - if plot_holidays: - panel_num += 1 + + def plot_holidays(self, fcst, ax=None, uncertainty=True): + artists = [] + if not ax: + ax = fig.add_subplot(111) + holiday_comps = self.holidays['holiday'].unique() + y_holiday = fcst[holiday_comps].sum(1) + y_holiday_l = fcst[[h + '_lower' for h in holiday_comps]].sum(1) + y_holiday_u = fcst[[h + '_upper' for h in holiday_comps]].sum(1) + # NOTE the above CI calculation is incorrect if holidays overlap + # in time. Since it is just for the visualization we will not + # worry about it now. + artists += ax.plot(fcst['ds'].values, y_holiday, ls='-', c=self.forecast_color) + if uncertainty: + artists += [ax.fill_between(fcst['ds'].values, y_holiday_l, y_holiday_u, + color=self.forecast_color, alpha=0.2)] + ax.grid(True, which='major', c='gray', ls='-', lw=1, alpha=0.2) + ax.xaxis.set_major_locator(MaxNLocator(nbins=7)) + ax.set_xlabel('ds') + ax.set_ylabel('holidays') + return artists + + def plot_weekly(self, fcst, ax=None, uncertainty=True): + artists = [] + if not ax: + ax = fig.add_subplot(111) + df_s = fcst.copy() + df_s['dow'] = df_s['ds'].dt.weekday_name + df_s = df_s.groupby('dow').first() + days = pd.date_range(start='2017-01-01', periods=7).weekday_name + y_weekly = [df_s.loc[d]['weekly'] for d in days] + y_weekly_l = [df_s.loc[d]['weekly_lower'] for d in days] + y_weekly_u = [df_s.loc[d]['weekly_upper'] for d in days] + artists += ax.plot(range(len(days)), y_weekly, ls='-', c=self.forecast_color) + if uncertainty: + artists += [ax.fill_between(range(len(days)), y_weekly_l, y_weekly_u, + color=self.forecast_color, alpha=0.2)] + ax.grid(True, which='major', c='gray', ls='-', lw=1, alpha=0.2) + ax.set_xticks(range(len(days))) + ax.set_xticklabels(days) + ax.set_xlabel('Day of week') + ax.set_ylabel('weekly') + return artists + + def plot_yearly(self, fcst, ax=None, uncertainty=True): + artists = [] + if not ax: ax = fig.add_subplot(npanel, 1, panel_num) - holiday_comps = self.holidays['holiday'].unique() - y_holiday = fcst[holiday_comps].sum(1) - y_holiday_l = fcst[[h + '_lower' for h in holiday_comps]].sum(1) - y_holiday_u = fcst[[h + '_upper' for h in holiday_comps]].sum(1) - # NOTE the above CI calculation is incorrect if holidays overlap - # in time. Since it is just for the visualization we will not - # worry about it now. - ax.plot(fcst['ds'].values, y_holiday, ls='-', c=forecast_color) - if uncertainty: - ax.fill_between(fcst['ds'].values, y_holiday_l, y_holiday_u, - color=forecast_color, alpha=0.2) - ax.grid(True, which='major', c='gray', ls='-', lw=1, alpha=0.2) - ax.xaxis.set_major_locator(MaxNLocator(nbins=7)) - ax.set_xlabel('ds') - ax.set_ylabel('holidays') + df_s = fcst.copy() + df_s['doy'] = df_s['ds'].map(lambda x: x.strftime('2000-%m-%d')) + df_s = df_s.groupby('doy').first().sort_index() + artists += ax.plot(pd.to_datetime(df_s.index), df_s['yearly'], ls='-', + c=self.forecast_color) + if uncertainty: + artists += [ax.fill_between( + pd.to_datetime(df_s.index), df_s['yearly_lower'], + df_s['yearly_upper'], color=self.forecast_color, alpha=0.2)] + ax.grid(True, which='major', c='gray', ls='-', lw=1, alpha=0.2) + months = MonthLocator(range(1, 13), bymonthday=1, interval=2) + ax.xaxis.set_major_formatter(DateFormatter('%B %-d')) + ax.xaxis.set_major_locator(months) + ax.set_xlabel('Day of year') + ax.set_ylabel('yearly') + return artists - if plot_weekly: - panel_num += 1 - ax = fig.add_subplot(npanel, 1, panel_num) - df_s = fcst.copy() - df_s['dow'] = df_s['ds'].dt.weekday_name - df_s = df_s.groupby('dow').first() - days = pd.date_range(start='2017-01-01', periods=7).weekday_name - y_weekly = [df_s.loc[d]['weekly'] for d in days] - y_weekly_l = [df_s.loc[d]['weekly_lower'] for d in days] - y_weekly_u = [df_s.loc[d]['weekly_upper'] for d in days] - ax.plot(range(len(days)), y_weekly, ls='-', c=forecast_color) - if uncertainty: - ax.fill_between(range(len(days)), y_weekly_l, y_weekly_u, - color=forecast_color, alpha=0.2) - ax.grid(True, which='major', c='gray', ls='-', lw=1, alpha=0.2) - ax.set_xticks(range(len(days))) - ax.set_xticklabels(days) - ax.set_xlabel('Day of week') - ax.set_ylabel('weekly') - if plot_yearly: - panel_num += 1 - ax = fig.add_subplot(npanel, 1, panel_num) - df_s = fcst.copy() - df_s['doy'] = df_s['ds'].map(lambda x: x.strftime('2000-%m-%d')) - df_s = df_s.groupby('doy').first().sort_index() - ax.plot(pd.to_datetime(df_s.index), df_s['yearly'], ls='-', - c=forecast_color) - if uncertainty: - ax.fill_between( - pd.to_datetime(df_s.index), df_s['yearly_lower'], - df_s['yearly_upper'], color=forecast_color, alpha=0.2) - ax.grid(True, which='major', c='gray', ls='-', lw=1, alpha=0.2) - months = MonthLocator(range(1, 13), bymonthday=1, interval=2) - ax.xaxis.set_major_formatter(DateFormatter('%B %-d')) - ax.xaxis.set_major_locator(months) - ax.set_xlabel('Day of year') - ax.set_ylabel('yearly') - - fig.tight_layout() - return fig # fb-block 9