nbs/core.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "524620c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| default_exp core"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15392f6f",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "12fa25a4",
   "metadata": {},
   "source": [
    "# Core\n",
    "> NeuralForecast contains two main components, PyTorch implementations deep learning predictive models, as well as parallelization and distributed computation utilities. The first component comprises low-level PyTorch model estimator classes like `models.NBEATS` and `models.RNN`. The second component is a high-level `core.NeuralForecast` wrapper class that operates with sets of time series data stored in pandas DataFrames."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2508f7a9-1433-4ad8-8f2f-0078c6ed6c3c",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "import shutil\n",
    "from fastcore.test import test_eq, test_fail\n",
    "from nbdev.showdoc import show_doc\n",
    "from neuralforecast.utils import generate_series"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44065066-e72a-431f-938f-1528adef9fe8",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "import os\n",
    "import pickle\n",
    "import warnings\n",
    "from copy import deepcopy\n",
    "from itertools import chain\n",
    "from typing import Any, Dict, List, Optional, Union\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import torch\n",
    "import utilsforecast.processing as ufp\n",
    "from coreforecast.grouped_array import GroupedArray\n",
    "from coreforecast.scalers import (\n",
    "    LocalMinMaxScaler,\n",
    "    LocalRobustScaler,\n",
    "    LocalStandardScaler,\n",
    ")\n",
    "from utilsforecast.compat import DataFrame, Series, pl_DataFrame, pl_Series\n",
    "from utilsforecast.validation import validate_freq\n",
    "\n",
    "from neuralforecast.tsdataset import TimeSeriesDataset\n",
    "from neuralforecast.models import (\n",
    "    GRU, LSTM, RNN, TCN, DeepAR, DilatedRNN,\n",
    "    MLP, NHITS, NBEATS, NBEATSx, DLinear,\n",
    "    TFT, VanillaTransformer,\n",
    "    Informer, Autoformer, FEDformer,\n",
    "    StemGNN, PatchTST, TimesNet\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb8b4b3c-04bf-4a92-9a1a-b60735997c36",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| exporti\n",
    "def _insample_times(\n",
    "    times: np.ndarray,\n",
    "    uids: Series,\n",
    "    indptr: np.ndarray,\n",
    "    h: int,\n",
    "    freq: Union[int, str, pd.offsets.BaseOffset],\n",
    "    step_size: int = 1,\n",
    "    id_col: str = 'unique_id',\n",
    "    time_col: str = 'ds',\n",
    ") -> DataFrame:\n",
    "    sizes = np.diff(indptr)\n",
    "    if (sizes < h).any():\n",
    "        raise ValueError('`sizes` should be greater or equal to `h`.')\n",
    "    # TODO: we can just truncate here instead of raising an error\n",
    "    ns, resids = np.divmod(sizes - h, step_size)\n",
    "    if (resids != 0).any():\n",
    "        raise ValueError('`sizes - h` should be multiples of `step_size`')\n",
    "    windows_per_serie = ns + 1\n",
    "    # determine the offsets for the cutoffs, e.g. 2 means the 3rd training date is a cutoff\n",
    "    cutoffs_offsets = step_size * np.hstack([np.arange(w) for w in windows_per_serie])\n",
    "    # start index of each serie, e.g. [0, 17] means the the second serie starts on the 18th entry\n",
    "    # we repeat each of these as many times as we have windows, e.g. windows_per_serie = [2, 3]\n",
    "    # would yield [0, 0, 17, 17, 17]\n",
    "    start_idxs = np.repeat(indptr[:-1], windows_per_serie)\n",
    "    # determine the actual indices of the cutoffs, we repeat the cutoff for the complete horizon\n",
    "    # e.g. if we have two series and h=2 this could be [0, 0, 1, 1, 17, 17, 18, 18]\n",
    "    # which would have the first two training dates from each serie as the cutoffs\n",
    "    cutoff_idxs = np.repeat(start_idxs + cutoffs_offsets, h)\n",
    "    cutoffs = times[cutoff_idxs]\n",
    "    total_windows = windows_per_serie.sum()\n",
    "    # determine the offsets for the actual dates. this is going to be [0, ..., h] repeated\n",
    "    ds_offsets = np.tile(np.arange(h), total_windows)\n",
    "    # determine the actual indices of the times\n",
    "    # e.g. if we have two series and h=2 this could be [0, 1, 1, 2, 17, 18, 18, 19]\n",
    "    ds_idxs = cutoff_idxs + ds_offsets\n",
    "    ds = times[ds_idxs]\n",
    "    if isinstance(uids, pl_Series):\n",
    "        df_constructor = pl_DataFrame\n",
    "    else:\n",
    "        df_constructor = pd.DataFrame\n",
    "    out = df_constructor(\n",
    "        {\n",
    "            id_col: ufp.repeat(uids, h * windows_per_serie),\n",
    "            time_col: ds,\n",
    "            'cutoff': cutoffs,\n",
    "        }\n",
    "    )\n",
    "    # the first cutoff is before the first train date\n",
    "    actual_cutoffs = ufp.offset_times(out['cutoff'], freq, -1)\n",
    "    out = ufp.assign_columns(out, 'cutoff', actual_cutoffs)\n",
    "    return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6aead6db-170a-4a74-baf3-7cbaf8b7468c",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "uids = pd.Series(['id_0', 'id_1'])\n",
    "indptr = np.array([0, 4, 10], dtype=np.int32)\n",
    "h = 2\n",
    "for step_size, freq, days in zip([1, 2], ['D', 'W-THU'], [1, 14]):\n",
    "    times = np.hstack([\n",
    "        pd.date_range('2000-01-01', freq=freq, periods=4),\n",
    "        pd.date_range('2000-10-10', freq=freq, periods=10),\n",
    "    ])    \n",
    "    times_df = _insample_times(times, uids, indptr, h, freq, step_size=step_size)\n",
    "    pd.testing.assert_frame_equal(\n",
    "        times_df.groupby('unique_id')['ds'].min().reset_index(),\n",
    "        pd.DataFrame({\n",
    "            'unique_id': uids,\n",
    "            'ds': times[indptr[:-1]],\n",
    "        })\n",
    "    )\n",
    "    pd.testing.assert_frame_equal(\n",
    "        times_df.groupby('unique_id')['ds'].max().reset_index(),\n",
    "        pd.DataFrame({\n",
    "            'unique_id': uids,\n",
    "            'ds': times[indptr[1:] - 1],\n",
    "        })\n",
    "    )\n",
    "    cutoff_deltas = (\n",
    "        times_df\n",
    "        .drop_duplicates(['unique_id', 'cutoff'])\n",
    "        .groupby('unique_id')\n",
    "        ['cutoff']\n",
    "        .diff()\n",
    "        .dropna()\n",
    "    )\n",
    "    assert cutoff_deltas.nunique() == 1\n",
    "    assert cutoff_deltas.unique()[0] == pd.Timedelta(f'{days}D')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c58a8a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| exporti\n",
    "MODEL_FILENAME_DICT = {\n",
    "    'autoformer': Autoformer, 'autoautoformer': Autoformer,\n",
    "    'deepar': DeepAR, 'autodeepar': DeepAR,\n",
    "    'dilatedrnn': DilatedRNN , 'autodilatedrnn': DilatedRNN,\n",
    "    'fedformer': FEDformer, 'autofedformer': FEDformer,\n",
    "    'gru': GRU, 'autogru': GRU,\n",
    "    'informer': Informer, 'autoinformer': Informer,\n",
    "    'lstm': LSTM, 'autolstm': LSTM,\n",
    "    'mlp': MLP, 'automlp': MLP,\n",
    "    'nbeats': NBEATS, 'autonbeats': NBEATS,\n",
    "    'nbeatsx': NBEATSx, 'autonbeatsx': NBEATSx,\n",
    "    'nhits': NHITS, 'autonhits': NHITS,\n",
    "    'dlinear': DLinear, 'autodlinear': DLinear,\n",
    "    'patchtst': PatchTST, 'autopatchtst': PatchTST,\n",
    "    'rnn': RNN, 'autornn': RNN,\n",
    "    'stemgnn': StemGNN, 'autostemgnn': StemGNN,\n",
    "    'tcn': TCN, 'autotcn': TCN, \n",
    "    'tft': TFT, 'autotft': TFT,\n",
    "    'timesnet': TimesNet, 'autotimesnet': TimesNet,\n",
    "    'vanillatransformer': VanillaTransformer, 'autovanillatransformer': VanillaTransformer,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c621a39-5658-4850-95c4-050eee97403d",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| exporti\n",
    "_type2scaler = {\n",
    "    'standard': LocalStandardScaler,\n",
    "    'robust': lambda: LocalRobustScaler(scale='mad'),\n",
    "    'robust-iqr': lambda: LocalRobustScaler(scale='iqr'),\n",
    "    'minmax': LocalMinMaxScaler,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4bf5a30-8378-40ac-920b-af757c228a9b",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| exporti\n",
    "def _id_as_idx() -> bool:\n",
    "    return not bool(os.getenv(\"NIXTLA_ID_AS_COL\", \"\"))\n",
    "\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 the `NIXTLA_ID_AS_COL` environment variable \"\n",
    "        \"to adopt the new behavior and to suppress this warning.\",\n",
    "        category=FutureWarning,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4dae43c-4d11-4bbc-a431-ac33b004859a",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "class NeuralForecast:\n",
    "    \n",
    "    def __init__(self, \n",
    "                 models: List[Any],\n",
    "                 freq: Union[str, int],\n",
    "                 local_scaler_type: Optional[str] = None):\n",
    "        \"\"\"\n",
    "        The `core.StatsForecast` class allows you to efficiently fit multiple `NeuralForecast` models \n",
    "        for large sets of time series. It operates with pandas DataFrame `df` that identifies series \n",
    "        and datestamps with the `unique_id` and `ds` columns. The `y` column denotes the target \n",
    "        time series variable.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        models : List[typing.Any]\n",
    "            Instantiated `neuralforecast.models` \n",
    "            see [collection here](https://nixtla.github.io/neuralforecast/models.html).\n",
    "        freq : str or int\n",
    "            Frequency of the data. Must be a valid pandas or polars offset alias, or an integer.\n",
    "        local_scaler_type : str, optional (default=None)\n",
    "            Scaler to apply per-serie to all features before fitting, which is inverted after predicting.\n",
    "            Can be 'standard', 'robust', 'robust-iqr' or 'minmax'\n",
    "        \n",
    "        Returns\n",
    "        -------\n",
    "        self : NeuralForecast\n",
    "            Returns instantiated `NeuralForecast` class.\n",
    "        \"\"\"\n",
    "        assert all(model.h == models[0].h for model in models), 'All models should have the same horizon'\n",
    "\n",
    "        self.h = models[0].h\n",
    "        self.models_init = models\n",
    "        self.freq = freq\n",
    "        if local_scaler_type is not None and local_scaler_type not in _type2scaler:\n",
    "            raise ValueError(f'scaler_type must be one of {_type2scaler.keys()}')\n",
    "        self.local_scaler_type = local_scaler_type\n",
    "        self.scalers_: Dict\n",
    "\n",
    "        # Flags and attributes\n",
    "        self._fitted = False\n",
    "        self._reset_models()\n",
    "\n",
    "    def _scalers_fit_transform(self, dataset: TimeSeriesDataset) -> None:\n",
    "        self.scalers_ = {}        \n",
    "        if self.local_scaler_type is None:\n",
    "            return None\n",
    "        for i, col in enumerate(dataset.temporal_cols):\n",
    "            if col == 'available_mask':\n",
    "                continue\n",
    "            ga = GroupedArray(dataset.temporal[:, i].numpy(), dataset.indptr)                \n",
    "            self.scalers_[col] = _type2scaler[self.local_scaler_type]().fit(ga)\n",
    "            dataset.temporal[:, i] = torch.from_numpy(self.scalers_[col].transform(ga))\n",
    "\n",
    "    def _scalers_transform(self, dataset: TimeSeriesDataset) -> None:\n",
    "        if not self.scalers_:\n",
    "            return None\n",
    "        for i, col in enumerate(dataset.temporal_cols):\n",
    "            scaler = self.scalers_.get(col, None)\n",
    "            if scaler is None:\n",
    "                continue\n",
    "            ga = GroupedArray(dataset.temporal[:, i].numpy(), dataset.indptr)\n",
    "            dataset.temporal[:, i] = torch.from_numpy(scaler.transform(ga))\n",
    "\n",
    "    def _scalers_target_inverse_transform(self, data: np.ndarray, indptr: np.ndarray) -> np.ndarray:\n",
    "        if not self.scalers_:\n",
    "            return data\n",
    "        for i in range(data.shape[1]):\n",
    "            ga = GroupedArray(data[:, i], indptr)\n",
    "            data[:, i] = self.scalers_[self.target_col].inverse_transform(ga)\n",
    "        return data\n",
    "\n",
    "    def _prepare_fit(self, df, static_df, sort_df, predict_only, id_col, time_col, target_col):\n",
    "        #TODO: uids, last_dates and ds should be properties of the dataset class. See github issue.\n",
    "        self.id_col = id_col\n",
    "        self.time_col = time_col\n",
    "        self.target_col = target_col\n",
    "        dataset, uids, last_dates, ds = TimeSeriesDataset.from_df(\n",
    "            df=df,\n",
    "            static_df=static_df,\n",
    "            sort_df=sort_df,\n",
    "            id_col=id_col,\n",
    "            time_col=time_col,\n",
    "            target_col=target_col,\n",
    "        )\n",
    "        if predict_only:\n",
    "            self._scalers_transform(dataset)\n",
    "        else:\n",
    "            self._scalers_fit_transform(dataset)\n",
    "        return dataset, uids, last_dates, ds\n",
    "\n",
    "    def fit(self,\n",
    "        df: Optional[DataFrame] = None,\n",
    "        static_df: Optional[DataFrame] = None,\n",
    "        val_size: Optional[int] = 0,\n",
    "        sort_df: bool = True,\n",
    "        use_init_models: bool = False,\n",
    "        verbose: bool = False,\n",
    "        id_col: str = 'unique_id',\n",
    "        time_col: str = 'ds',\n",
    "        target_col: str = 'y',\n",
    "    ) -> None:\n",
    "        \"\"\"Fit the core.NeuralForecast.\n",
    "\n",
    "        Fit `models` to a large set of time series from DataFrame `df`.\n",
    "        and store fitted models for later inspection.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        df : pandas or polars DataFrame, optional (default=None)\n",
    "            DataFrame with columns [`unique_id`, `ds`, `y`] and exogenous variables.\n",
    "            If None, a previously stored dataset is required.\n",
    "        static_df : pandas or polars DataFrame, optional (default=None)\n",
    "            DataFrame with columns [`unique_id`] and static exogenous.\n",
    "        val_size : int, optional (default=0)\n",
    "            Size of validation set.\n",
    "        sort_df : bool, optional (default=False)\n",
    "            Sort `df` before fitting.\n",
    "        use_init_models : bool, optional (default=False)\n",
    "            Use initial model passed when NeuralForecast object was instantiated.\n",
    "        verbose : bool (default=False)\n",
    "            Print processing steps.\n",
    "        id_col : str (default='unique_id')\n",
    "            Column that identifies each serie.\n",
    "        time_col : str (default='ds')\n",
    "            Column that identifies each timestep, its values can be timestamps or integers.\n",
    "        target_col : str (default='y')\n",
    "            Column that contains the target.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        self : NeuralForecast\n",
    "            Returns `NeuralForecast` class with fitted `models`.\n",
    "        \"\"\"\n",
    "        if (df is None) and not (hasattr(self, 'dataset')):\n",
    "            raise Exception('You must pass a DataFrame or have one stored.')\n",
    "\n",
    "        # Model and datasets interactions protections\n",
    "        if (any(model.early_stop_patience_steps>0 for model in self.models)) \\\n",
    "            and (val_size==0):\n",
    "                raise Exception('Set val_size>0 if early stopping is enabled.')\n",
    "\n",
    "        # Process and save new dataset (in self)\n",
    "        if df is not None:\n",
    "            validate_freq(df[time_col], self.freq)\n",
    "            self.dataset, self.uids, self.last_dates, self.ds = self._prepare_fit(\n",
    "                df=df,\n",
    "                static_df=static_df,\n",
    "                sort_df=sort_df,\n",
    "                predict_only=False,\n",
    "                id_col=id_col,\n",
    "                time_col=time_col,\n",
    "                target_col=target_col,\n",
    "            )\n",
    "            self.sort_df = sort_df\n",
    "        else:\n",
    "            if verbose: print('Using stored dataset.')\n",
    "\n",
    "        if val_size is not None:\n",
    "            if self.dataset.min_size < val_size:\n",
    "                warnings.warn('Validation set size is larger than the shorter time-series.')\n",
    "\n",
    "        # Recover initial model if use_init_models\n",
    "        if use_init_models:\n",
    "            self._reset_models()\n",
    "\n",
    "        for model in self.models:\n",
    "            model.fit(self.dataset, val_size=val_size)\n",
    "\n",
    "        self._fitted = True\n",
    "\n",
    "    def make_future_dataframe(self, df: Optional[DataFrame] = None) -> DataFrame:\n",
    "        \"\"\"Create a dataframe with all ids and future times in the forecasting horizon.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        df : pandas or polars DataFrame, optional (default=None)\n",
    "            DataFrame with columns [`unique_id`, `ds`, `y`] and exogenous variables.\n",
    "            Only required if this is different than the one used in the fit step.\n",
    "        \"\"\"\n",
    "        if not self._fitted:\n",
    "            raise Exception('You must fit the model first.')\n",
    "        if df is not None:\n",
    "            df = ufp.sort(df, by=[self.id_col, self.time_col])\n",
    "            last_times_by_id = ufp.group_by_agg(\n",
    "                df,\n",
    "                by=self.id_col,\n",
    "                aggs={self.time_col: 'max'},\n",
    "                maintain_order=True,\n",
    "            )\n",
    "            uids = last_times_by_id[self.id_col]\n",
    "            last_times = last_times_by_id[self.time_col]\n",
    "        else:\n",
    "            uids = self.uids\n",
    "            last_times = self.last_dates\n",
    "        return ufp.make_future_dataframe(\n",
    "            uids=uids,\n",
    "            last_times=last_times,\n",
    "            freq=self.freq,\n",
    "            h=self.h,\n",
    "            id_col=self.id_col,\n",
    "            time_col=self.time_col,\n",
    "        )\n",
    "\n",
    "    def get_missing_future(\n",
    "        self, futr_df: DataFrame, df: Optional[DataFrame] = None\n",
    "    ) -> DataFrame:\n",
    "        \"\"\"Get the missing ids and times combinations in `futr_df`.\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        futr_df : pandas or polars DataFrame\n",
    "            DataFrame with [`unique_id`, `ds`] columns and `df`'s future exogenous.\n",
    "        df : pandas or polars DataFrame, optional (default=None)\n",
    "            DataFrame with columns [`unique_id`, `ds`, `y`] and exogenous variables.\n",
    "            Only required if this is different than the one used in the fit step.\n",
    "        \"\"\"\n",
    "        expected = self.make_future_dataframe(df)\n",
    "        ids = [self.id_col, self.time_col]\n",
    "        return ufp.anti_join(expected, futr_df[ids], on=ids)\n",
    "\n",
    "    def _get_needed_futr_exog(self):\n",
    "        return set(chain.from_iterable(getattr(m, 'futr_exog_list', []) for m in self.models))\n",
    "\n",
    "    def predict(self,\n",
    "                df: Optional[DataFrame] = None,\n",
    "                static_df: Optional[DataFrame] = None,\n",
    "                futr_df: Optional[DataFrame] = None,\n",
    "                sort_df: bool = True,\n",
    "                verbose: bool = False,\n",
    "                **data_kwargs):\n",
    "        \"\"\"Predict with core.NeuralForecast.\n",
    "\n",
    "        Use stored fitted `models` to predict large set of time series from DataFrame `df`.        \n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        df : pandas or polars DataFrame, optional (default=None)\n",
    "            DataFrame with columns [`unique_id`, `ds`, `y`] and exogenous variables.\n",
    "            If a DataFrame is passed, it is used to generate forecasts.\n",
    "        static_df : pandas or polars DataFrame, optional (default=None)\n",
    "            DataFrame with columns [`unique_id`] and static exogenous.\n",
    "        futr_df : pandas or polars DataFrame, optional (default=None)\n",
    "            DataFrame with [`unique_id`, `ds`] columns and `df`'s future exogenous.\n",
    "        sort_df : bool (default=True)\n",
    "            Sort `df` before fitting.\n",
    "        verbose : bool (default=False)\n",
    "            Print processing steps.\n",
    "        data_kwargs : kwargs\n",
    "            Extra arguments to be passed to the dataset within each model.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        fcsts_df : pandas or polars DataFrame\n",
    "            DataFrame with insample `models` columns for point predictions and probabilistic\n",
    "            predictions for all fitted `models`.    \n",
    "        \"\"\"\n",
    "        if (df is None) and not (hasattr(self, 'dataset')):\n",
    "            raise Exception('You must pass a DataFrame or have one stored.')\n",
    "\n",
    "        if not self._fitted:\n",
    "            raise Exception(\"You must fit the model before predicting.\")\n",
    "\n",
    "        needed_futr_exog = self._get_needed_futr_exog()\n",
    "        if needed_futr_exog:\n",
    "            if futr_df is None:\n",
    "                raise ValueError(\n",
    "                    f'Models require the following future exogenous features: {needed_futr_exog}. '\n",
    "                    'Please provide them through the `futr_df` argument.'\n",
    "                )\n",
    "            else:\n",
    "                missing = needed_futr_exog - set(futr_df.columns)\n",
    "                if missing:\n",
    "                    raise ValueError(f'The following features are missing from `futr_df`: {missing}')\n",
    "        \n",
    "        # Process new dataset but does not store it.\n",
    "        if df is not None:\n",
    "            validate_freq(df[self.time_col], self.freq)\n",
    "            dataset, uids, last_dates, _ = self._prepare_fit(\n",
    "                df=df,\n",
    "                static_df=static_df,\n",
    "                sort_df=sort_df,\n",
    "                predict_only=True,\n",
    "                id_col=self.id_col,\n",
    "                time_col=self.time_col,\n",
    "                target_col=self.target_col,\n",
    "            )\n",
    "        else:\n",
    "            dataset = self.dataset\n",
    "            uids = self.uids\n",
    "            last_dates = self.last_dates\n",
    "            if verbose: print('Using stored dataset.')\n",
    "\n",
    "        cols = []\n",
    "        count_names = {'model': 0}\n",
    "        for model in self.models:\n",
    "            model_name = repr(model)\n",
    "            count_names[model_name] = count_names.get(model_name, -1) + 1\n",
    "            if count_names[model_name] > 0:\n",
    "                model_name += str(count_names[model_name])\n",
    "            cols += [model_name + n for n in model.loss.output_names]\n",
    "\n",
    "        # Placeholder dataframe for predictions with unique_id and ds\n",
    "        fcsts_df = ufp.make_future_dataframe(\n",
    "            uids=uids,\n",
    "            last_times=last_dates,\n",
    "            freq=self.freq,\n",
    "            h=self.h,\n",
    "            id_col=self.id_col,\n",
    "            time_col=self.time_col,\n",
    "        )\n",
    "\n",
    "        # Update and define new forecasting dataset\n",
    "        if futr_df is None:\n",
    "            futr_df = fcsts_df\n",
    "        else:\n",
    "            futr_orig_rows = futr_df.shape[0]\n",
    "            futr_df = ufp.join(futr_df, fcsts_df, on=[self.id_col, self.time_col])\n",
    "            if futr_df.shape[0] < fcsts_df.shape[0]:\n",
    "                if df is None:\n",
    "                    expected_cmd = 'make_future_dataframe()'\n",
    "                    missing_cmd = 'get_missing_future(futr_df)'\n",
    "                else:\n",
    "                    expected_cmd = 'make_future_dataframe(df)'\n",
    "                    missing_cmd = 'get_missing_future(futr_df, df)'\n",
    "                raise ValueError(\n",
    "                    'There are missing combinations of ids and times in `futr_df`.\\n'\n",
    "                    f'You can run the `{expected_cmd}` method to get the expected combinations or '\n",
    "                    f'the `{missing_cmd}` method to get the missing combinations.'\n",
    "                )\n",
    "            if futr_orig_rows > futr_df.shape[0]:\n",
    "                dropped_rows = futr_orig_rows - futr_df.shape[0]\n",
    "                warnings.warn(\n",
    "                    f'Dropped {dropped_rows:,} unused rows from `futr_df`.'\n",
    "                )\n",
    "            if any(ufp.is_none(futr_df[col]).any() for col in needed_futr_exog):\n",
    "                raise ValueError('Found null values in `futr_df`')\n",
    "        futr_dataset = dataset.align(\n",
    "            futr_df,\n",
    "            id_col=self.id_col,\n",
    "            time_col=self.time_col,\n",
    "            target_col=self.target_col,\n",
    "        )\n",
    "        self._scalers_transform(futr_dataset)\n",
    "        dataset = dataset.append(futr_dataset)\n",
    "\n",
    "        col_idx = 0\n",
    "        fcsts = np.full((self.h * len(uids), len(cols)), fill_value=np.nan, dtype=np.float32)\n",
    "        for model in self.models:\n",
    "            old_test_size = model.get_test_size()\n",
    "            model.set_test_size(self.h) # To predict h steps ahead\n",
    "            model_fcsts = model.predict(dataset=dataset, **data_kwargs)\n",
    "            # Append predictions in memory placeholder\n",
    "            output_length = len(model.loss.output_names)\n",
    "            fcsts[:, col_idx : col_idx + output_length] = model_fcsts\n",
    "            col_idx += output_length\n",
    "            model.set_test_size(old_test_size) # Set back to original value\n",
    "        if self.scalers_:\n",
    "            indptr = np.append(0, np.full(len(uids), self.h).cumsum())\n",
    "            fcsts = self._scalers_target_inverse_transform(fcsts, indptr)\n",
    "\n",
    "        # Declare predictions pd.DataFrame\n",
    "        if isinstance(fcsts_df, pl_DataFrame):\n",
    "            fcsts = pl_DataFrame(dict(zip(cols, fcsts.T)))\n",
    "        else:\n",
    "            fcsts = pd.DataFrame(fcsts, columns=cols)\n",
    "        fcsts_df = ufp.horizontal_concat([fcsts_df, fcsts])\n",
    "        if isinstance(fcsts_df, pd.DataFrame) and _id_as_idx():\n",
    "            _warn_id_as_idx()\n",
    "            fcsts_df = fcsts_df.set_index(self.id_col)\n",
    "        return fcsts_df\n",
    "\n",
    "    def _reset_models(self):\n",
    "        self.models = [deepcopy(model) for model in self.models_init]\n",
    "        if self._fitted:\n",
    "            print('WARNING: Deleting previously fitted models.')        \n",
    "    \n",
    "    def _no_refit_cross_validation(\n",
    "        self,\n",
    "        df: Optional[DataFrame],\n",
    "        static_df: Optional[DataFrame],\n",
    "        n_windows: int,\n",
    "        step_size: int,\n",
    "        val_size: Optional[int], \n",
    "        test_size: int,\n",
    "        sort_df: bool,\n",
    "        verbose: bool,\n",
    "        id_col: str,\n",
    "        time_col: str,\n",
    "        target_col: str,\n",
    "        **data_kwargs\n",
    "    ) -> DataFrame:\n",
    "        if (df is None) and not (hasattr(self, 'dataset')):\n",
    "            raise Exception('You must pass a DataFrame or have one stored.')\n",
    "\n",
    "        # Process and save new dataset (in self)\n",
    "        if df is not None:\n",
    "            validate_freq(df[time_col], self.freq)\n",
    "            self.dataset, self.uids, self.last_dates, self.ds = self._prepare_fit(\n",
    "                df=df,\n",
    "                static_df=static_df,\n",
    "                sort_df=sort_df,\n",
    "                predict_only=False,\n",
    "                id_col=id_col,\n",
    "                time_col=time_col,\n",
    "                target_col=target_col,\n",
    "            )\n",
    "            self.sort_df = sort_df\n",
    "        else:\n",
    "            if verbose: print('Using stored dataset.')\n",
    "\n",
    "        if val_size is not None:\n",
    "            if self.dataset.min_size < (val_size+test_size):\n",
    "                warnings.warn('Validation and test sets are larger than the shorter time-series.')\n",
    "\n",
    "        cols = []\n",
    "        count_names = {'model': 0}\n",
    "        for model in self.models:\n",
    "            model_name = repr(model)\n",
    "            count_names[model_name] = count_names.get(model_name, -1) + 1\n",
    "            if count_names[model_name] > 0:\n",
    "                model_name += str(count_names[model_name])\n",
    "            cols += [model_name + n for n in model.loss.output_names]\n",
    "\n",
    "        fcsts_df = ufp.cv_times(\n",
    "            times=self.ds,\n",
    "            uids=self.uids,\n",
    "            indptr=self.dataset.indptr,\n",
    "            h=self.h,\n",
    "            test_size=test_size,\n",
    "            step_size=step_size,\n",
    "            id_col=id_col,\n",
    "            time_col=time_col,\n",
    "        )\n",
    "        # the cv_times is sorted by window and then id\n",
    "        fcsts_df = ufp.sort(fcsts_df, [id_col, 'cutoff', time_col])\n",
    "\n",
    "        col_idx = 0\n",
    "        fcsts = np.full((self.dataset.n_groups * self.h * n_windows, len(cols)),\n",
    "                         np.nan, dtype=np.float32)\n",
    "        \n",
    "        for model in self.models:\n",
    "            model.fit(dataset=self.dataset,\n",
    "                        val_size=val_size, \n",
    "                        test_size=test_size)\n",
    "            model_fcsts = model.predict(self.dataset, step_size=step_size, **data_kwargs)\n",
    "\n",
    "            # Append predictions in memory placeholder\n",
    "            output_length = len(model.loss.output_names)\n",
    "            fcsts[:,col_idx:(col_idx + output_length)] = model_fcsts\n",
    "            col_idx += output_length\n",
    "        if self.scalers_:            \n",
    "            indptr = np.append(0, np.full(self.dataset.n_groups, self.h * n_windows).cumsum())\n",
    "            fcsts = self._scalers_target_inverse_transform(fcsts, indptr)\n",
    "\n",
    "        self._fitted = True\n",
    "\n",
    "        # Add predictions to forecasts DataFrame\n",
    "        if isinstance(self.uids, pl_Series):\n",
    "            fcsts = pl_DataFrame(dict(zip(cols, fcsts.T)))\n",
    "        else:\n",
    "            fcsts = pd.DataFrame(fcsts, columns=cols)\n",
    "        fcsts_df = ufp.horizontal_concat([fcsts_df, fcsts])\n",
    "\n",
    "        # Add original input df's y to forecasts DataFrame    \n",
    "        fcsts_df = ufp.join(\n",
    "            fcsts_df,\n",
    "            df[[id_col, time_col, target_col]],\n",
    "            how='left',\n",
    "            on=[id_col, time_col],\n",
    "        )\n",
    "        if isinstance(fcsts_df, pd.DataFrame) and _id_as_idx():\n",
    "            _warn_id_as_idx()\n",
    "            fcsts_df = fcsts_df.set_index(id_col)\n",
    "        return fcsts_df\n",
    "\n",
    "    def cross_validation(\n",
    "        self,\n",
    "        df: Optional[DataFrame] = None,\n",
    "        static_df: Optional[DataFrame] = None,\n",
    "        n_windows: int = 1,\n",
    "        step_size: int = 1,\n",
    "        val_size: Optional[int] = 0, \n",
    "        test_size: Optional[int] = None,\n",
    "        sort_df: bool = True,\n",
    "        use_init_models: bool = False,\n",
    "        verbose: bool = False,\n",
    "        refit: Union[bool, int] = False,\n",
    "        id_col: str = 'unique_id',\n",
    "        time_col: str = 'ds',\n",
    "        target_col: str = 'y',\n",
    "        **data_kwargs\n",
    "    ) -> DataFrame:\n",
    "        \"\"\"Temporal Cross-Validation with core.NeuralForecast.\n",
    "\n",
    "        `core.NeuralForecast`'s cross-validation efficiently fits a list of NeuralForecast \n",
    "        models through multiple windows, in either chained or rolled manner.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        df : pandas or polars DataFrame, optional (default=None)\n",
    "            DataFrame with columns [`unique_id`, `ds`, `y`] and exogenous variables.\n",
    "            If None, a previously stored dataset is required.\n",
    "        static_df : pandas or polars DataFrame, optional (default=None)\n",
    "            DataFrame with columns [`unique_id`] and static exogenous.\n",
    "        n_windows : int (default=1)\n",
    "            Number of windows used for cross validation.\n",
    "        step_size : int (default=1)\n",
    "            Step size between each window.\n",
    "        val_size : int, optional (default=None)\n",
    "            Length of validation size. If passed, set `n_windows=None`.\n",
    "        test_size : int, optional (default=None)\n",
    "            Length of test size. If passed, set `n_windows=None`.\n",
    "        sort_df : bool (default=True)\n",
    "            Sort `df` before fitting.\n",
    "        use_init_models : bool, option (default=False)\n",
    "            Use initial model passed when object was instantiated.\n",
    "        verbose : bool (default=False)\n",
    "            Print processing steps.\n",
    "        refit : bool or int (default=False)\n",
    "            Retrain model for each cross validation window.\n",
    "            If False, the models are trained at the beginning and then used to predict each window.\n",
    "            If positive int, the models are retrained every `refit` windows.\n",
    "        id_col : str (default='unique_id')\n",
    "            Column that identifies each serie.\n",
    "        time_col : str (default='ds')\n",
    "            Column that identifies each timestep, its values can be timestamps or integers.\n",
    "        target_col : str (default='y')\n",
    "            Column that contains the target.            \n",
    "        data_kwargs : kwargs\n",
    "            Extra arguments to be passed to the dataset within each model.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        fcsts_df : pandas or polars DataFrame\n",
    "            DataFrame with insample `models` columns for point predictions and probabilistic\n",
    "            predictions for all fitted `models`.    \n",
    "        \"\"\"\n",
    "        h = self.h\n",
    "        if n_windows is None and test_size is None:\n",
    "            raise Exception('you must define `n_windows` or `test_size`.')            \n",
    "        if test_size is None:\n",
    "            test_size = h + step_size * (n_windows - 1)\n",
    "        elif n_windows is None:\n",
    "            if (test_size - h) % step_size:\n",
    "                raise Exception('`test_size - h` should be module `step_size`')\n",
    "            n_windows = int((test_size - h) / step_size) + 1\n",
    "        else:\n",
    "            raise Exception('you must define `n_windows` or `test_size` but not both')       \n",
    "        # Recover initial model if use_init_models.\n",
    "        if use_init_models:\n",
    "            self._reset_models()        \n",
    "        if not refit:\n",
    "            return self._no_refit_cross_validation(\n",
    "                df=df,\n",
    "                static_df=static_df,\n",
    "                n_windows=n_windows,\n",
    "                step_size=step_size,\n",
    "                val_size=val_size,\n",
    "                test_size=test_size,\n",
    "                sort_df=sort_df,\n",
    "                verbose=verbose,\n",
    "                id_col=id_col,\n",
    "                time_col=time_col,\n",
    "                target_col=target_col,\n",
    "                **data_kwargs\n",
    "            )\n",
    "        if df is None:\n",
    "            raise ValueError('Must specify `df` with `refit!=False`.')\n",
    "        validate_freq(df[time_col], self.freq)\n",
    "        splits = ufp.backtest_splits(\n",
    "            df,\n",
    "            n_windows=n_windows,\n",
    "            h=self.h,\n",
    "            id_col=id_col,\n",
    "            time_col=time_col,\n",
    "            freq=self.freq,\n",
    "            step_size=step_size,\n",
    "            input_size=None,\n",
    "        )\n",
    "        results = []\n",
    "        for i_window, (cutoffs, train, test) in enumerate(splits):\n",
    "            should_fit = i_window == 0 or (refit > 0 and i_window % refit == 0)\n",
    "            if should_fit:\n",
    "                self.fit(\n",
    "                    df=train,\n",
    "                    static_df=static_df,\n",
    "                    val_size=val_size,\n",
    "                    sort_df=sort_df,\n",
    "                    use_init_models=False,\n",
    "                    verbose=verbose,                 \n",
    "                )\n",
    "                predict_df: Optional[DataFrame] = None\n",
    "            else:\n",
    "                predict_df = train\n",
    "            needed_futr_exog = self._get_needed_futr_exog()\n",
    "            if needed_futr_exog:\n",
    "                futr_df: Optional[DataFrame] = test\n",
    "            else:\n",
    "                futr_df = None\n",
    "            preds = self.predict(\n",
    "                df=predict_df,\n",
    "                static_df=static_df,\n",
    "                futr_df=futr_df,\n",
    "                sort_df=sort_df,\n",
    "                verbose=verbose,\n",
    "                **data_kwargs\n",
    "            )\n",
    "            preds = ufp.join(preds, cutoffs, on=id_col, how='left')\n",
    "            fold_result = ufp.join(\n",
    "                preds, test[[id_col, time_col, target_col]], on=[id_col, time_col]\n",
    "            )\n",
    "            results.append(fold_result)\n",
    "        out = ufp.vertical_concat(results, match_categories=False)\n",
    "        out = ufp.drop_index_if_pandas(out)\n",
    "        # match order of cv with no refit\n",
    "        first_out_cols = [id_col, time_col, \"cutoff\"]\n",
    "        remaining_cols = [\n",
    "            c for c in out.columns if c not in first_out_cols + [target_col]\n",
    "        ]\n",
    "        cols_order = first_out_cols + remaining_cols + [target_col]\n",
    "        out = ufp.sort(out[cols_order], by=[id_col, 'cutoff', time_col])\n",
    "        if isinstance(out, pd.DataFrame) and _id_as_idx():\n",
    "            _warn_id_as_idx()\n",
    "            out = out.set_index(id_col)\n",
    "        return out\n",
    "\n",
    "    def predict_insample(self, step_size: int = 1):\n",
    "        \"\"\"Predict insample with core.NeuralForecast.\n",
    "\n",
    "        `core.NeuralForecast`'s `predict_insample` uses stored fitted `models`\n",
    "        to predict historic values of a time series from the stored dataframe.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        step_size : int (default=1)\n",
    "            Step size between each window.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        fcsts_df : pandas.DataFrame\n",
    "            DataFrame with insample predictions for all fitted `models`.    \n",
    "        \"\"\"\n",
    "        if not self._fitted:\n",
    "            raise Exception('The models must be fitted first with `fit` or `cross_validation`.')\n",
    "\n",
    "        for model in self.models:\n",
    "            if model.SAMPLING_TYPE == 'recurrent':\n",
    "                warnings.warn(f'Predict insample might not provide accurate predictions for \\\n",
    "                       recurrent model {repr(model)} class yet due to scaling.')\n",
    "                print(f'WARNING: Predict insample might not provide accurate predictions for \\\n",
    "                      recurrent model {repr(model)} class yet due to scaling.')\n",
    "        \n",
    "        cols = []\n",
    "        count_names = {'model': 0}\n",
    "        for model in self.models:\n",
    "            model_name = repr(model)\n",
    "            count_names[model_name] = count_names.get(model_name, -1) + 1\n",
    "            if count_names[model_name] > 0:\n",
    "                model_name += str(count_names[model_name])\n",
    "            cols += [model_name + n for n in model.loss.output_names]\n",
    "\n",
    "        # Remove test set from dataset and last dates\n",
    "        test_size = self.models[0].get_test_size()\n",
    "        if test_size>0:\n",
    "            trimmed_dataset = TimeSeriesDataset.trim_dataset(dataset=self.dataset,\n",
    "                                                     right_trim=test_size,\n",
    "                                                     left_trim=0)\n",
    "            new_idxs = np.hstack(\n",
    "                [\n",
    "                    np.arange(self.dataset.indptr[i], self.dataset.indptr[i + 1] - test_size)\n",
    "                    for i in range(self.dataset.n_groups)\n",
    "                ]\n",
    "            )\n",
    "            times = self.ds[new_idxs]\n",
    "        else:\n",
    "            trimmed_dataset = self.dataset\n",
    "            times = self.ds\n",
    "\n",
    "        # Generate dates\n",
    "        fcsts_df = _insample_times(\n",
    "            times=times,\n",
    "            uids=self.uids,\n",
    "            indptr=trimmed_dataset.indptr,\n",
    "            h=self.h,\n",
    "            freq=self.freq,\n",
    "            step_size=step_size,\n",
    "            id_col=self.id_col,\n",
    "            time_col=self.time_col,\n",
    "        )\n",
    "\n",
    "        col_idx = 0\n",
    "        fcsts = np.full((len(fcsts_df), len(cols)), np.nan, dtype=np.float32)\n",
    "\n",
    "        for model in self.models:\n",
    "            # Test size is the number of periods to forecast (full size of trimmed dataset)\n",
    "            model.set_test_size(test_size=trimmed_dataset.max_size)\n",
    "\n",
    "            # Predict\n",
    "            model_fcsts = model.predict(trimmed_dataset, step_size=step_size)\n",
    "            # Append predictions in memory placeholder\n",
    "            output_length = len(model.loss.output_names)\n",
    "            fcsts[:,col_idx:(col_idx + output_length)] = model_fcsts\n",
    "            col_idx += output_length          \n",
    "            model.set_test_size(test_size=test_size) # Set original test_size\n",
    "\n",
    "        # original y\n",
    "        original_y = {\n",
    "            self.id_col: ufp.repeat(self.uids, np.diff(self.dataset.indptr)),\n",
    "            self.time_col: self.ds,\n",
    "            self.target_col: self.dataset.temporal[:, 0].numpy(),\n",
    "        }\n",
    "\n",
    "        # Add predictions to forecasts DataFrame\n",
    "        if isinstance(self.uids, pl_Series):\n",
    "            fcsts = pl_DataFrame(dict(zip(cols, fcsts.T)))\n",
    "            Y_df = pl_DataFrame(original_y)\n",
    "        else:\n",
    "            fcsts = pd.DataFrame(fcsts, columns=cols)\n",
    "            Y_df = pd.DataFrame(original_y).reset_index(drop=True)\n",
    "        fcsts_df = ufp.horizontal_concat([fcsts_df, fcsts])\n",
    "\n",
    "        # Add original input df's y to forecasts DataFrame\n",
    "        fcsts_df = ufp.join(fcsts_df, Y_df, how='left', on=[self.id_col, self.time_col])\n",
    "        if self.scalers_:\n",
    "            sizes = ufp.counts_by_id(fcsts_df, self.id_col)['counts'].to_numpy()\n",
    "            indptr = np.append(0, sizes.cumsum())\n",
    "            invert_cols = cols + [self.target_col]\n",
    "            fcsts_df[invert_cols] = self._scalers_target_inverse_transform(\n",
    "                fcsts_df[invert_cols].to_numpy(),\n",
    "                indptr\n",
    "            )\n",
    "        if isinstance(fcsts_df, pd.DataFrame) and _id_as_idx():\n",
    "            _warn_id_as_idx()\n",
    "            fcsts_df = fcsts_df.set_index(self.id_col)            \n",
    "        return fcsts_df\n",
    "        \n",
    "    # Save list of models with pytorch lightning save_checkpoint function\n",
    "    def save(self, path: str, model_index: Optional[List]=None, save_dataset: bool=True, overwrite: bool=False):\n",
    "        \"\"\"Save NeuralForecast core class.\n",
    "\n",
    "        `core.NeuralForecast`'s method to save current status of models, dataset, and configuration.\n",
    "        Note that by default the `models` are not saving training checkpoints to save disk memory,\n",
    "        to get them change the individual model `**trainer_kwargs` to include `enable_checkpointing=True`.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        path : str\n",
    "            Directory to save current status.\n",
    "        model_index : list, optional (default=None)\n",
    "            List to specify which models from list of self.models to save.\n",
    "        save_dataset : bool (default=True)\n",
    "            Whether to save dataset or not.\n",
    "        overwrite : bool (default=False)\n",
    "            Whether to overwrite files or not.\n",
    "        \"\"\"\n",
    "        # Standarize path without '/'\n",
    "        if path[-1] == '/':\n",
    "            path = path[:-1]\n",
    "\n",
    "        # Model index list\n",
    "        if model_index is None:\n",
    "            model_index = list(range(len(self.models)))\n",
    "\n",
    "        # Create directory if not exists\n",
    "        os.makedirs(path, exist_ok = True)\n",
    "\n",
    "        # Check if directory is empty to protect overwriting files\n",
    "        dir = os.listdir(path)\n",
    "        \n",
    "        # Checking if the list is empty or not\n",
    "        if (len(dir) > 0) and (not overwrite):\n",
    "            raise Exception('Directory is not empty. Set `overwrite=True` to overwrite files.')\n",
    "\n",
    "        # Save models\n",
    "        count_names = {'model': 0}\n",
    "        alias_to_model = {}\n",
    "        for i, model in enumerate(self.models):\n",
    "            # Skip model if not in list\n",
    "            if i not in model_index:\n",
    "                continue\n",
    "\n",
    "            model_name = repr(model).lower().replace('_', '')\n",
    "            model_class_name = model.__class__.__name__.lower()\n",
    "            if model_name != model_class_name:\n",
    "                alias_to_model[model_name] = model_class_name\n",
    "            count_names[model_name] = count_names.get(model_name, -1) + 1\n",
    "            model.save(f\"{path}/{model_name}_{count_names[model_name]}.ckpt\")\n",
    "        if alias_to_model:\n",
    "            with open(f\"{path}/alias_to_model.pkl\", \"wb\") as f:\n",
    "                pickle.dump(alias_to_model, f)\n",
    "\n",
    "        # Save dataset\n",
    "        if (save_dataset) and (hasattr(self, 'dataset')):\n",
    "            with open(f\"{path}/dataset.pkl\", \"wb\") as f:\n",
    "                pickle.dump(self.dataset, f)\n",
    "        elif save_dataset:\n",
    "            raise Exception('You need to have a stored dataset to save it, \\\n",
    "                             set `save_dataset=False` to skip saving dataset.')\n",
    "\n",
    "        # Save configuration and parameters\n",
    "        config_dict = {\n",
    "            \"h\": self.h,\n",
    "            \"freq\": self.freq,\n",
    "            \"uids\": self.uids,\n",
    "            \"last_dates\": self.last_dates,\n",
    "            \"ds\": self.ds,\n",
    "            \"sort_df\": self.sort_df,\n",
    "            \"_fitted\": self._fitted,\n",
    "            \"local_scaler_type\": self.local_scaler_type,\n",
    "            \"scalers_\": self.scalers_,\n",
    "            \"id_col\": self.id_col,\n",
    "            \"time_col\": self.time_col,\n",
    "            \"target_col\": self.target_col,\n",
    "        }\n",
    "\n",
    "        with open(f\"{path}/configuration.pkl\", \"wb\") as f:\n",
    "            pickle.dump(config_dict, f)\n",
    "\n",
    "    @staticmethod\n",
    "    def load(path, verbose=False, **kwargs):\n",
    "        \"\"\"Load NeuralForecast\n",
    "\n",
    "        `core.NeuralForecast`'s method to load checkpoint from path.\n",
    "\n",
    "        Parameters\n",
    "        -----------\n",
    "        path : str\n",
    "            Directory to save current status.\n",
    "        kwargs\n",
    "            Additional keyword arguments to be passed to the function\n",
    "            `load_from_checkpoint`.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        result : NeuralForecast\n",
    "            Instantiated `NeuralForecast` class.\n",
    "        \"\"\"\n",
    "        files = [f for f in os.listdir(path) if os.path.isfile(os.path.join(path, f))]\n",
    "\n",
    "        # Load models\n",
    "        models_ckpt = [f for f in files if f.endswith('.ckpt')]\n",
    "        if len(models_ckpt) == 0:\n",
    "            raise Exception('No model found in directory.') \n",
    "        \n",
    "        if verbose: print(10 * '-' + ' Loading models ' + 10 * '-')\n",
    "        models = []\n",
    "        alias_file = 'alias_to_model.pkl'\n",
    "        if os.path.isfile(f'{path}/{alias_file}'):\n",
    "            with open(f'{path}/{alias_file}', 'rb') as f:\n",
    "                alias_to_model = pickle.load(f)\n",
    "        else:\n",
    "            alias_to_model = {}\n",
    "        for model in models_ckpt:\n",
    "            model_name = model.split('_')[0]\n",
    "            model_class_name = alias_to_model.get(model_name, model_name)\n",
    "            models.append(MODEL_FILENAME_DICT[model_class_name].load_from_checkpoint(f\"{path}/{model}\", **kwargs))\n",
    "            if verbose: print(f\"Model {model_name} loaded.\")\n",
    "\n",
    "        if verbose: print(10*'-' + ' Loading dataset ' + 10*'-')\n",
    "        # Load dataset\n",
    "        if 'dataset.pkl' in files:\n",
    "            with open(f\"{path}/dataset.pkl\", \"rb\") as f:\n",
    "                dataset = pickle.load(f)\n",
    "            if verbose: print('Dataset loaded.')\n",
    "        else:\n",
    "            dataset = None\n",
    "            if verbose: print('No dataset found in directory.')\n",
    "        \n",
    "        if verbose: print(10*'-' + ' Loading configuration ' + 10*'-')\n",
    "        # Load configuration\n",
    "        if 'configuration.pkl' in files:\n",
    "            with open(f\"{path}/configuration.pkl\", \"rb\") as f:\n",
    "                config_dict = pickle.load(f)\n",
    "            if verbose: print('Configuration loaded.')\n",
    "        else:\n",
    "            raise Exception('No configuration found in directory.')\n",
    "\n",
    "        # Create NeuralForecast object\n",
    "        neuralforecast = NeuralForecast(\n",
    "            models=models,\n",
    "            freq=config_dict['freq'],\n",
    "            local_scaler_type=config_dict['local_scaler_type'],\n",
    "        )\n",
    "\n",
    "        for attr in ['id_col', 'time_col', 'target_col']:\n",
    "            setattr(neuralforecast, attr, config_dict[attr])\n",
    "\n",
    "        # Dataset\n",
    "        if dataset is not None:\n",
    "            neuralforecast.dataset = dataset\n",
    "            restore_attrs = [\n",
    "                'uids',\n",
    "                'last_dates',\n",
    "                'ds',\n",
    "                'sort_df',\n",
    "            ]\n",
    "            for attr in restore_attrs:\n",
    "                setattr(neuralforecast, attr, config_dict[attr])\n",
    "\n",
    "        # Fitted flag\n",
    "        neuralforecast._fitted = config_dict['_fitted']\n",
    "\n",
    "        neuralforecast.scalers_ = config_dict['scalers_']\n",
    "\n",
    "        return neuralforecast"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d6ef366-daec-4ec6-a2ae-199c6ea39a51",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "import logging\n",
    "import warnings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ac1aa65-40a4-4909-bdfb-1439c30439b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "logging.getLogger(\"pytorch_lightning\").setLevel(logging.ERROR)\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4bede563-78c0-40ee-ba76-f06f329cd772",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(NeuralForecast.fit, title_level=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f90209f6-16da-40a6-8302-1c5c2f66c619",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(NeuralForecast.predict, title_level=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19a8923a-f4f3-4e60-b9b9-a7088fc9bff5",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(NeuralForecast.cross_validation, title_level=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "355df52b",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(NeuralForecast.predict_insample, title_level=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93155738-b40f-43d3-ba76-d345bf2583d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(NeuralForecast.save, title_level=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e915796-173c-4400-812f-c6351d5df3be",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(NeuralForecast.load, title_level=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b534d29d-eecc-43ba-8468-c23305fa24a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "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",
    "    AutoMLP, AutoNBEATS, \n",
    "    AutoRNN, AutoTCN, AutoDilatedRNN,\n",
    ")\n",
    "\n",
    "from neuralforecast.models.rnn import RNN\n",
    "from neuralforecast.models.tcn import TCN\n",
    "from neuralforecast.models.deepar import DeepAR\n",
    "from neuralforecast.models.dilated_rnn import DilatedRNN\n",
    "\n",
    "from neuralforecast.models.mlp import MLP\n",
    "from neuralforecast.models.nhits import NHITS\n",
    "from neuralforecast.models.nbeats import NBEATS\n",
    "from neuralforecast.models.nbeatsx import NBEATSx\n",
    "\n",
    "from neuralforecast.models.tft import TFT\n",
    "from neuralforecast.models.vanillatransformer import VanillaTransformer\n",
    "from neuralforecast.models.informer import Informer\n",
    "from neuralforecast.models.autoformer import Autoformer\n",
    "\n",
    "from neuralforecast.models.stemgnn import StemGNN\n",
    "\n",
    "from neuralforecast.losses.pytorch import MQLoss, MAE, MSE\n",
    "from neuralforecast.utils import AirPassengersDF, AirPassengersPanel, AirPassengersStatic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6fd1507c",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "AirPassengersPanel_train = AirPassengersPanel[AirPassengersPanel['ds'] < AirPassengersPanel['ds'].values[-12]].reset_index(drop=True)\n",
    "AirPassengersPanel_test = AirPassengersPanel[AirPassengersPanel['ds'] >= AirPassengersPanel['ds'].values[-12]].reset_index(drop=True)\n",
    "AirPassengersPanel_test['y'] = np.nan\n",
    "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=FutureWarning)\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": [
    "os.environ['NIXTLA_ID_AS_COL'] = '1'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2e35f8f",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# Unitest for early stopping without val_size protection\n",
    "models = [\n",
    "    NHITS(h=12, input_size=12, max_steps=1, early_stop_patience_steps=5)\n",
    "]\n",
    "nf = NeuralForecast(models=models, freq='M')\n",
    "test_fail(nf.fit,\n",
    "          contains='Set val_size>0 if early stopping is enabled.',\n",
    "          args=(AirPassengersPanel_train,))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14ae4692",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test fit+cross_validation behaviour\n",
    "models = [NHITS(h=12, input_size=24, max_steps=10)]\n",
    "nf = NeuralForecast(models=models, freq='M')\n",
    "nf.fit(AirPassengersPanel_train)\n",
    "init_fcst = nf.predict()\n",
    "init_cv = nf.cross_validation(AirPassengersPanel_train, use_init_models=True)\n",
    "after_cv = nf.cross_validation(AirPassengersPanel_train, use_init_models=True)\n",
    "nf.fit(AirPassengersPanel_train, use_init_models=True)\n",
    "after_fcst = nf.predict()\n",
    "test_eq(init_cv, after_cv)\n",
    "test_eq(init_fcst, after_fcst)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fddd5459-cc9b-4fd3-b50c-493b969b83f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test cross_validation with refit\n",
    "models = [\n",
    "    NHITS(\n",
    "        h=12,\n",
    "        input_size=24,\n",
    "        max_steps=2,\n",
    "        futr_exog_list=['trend'],\n",
    "        stat_exog_list=['airline1', 'airline2']\n",
    "    )\n",
    "]\n",
    "nf = NeuralForecast(models=models, freq='M')\n",
    "cv_kwargs = dict(\n",
    "    df=AirPassengersPanel_train,\n",
    "    static_df=AirPassengersStatic,\n",
    "    n_windows=4,\n",
    "    use_init_models=True,\n",
    ")\n",
    "cv_res_norefit = nf.cross_validation(refit=False, **cv_kwargs)\n",
    "cutoffs = cv_res_norefit['cutoff'].unique()\n",
    "for refit in [True, 2]:\n",
    "    cv_res = nf.cross_validation(refit=refit, **cv_kwargs)\n",
    "    refit = int(refit)\n",
    "    fltr = lambda df: df['cutoff'].isin(cutoffs[:refit])\n",
    "    expected = cv_res_norefit[fltr]\n",
    "    actual = cv_res[fltr]\n",
    "    # predictions for the no-refit windows should be the same\n",
    "    pd.testing.assert_frame_equal(\n",
    "        actual.reset_index(drop=True),\n",
    "        expected.reset_index(drop=True)\n",
    "    )\n",
    "    # predictions after refit should be different\n",
    "    test_fail(\n",
    "        lambda: pd.testing.assert_frame_equal(\n",
    "            cv_res_norefit.drop(expected.index).reset_index(drop=True),\n",
    "            cv_res.drop(actual.index).reset_index(drop=True),\n",
    "        ),\n",
    "        contains='(column name=\"NHITS\") are different',\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1161197c-c0c2-4d71-9b39-701777db26e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test scaling\n",
    "models = [NHITS(h=12, input_size=24, max_steps=10)]\n",
    "models_exog = [NHITS(h=12, input_size=12, max_steps=10, hist_exog_list=['trend'], futr_exog_list=['trend'])]\n",
    "\n",
    "# fit+predict\n",
    "nf = NeuralForecast(models=models, freq='M', local_scaler_type='standard')\n",
    "nf.fit(AirPassengersPanel_train)\n",
    "scaled_fcst = nf.predict()\n",
    "# check that the forecasts are similar to the one without scaling\n",
    "np.testing.assert_allclose(\n",
    "    init_fcst['NHITS'].values,\n",
    "    scaled_fcst['NHITS'].values,\n",
    "    rtol=0.3,\n",
    ")\n",
    "# with exog\n",
    "nf = NeuralForecast(models=models_exog, freq='M', local_scaler_type='standard')\n",
    "nf.fit(AirPassengersPanel_train)\n",
    "scaled_exog_fcst = nf.predict(futr_df=AirPassengersPanel_test)\n",
    "# check that the forecasts are similar to the one without exog\n",
    "np.testing.assert_allclose(\n",
    "    scaled_fcst['NHITS'].values,\n",
    "    scaled_exog_fcst['NHITS'].values,\n",
    "    rtol=0.3,\n",
    ")\n",
    "\n",
    "# CV\n",
    "nf = NeuralForecast(models=models, freq='M', local_scaler_type='robust')\n",
    "cv_res = nf.cross_validation(AirPassengersPanel)\n",
    "# check that the forecasts are similar to the original values (originals are restored directly from the df)\n",
    "np.testing.assert_allclose(\n",
    "    cv_res['NHITS'].values,\n",
    "    cv_res['y'].values,\n",
    "    rtol=0.3,\n",
    ")\n",
    "# with exog\n",
    "nf = NeuralForecast(models=models_exog, freq='M', local_scaler_type='robust-iqr')\n",
    "cv_res_exog = nf.cross_validation(AirPassengersPanel)\n",
    "# check that the forecasts are similar to the original values (originals are restored directly from the df)\n",
    "np.testing.assert_allclose(\n",
    "    cv_res_exog['NHITS'].values,\n",
    "    cv_res_exog['y'].values,\n",
    "    rtol=0.2,\n",
    ")\n",
    "\n",
    "# fit+predict_insample\n",
    "nf = NeuralForecast(models=models, freq='M', local_scaler_type='minmax')\n",
    "nf.fit(AirPassengersPanel_train)\n",
    "insample_res = (\n",
    "    nf.predict_insample()\n",
    "    .groupby('unique_id').tail(-12) # first values aren't reliable\n",
    "    .merge(\n",
    "        AirPassengersPanel_train[['unique_id', 'ds', 'y']],\n",
    "        on=['unique_id', 'ds'],\n",
    "        how='left',\n",
    "        suffixes=('_actual', '_expected'),\n",
    "    )\n",
    ")\n",
    "# y is inverted correctly\n",
    "np.testing.assert_allclose(\n",
    "    insample_res['y_actual'].values,\n",
    "    insample_res['y_expected'].values,\n",
    "    rtol=1e-5,\n",
    ")\n",
    "# predictions are in the same scale\n",
    "np.testing.assert_allclose(\n",
    "    insample_res['NHITS'].values,\n",
    "    insample_res['y_expected'].values,\n",
    "    rtol=0.7,\n",
    ")\n",
    "# with exog\n",
    "nf = NeuralForecast(models=models_exog, freq='M', local_scaler_type='minmax')\n",
    "nf.fit(AirPassengersPanel_train)\n",
    "insample_res_exog = (\n",
    "    nf.predict_insample()\n",
    "    .groupby('unique_id').tail(-12) # first values aren't reliable\n",
    "    .merge(\n",
    "        AirPassengersPanel_train[['unique_id', 'ds', 'y']],\n",
    "        on=['unique_id', 'ds'],\n",
    "        how='left',\n",
    "        suffixes=('_actual', '_expected'),\n",
    "    )\n",
    ")\n",
    "# y is inverted correctly\n",
    "np.testing.assert_allclose(\n",
    "    insample_res_exog['y_actual'].values,\n",
    "    insample_res_exog['y_expected'].values,\n",
    "    rtol=1e-5,\n",
    ")\n",
    "# predictions are similar than without exog\n",
    "np.testing.assert_allclose(\n",
    "    insample_res['NHITS'].values,\n",
    "    insample_res_exog['NHITS'].values,\n",
    "    rtol=0.2,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1b34d37-29b0-49d4-9c7b-cb8add7ce9cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test futr_df contents\n",
    "models = [NHITS(h=6, input_size=24, max_steps=10, hist_exog_list=['trend'], futr_exog_list=['trend'])]\n",
    "nf = NeuralForecast(models=models, freq='M')\n",
    "nf.fit(AirPassengersPanel_train)\n",
    "# not enough rows in futr_df raises an error\n",
    "test_fail(lambda: nf.predict(futr_df=AirPassengersPanel_test.head()), contains='There are missing combinations')\n",
    "# extra rows issues a warning\n",
    "with warnings.catch_warnings(record=True) as issued_warnings:\n",
    "    warnings.simplefilter('always', UserWarning)\n",
    "    nf.predict(futr_df=AirPassengersPanel_test)\n",
    "assert any('Dropped 12 unused rows' in str(w.message) for w in issued_warnings)\n",
    "# models require futr_df and not provided raises an error\n",
    "test_fail(lambda: nf.predict(), contains=\"Models require the following future exogenous features: {'trend'}\") \n",
    "# missing feature in futr_df raises an error\n",
    "test_fail(lambda: nf.predict(futr_df=AirPassengersPanel_test.drop(columns='trend')), contains=\"missing from `futr_df`: {'trend'}\")\n",
    "# null values in futr_df raises an error\n",
    "test_fail(lambda: nf.predict(futr_df=AirPassengersPanel_test.assign(trend=np.nan)), contains='Found null values in `futr_df`')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e78b113",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# Test inplace model fitting\n",
    "models = [MLP(h=12, input_size=12, max_steps=1, scaler_type='robust')]\n",
    "initial_weights = models[0].mlp[0].weight.detach().clone()\n",
    "fcst = NeuralForecast(models=models, freq='M')\n",
    "fcst.fit(df=AirPassengersPanel_train, static_df=AirPassengersStatic, use_init_models=True)\n",
    "after_weights = fcst.models_init[0].mlp[0].weight.detach().clone()\n",
    "assert np.allclose(initial_weights, after_weights), 'init models should not be modified'\n",
    "assert len(fcst.models[0].train_trajectories)>0, 'models stored trajectories should not be empty'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "340dd8a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# Test predict_insample\n",
    "test_size = 12\n",
    "n_series = 2\n",
    "h = 12\n",
    "\n",
    "config = {'input_size': tune.choice([12, 24]), \n",
    "          'hidden_size': 128,\n",
    "          'max_steps': 1,\n",
    "          'val_check_steps': 1,\n",
    "          'step_size': 12}\n",
    "\n",
    "models = [\n",
    "    NHITS(h=h, input_size=24, loss=MQLoss(level=[80]), max_steps=1, alias='NHITS', scaler_type=None),\n",
    "    AutoMLP(h=12, config=config, cpus=1, num_samples=1),\n",
    "    RNN(h=h, input_size=-1, loss=MAE(), max_steps=1, alias='RNN', scaler_type=None),\n",
    "    ]\n",
    "\n",
    "nf = NeuralForecast(models=models, freq='M')\n",
    "cv = nf.cross_validation(df=AirPassengersPanel_train, static_df=AirPassengersStatic, val_size=0, test_size=test_size, n_windows=None)\n",
    "\n",
    "forecasts = nf.predict_insample(step_size=1)\n",
    "\n",
    "expected_size = n_series*((len(AirPassengersPanel_train)//n_series-test_size)-h+1)*h\n",
    "assert len(forecasts) == expected_size, 'Shape mistmach in predict_insample'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d6bbc2c-d38f-4cec-a3ef-15164852479f",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# tests aliases\n",
    "config_drnn = {'input_size': tune.choice([-1]), \n",
    "               'encoder_hidden_size': tune.choice([5, 10]),\n",
    "               'max_steps': 1,\n",
    "               'val_check_steps': 1,\n",
    "               'step_size': 1}\n",
    "models = [\n",
    "    # test Auto\n",
    "    AutoDilatedRNN(h=12, config=config_drnn, cpus=1, num_samples=2, alias='AutoDIL'),\n",
    "    # test BaseWindows\n",
    "    NHITS(h=12, input_size=24, loss=MQLoss(level=[80]), max_steps=1, alias='NHITSMQ'),\n",
    "    # test BaseRecurrent\n",
    "    RNN(h=12, input_size=-1, encoder_hidden_size=10, max_steps=1,\n",
    "            stat_exog_list=['airline1'],\n",
    "            futr_exog_list=['trend'], hist_exog_list=['y_[lag12]'], alias='MyRNN'),\n",
    "    # test BaseMultivariate\n",
    "    StemGNN(h=12, input_size=24, n_series=2, max_steps=1, scaler_type='robust', alias='StemMulti'),\n",
    "    # test model without alias\n",
    "    NHITS(h=12, input_size=24, max_steps=1),\n",
    "]\n",
    "nf = NeuralForecast(models=models, freq='M')\n",
    "nf.fit(df=AirPassengersPanel_train, static_df=AirPassengersStatic)\n",
    "forecasts = nf.predict(futr_df=AirPassengersPanel_test)\n",
    "test_eq(\n",
    "    forecasts.columns.to_list(),\n",
    "    ['unique_id', 'ds', 'AutoDIL', 'NHITSMQ-median', 'NHITSMQ-lo-80', 'NHITSMQ-hi-80', 'MyRNN', 'StemMulti', 'NHITS']\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d3779a6-2d03-4ac3-9f01-8bd5cb306845",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# Unit test for core/model interactions\n",
    "config = {'input_size': tune.choice([12, 24]), \n",
    "          'hidden_size': 256,\n",
    "          'max_steps': 1,\n",
    "          'val_check_steps': 1,\n",
    "          'step_size': 12}\n",
    "\n",
    "config_drnn = {'input_size': tune.choice([-1]), \n",
    "               'encoder_hidden_size': tune.choice([5, 10]),\n",
    "               'max_steps': 1,\n",
    "               'val_check_steps': 1,\n",
    "               'step_size': 1}\n",
    "\n",
    "fcst = NeuralForecast(\n",
    "    models=[\n",
    "        AutoDilatedRNN(h=12, config=config_drnn, cpus=1, num_samples=2),\n",
    "        DeepAR(h=12, input_size=24, max_steps=1,\n",
    "               stat_exog_list=['airline1'], futr_exog_list=['trend']),\n",
    "        DilatedRNN(h=12, input_size=-1, encoder_hidden_size=10, max_steps=1,\n",
    "                   stat_exog_list=['airline1'],\n",
    "                   futr_exog_list=['trend'], hist_exog_list=['y_[lag12]']),\n",
    "        RNN(h=12, input_size=-1, encoder_hidden_size=10, max_steps=1,\n",
    "            inference_input_size=24,\n",
    "            stat_exog_list=['airline1'],\n",
    "            futr_exog_list=['trend'], hist_exog_list=['y_[lag12]']),\n",
    "        TCN(h=12, input_size=-1, encoder_hidden_size=10, max_steps=1,\n",
    "            stat_exog_list=['airline1'],\n",
    "            futr_exog_list=['trend'], hist_exog_list=['y_[lag12]']),\n",
    "        AutoMLP(h=12, config=config, cpus=1, num_samples=2),\n",
    "        NBEATSx(h=12, input_size=12, max_steps=1,\n",
    "                stat_exog_list=['airline1'],\n",
    "                futr_exog_list=['trend'], hist_exog_list=['y_[lag12]']),\n",
    "        NHITS(h=12, input_size=24, loss=MQLoss(level=[80]), max_steps=1),\n",
    "        NHITS(h=12, input_size=12, max_steps=1,\n",
    "              stat_exog_list=['airline1'],\n",
    "              futr_exog_list=['trend'], hist_exog_list=['y_[lag12]']),\n",
    "        DLinear(h=12, input_size=24, max_steps=1),\n",
    "        MLP(h=12, input_size=12, max_steps=1,\n",
    "            stat_exog_list=['airline1'],\n",
    "            futr_exog_list=['trend'], hist_exog_list=['y_[lag12]']),\n",
    "        TFT(h=12, input_size=24, max_steps=1),\n",
    "        VanillaTransformer(h=12, input_size=24, max_steps=1),\n",
    "        Informer(h=12, input_size=24, max_steps=1),\n",
    "        Autoformer(h=12, input_size=24, max_steps=1),\n",
    "        FEDformer(h=12, input_size=24, max_steps=1),\n",
    "        PatchTST(h=12, input_size=24, max_steps=1),\n",
    "        TimesNet(h=12, input_size=24, max_steps=1),\n",
    "        StemGNN(h=12, input_size=24, n_series=2, max_steps=1, scaler_type='robust'),\n",
    "    ],\n",
    "    freq='M'\n",
    ")\n",
    "fcst.fit(df=AirPassengersPanel_train, static_df=AirPassengersStatic)\n",
    "forecasts = fcst.predict(futr_df=AirPassengersPanel_test)\n",
    "forecasts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40038532-fd68-4375-b7da-ba5bc2491c5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "fig, ax = plt.subplots(1, 1, figsize = (20, 7))\n",
    "plot_df = pd.concat([AirPassengersPanel_train, forecasts.reset_index()]).set_index('ds')\n",
    "\n",
    "plot_df[plot_df['unique_id']=='Airline1'].drop(['unique_id','trend','y_[lag12]'], axis=1).plot(ax=ax, linewidth=2)\n",
    "\n",
    "ax.set_title('AirPassengers Forecast', fontsize=22)\n",
    "ax.set_ylabel('Monthly Passengers', fontsize=20)\n",
    "ax.set_xlabel('Timestamp [t]', fontsize=20)\n",
    "ax.legend(prop={'size': 15})\n",
    "ax.grid()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d61909b",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "fig, ax = plt.subplots(1, 1, figsize = (20, 7))\n",
    "plot_df = pd.concat([AirPassengersPanel_train, forecasts.reset_index()]).set_index('ds')\n",
    "\n",
    "plot_df[plot_df['unique_id']=='Airline2'].drop(['unique_id','trend','y_[lag12]'], axis=1).plot(ax=ax, linewidth=2)\n",
    "\n",
    "ax.set_title('AirPassengers Forecast', fontsize=22)\n",
    "ax.set_ylabel('Monthly Passengers', fontsize=20)\n",
    "ax.set_xlabel('Timestamp [t]', fontsize=20)\n",
    "ax.legend(prop={'size': 15})\n",
    "ax.grid()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4a8162a-3d9d-48df-a314-3a2ce0377e36",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "config = {'input_size': tune.choice([12, 24]), \n",
    "          'hidden_size': 256,\n",
    "          'max_steps': 1,\n",
    "          'val_check_steps': 1,\n",
    "          'step_size': 12}\n",
    "\n",
    "config_drnn = {'input_size': tune.choice([-1]), \n",
    "               'encoder_hidden_size': tune.choice([5, 10]),\n",
    "               'max_steps': 1,\n",
    "               'val_check_steps': 1,\n",
    "               'step_size': 1}\n",
    "\n",
    "fcst = NeuralForecast(\n",
    "    models=[\n",
    "        DilatedRNN(h=12, input_size=-1, encoder_hidden_size=10, max_steps=1),\n",
    "        AutoMLP(h=12, config=config, cpus=1, num_samples=1),\n",
    "        NHITS(h=12, input_size=12, max_steps=1)\n",
    "    ],\n",
    "    freq='M'\n",
    ")\n",
    "cv_df = fcst.cross_validation(df=AirPassengersPanel, static_df=AirPassengersStatic, n_windows=3, step_size=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17c5ea12-ed87-4e46-ad04-3088e7167dfd",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "#test cross validation no leakage\n",
    "def test_cross_validation(df, static_df, h, test_size):\n",
    "    if (test_size - h) % 1:\n",
    "        raise Exception(\"`test_size - h` should be module `step_size`\")\n",
    "    \n",
    "    n_windows = int((test_size - h) / 1) + 1\n",
    "    Y_test_df = df.groupby('unique_id').tail(test_size)\n",
    "    Y_train_df = df.drop(Y_test_df.index)\n",
    "    config = {'input_size': tune.choice([12, 24]),\n",
    "              'step_size': 12, 'hidden_size': 256, 'max_steps': 1, 'val_check_steps': 1}\n",
    "    config_drnn = {'input_size': tune.choice([-1]), 'encoder_hidden_size': tune.choice([5, 10]),\n",
    "                   'max_steps': 1, 'val_check_steps': 1}\n",
    "    fcst = NeuralForecast(\n",
    "        models=[\n",
    "            AutoDilatedRNN(h=12, config=config_drnn, cpus=1, num_samples=1),\n",
    "            DilatedRNN(h=12, input_size=-1, encoder_hidden_size=5, max_steps=1),\n",
    "            RNN(h=12, input_size=-1, encoder_hidden_size=5, max_steps=1,\n",
    "                stat_exog_list=['airline1'], futr_exog_list=['trend'], hist_exog_list=['y_[lag12]']),\n",
    "            TCN(h=12, input_size=-1, encoder_hidden_size=5, max_steps=1,\n",
    "                stat_exog_list=['airline1'], futr_exog_list=['trend'], hist_exog_list=['y_[lag12]']),\n",
    "            AutoMLP(h=12, config=config, cpus=1, num_samples=1),\n",
    "            MLP(h=12, input_size=12, max_steps=1, scaler_type='robust'),\n",
    "            NBEATSx(h=12, input_size=12, max_steps=1,\n",
    "                    stat_exog_list=['airline1'], futr_exog_list=['trend'], hist_exog_list=['y_[lag12]']),\n",
    "            NHITS(h=12, input_size=12, max_steps=1, scaler_type='robust'),\n",
    "            NHITS(h=12, input_size=12, loss=MQLoss(level=[80]), max_steps=1),\n",
    "            TFT(h=12, input_size=24, max_steps=1, scaler_type='robust'),\n",
    "            DLinear(h=12, input_size=24, max_steps=1),\n",
    "            VanillaTransformer(h=12, input_size=12, max_steps=1, scaler_type=None),\n",
    "            Informer(h=12, input_size=12, max_steps=1, scaler_type=None),\n",
    "            Autoformer(h=12, input_size=12, max_steps=1, scaler_type=None),\n",
    "            FEDformer(h=12, input_size=12, max_steps=1, scaler_type=None),\n",
    "            PatchTST(h=12, input_size=24, max_steps=1, scaler_type=None),\n",
    "            TimesNet(h=12, input_size=24, max_steps=1, scaler_type='standard'),\n",
    "            StemGNN(h=12, input_size=12, n_series=2, max_steps=1, scaler_type='robust'),\n",
    "            DeepAR(h=12, input_size=24, max_steps=1,\n",
    "               stat_exog_list=['airline1'], futr_exog_list=['trend']),\n",
    "        ],\n",
    "        freq='M'\n",
    "    )\n",
    "    fcst.fit(df=Y_train_df, static_df=static_df)\n",
    "    Y_hat_df = fcst.predict(futr_df=Y_test_df)\n",
    "    Y_hat_df = Y_hat_df.merge(Y_test_df, how='left', on=['unique_id', 'ds'])\n",
    "    last_dates = Y_train_df.groupby('unique_id').tail(1)\n",
    "    last_dates = last_dates[['unique_id', 'ds']].rename(columns={'ds': 'cutoff'})\n",
    "    Y_hat_df = Y_hat_df.merge(last_dates, how='left', on='unique_id')\n",
    "    \n",
    "    #cross validation\n",
    "    fcst = NeuralForecast(\n",
    "        models=[\n",
    "            AutoDilatedRNN(h=12, config=config_drnn, cpus=1, num_samples=1),\n",
    "            DilatedRNN(h=12, input_size=-1, encoder_hidden_size=5, max_steps=1),\n",
    "            RNN(h=12, input_size=-1, encoder_hidden_size=5, max_steps=1,\n",
    "                stat_exog_list=['airline1'], futr_exog_list=['trend'], hist_exog_list=['y_[lag12]']),\n",
    "            TCN(h=12, input_size=-1, encoder_hidden_size=5, max_steps=1,\n",
    "                stat_exog_list=['airline1'], futr_exog_list=['trend'], hist_exog_list=['y_[lag12]']),\n",
    "            AutoMLP(h=12, config=config, cpus=1, num_samples=1),\n",
    "            MLP(h=12, input_size=12, max_steps=1, scaler_type='robust'),\n",
    "            NBEATSx(h=12, input_size=12, max_steps=1,\n",
    "                    stat_exog_list=['airline1'], futr_exog_list=['trend'], hist_exog_list=['y_[lag12]']),\n",
    "            NHITS(h=12, input_size=12, max_steps=1, scaler_type='robust'),\n",
    "            NHITS(h=12, input_size=12, loss=MQLoss(level=[80]), max_steps=1),\n",
    "            TFT(h=12, input_size=24, max_steps=1, scaler_type='robust'),\n",
    "            DLinear(h=12, input_size=24, max_steps=1),\n",
    "            VanillaTransformer(h=12, input_size=12, max_steps=1, scaler_type=None),\n",
    "            Informer(h=12, input_size=12, max_steps=1, scaler_type=None),\n",
    "            Autoformer(h=12, input_size=12, max_steps=1, scaler_type=None),\n",
    "            FEDformer(h=12, input_size=12, max_steps=1, scaler_type=None),\n",
    "            PatchTST(h=12, input_size=24, max_steps=1, scaler_type=None),\n",
    "            TimesNet(h=12, input_size=24, max_steps=1, scaler_type='standard'),\n",
    "            StemGNN(h=12, input_size=12, n_series=2, max_steps=1, scaler_type='robust'),\n",
    "            DeepAR(h=12, input_size=24, max_steps=1,\n",
    "               stat_exog_list=['airline1'], futr_exog_list=['trend']),\n",
    "        ],\n",
    "        freq='M'\n",
    "    )\n",
    "    Y_hat_df_cv = fcst.cross_validation(df, static_df=static_df, test_size=test_size, \n",
    "                                        n_windows=None)\n",
    "    for col in ['ds', 'cutoff']:\n",
    "        Y_hat_df_cv[col] = pd.to_datetime(Y_hat_df_cv[col].astype(str))\n",
    "        Y_hat_df[col] = pd.to_datetime(Y_hat_df[col].astype(str))\n",
    "    pd.testing.assert_frame_equal(\n",
    "        Y_hat_df[Y_hat_df_cv.columns],\n",
    "        Y_hat_df_cv,\n",
    "        check_dtype=False,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0467904-748e-42ec-99cc-bac514626304",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "test_cross_validation(AirPassengersPanel, AirPassengersStatic, h=12, test_size=12)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f61d030-a51e-49f8-a16e-b97bccb62401",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test save and load\n",
    "config = {'input_size': tune.choice([12, 24]), \n",
    "          'hidden_size': 256,\n",
    "          'max_steps': 1,\n",
    "          'val_check_steps': 1,\n",
    "          'step_size': 12}\n",
    "\n",
    "config_drnn = {'input_size': tune.choice([-1]),\n",
    "               'encoder_hidden_size': tune.choice([5, 10]),\n",
    "               'max_steps': 1,\n",
    "               'val_check_steps': 1}\n",
    "\n",
    "fcst = NeuralForecast(\n",
    "    models=[\n",
    "        AutoRNN(h=12, config=config_drnn, cpus=1, num_samples=2, refit_with_val=True),\n",
    "        DilatedRNN(h=12, input_size=-1, encoder_hidden_size=5, max_steps=1),\n",
    "        AutoMLP(h=12, config=config, cpus=1, num_samples=2),\n",
    "        NHITS(h=12, input_size=12, max_steps=1,\n",
    "              futr_exog_list=['trend'], hist_exog_list=['y_[lag12]'], alias='Model1'),\n",
    "        StemGNN(h=12, input_size=12, n_series=2, max_steps=1, scaler_type='robust')\n",
    "    ],\n",
    "    freq='M'\n",
    ")\n",
    "fcst.fit(AirPassengersPanel_train)\n",
    "forecasts1 = fcst.predict(futr_df=AirPassengersPanel_test)\n",
    "fcst.save(path='./examples/debug_run/', model_index=None, overwrite=True, save_dataset=True)\n",
    "\n",
    "fcst2 = NeuralForecast.load(path='./examples/debug_run/')\n",
    "forecasts2 = fcst2.predict(futr_df=AirPassengersPanel_test)\n",
    "pairwise_tuples = [('AutoRNN', 'RNN'), ('DilatedRNN','DilatedRNN'), ('AutoMLP','MLP'), ('Model1','Model1'), ('StemGNN','StemGNN')]\n",
    "for model1, model2 in pairwise_tuples:\n",
    "    np.testing.assert_allclose(forecasts1[model1], forecasts2[model2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d221d90-7a89-4338-96b8-09543f3d5554",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test save and load without dataset\n",
    "shutil.rmtree('examples/debug_run')\n",
    "fcst = NeuralForecast(\n",
    "    models=[DilatedRNN(h=12, input_size=-1, encoder_hidden_size=5, max_steps=1)],\n",
    "    freq='M',\n",
    ")\n",
    "fcst.fit(AirPassengersPanel_train)\n",
    "forecasts1 = fcst.predict(futr_df=AirPassengersPanel_test)\n",
    "fcst.save(path='./examples/debug_run/', model_index=None, overwrite=True, save_dataset=False)\n",
    "fcst2 = NeuralForecast.load(path='./examples/debug_run/')\n",
    "forecasts2 = fcst2.predict(df=AirPassengersPanel_train, futr_df=AirPassengersPanel_test)\n",
    "np.testing.assert_allclose(forecasts1['DilatedRNN'], forecasts2['DilatedRNN'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c22ad495",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test `enable_checkpointing=True` should generate chkpt\n",
    "shutil.rmtree('lightning_logs')\n",
    "fcst = NeuralForecast(\n",
    "    models=[\n",
    "        MLP(h=12, input_size=12, max_steps=10, val_check_steps=5, enable_checkpointing=True),\n",
    "        RNN(h=12, input_size=-1, max_steps=10, val_check_steps=5, enable_checkpointing=True)\n",
    "    ],\n",
    "    freq='M'\n",
    ")\n",
    "fcst.fit(AirPassengersPanel_train)\n",
    "last_log = f\"lightning_logs/{os.listdir('lightning_logs')[-1]}\"\n",
    "no_chkpt_found = ~np.any([file.endswith('checkpoints') for file in os.listdir(last_log)])\n",
    "test_eq(no_chkpt_found, False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ac7a0b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test `enable_checkpointing=False` should not generate chkpt\n",
    "shutil.rmtree('lightning_logs')\n",
    "fcst = NeuralForecast(\n",
    "    models=[\n",
    "        MLP(h=12, input_size=12, max_steps=10, val_check_steps=5),\n",
    "        RNN(h=12, input_size=-1, max_steps=10, val_check_steps=5)\n",
    "    ],\n",
    "    freq='M'\n",
    ")\n",
    "fcst.fit(AirPassengersPanel_train)\n",
    "last_log = f\"lightning_logs/{os.listdir('lightning_logs')[-1]}\"\n",
    "no_chkpt_found = ~np.any([file.endswith('checkpoints') for file in os.listdir(last_log)])\n",
    "test_eq(no_chkpt_found, True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8602941e",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test short time series\n",
    "config = {'input_size': tune.choice([12, 24]), \n",
    "          'max_steps': 1,\n",
    "          'val_check_steps': 1}\n",
    "\n",
    "fcst = NeuralForecast(\n",
    "    models=[\n",
    "        AutoNBEATS(h=12, config=config, cpus=1, num_samples=2)],\n",
    "    freq='M'\n",
    ")\n",
    "\n",
    "AirPassengersShort = AirPassengersPanel.tail(36+144).reset_index(drop=True)\n",
    "forecasts = fcst.cross_validation(AirPassengersShort, val_size=48, n_windows=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cadac88d",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test validation scale BaseWindows\n",
    "\n",
    "models = [NHITS(h=12, input_size=24, max_steps=50, scaler_type='robust')]\n",
    "nf = NeuralForecast(models=models, freq='M')\n",
    "nf.fit(AirPassengersPanel_train,val_size=12)\n",
    "valid_losses = nf.models[0].valid_trajectories\n",
    "assert valid_losses[-1][1] < 40, 'Validation loss is too high'\n",
    "assert valid_losses[-1][1] > 10, 'Validation loss is too low'\n",
    "\n",
    "models = [NHITS(h=12, input_size=24, max_steps=50, scaler_type=None)]\n",
    "nf = NeuralForecast(models=models, freq='M')\n",
    "nf.fit(AirPassengersPanel_train,val_size=12)\n",
    "valid_losses = nf.models[0].valid_trajectories\n",
    "assert valid_losses[-1][1] < 40, 'Validation loss is too high'\n",
    "assert valid_losses[-1][1] > 10, 'Validation loss is too low'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee083d85",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test validation scale BaseRecurrent\n",
    "\n",
    "nf = NeuralForecast(\n",
    "    models=[LSTM(h=12,\n",
    "                 input_size=-1,\n",
    "                 loss=MAE(),\n",
    "                 scaler_type='robust',\n",
    "                 encoder_n_layers=2,\n",
    "                 encoder_hidden_size=128,\n",
    "                 context_size=10,\n",
    "                 decoder_hidden_size=128,\n",
    "                 decoder_layers=2,\n",
    "                 max_steps=50,\n",
    "                 val_check_steps=10,\n",
    "                 )\n",
    "    ],\n",
    "    freq='M'\n",
    ")\n",
    "nf.fit(AirPassengersPanel_train,val_size=12)\n",
    "valid_losses = nf.models[0].valid_trajectories\n",
    "assert valid_losses[-1][1] < 100, 'Validation loss is too high'\n",
    "assert valid_losses[-1][1] > 30, 'Validation loss is too low'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02ee53b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# Test order of variables does not affect validation loss\n",
    "\n",
    "AirPassengersPanel_train['zeros'] = 0\n",
    "AirPassengersPanel_train['large_number'] = 100000\n",
    "AirPassengersPanel_train['available_mask'] = 1\n",
    "AirPassengersPanel_train = AirPassengersPanel_train[['unique_id','ds','zeros','y','available_mask','large_number']]\n",
    "\n",
    "models = [NHITS(h=12, input_size=24, max_steps=50, scaler_type='robust')]\n",
    "nf = NeuralForecast(models=models, freq='M')\n",
    "nf.fit(AirPassengersPanel_train,val_size=12)\n",
    "valid_losses = nf.models[0].valid_trajectories\n",
    "assert valid_losses[-1][1] < 40, 'Validation loss is too high'\n",
    "assert valid_losses[-1][1] > 10, 'Validation loss is too low'\n",
    "\n",
    "models = [NHITS(h=12, input_size=24, max_steps=50, scaler_type=None)]\n",
    "nf = NeuralForecast(models=models, freq='M')\n",
    "nf.fit(AirPassengersPanel_train,val_size=12)\n",
    "valid_losses = nf.models[0].valid_trajectories\n",
    "assert valid_losses[-1][1] < 40, 'Validation loss is too high'\n",
    "assert valid_losses[-1][1] > 10, 'Validation loss is too low'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ba31378",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# Test fit fails if variable not in dataframe\n",
    "\n",
    "# Base Windows\n",
    "models = [NHITS(h=12, input_size=24, max_steps=1, hist_exog_list=['not_included'], scaler_type='robust')]\n",
    "nf = NeuralForecast(models=models, freq='M')\n",
    "test_fail(nf.fit,\n",
    "          contains='historical exogenous variables not found in input dataset',\n",
    "          args=(AirPassengersPanel_train,))\n",
    "\n",
    "models = [NHITS(h=12, input_size=24, max_steps=1, futr_exog_list=['not_included'], scaler_type='robust')]\n",
    "nf = NeuralForecast(models=models, freq='M')\n",
    "test_fail(nf.fit,\n",
    "          contains='future exogenous variables not found in input dataset',\n",
    "          args=(AirPassengersPanel_train,))\n",
    "\n",
    "models = [NHITS(h=12, input_size=24, max_steps=1, stat_exog_list=['not_included'], scaler_type='robust')]\n",
    "nf = NeuralForecast(models=models, freq='M')\n",
    "test_fail(nf.fit,\n",
    "          contains='static exogenous variables not found in input dataset',\n",
    "          args=(AirPassengersPanel_train,))\n",
    "\n",
    "# Base Recurrent\n",
    "models = [LSTM(h=12, input_size=24, max_steps=1, hist_exog_list=['not_included'], scaler_type='robust')]\n",
    "nf = NeuralForecast(models=models, freq='M')\n",
    "test_fail(nf.fit,\n",
    "          contains='historical exogenous variables not found in input dataset',\n",
    "          args=(AirPassengersPanel_train,))\n",
    "\n",
    "models = [LSTM(h=12, input_size=24, max_steps=1, futr_exog_list=['not_included'], scaler_type='robust')]\n",
    "nf = NeuralForecast(models=models, freq='M')\n",
    "test_fail(nf.fit,\n",
    "          contains='future exogenous variables not found in input dataset',\n",
    "          args=(AirPassengersPanel_train,))\n",
    "\n",
    "models = [LSTM(h=12, input_size=24, max_steps=1, stat_exog_list=['not_included'], scaler_type='robust')]\n",
    "nf = NeuralForecast(models=models, freq='M')\n",
    "test_fail(nf.fit,\n",
    "          contains='static exogenous variables not found in input dataset',\n",
    "          args=(AirPassengersPanel_train,))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d221479c",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# Test passing unused variables in dataframe does not affect forecasts  \n",
    "\n",
    "models = [NHITS(h=12, input_size=24, max_steps=5, hist_exog_list=['zeros'], scaler_type='robust')]\n",
    "nf = NeuralForecast(models=models, freq='M')\n",
    "nf.fit(AirPassengersPanel_train)\n",
    "\n",
    "Y_hat1 = nf.predict(df=AirPassengersPanel_train[['unique_id','ds','y','zeros','large_number']])\n",
    "Y_hat2 = nf.predict(df=AirPassengersPanel_train[['unique_id','ds','y','zeros']])\n",
    "\n",
    "pd.testing.assert_frame_equal(\n",
    "    Y_hat1,\n",
    "    Y_hat2,\n",
    "    check_dtype=False,\n",
    ")\n",
    "\n",
    "models = [LSTM(h=12, input_size=24, max_steps=5, hist_exog_list=['zeros'], scaler_type='robust')]\n",
    "nf = NeuralForecast(models=models, freq='M')\n",
    "nf.fit(AirPassengersPanel_train)\n",
    "\n",
    "Y_hat1 = nf.predict(df=AirPassengersPanel_train[['unique_id','ds','y','zeros','large_number']])\n",
    "Y_hat2 = nf.predict(df=AirPassengersPanel_train[['unique_id','ds','y','zeros']])\n",
    "\n",
    "pd.testing.assert_frame_equal(\n",
    "    Y_hat1,\n",
    "    Y_hat2,\n",
    "    check_dtype=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30890c07-1763-4795-afba-f5ed916245be",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "#| polars\n",
    "import polars"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7502d18c-62a5-4381-bfdb-bba5300c5290",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "#| polars\n",
    "models = [LSTM(h=12, input_size=24, max_steps=5, hist_exog_list=['zeros'], scaler_type='robust')]\n",
    "nf = NeuralForecast(models=models, freq='M')\n",
    "nf.fit(AirPassengersPanel_train, static_df=AirPassengersStatic)\n",
    "insample_preds = nf.predict_insample()\n",
    "preds = nf.predict()\n",
    "cv_res = nf.cross_validation(df=AirPassengersPanel_train, static_df=AirPassengersStatic)\n",
    "\n",
    "renamer = {'unique_id': 'uid', 'ds': 'time', 'y': 'target'}\n",
    "inverse_renamer = {v: k for k, v in renamer.items()}\n",
    "AirPassengers_pl = polars.from_pandas(AirPassengersPanel_train)\n",
    "AirPassengers_pl = AirPassengers_pl.rename(renamer)\n",
    "AirPassengersStatic_pl = polars.from_pandas(AirPassengersStatic)\n",
    "AirPassengersStatic_pl = AirPassengersStatic_pl.rename({'unique_id': 'uid'})\n",
    "nf = NeuralForecast(models=models, freq='1mo')\n",
    "nf.fit(\n",
    "    AirPassengers_pl,\n",
    "    static_df=AirPassengersStatic_pl,\n",
    "    id_col='uid',\n",
    "    time_col='time',\n",
    "    target_col='target',\n",
    ")\n",
    "insample_preds_pl = nf.predict_insample()\n",
    "preds_pl = nf.predict()\n",
    "cv_res_pl = nf.cross_validation(\n",
    "    df=AirPassengers_pl,\n",
    "    static_df=AirPassengersStatic_pl,\n",
    "    id_col='uid',\n",
    "    time_col='time',\n",
    "    target_col='target',\n",
    ")\n",
    "\n",
    "def assert_equal_dfs(pandas_df, polars_df):\n",
    "    mapping = {k: v for k, v in inverse_renamer.items() if k in polars_df}\n",
    "    pd.testing.assert_frame_equal(\n",
    "        pandas_df,\n",
    "        polars_df.rename(mapping).to_pandas(),\n",
    "    )\n",
    "\n",
    "assert_equal_dfs(preds, preds_pl)\n",
    "assert_equal_dfs(insample_preds, insample_preds_pl)\n",
    "assert_equal_dfs(cv_res, cv_res_pl)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}