add config module and deprecation warnings for id as index

Nixtla · Nov 23, 2023 · a543d82 · a543d82
1 parent 690e8b6
commit a543d82
Show file tree

Hide file tree

Showing 8 changed files with 137 additions and 5 deletions.
diff --git a/action_files/test_models/src/models.py b/action_files/test_models/src/models.py
@@ -7,6 +7,7 @@
 import pytorch_lightning as pl
 import torch
 
+import neuralforecast
 from neuralforecast.core import NeuralForecast
 
 from neuralforecast.models.gru import GRU
@@ -34,6 +35,8 @@
 
 from src.data import get_data
 
+neuralforecast.config.id_as_index = False
+
 
 def main(dataset: str = 'M3', group: str = 'Monthly') -> None:
     train, horizon, freq, seasonality = get_data('data/', dataset, group)

diff --git a/action_files/test_models/src/models2.py b/action_files/test_models/src/models2.py
@@ -7,6 +7,7 @@
 import pytorch_lightning as pl
 import torch
 
+import neuralforecast
 from neuralforecast.core import NeuralForecast
 
 from neuralforecast.models.gru import GRU
@@ -34,6 +35,8 @@
 
 from src.data import get_data
 
+neuralforecast.config.id_as_index = False
+
 
 def main(dataset: str = 'M3', group: str = 'Monthly') -> None:
     train, horizon, freq, seasonality = get_data('data/', dataset, group)

diff --git a/nbs/config.ipynb b/nbs/config.ipynb
@@ -0,0 +1,34 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "06dabc05-6dc5-4cb2-afde-62a1dd58f606",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#| default_exp config"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e84cc420-9d82-4e69-abe2-87564faeb159",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#| export\n",
+    "id_as_index = True"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "python3",
+   "language": "python",
+   "name": "python3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/nbs/core.ipynb b/nbs/core.ipynb
@@ -75,6 +75,7 @@
     "    LocalStandardScaler,\n",
     ")\n",
     "\n",
+    "import neuralforecast.config as nf_config\n",
     "from neuralforecast.tsdataset import TimeSeriesDataset\n",
     "from neuralforecast.models import (\n",
     "    GRU, LSTM, RNN, TCN, DeepAR, DilatedRNN,\n",
@@ -229,6 +230,23 @@
     "}"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d4bf5a30-8378-40ac-920b-af757c228a9b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#| exporti\n",
+    "def _warn_id_as_idx():\n",
+    "    warnings.warn(\n",
+    "        \"In a future version the predictions will have the id as a column. \"\n",
+    "        \"You can set `neuralforecast.config.id_as_index = False` \"\n",
+    "        \"to adopt the new behavior and to suppress this warning.\",\n",
+    "        category=DeprecationWarning,\n",
+    "    )"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -509,7 +527,9 @@
     "        else:\n",
     "            fcsts = pd.DataFrame(fcsts, columns=cols)\n",
     "        fcsts_df = ufp.horizontal_concat([fcsts_df, fcsts])\n",
-    "\n",
+    "        if isinstance(fcsts_df, pd.DataFrame) and nf_config.id_as_index:\n",
+    "            _warn_id_as_idx()\n",
+    "            fcsts_df = fcsts_df.set_index('unique_id')\n",
     "        return fcsts_df\n",
     "    \n",
     "    def cross_validation(self,\n",
@@ -640,7 +660,11 @@
     "        fcsts_df = ufp.horizontal_concat([fcsts_df, fcsts])\n",
     "\n",
     "        # Add original input df's y to forecasts DataFrame    \n",
-    "        return ufp.join(fcsts_df, df, how='left', on=['unique_id', 'ds'])\n",
+    "        fcsts_df = ufp.join(fcsts_df, df, how='left', on=['unique_id', 'ds'])\n",
+    "        if isinstance(fcsts_df, pd.DataFrame) and nf_config.id_as_index:\n",
+    "            _warn_id_as_idx()\n",
+    "            fcsts_df = fcsts_df.set_index('unique_id')\n",
+    "        return fcsts_df\n",
     "\n",
     "    def predict_insample(self, step_size: int = 1):\n",
     "        \"\"\"Predict insample with core.NeuralForecast.\n",
@@ -745,6 +769,9 @@
     "                fcsts_df[invert_cols].to_numpy(),\n",
     "                indptr\n",
     "            )\n",
+    "        if isinstance(fcsts_df, pd.DataFrame) and nf_config.id_as_index:\n",
+    "            _warn_id_as_idx()\n",
+    "            fcsts_df = fcsts_df.set_index('unique_id')            \n",
     "        return fcsts_df\n",
     "        \n",
     "    # Save list of models with pytorch lightning save_checkpoint function\n",
@@ -989,6 +1016,7 @@
     "import matplotlib.pyplot as plt\n",
     "import pytorch_lightning as pl\n",
     "\n",
+    "import neuralforecast\n",
     "from ray import tune\n",
     "\n",
     "from neuralforecast.auto import (\n",
@@ -1031,6 +1059,41 @@
     "AirPassengersPanel_test['y_[lag12]'] = np.nan"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c596acd4-c95a-41f3-a710-cb9b2c27459d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#| hide\n",
