{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TimesFM with Covariates\n",
    "\n",
    "This toturial notebook demonstrates how to utilize exogenous covariates with TimesFM when making forecasts. Before running this notebook, make sure:\n",
    "\n",
    "- You've read through the README of TimesFM.\n",
    "- A local kernel with Python 3.10 is up and running, for the jax version.\n",
    "- Install the JAX version following the installation instructions."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup the environment and install TimesFM."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load the checkpoint\n",
    "\n",
    "**Notice:** Please set up the backend as per your machine (\"cpu\", \"gpu\" or \"tpu\"). This notebook will run by default on GPU.\n",
    "\n",
    "We load the 2.0-500m model checkpoint from HuggingFace."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import timesfm\n",
    "timesfm_backend = \"gpu\"  # @param\n",
    "\n",
    "model = timesfm.TimesFm(\n",
    "      hparams=timesfm.TimesFmHparams(\n",
    "          backend=timesfm_backend,\n",
    "          per_core_batch_size=32,\n",
    "          horizon_len=128,\n",
    "          num_layers=50,\n",
    "          use_positional_embedding=False,\n",
    "          context_len=2048,\n",
    "      ),\n",
    "      checkpoint=timesfm.TimesFmCheckpoint(\n",
    "          huggingface_repo_id=\"google/timesfm-2.0-500m-jax\"),\n",
    "  )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Covariates\n",
    "\n",
    "Let's take a toy example of forecasting sales for a grocery store: \n",
    "\n",
    "**Task:** Given the observed the daily sales of this week (7 days), forecast the daily sales of next week (7 days).\n",
    "\n",
    "```\n",
    "Product: ice cream\n",
    "Daily_sales: [30, 30, 4, 5, 7, 8, 10]\n",
    "Category: food\n",
    "Base_price: 1.99\n",
    "Weekday: [0, 1, 2, 3, 4, 5, 6, 0, 1, 2, 3, 4, 5, 6]\n",
    "Has_promotion: [Yes, Yes, No, No, No, Yes, Yes, No, No, No, No, No, No, No]\n",
    "Daily_temperature: [31.0, 24.3, 19.4, 26.2, 24.6, 30.0, 31.1, 32.4, 30.9, 26.0, 25.0, 27.8, 29.5, 31.2]\n",
    "```\n",
    "\n",
    "```\n",
    "Product: sunscreen\n",
    "Daily_sales: [5, 7, 12, 13, 5, 6, 10]\n",
    "Category: skin product\n",
    "Base_price: 29.99\n",
    "Weekday: [0, 1, 2, 3, 4, 5, 6, 0, 1, 2, 3, 4, 5, 6]\n",
    "Has_promotion: [No, No, Yes, Yes, No, No, No, Yes, Yes, Yes, Yes, Yes, Yes, Yes]\n",
    "Daily_temperature: [31.0, 24.3, 19.4, 26.2, 24.6, 30.0, 31.1, 32.4, 30.9, 26.0, 25.0, 27.8, 29.5, 31.2]\n",
    "```\n",
    "\n",
    "In this example, besides the `Daily_sales`, we also have covariates `Category`, `Base_price`, `Weekday`, `Has_promotion`, `Daily_temperature`. Let's introduce some concepts:\n",
    "\n",
    "**Static covariates** are covariates for each time series. \n",
    "- In our example, `Category` is a **static categorical covariate**, \n",
    "- `Base_price` is a **static numerical covariates**.\n",
    "\n",
    "**Dynamic covariates** are covaraites for each time stamps.\n",
    "- Date / time related features can be usually treated as dynamic covariates.\n",
    "- In our example, `Weekday` and `Has_promotion` are **dynamic categorical covariates**.\n",
    "- `Daily_temperate` is a **dynamic numerical covariate**.\n",
    "\n",
    "**Notice:** Here we make it mandatory that the dynamic covariates need to cover both the forecasting context and horizon. For example, all dynamic covariates in the example have 14 values: the first 7 correspond to the observed 7 days, and the last 7 correspond to the next 7 days."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TimesFM with Covariates\n",
    "\n",
    "\n",
    "The strategy we take here is to treat covariates as batched in-context exogenous regressors (XReg) and fit linear models on them outside of TimesFM. The final forecast will be the sum of the TimesFM forecast and the linear model forecast.\n",
    "\n",
    " In simple words, we consider these two options.\n",
    "\n",
    "**Option 1:** Get the TimesFM forecast, and fit the linear model regressing the residuals on the covariates (\"timesfm + xreg\").\n",
    "\n",
    "**Option 2:** Fit the linear model of the time series itself on the covariates, then forecast the residuals using TimesFM  (\"xreg + timesfm\").\n",
    "\n",
    "Let's take a code at the example of Electricity Price Forecasting (EPF). \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "from collections import defaultdict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.read_csv('https://datasets-nixtla.s3.amazonaws.com/EPF_FR_BE.csv')\n",
    "df['ds'] = pd.to_datetime(df['ds'])\n",
    "df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This dataset has a few covariates beside the hourly target `y`:\n",
    "\n",
    "- `unique_id`: a static categorical covariate indicating the country.\n",
    "- `gen_forecast`: a dynamic numerical covariate indicating the estimated electricity to be generated.\n",
    "- `system_load`: the observed system load. Notice that this **CANNOT** be considered as a dynamic numerical covariate because we cannot know its values over the forecasting horizon in advance.\n",
    "- `weekday`: a dynamic categorical covariate.\\\n",
    "\n",
    "Let's now make some forecasting tasks for TimesFM based on this dataset. For simplicity we create forecast contexts of 120 time points (hours) and forecast horizons of 24 time points."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Data pipelining\n",
    "def get_batched_data_fn(\n",
    "    batch_size: int = 128, \n",
    "    context_len: int = 120, \n",
    "    horizon_len: int = 24,\n",
    "):\n",
    "  examples = defaultdict(list)\n",
    "\n",
    "  num_examples = 0\n",
    "  for country in (\"FR\", \"BE\"):\n",
    "    sub_df = df[df[\"unique_id\"] == country]\n",
    "    for start in range(0, len(sub_df) - (context_len + horizon_len), horizon_len):\n",
    "      num_examples += 1\n",
    "      examples[\"country\"].append(country)\n",
    "      examples[\"inputs\"].append(sub_df[\"y\"][start:(context_end := start + context_len)].tolist())\n",
    "      examples[\"gen_forecast\"].append(sub_df[\"gen_forecast\"][start:context_end + horizon_len].tolist())\n",
    "      examples[\"week_day\"].append(sub_df[\"week_day\"][start:context_end + horizon_len].tolist())\n",
    "      examples[\"outputs\"].append(sub_df[\"y\"][context_end:(context_end + horizon_len)].tolist())\n",
    "  \n",
    "  def data_fn():\n",
    "    for i in range(1 + (num_examples - 1) // batch_size):\n",
    "      yield {k: v[(i * batch_size) : ((i + 1) * batch_size)] for k, v in examples.items()}\n",
    "  \n",
    "  return data_fn\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define metrics\n",
    "def mse(y_pred, y_true):\n",
    "  y_pred = np.array(y_pred)\n",
    "  y_true = np.array(y_true)\n",
    "  return np.mean(np.square(y_pred - y_true), axis=1, keepdims=True)\n",
    "\n",
    "def mae(y_pred, y_true):\n",
    "  y_pred = np.array(y_pred)\n",
    "  y_true = np.array(y_true)\n",
    "  return np.mean(np.abs(y_pred - y_true), axis=1, keepdims=True)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's try `model.forecast_with_covariates`. \n",
    "\n",
    "In particular, the output is a tuple whose first element is the new forecast."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "# Benchmark\n",
    "batch_size = 128\n",
    "context_len = 120\n",
    "horizon_len = 24\n",
    "input_data = get_batched_data_fn(batch_size = 128)\n",
    "metrics = defaultdict(list)\n",
    "\n",
    "\n",
    "for i, example in enumerate(input_data()):\n",
    "  raw_forecast, _ = model.forecast(\n",
    "      inputs=example[\"inputs\"], freq=[0] * len(example[\"inputs\"])\n",
    "  )\n",
    "  start_time = time.time()\n",
    "  # Forecast with covariates\n",
    "  # Output: new forecast, forecast by the xreg\n",
    "  cov_forecast, ols_forecast = model.forecast_with_covariates(  \n",
    "      inputs=example[\"inputs\"],\n",
    "      dynamic_numerical_covariates={\n",
    "          \"gen_forecast\": example[\"gen_forecast\"],\n",
    "      },\n",
    "      dynamic_categorical_covariates={\n",
    "          \"week_day\": example[\"week_day\"],\n",
    "      },\n",
    "      static_numerical_covariates={},\n",
    "      static_categorical_covariates={\n",
    "          \"country\": example[\"country\"]\n",
    "      },\n",
    "      freq=[0] * len(example[\"inputs\"]),\n",
    "      xreg_mode=\"xreg + timesfm\",              # default\n",
    "      ridge=0.0,\n",
    "      force_on_cpu=False,\n",
    "      normalize_xreg_target_per_input=True,    # default\n",
    "  )\n",
    "  print(\n",
    "      f\"\\rFinished batch {i} linear in {time.time() - start_time} seconds\",\n",
    "      end=\"\",\n",
    "  )\n",
    "  metrics[\"eval_mae_timesfm\"].extend(\n",
    "      mae(raw_forecast[:, :horizon_len], example[\"outputs\"])\n",
    "  )\n",
    "  metrics[\"eval_mae_xreg_timesfm\"].extend(mae(cov_forecast, example[\"outputs\"]))\n",
    "  metrics[\"eval_mae_xreg\"].extend(mae(ols_forecast, example[\"outputs\"]))\n",
    "  metrics[\"eval_mse_timesfm\"].extend(\n",
    "      mse(raw_forecast[:, :horizon_len], example[\"outputs\"])\n",
    "  )\n",
    "  metrics[\"eval_mse_xreg_timesfm\"].extend(mse(cov_forecast, example[\"outputs\"]))\n",
    "  metrics[\"eval_mse_xreg\"].extend(mse(ols_forecast, example[\"outputs\"]))\n",
    "\n",
    "print()\n",
    "\n",
    "for k, v in metrics.items():\n",
    "  print(f\"{k}: {np.mean(v)}\")\n",
    "\n",
    "# My output:\n",
    "# eval_mae_timesfm: 6.762283045916956\n",
    "# eval_mae_xreg_timesfm: 5.39219617611074\n",
    "# eval_mae_xreg: 37.15275842572484\n",
    "# eval_mse_timesfm: 166.7771466306823\n",
    "# eval_mse_xreg_timesfm: 120.64757721021306\n",
    "# eval_mse_xreg: 1672.2116821201796"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You should see results close to \n",
    "```\n",
    "eval_mae_timesfm: 6.729583250571446\n",
    "eval_mae_xreg_timesfm: 5.3375301110158\n",
    "eval_mae_xreg: 37.152760709266\n",
    "eval_mse_timesfm: 162.3132151851567\n",
    "eval_mse_xreg_timesfm: 120.9900627409689\n",
    "eval_mse_xreg: 1672.208769045399\n",
    "```\n",
    "\n",
    "With the covariates, the TimesFM forecast Mean Absolute Error improves from 6.73 to 5.34, and Mean Squred Error from 162.31 to 120.99. The results of purely fitting the linear model are also provided for reference."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Formatting Your Request\n",
    "\n",
    "It is quite crucial to get the covariates properly formatted so that we can call this `model.forecast_with_covariates`. Please see its docstring for details. Here let's also grab a batch from a toy data input pipeline for quick explanations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "toy_input_pipeline = get_batched_data_fn(batch_size=2, context_len=5, horizon_len=2)\n",
    "print(next(toy_input_pipeline()))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You should see something similar to this\n",
    "```\n",
    "{\n",
    "    'country': ['FR', 'FR'], \n",
    "    'inputs': [[53.48, 51.93, 48.76, 42.27, 38.41], [48.76, 42.27, 38.41, 35.72, 32.66]], \n",
    "    'gen_forecast': [[76905.0, 75492.0, 74394.0, 72639.0, 69347.0, 67960.0, 67564.0], [74394.0, 72639.0, 69347.0, 67960.0, 67564.0, 67277.0, 67019.0]], \n",
    "    'week_day': [[3, 3, 3, 3, 3, 3, 3], [3, 3, 3, 3, 3, 3, 3]], \n",
    "    'outputs': [[35.72, 32.66], [32.83, 30.06]],\n",
    "}\n",
    "```\n",
    "\n",
    "Notice:\n",
    "- We have two examples in this batch.\n",
    "- For each example we support different context lengths and horizon lengths just as `model.forecast`. Although it is not demonstrated in this dataset.\n",
    "- If dynamic covariates are present, the horizon lengths will be inferred from them, e.g. how many values are provided in additional to the ones corresponding to the inputs. Make sure all your dynamic covariates have the same length per example.\n",
    "- The static covariates are one per example.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## More Applications\n",
    "\n",
    "### Past Dynamic Covariates\n",
    "\n",
    "Past dynamic covariates are covariates that are only available for the context. For instance in our example `system_load` is a past dynamic covariate. Time series models generally can handle this, however it is something the batched in context regression cannot address, because these regressors are not available in the future. If you do have those covariates and consider them very meaningful, there are two hacky options to try immediately:\n",
    "\n",
    "1. Shift and repeat these past dynamic covariates to use their delayed version. For example, if you think the `system_load` for this week is meaningful for forecasting next week, you can create a `delay_7_system_load` by shifting 7 timestamps and use this as one dynamic numerical covariate for TimesFM.\n",
    "2. Bootstrap, that is to run TimesFM once to forecast these past dynamic covariates into the horizon, then call TimesFM again using these forecasts as the future part for these dynamic covariates.\n",
    "\n",
    "### Multivariate Time Series\n",
    "\n",
    "For multivariate time series, if we need univariate forecast, we can try treating the main time series as the target and use the rest as the dynamic covariates."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "chronos-v2",
   "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.10.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
