Merge pull request #52 from Nixtla/feat/plots

[FEAT] Add HierarchicalPlot class

environment.yml

@@ -11,3 +11,4 @@ dependencies:
   - pip
   - pip:
     - nbdev
+    - mycolorpy

hierarchicalforecast/_modidx.py

@@ -27,7 +27,7 @@
                 'nbs_path': 'nbs',
                 'readme_nb': 'index.ipynb',
                 'recursive': 'False',
-                'requirements': 'numpy numba pandas scikit-learn statsmodels',
+                'requirements': 'numpy numba pandas scikit-learn statsmodels mycolorpy',
                 'status': '2',
                 'title': 'hierarchicalforecast',
                 'tst_flags': '',
@@ -49,4 +49,8 @@
                                               'hierarchicalforecast.methods.OptimalCombination.reconcile': 'https://Nixtla.github.io/hierarchicalforecast/methods.html#optimalcombination.reconcile',
                                               'hierarchicalforecast.methods.TopDown': 'https://Nixtla.github.io/hierarchicalforecast/methods.html#topdown',
                                               'hierarchicalforecast.methods.TopDown.reconcile': 'https://Nixtla.github.io/hierarchicalforecast/methods.html#topdown.reconcile'},
-            'hierarchicalforecast.utils': {'hierarchicalforecast.utils.aggregate': 'https://Nixtla.github.io/hierarchicalforecast/utils.html#aggregate'}}}
+            'hierarchicalforecast.utils': { 'hierarchicalforecast.utils.HierarchicalPlot': 'https://Nixtla.github.io/hierarchicalforecast/utils.html#hierarchicalplot',
+                                            'hierarchicalforecast.utils.HierarchicalPlot.plot_hierarchically_linked_series': 'https://Nixtla.github.io/hierarchicalforecast/utils.html#hierarchicalplot.plot_hierarchically_linked_series',
+                                            'hierarchicalforecast.utils.HierarchicalPlot.plot_series': 'https://Nixtla.github.io/hierarchicalforecast/utils.html#hierarchicalplot.plot_series',
+                                            'hierarchicalforecast.utils.HierarchicalPlot.plot_summing_matrix': 'https://Nixtla.github.io/hierarchicalforecast/utils.html#hierarchicalplot.plot_summing_matrix',
+                                            'hierarchicalforecast.utils.aggregate': 'https://Nixtla.github.io/hierarchicalforecast/utils.html#aggregate'}}}

hierarchicalforecast/utils.py

@@ -1,15 +1,21 @@
 # AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/utils.ipynb.
 
 # %% auto 0
-__all__ = ['aggregate']
+__all__ = ['aggregate', 'HierarchicalPlot']
 
 # %% ../nbs/utils.ipynb 2
+import sys
 from itertools import chain
-from typing import Callable, List
+from typing import Callable, Dict, List, Optional
 
+import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
 from sklearn.preprocessing import OneHotEncoder
+from matplotlib import rcParams
+from mycolorpy import colorlist as mcp
+
+plt.rcParams['font.family'] = 'serif'
 
 # %% ../nbs/utils.ipynb 4
 def _to_summing_matrix(S_df: pd.DataFrame):
