fix readme and index examples (#237)

Nixtla · Oct 4, 2023 · 9dce084 · 9dce084
1 parent e0a4f8d
commit 9dce084
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 35 deletions.
diff --git a/README.md b/README.md
@@ -65,30 +65,35 @@ You can open a complete example in Colab [![Open In Colab](https://colab.researc
 Minimal Example:
 ```python
 # !pip install -U numba statsforecast datasetsforecast
+import numpy as np
 import pandas as pd
 
+#obtain hierarchical dataset
+from datasetsforecast.hierarchical import HierarchicalData
+
 # compute base forecast no coherent
 from statsforecast.core import StatsForecast
 from statsforecast.models import AutoARIMA, Naive
 
 #obtain hierarchical reconciliation methods and evaluation
 from hierarchicalforecast.core import HierarchicalReconciliation
+from hierarchicalforecast.evaluation import HierarchicalEvaluation
 from hierarchicalforecast.methods import BottomUp, TopDown, MiddleOut
 
-#obtain hierarchical datasets
-from datasetsforecast.hierarchical import HierarchicalData
 
 # Load TourismSmall dataset
 Y_df, S, tags = HierarchicalData.load('./data', 'TourismSmall')
 Y_df['ds'] = pd.to_datetime(Y_df['ds'])
 
+#split train/test sets
+Y_test_df  = Y_df.groupby('unique_id').tail(4)
+Y_train_df = Y_df.drop(Y_test_df.index)
 
-# Compute base level predictions 
-sf = StatsForecast(df=Y_df, 
-                   models=[AutoARIMA(season_length=12), Naive()], 
-                   freq='M', n_jobs=-1)
-
-forecasts_df = sf.forecast(h=12)
+# Compute base auto-ARIMA predictions
+fcst = StatsForecast(df=Y_train_df,
+                     models=[AutoARIMA(season_length=4), Naive()],
+                     freq='Q', n_jobs=-1)
+Y_hat_df = fcst.forecast(h=4)
 
 # Reconcile the base predictions
 reconcilers = [
@@ -97,23 +102,20 @@ reconcilers = [
     MiddleOut(middle_level='Country/Purpose/State',
               top_down_method='forecast_proportions')
 ]
-
 hrec = HierarchicalReconciliation(reconcilers=reconcilers)
-
-reconciled_forecasts = hrec.reconcile(Y_hat_df=forecasts_df, S=S, tags=tags)
+Y_rec_df = hrec.reconcile(Y_hat_df=Y_hat_df, Y_df=Y_train_df,
+                          S=S, tags=tags)
 ```
 
 ### Evaluation
 Assumes you have a test dataframe.
 
 ```python
-from hierarchicalforecast.core import HierarchicalEvaluation
-
 def mse(y, y_hat):
     return np.mean((y-y_hat)**2)
 
 evaluator = HierarchicalEvaluation(evaluators=[mse])
-evaluator.evaluate(Y_hat_df=Y_rec_df, Y_test=Y_test_df, 
+evaluator.evaluate(Y_hat_df=Y_rec_df, Y_test=Y_test_df.set_index('unique_id'),
                    tags=tags, benchmark='Naive')
 ```
 

diff --git a/nbs/examples/TourismSmall.ipynb b/nbs/examples/TourismSmall.ipynb
@@ -352,17 +352,6 @@
     "Y_train_df = Y_df.drop(Y_test_df.index)"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "ab321103-cca1-4931-8f5d-cd8026ab63fe",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "Y_test_df = Y_test_df.set_index('unique_id')\n",
-    "Y_train_df = Y_train_df.set_index('unique_id')"
-   ]
-  },
   {
    "attachments": {},
    "cell_type": "markdown",
@@ -604,7 +593,7 @@
     "\n",
     "evaluator = HierarchicalEvaluation(evaluators=[mse])\n",
     "evaluation = evaluator.evaluate(\n",
-    "        Y_hat_df=Y_rec_df, Y_test_df=Y_test_df, \n",
+    "        Y_hat_df=Y_rec_df, Y_test_df=Y_test_df.set_index('unique_id'),\n",
     "        tags=tags, benchmark='Naive'\n",
     ")\n",
     "evaluation.filter(like='ARIMA', axis=1).T"

diff --git a/nbs/index.ipynb b/nbs/index.ipynb
@@ -125,16 +125,14 @@
     "Y_df['ds'] = pd.to_datetime(Y_df['ds'])\n",
     "\n",
     "#split train/test sets\n",
-    "Y_test_df  = Y_df.groupby('unique_id').tail(12)\n",
+    "Y_test_df  = Y_df.groupby('unique_id').tail(4)\n",
     "Y_train_df = Y_df.drop(Y_test_df.index)\n",
-    "Y_test_df  = Y_test_df.set_index('unique_id')\n",
-    "Y_train_df = Y_train_df.set_index('unique_id')\n",
     "\n",
     "# Compute base auto-ARIMA predictions\n",
-    "fcst = StatsForecast(df=Y_train_df, \n",
-    "                     models=[AutoARIMA(season_length=12), Naive()], \n",
-    "                     freq='M', n_jobs=-1)\n",
-    "Y_hat_df = fcst.forecast(h=12)\n",
+    "fcst = StatsForecast(df=Y_train_df,\n",
+    "                     models=[AutoARIMA(season_length=4), Naive()],\n",
+    "                     freq='Q', n_jobs=-1)\n",
+    "Y_hat_df = fcst.forecast(h=4)\n",
     "\n",
     "# Reconcile the base predictions\n",
     "reconcilers = [\n",
@@ -144,7 +142,7 @@
     "              top_down_method='forecast_proportions')\n",
     "]\n",
     "hrec = HierarchicalReconciliation(reconcilers=reconcilers)\n",
-    "Y_rec_df = hrec.reconcile(Y_hat_df=Y_hat_df, Y_df=Y_train_df, \n",
+    "Y_rec_df = hrec.reconcile(Y_hat_df=Y_hat_df, Y_df=Y_train_df,\n",
     "                          S=S, tags=tags)\n",
     "```"
    ]
@@ -161,7 +159,7 @@
     "    return np.mean((y-y_hat)**2)\n",
     "\n",
     "evaluator = HierarchicalEvaluation(evaluators=[mse])\n",
-    "evaluator.evaluate(Y_h=Y_rec_df, Y_test=Y_df_test, \n",
+    "evaluator.evaluate(Y_hat_df=Y_rec_df, Y_test=Y_test_df.set_index('unique_id'),\n",
     "                   tags=tags, benchmark='Naive')\n",
     "```"
    ]