+    "# id as index warnings\n",
+    "models = [\n",
+    "    NHITS(h=12, input_size=12, max_steps=1)\n",
+    "]\n",
+    "nf = NeuralForecast(models=models, freq='M')\n",
+    "nf.fit(df=AirPassengersPanel_train)\n",
+    "with warnings.catch_warnings(record=True) as issued_warnings:\n",
+    "    warnings.simplefilter('always', category=DeprecationWarning)\n",
+    "    nf.predict()\n",
+    "    nf.predict_insample()\n",
+    "    nf.cross_validation(df=AirPassengersPanel_train)\n",
+    "id_warnings = [\n",
+    "    w for w in issued_warnings if 'the predictions will have the id as a column' in str(w.message)\n",
+    "]\n",
+    "assert len(id_warnings) == 3"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "51ecf21b-1919-44b2-843d-9059fb30f1df",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "neuralforecast.config.id_as_index = False"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,

diff --git a/neuralforecast/__init__.py b/neuralforecast/__init__.py
@@ -1,3 +1,4 @@
 __version__ = "1.6.4"
 __all__ = ['NeuralForecast']
+import neuralforecast.config
 from .core import NeuralForecast
diff --git a/neuralforecast/_modidx.py b/neuralforecast/_modidx.py
@@ -57,6 +57,7 @@
                                                                                      'neuralforecast/auto.py'),
                                      'neuralforecast.auto.AutoVanillaTransformer.__init__': ( 'models.html#autovanillatransformer.__init__',
                                                                                               'neuralforecast/auto.py')},
+            'neuralforecast.config': {},
             'neuralforecast.core': { 'neuralforecast.core.NeuralForecast': ('core.html#neuralforecast', 'neuralforecast/core.py'),
                                      'neuralforecast.core.NeuralForecast.__init__': ( 'core.html#neuralforecast.__init__',
                                                                                       'neuralforecast/core.py'),
@@ -77,7 +78,8 @@
                                      'neuralforecast.core.NeuralForecast.predict_insample': ( 'core.html#neuralforecast.predict_insample',
                                                                                               'neuralforecast/core.py'),
                                      'neuralforecast.core.NeuralForecast.save': ('core.html#neuralforecast.save', 'neuralforecast/core.py'),
-                                     'neuralforecast.core._insample_times': ('core.html#_insample_times', 'neuralforecast/core.py')},
+                                     'neuralforecast.core._insample_times': ('core.html#_insample_times', 'neuralforecast/core.py'),
+                                     'neuralforecast.core._warn_id_as_idx': ('core.html#_warn_id_as_idx', 'neuralforecast/core.py')},
             'neuralforecast.losses.numpy': { 'neuralforecast.losses.numpy._divide_no_nan': ( 'losses.numpy.html#_divide_no_nan',
                                                                                              'neuralforecast/losses/numpy.py'),
                                              'neuralforecast.losses.numpy._metric_protections': ( 'losses.numpy.html#_metric_protections',

diff --git a/neuralforecast/config.py b/neuralforecast/config.py
@@ -0,0 +1,7 @@
+# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/config.ipynb.
+
+# %% auto 0
+__all__ = ['id_as_index']
+
+# %% ../nbs/config.ipynb 1
+id_as_index = True
diff --git a/neuralforecast/core.py b/neuralforecast/core.py
@@ -25,6 +25,7 @@
     LocalStandardScaler,
 )
 
+import neuralforecast.config as nf_config
 from .tsdataset import TimeSeriesDataset
 from neuralforecast.models import (
     GRU,
@@ -146,6 +147,15 @@ def _insample_times(
 }
 
 # %% ../nbs/core.ipynb 9
+def _warn_id_as_idx():
+    warnings.warn(
+        "In a future version the predictions will have the id as a column. "
+        "You can set `neuralforecast.config.id_as_index = False` "
+        "to adopt the new behavior and to suppress this warning.",
+        category=DeprecationWarning,
+    )
+
+# %% ../nbs/core.ipynb 10
 class NeuralForecast:
     def __init__(
         self, models: List[Any], freq: str, local_scaler_type: Optional[str] = None
@@ -437,7 +447,9 @@ def predict(
         else:
             fcsts = pd.DataFrame(fcsts, columns=cols)
         fcsts_df = ufp.horizontal_concat([fcsts_df, fcsts])
-
+        if isinstance(fcsts_df, pd.DataFrame) and nf_config.id_as_index:
+            _warn_id_as_idx()
+            fcsts_df = fcsts_df.set_index("unique_id")
         return fcsts_df
 
     def cross_validation(
@@ -576,7 +588,11 @@ def cross_validation(
         fcsts_df = ufp.horizontal_concat([fcsts_df, fcsts])
 
         # Add original input df's y to forecasts DataFrame
-        return ufp.join(fcsts_df, df, how="left", on=["unique_id", "ds"])
+        fcsts_df = ufp.join(fcsts_df, df, how="left", on=["unique_id", "ds"])
+        if isinstance(fcsts_df, pd.DataFrame) and nf_config.id_as_index:
+            _warn_id_as_idx()
+            fcsts_df = fcsts_df.set_index("unique_id")
+        return fcsts_df
 
     def predict_insample(self, step_size: int = 1):
         """Predict insample with core.NeuralForecast.
@@ -688,6 +704,9 @@ def predict_insample(self, step_size: int = 1):
             fcsts_df[invert_cols] = self._scalers_target_inverse_transform(
                 fcsts_df[invert_cols].to_numpy(), indptr
             )
+        if isinstance(fcsts_df, pd.DataFrame) and nf_config.id_as_index:
+            _warn_id_as_idx()
+            fcsts_df = fcsts_df.set_index("unique_id")
         return fcsts_df
 
     # Save list of models with pytorch lightning save_checkpoint function