@@ -54,3 +60,116 @@ def aggregate(
     #S definition
     S, tags = _to_summing_matrix(S_df.loc[bottom_hier, hiers_cols])
     return y_df, S, tags
+
+# %% ../nbs/utils.ipynb 9
+class HierarchicalPlot:
+
+    def __init__(
+        self,
+        S: pd.DataFrame,    #  Summing matrix of size `(base, bottom)`.
+        tags: Dict[str, np.ndarray], # Each key is a level and its value contains tags associated to that level.
+    ):
+        self.S = S
+        self.tags = tags
+
+    def plot_summing_matrix(self):
+        plt.figure(num=1, figsize=(4, 6), dpi=80, facecolor='w')
+        plt.spy(self.S)
+        plt.show()
+        plt.close()
+
+    def plot_series(
+            self, 
+            series: str, # Time series to plot
+            Y_df: Optional[pd.DataFrame] = None, # Dataframe with columns `ds` and models to plot indexed by `unique_id`.
+            models: Optional[List[str]] = None, # Models to plot. 
+            level: Optional[List[int]] = None # Levels for probabilistic intervals
+        ):
+        if series not in self.S.index:
+            raise Exception(f'time series {series} not found')
+        fig, ax = plt.subplots(1, 1, figsize = (20, 7))
+        df_plot = Y_df.loc[series].set_index('ds')
+        cols = models if models is not None else df_plot.columns
+        cols_wo_levels = [col for col in cols if ('lo' not in col and 'hi' not in col)]
+        cmap = mcp.gen_color('tab10', 10)[:len(cols_wo_levels)]
+        cmap_dict = dict(zip(cols_wo_levels, cmap))
+        df_plot[cols_wo_levels].plot(ax=ax, linewidth=2, color=cmap)
+        if level is not None:
+            for lv in level:
+                for col in cols_wo_levels:
+                    if col == 'y':
+                        # we dont need intervals
+                        # for the actual value
+                        continue
+                    if f'{col}-lo-{lv}' not in df_plot.columns:
+                        # if model
+                        # doesnt have levels
+                        continue
+                    ax.fill_between(
+                        df_plot.index, 
+                        df_plot[f'{col}-lo-{lv}'], 
+                        df_plot[f'{col}-hi-{lv}'],
+                        alpha=-lv/50 + 2,
+                        color=cmap_dict[col],
+                        label=f'{col}_level_{lv}'
+                    )
+        ax.set_title(f'{series} Forecast', fontsize=22)
+        ax.set_xlabel('Timestamp [t]', fontsize=20)
+        ax.legend(prop={'size': 15})
+        ax.grid()
+        for label in (ax.get_xticklabels() + ax.get_yticklabels()):
+            label.set_fontsize(20)
+
+    def plot_hierarchically_linked_series(
+        self,
+        bottom_series: str, # Bottom time series to plot
+        Y_df: Optional[pd.DataFrame] = None, # Dataframe with columns `ds` and models to plot indexed by `unique_id`.
+        models: Optional[List[str]] = None, # Models to plot. 
+        level: Optional[List[int]] = None # Levels for probabilistic intervals
+    ):
+        if bottom_series not in self.S.columns:
+            raise Exception(f'bottom time series {bottom_series} not found')
+        linked_series = self.S[bottom_series].loc[lambda x: x == 1.].index
+        fig, axs = plt.subplots(len(linked_series), 1, figsize=(20, 2 * len(linked_series)))
+        cols = models if models is not None else Y_df.drop(['ds'], axis=1)
+        cols_wo_levels = [col for col in cols if ('lo' not in col and 'hi' not in col)]
+        cmap = mcp.gen_color('tab10', 10)[:len(cols_wo_levels)]
+        cmap_dict = dict(zip(cols_wo_levels, cmap))
+        for idx, series in enumerate(linked_series):
+            df_plot = Y_df.loc[series].set_index('ds')
+            df_plot[cols_wo_levels].plot(ax=axs[idx], linewidth=2, color=cmap)
+            if level is not None:
+                for lv in level:
+                    for col in cols_wo_levels:
+                        if col == 'y':
+                            # we dont need intervals
+                            # for the actual value
+                            continue
+                        if f'{col}-lo-{lv}' not in df_plot.columns:
+                            # if model
+                            # doesnt have levels
+                            continue
+                        axs[idx].fill_between(
+                            df_plot.index, 
+                            df_plot[f'{col}-lo-{lv}'], 
+                            df_plot[f'{col}-hi-{lv}'],
+                            alpha=-lv/50 + 2,
+                            color=cmap_dict[col],
+                            label=f'{col}_level_{lv}'
+                        )
+            axs[idx].set_title(f'{series}', fontsize=10)
+            axs[idx].grid()
+            axs[idx].get_xaxis().label.set_visible(False)
+            axs[idx].legend().set_visible(False)
+            for label in (axs[idx].get_xticklabels() + axs[idx].get_yticklabels()):
+                label.set_fontsize(10)
+        plt.subplots_adjust(hspace=0.4)
+        handles, labels = axs[0].get_legend_handles_labels()
+        kwargs = dict(
+            loc='lower center', 
+            prop={'size': 10}, 
+            bbox_to_anchor=(0, 0.05, 1, 1)
+        )
+        if sys.version_info.minor > 7:
+            kwargs['ncols'] = np.max([2, np.ceil(len(labels) / 2)])
+        fig.legend(handles, labels, **kwargs)