<commit_before>{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4dba3df0-5f5f-4645-80cf-acb6f2ff7130",
   "metadata": {},
   "outputs": [],
   "source": [
    "#default_exp hierarchical"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1045633d-6ddd-4e98-a316-b07d333206a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c6faac6-8bf6-4173-8ca8-3acd39e3d4d4",
   "metadata": {},
   "source": [
    "# Hierarchical Time Series Forecasting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18f252aa-88e4-41f2-93eb-ccf0eeac9c93",
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore', category=FutureWarning)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4c06e3b-5315-4fcf-8642-618e2970a191",
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "from fastcore.test import test_eq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 109,
   "id": "c915637b-ae73-4f2c-ae4b-8b3b05b30536",
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "from typing import Callable, List, Optional, Tuple\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "from statsforecast.core import StatsForecast"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 141,
   "id": "2422c143-7ffe-45ec-84c4-b53ea0badb89",
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def _aggregate_key(df: pd.DataFrame, keys: List[List[str]], agg_fn: Callable = np.sum):\n",
    "    \"\"\"Aggregates `df` according to `keys` using `agg_fn`.\"\"\"\n",
    "    df = df.copy()\n",
    "    max_len_idx = np.argmax([len(key) for key in keys])\n",
    "    bottom_comb = keys[max_len_idx]\n",
    "    orig_cols = df.drop(labels=['ds', 'y'], axis=1).columns.to_list()\n",
    "    df_keys = []\n",
    "    for key in keys:\n",
    "        if key == ['total']:\n",
    "            df = df.assign(total='total')\n",
    "        df_key = df.groupby(key + ['ds'])['y'].apply(agg_fn).reset_index()\n",
    "        df_key['unique_id'] = df_key[key].agg('_'.join, axis=1)\n",
    "        if key == bottom_comb:\n",
    "            bottom_keys = df_key['unique_id'].unique()\n",
    "        df_keys.append(df_key)\n",
    "    df_keys = pd.concat(df_keys)\n",
    "    s_df = df_keys[['unique_id'] + orig_cols].drop_duplicates().reset_index(drop=True)\n",
    "    y_df = df_keys[['unique_id', 'ds', 'y']].set_index('unique_id')\n",
    "    #s_mat definition\n",
    "    s_mat = np.zeros((len(s_df), len(bottom_keys)))\n",
    "    for idx, label in enumerate(s_df['unique_id'], start=0):\n",
    "        if label == 'total':\n",
    "            s_mat[idx] = 1\n",
    "        else:\n",
    "            s_mat[idx, [label in bt for bt in bottom_keys]] = 1\n",
    "    return s_df, s_mat, y_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 150,
   "id": "17e32eb4-b464-4023-a6f1-a84aa6a7e61f",
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "df = pd.read_csv('https://raw.githubusercontent.com/Nixtla/transfer-learning-time-series/main/datasets/tourism.csv')\n",
    "df = df.rename({'Trips': 'y', 'Quarter': 'ds'}, axis=1)\n",
    "s_df, s_mat, y_df = _aggregate_key(df, [['total'],\n",
    "                                         ['State'], \n",
    "                                         ['Purpose'], \n",
    "                                         ['State', 'Region'], \n",
    "                                         ['State', 'Purpose'], \n",
    "                                         ['State', 'Region', 'Purpose']])\n",
    "test_eq(len(y_df), 34_000)\n",
    "test_eq(y_df.index.nunique(), 425)\n",
    "test_eq(s_mat.shape, (425, 304))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 114,
   "id": "80b6901b-4710-4757-b94a-2ad2599b261d",
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def bottom_up(y_hat: np.ndarray, s_mat: np.ndarray):\n",
    "    ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 115,
   "id": "d938b6ae-1d8a-4c88-a152-6f2541a6217f",
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class HierarchicalStatsForecast:\n",
    "    \n",
    "    def __init__(self, df: pd.DataFrame, keys: List[List[str]], \n",
    "                 models: List, freq: str, \n",
    "                 n_jobs: int = 1, ray_address: Optional[str] = None):\n",
    "        self.s_df, y_df = _aggregate_key(df, keys=keys)\n",
    "        self.fcst = StatsForecast(df=y_df, models=models, freq=freq, \n",
    "                                  n_jobs=n_jobs, ray_address=ray_address)\n",
    "    \n",
    "    def forecast(self, h: int, reconcile_fns: List[Callable] = [bottom_up],\n",
    "                 xreg: Optional[pd.DataFrame] = None, \n",
    "                 level: Optional[Tuple] = None):\n",
    "        fcsts = self.fcst.forecast(h=h, xreg=xreg, level=level)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8334ae8-c49e-4c37-a3d6-83592be5ec69",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {<commit_msg>feat: add forecast hierarchicalstatsforecast<commit_after>{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "4dba3df0-5f5f-4645-80cf-acb6f2ff7130",
   "metadata": {},
   "outputs": [],
   "source": [
    "#default_exp hierarchical"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "1045633d-6ddd-4e98-a316-b07d333206a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c6faac6-8bf6-4173-8ca8-3acd39e3d4d4",
   "metadata": {},
   "source": [
    "# Hierarchical Time Series Forecasting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "18f252aa-88e4-41f2-93eb-ccf0eeac9c93",
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore', category=FutureWarning)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "a4c06e3b-5315-4fcf-8642-618e2970a191",
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "from fastcore.test import test_eq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c915637b-ae73-4f2c-ae4b-8b3b05b30536",
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "from typing import Callable, List, Optional, Tuple\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "from statsforecast.core import StatsForecast"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "2422c143-7ffe-45ec-84c4-b53ea0badb89",
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def _aggregate_key(df: pd.DataFrame, keys: List[List[str]], agg_fn: Callable = np.sum):\n",
    "    \"\"\"Aggregates `df` according to `keys` using `agg_fn`.\"\"\"\n",
    "    df = df.copy()\n",
    "    max_len_idx = np.argmax([len(key) for key in keys])\n",
    "    bottom_comb = keys[max_len_idx]\n",
    "    orig_cols = df.drop(labels=['ds', 'y'], axis=1).columns.to_list()\n",
    "    df_keys = []\n",
    "    for key in keys:\n",
    "        if key == ['total']:\n",
    "            df = df.assign(total='total')\n",
    "        df_key = df.groupby(key + ['ds'])['y'].apply(agg_fn).reset_index()\n",
    "        df_key['unique_id'] = df_key[key].agg('_'.join, axis=1)\n",
    "        if key == bottom_comb:\n",
    "            bottom_keys = df_key['unique_id'].unique()\n",
    "        df_keys.append(df_key)\n",
    "    df_keys = pd.concat(df_keys)\n",
    "    s_df = df_keys[['unique_id'] + orig_cols].drop_duplicates().reset_index(drop=True)\n",
    "    s_df = s_df.set_index('unique_id')\n",
    "    y_df = df_keys[['unique_id', 'ds', 'y']].set_index('unique_id')\n",
    "    #s_mat definition\n",
    "    s_mat = np.zeros((len(s_df), len(bottom_keys)))\n",
    "    for idx, label in enumerate(s_df.index, start=0):\n",
    "        if label == 'total':\n",
    "            s_mat[idx] = 1\n",
    "        else:\n",
    "            s_mat[idx, [label in bt for bt in bottom_keys]] = 1\n",
    "    return s_df, s_mat, y_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "17e32eb4-b464-4023-a6f1-a84aa6a7e61f",
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "df = pd.read_csv('https://raw.githubusercontent.com/Nixtla/transfer-learning-time-series/main/datasets/tourism.csv')\n",
    "df = df.rename({'Trips': 'y', 'Quarter': 'ds'}, axis=1)\n",
    "s_df, s_mat, y_df = _aggregate_key(df, [['total'],\n",
    "                                         ['State'], \n",
    "                                         ['Purpose'], \n",
    "                                         ['State', 'Region'], \n",
    "                                         ['State', 'Purpose'], \n",
    "                                         ['State', 'Region', 'Purpose']])\n",
    "test_eq(len(y_df), 34_000)\n",
    "test_eq(y_df.index.nunique(), 425)\n",
    "test_eq(s_mat.shape, (425, 304))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "80b6901b-4710-4757-b94a-2ad2599b261d",
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def bottom_up(y_hat: np.ndarray, s_mat: np.ndarray):\n",
    "    # size of y_hat = n_levels, horizon\n",
    "    _, h = y_hat.shape\n",
    "    n_levels, n_bottom = s_mat.shape\n",
    "    g_mat = np.eye(n_bottom, n_levels, k=(n_levels - n_bottom))\n",
    "    s_g_mat = s_mat @ g_mat\n",
    "    return np.vstack([s_g_mat @ y_hat[:,i] for i in range(h)]).T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "d938b6ae-1d8a-4c88-a152-6f2541a6217f",
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class HierarchicalStatsForecast:\n",
    "    \n",
    "    def __init__(self, df: pd.DataFrame, keys: List[List[str]], \n",
    "                 models: List, freq: str, \n",
    "                 n_jobs: int = 1, ray_address: Optional[str] = None):\n",
    "        self.s_df, self.s_mat, y_df = _aggregate_key(df, keys=keys)\n",
    "        self.fcst = StatsForecast(df=y_df, models=models, freq=freq, \n",
    "                                  n_jobs=n_jobs, ray_address=ray_address,\n",
    "                                  sort_df=False)\n",
    "    \n",
    "    def forecast(self, h: int, reconcile_fns: List[Callable] = [bottom_up],\n",
    "                 xreg: Optional[pd.DataFrame] = None, \n",
    "                 level: Optional[Tuple] = None):\n",
    "        fcsts = self.fcst.forecast(h=h, xreg=xreg, level=level)\n",
    "        model_names = fcsts.drop(columns=['ds'], axis=1).columns.to_list()\n",
    "        for model_name in model_names:\n",
    "            fcsts_model = fcsts[model_name].values.reshape(-1, h)\n",
    "            for reconcile_fn in reconcile_fns:\n",
    "                reconcile_fn_name = reconcile_fn.__name__\n",
    "                fcsts_model = reconcile_fn(fcsts_model, self.s_mat)\n",
    "                fcsts[f'{reconcile_fn_name}_{model_name}'] = fcsts_model.flatten()\n",
    "        return fcsts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "0a5b4c24-cfda-4c7d-a290-821e83b65c0c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from statsforecast.models import naive"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "770fe257-3279-4aca-b296-2e367fe76d27",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_int = df[['ds']].drop_duplicates().assign(ds_int = lambda df: np.arange(len(df)) + 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "9aad79dc-3b57-4a19-998b-04a9b3c1504c",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = df.merge(ds_int, how='left', on=['ds']).drop('ds', axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "21e06081-62a5-40c2-98f4-9bf040a19a38",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = df.rename(columns={'ds_int': 'ds'})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "f8334ae8-c49e-4c37-a3d6-83592be5ec69",
   "metadata": {},
   "outputs": [],
   "source": [
    "hier_fcst = HierarchicalStatsForecast(df, \n",
    "                                      keys=[['total'],\n",
    "                                            ['State'], \n",
    "                                            ['Purpose'], \n",
    "                                            ['State', 'Region'], \n",
    "                                            ['State', 'Purpose'], \n",
    "                                            ['State', 'Region', 'Purpose']],\n",
    "                                      models=[naive],\n",
    "                                      freq='D', \n",
    "                                      n_jobs=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "5f4ecb31-699a-4417-b3aa-dd012e15d6ed",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>ds</th>\n",
       "      <th>naive</th>\n",
       "      <th>bottom_up_naive</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>unique_id</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>total</th>\n",
       "      <td>81</td>\n",
       "      <td>27593.554688</td>\n",
       "      <td>27593.554198</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>total</th>\n",
       "      <td>82</td>\n",
       "      <td>27593.554688</td>\n",
       "      <td>27593.554198</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>total</th>\n",
       "      <td>83</td>\n",
       "      <td>27593.554688</td>\n",
       "      <td>27593.554198</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>total</th>\n",
       "      <td>84</td>\n",
       "      <td>27593.554688</td>\n",
       "      <td>27593.554198</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>total</th>\n",
       "      <td>85</td>\n",
       "      <td>27593.554688</td>\n",
       "      <td>27593.554198</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Western Australia_Experience Perth_Visiting</th>\n",
       "      <td>86</td>\n",
       "      <td>455.316711</td>\n",
       "      <td>455.316711</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Western Australia_Experience Perth_Visiting</th>\n",
       "      <td>87</td>\n",
       "      <td>455.316711</td>\n",
       "      <td>455.316711</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Western Australia_Experience Perth_Visiting</th>\n",
       "      <td>88</td>\n",
       "      <td>455.316711</td>\n",
       "      <td>455.316711</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Western Australia_Experience Perth_Visiting</th>\n",
       "      <td>89</td>\n",
       "      <td>455.316711</td>\n",
       "      <td>455.316711</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Western Australia_Experience Perth_Visiting</th>\n",
       "      <td>90</td>\n",
       "      <td>455.316711</td>\n",
       "      <td>455.316711</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>4250 rows × 3 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "                                             ds         naive  bottom_up_naive\n",
       "unique_id                                                                     \n",
       "total                                        81  27593.554688     27593.554198\n",
       "total                                        82  27593.554688     27593.554198\n",
       "total                                        83  27593.554688     27593.554198\n",
       "total                                        84  27593.554688     27593.554198\n",
       "total                                        85  27593.554688     27593.554198\n",
       "...                                          ..           ...              ...\n",
       "Western Australia_Experience Perth_Visiting  86    455.316711       455.316711\n",
       "Western Australia_Experience Perth_Visiting  87    455.316711       455.316711\n",
       "Western Australia_Experience Perth_Visiting  88    455.316711       455.316711\n",
       "Western Australia_Experience Perth_Visiting  89    455.316711       455.316711\n",
       "Western Australia_Experience Perth_Visiting  90    455.316711       455.316711\n",
       "\n",
       "[4250 rows x 3 columns]"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hier_fcst.forecast(10)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
{'fraction_non_alphanumeric': 0.20671119074767877, 'fraction_numerical': 0.05902155617092903, 'mean_word_length': 1.7522414823670054, 'pattern_counts': {'":': 178, '<': 145, '<?xml version=': 0, '>': 145, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 58, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6474364', 'n_tokens_mistral': 7725, 'n_tokens_neox': 7063, 'n_words': 1015}