{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Adding a custom time series forecasting model\n",
    "\n",
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/autogluon/autogluon/blob/master/docs/tutorials/timeseries/advanced/forecasting-custom-model.ipynb)\n",
    "[![Open In SageMaker Studio Lab](https://studiolab.sagemaker.aws/studiolab.svg)](https://studiolab.sagemaker.aws/import/github/autogluon/autogluon/blob/master/docs/tutorials/timeseries/advanced/forecasting-custom-model.ipynb)\n",
    "\n",
    "This tutorial describes how to add a custom forecasting model that can be trained, hyperparameter-tuned, and ensembled alongside the [default forecasting models](../forecasting-model-zoo.html).\n",
    "\n",
    "As an example, we will implement an AutoGluon wrapper for the [NHITS](https://nixtlaverse.nixtla.io/neuralforecast/models.nhits.html#nhits) model from the [NeuralForecast](https://github.com/Nixtla/NeuralForecast) library.\n",
    "\n",
    "This tutorial consists of the following sections:\n",
    "1. Implementing the model wrapper.\n",
    "2. Loading & preprocessing the dataset used for model development.\n",
    "3. Using the custom model in standalone mode.\n",
    "4. Using the custom model inside the `TimeSeriesPredictor`.\n",
    "\n",
    "\n",
    "```{warning}\n",
    "\n",
    "This tutorial is designed for advanced AutoGluon users.\n",
    "\n",
    "Custom model implementations rely heavily on the private of API of AutoGluon that might change over time. For this reason, it might be necessary to update your custom model implementations as you upgrade to new versions of AutoGluon.\n",
    "\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-cell",
     "skip-execution"
    ]
   },
   "outputs": [],
   "source": [
    "# We use uv for faster installation\n",
    "!pip install uv\n",
    "!uv pip install -q autogluon.timeseries --system\n",
    "!uv pip uninstall -q torchaudio torchvision torchtext --system # fix incompatible package versions on Colab"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we install the NeuralForecast library that contains the implementation of the custom model used in this tutorial."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pip install -q neuralforecast==2.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Implement the custom model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "To implement a custom model we need to create a subclass of the [`AbstractTimeSeriesModel`](https://github.com/autogluon/autogluon/blob/master/timeseries/src/autogluon/timeseries/models/abstract/abstract_timeseries_model.py) class. This subclass must implement two methods: `_fit` and `_predict`. For models that require a custom preprocessing logic (e.g., to handle missing values), we also need to implement the `preprocess` method.\n",
    "\n",
    "Please have a look at the following code and read the comments to understand the different components of the custom model wrapper."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "import pprint\n",
    "from typing import Optional, Tuple\n",
    "\n",
    "import pandas as pd\n",
    "\n",
    "from autogluon.timeseries import TimeSeriesDataFrame\n",
    "from autogluon.timeseries.models.abstract import AbstractTimeSeriesModel\n",
    "from autogluon.timeseries.utils.warning_filters import warning_filter\n",
    "\n",
    "# Optional - disable annoying PyTorch-Lightning loggers\n",
    "for logger_name in [\n",
    "    \"lightning.pytorch.utilities.rank_zero\",\n",
    "    \"pytorch_lightning.accelerators.cuda\",\n",
    "    \"lightning_fabric.utilities.seed\",\n",
    "]:\n",
    "    logging.getLogger(logger_name).setLevel(logging.ERROR)\n",
    "\n",
    "\n",
    "class NHITSModel(AbstractTimeSeriesModel):\n",
    "    \"\"\"AutoGluon-compatible wrapper for the NHITS model from NeuralForecast.\"\"\"\n",
    "\n",
    "    # Set these attributes to ensure that AutoGluon passes correct features to the model\n",
    "    _supports_known_covariates: bool = True\n",
    "    _supports_past_covariates: bool = True\n",
    "    _supports_static_features: bool = True\n",
    "\n",
    "    def preprocess(\n",
    "        self,\n",
    "        data: TimeSeriesDataFrame,\n",
    "        known_covariates: Optional[TimeSeriesDataFrame] = None,\n",
    "        is_train: bool = False,\n",
    "        **kwargs,\n",
    "    ) -> Tuple[TimeSeriesDataFrame, Optional[TimeSeriesDataFrame]]:\n",
    "        \"\"\"Method that implements model-specific preprocessing logic.\n",
    "\n",
    "        This method is called on all data that is passed to `_fit` and `_predict` methods.\n",
    "        \"\"\"\n",
    "        # NeuralForecast cannot handle missing values represented by NaN. Therefore, we\n",
    "        # need to impute them before the data is passed to the model. First, we\n",
    "        # forward-fill and backward-fill all time series\n",
    "        data = data.fill_missing_values()\n",
    "        # Some time series might consist completely of missing values, so the previous\n",
    "        # line has no effect on them. We fill them with 0.0\n",
    "        data = data.fill_missing_values(method=\"constant\", value=0.0)\n",
    "        # Some models (e.g., Chronos) can natively handle NaNs - for them we don't need\n",
    "        # to define a custom preprocessing logic\n",
    "        return data, known_covariates\n",
    "\n",
    "    def _get_default_hyperparameters(self) -> dict:\n",
    "        \"\"\"Default hyperparameters that will be provided to the inner model, i.e., the\n",
    "        NHITS implementation in neuralforecast. \"\"\"\n",
    "        import torch\n",
    "        from neuralforecast.losses.pytorch import MQLoss\n",
    "\n",
    "        default_hyperparameters = dict(\n",
    "            loss=MQLoss(quantiles=self.quantile_levels),\n",
    "            input_size=2 * self.prediction_length,\n",
    "            scaler_type=\"standard\",\n",
    "            enable_progress_bar=False,\n",
    "            enable_model_summary=False,\n",
    "            logger=False,\n",
    "            accelerator=\"cpu\",\n",
    "            # The model wrapper should handle any time series length - even time series\n",
    "            # with 1 observation\n",
    "            start_padding_enabled=True,\n",
    "            # NeuralForecast requires that names of the past/future/static covariates are\n",
    "            # passed as model arguments. AutoGluon models have access to this information\n",
    "            # using the `metadata` attribute that is set automatically at model creation.\n",
    "            #\n",
    "            # Note that NeuralForecast does not support categorical covariates, so we\n",
    "            # only use the real-valued covariates here. To use categorical features in\n",
    "            # you wrapper, you need to either use techniques like one-hot-encoding, or\n",
    "            # rely on models that natively handle categorical features.\n",
    "            futr_exog_list=self.covariate_metadata.known_covariates_real,\n",
    "            hist_exog_list=self.covariate_metadata.past_covariates_real,\n",
    "            stat_exog_list=self.covariate_metadata.static_features_real,\n",
    "        )\n",
    "\n",
    "        if torch.cuda.is_available():\n",
    "            default_hyperparameters[\"accelerator\"] = \"gpu\"\n",
    "            default_hyperparameters[\"devices\"] = 1\n",
    "\n",
    "        return default_hyperparameters\n",
    "\n",
    "    def _fit(\n",
    "        self,\n",
    "        train_data: TimeSeriesDataFrame,\n",
    "        val_data: Optional[TimeSeriesDataFrame] = None,\n",
    "        time_limit: Optional[float] = None,\n",
    "        **kwargs,\n",
    "    ) -> None:\n",
    "        \"\"\"Fit the model on the available training data.\"\"\"\n",
    "        print(\"Entering the `_fit` method\")\n",
    "\n",
    "        # We lazily import other libraries inside the _fit method. This reduces the\n",
    "        # import time for autogluon and ensures that even if one model has some problems\n",
    "        # with dependencies, the training process won't crash\n",
    "        from neuralforecast import NeuralForecast\n",
    "        from neuralforecast.models import NHITS\n",
    "\n",
    "        # It's important to ensure that the model respects the time_limit during `fit`.\n",
    "        # Since NeuralForecast is based on PyTorch-Lightning, this can be easily enforced\n",
    "        # using the `max_time` argument to `pl.Trainer`. For other model types such as\n",
    "        # ARIMA implementing the time_limit logic may require a lot of work.\n",
    "        hyperparameter_overrides = {}\n",
    "        if time_limit is not None:\n",
    "            hyperparameter_overrides = {\"max_time\": {\"seconds\": time_limit}}\n",
    "\n",
    "        # The method `get_hyperparameters()` returns the model hyperparameters in\n",
    "        # `_get_default_hyperparameters` overridden with the hyperparameters provided by the user in\n",
    "        # `predictor.fit(..., hyperparameters={NHITSModel: {}})`. We override these with other\n",
    "        # hyperparameters available at training time.\n",
    "        model_params = self.get_hyperparameters() | hyperparameter_overrides\n",
    "        print(f\"Hyperparameters:\\n{pprint.pformat(model_params, sort_dicts=False)}\")\n",
    "\n",
    "        model = NHITS(h=self.prediction_length, **model_params)\n",
    "        self.nf = NeuralForecast(models=[model], freq=self.freq)\n",
    "\n",
    "        # Convert data into a format expected by the model. NeuralForecast expects time\n",
    "        # series data in pandas.DataFrame format that is quite similar to AutoGluon, so\n",
    "        # the transformation is very easy.\n",
    "        #\n",
    "        # Note that the `preprocess` method was already applied to train_data and val_data.\n",
    "        train_df, static_df = self._to_neuralforecast_format(train_data)\n",
    "        self.nf.fit(\n",
    "            train_df,\n",
    "            static_df=static_df,\n",
    "            id_col=\"item_id\",\n",
    "            time_col=\"timestamp\",\n",
    "            target_col=self.target,\n",
    "        )\n",
    "        print(\"Exiting the `_fit` method\")\n",
    "\n",
    "    def _to_neuralforecast_format(self, data: TimeSeriesDataFrame) -> Tuple[pd.DataFrame, Optional[pd.DataFrame]]:\n",
    "        \"\"\"Convert a TimeSeriesDataFrame to the format expected by NeuralForecast.\"\"\"\n",
    "        df = data.to_data_frame().reset_index()\n",
    "        # Drop the categorical covariates to avoid NeuralForecast errors\n",
    "        df = df.drop(columns=self.covariate_metadata.covariates_cat)\n",
    "        static_df = data.static_features\n",
    "        if len(self.covariate_metadata.static_features_real) > 0:\n",
    "            static_df = static_df.reset_index()\n",
    "            static_df = static_df.drop(columns=self.covariate_metadata.static_features_cat)\n",
    "        return df, static_df\n",
    "\n",
    "    def _predict(\n",
    "        self,\n",
    "        data: TimeSeriesDataFrame,\n",
    "        known_covariates: Optional[TimeSeriesDataFrame] = None,\n",
    "        **kwargs,\n",
    "    ) -> TimeSeriesDataFrame:\n",
    "        \"\"\"Predict future target given the historical time series data and the future values of known_covariates.\"\"\"\n",
    "        print(\"Entering the `_predict` method\")\n",
    "\n",
    "        from neuralforecast.losses.pytorch import quantiles_to_outputs\n",
    "\n",
    "        df, static_df = self._to_neuralforecast_format(data)\n",
    "        if len(self.covariate_metadata.known_covariates_real) > 0:\n",
    "            futr_df, _ = self._to_neuralforecast_format(known_covariates)\n",
    "        else:\n",
    "            futr_df = None\n",
    "\n",
    "        with warning_filter():\n",
    "            predictions = self.nf.predict(df, static_df=static_df, futr_df=futr_df)\n",
    "\n",
    "        # predictions must be a TimeSeriesDataFrame with columns\n",
    "        # [\"mean\"] + [str(q) for q in self.quantile_levels]\n",
    "        model_name = str(self.nf.models[0])\n",
    "        rename_columns = {\n",
    "            f\"{model_name}{suffix}\": str(quantile)\n",
    "            for quantile, suffix in zip(*quantiles_to_outputs(self.quantile_levels))\n",
    "        }\n",
    "        predictions = predictions.rename(columns=rename_columns)\n",
    "        predictions[\"mean\"] = predictions[\"0.5\"]\n",
    "        predictions = TimeSeriesDataFrame(predictions)\n",
    "        return predictions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For convenience, here is an overview of the main constraints on the inputs and outputs of different methods.\n",
    "\n",
    "- Input data received by `_fit` and `_predict` methods satisfies\n",
    "    - the index is sorted by `(item_id, timestamp)`\n",
    "    - timestamps of observations have a regular frequency corresponding to `self.freq`\n",
    "    - column `self.target` contains the target values of the time series\n",
    "    - target column might contain missing values represented by `NaN`\n",
    "    - data may contain covariates (incl. static features) with schema described in `self.covariate_metadata`\n",
    "        - real-valued covariates have dtype `float32`\n",
    "        - categorical covariates have dtype `category`\n",
    "        - covariates do not contain any missing values\n",
    "    - static features, if present, are available as `data.static_features`\n",
    "- Predictions returned by `_predict` must satisfy:\n",
    "    - returns predictions as a `TimeSeriesDataFrame` object\n",
    "    - predictions contain columns `[\"mean\"] + [str(q) for q in self.quantile_levels]` containing the point and quantile forecasts, respectively\n",
    "    - the index of predictions contains exactly `self.prediction_length` future time steps of each time series present in `data`\n",
    "    - the frequency of the prediction timestamps matches `self.freq`\n",
    "    - the index of predictions is sorted by `(item_id, timestamp)`\n",
    "    - predictions contain no missing values represented by `NaN` and no gaps\n",
    "- The runtime of `_fit` method should not exceed `time_limit` seconds, if `time_limit` is provided.\n",
    "- None of the methods should modify the data in-place. If modifications are needed, create a copy of the data first.\n",
    "- All methods should work even if some time series consist of all NaNs, or only have a single observation.\n",
    "\n",
    "-----"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will now use this wrapper in two modes:\n",
    "1. Standalone mode (outside the `TimeSeriesPredictor`).\n",
    "    - This mode should be used for development and debugging. In this case, we need to take manually take care of preprocessing and model configuration.\n",
    "2. Inside the `TimeSeriesPredictor`.\n",
    "    - This mode makes it easy to combine & compare the custom model with other models available in AutoGluon. The main purpose of writing a custom model wrapper is to use it in this mode."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load and preprocess the data\n",
    "\n",
    "First, we load the Grocery Sales dataset that we will use for development and evaluation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from autogluon.timeseries import TimeSeriesDataFrame\n",
    "\n",
    "raw_data = TimeSeriesDataFrame.from_path(\n",
    "    \"https://autogluon.s3.amazonaws.com/datasets/timeseries/grocery_sales/test.csv\",\n",
    "    static_features_path=\"https://autogluon.s3.amazonaws.com/datasets/timeseries/grocery_sales/static.csv\",\n",
    ")\n",
    "raw_data.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "raw_data.static_features.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Types of the columns in raw data:\")\n",
    "print(raw_data.dtypes)\n",
    "print(\"\\nTypes of the columns in raw static features:\")\n",
    "print(raw_data.static_features.dtypes)\n",
    "\n",
    "print(\"\\nNumber of missing values per column:\")\n",
    "print(raw_data.isna().sum())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define the forecasting task"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prediction_length = 7  # number of future steps to predict\n",
    "target = \"unit_sales\"  # target column\n",
    "known_covariates_names = [\"promotion_email\", \"promotion_homepage\"]  # covariates known in the future"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before we use the model in standalone mode, we need to apply the general AutoGluon preprocessing to the data.\n",
    "\n",
    "The `TimeSeriesFeatureGenerator` captures preprocessing steps like normalizing the data types and imputing the missing values in the covariates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from autogluon.timeseries.utils.features import TimeSeriesFeatureGenerator\n",
    "\n",
    "feature_generator = TimeSeriesFeatureGenerator(target=target, known_covariates_names=known_covariates_names)\n",
    "data = feature_generator.fit_transform(raw_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Types of the columns in preprocessed data:\")\n",
    "print(data.dtypes)\n",
    "print(\"\\nTypes of the columns in preprocessed static features:\")\n",
    "print(data.static_features.dtypes)\n",
    "\n",
    "print(\"\\nNumber of missing values per column:\")\n",
    "print(data.isna().sum())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using the custom model in standalone mode\n",
    "Using the model in standalone mode is useful for debugging our implementation. Once we make sure that all methods work as expected, we will use the model inside the `TimeSeriesPredictor`.\n",
    "\n",
    "### Training\n",
    "We are now ready to train the custom model on the preprocessed data.\n",
    "\n",
    "When using the model in standalone mode, we need to manually configure its parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = NHITSModel(\n",
    "    prediction_length=prediction_length,\n",
    "    target=target,\n",
    "    covariate_metadata=feature_generator.covariate_metadata,\n",
    "    freq=data.freq,\n",
    "    quantile_levels=[0.1, 0.5, 0.9],\n",
    ")\n",
    "model.fit(train_data=data, time_limit=20)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Predicting and scoring"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "past_data, known_covariates = data.get_model_inputs_for_scoring(\n",
    "    prediction_length=prediction_length,\n",
    "    known_covariates_names=known_covariates_names,\n",
    ")\n",
    "predictions = model.predict(past_data, known_covariates)\n",
    "predictions.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.score(data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using the custom model inside the `TimeSeriesPredictor`\n",
    "After we made sure that our custom model works in standalone mode, we can pass it to the TimeSeriesPredictor alongside other models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from autogluon.timeseries import TimeSeriesPredictor\n",
    "\n",
    "train_data, test_data = raw_data.train_test_split(prediction_length)\n",
    "\n",
    "predictor = TimeSeriesPredictor(\n",
    "    prediction_length=prediction_length,\n",
    "    target=target,\n",
    "    known_covariates_names=known_covariates_names,\n",
    ")\n",
    "\n",
    "predictor.fit(\n",
    "    train_data,\n",
    "    hyperparameters={\n",
    "        \"Naive\": {},\n",
    "        \"Chronos\": {\"model_path\": \"bolt_small\"},\n",
    "        \"ETS\": {},\n",
    "        NHITSModel: {},\n",
    "    },\n",
    "    time_limit=120,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that when we use the custom model inside the predictor, we don't need to worry about:\n",
    "- manually configuring the model (setting `freq`, `prediction_length`)\n",
    "- preprocessing the data using `TimeSeriesFeatureGenerator`\n",
    "- setting the time limits\n",
    "\n",
    "The `TimeSeriesPredictor` automatically takes care of all above aspects.\n",
    "\n",
    "We can also easily compare our custom model with other model trained by the predictor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "predictor.leaderboard(test_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also take advantage of other predictor functionality such as `feature_importance`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "predictor.feature_importance(test_data, model=\"NHITS\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As expected, features `product_category` and `product_subcategory` have zero importance because our implementation ignores categorical features.\n",
    "\n",
    "-----"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is how we can train multiple versions of the custom model with different hyperparameter configurations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "predictor = TimeSeriesPredictor(\n",
    "    prediction_length=prediction_length,\n",
    "    target=target,\n",
    "    known_covariates_names=known_covariates_names,\n",
    ")\n",
    "predictor.fit(\n",
    "    train_data,\n",
    "    hyperparameters={\n",
    "        NHITSModel: [\n",
    "            {},  # default hyperparameters\n",
    "            {\"input_size\": 20},  # custom input_size\n",
    "            {\"scaler_type\": \"robust\"},  # custom scaler_type\n",
    "        ]\n",
    "    },\n",
    "    time_limit=60,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "predictor.leaderboard(test_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Wrapping up\n",
    "\n",
    "That's all it takes to add a custom forecasting model to AutoGluon. If you create a custom model, consider [submitting a PR](https://github.com/autogluon/autogluon/pulls) so that we can add it officially to AutoGluon!\n",
    "\n",
    "For more tutorials, refer to [Forecasting Time Series - Quick Start](../forecasting-quick-start.ipynb) and [Forecasting Time Series - In Depth](../forecasting-indepth.ipynb)."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ag",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}