Merge pull request #64 from google-research/benchmarking

Adding TimeGPT-1 to extended benchmarks

experiments/baselines/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

experiments/baselines/timegpt_pipeline.py

@@ -0,0 +1,251 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import List, Optional, Tuple
+import os
+import pandas as pd
+from gluonts.time_feature.seasonality import get_seasonality as _get_seasonality
+from tqdm import tqdm
+from utilsforecast.processing import (
+    backtest_splits,
+    drop_index_if_pandas,
+    join,
+    maybe_compute_sort_indices,
+    take_rows,
+    vertical_concat,
+)
+from time import time
+from dotenv import load_dotenv
+from nixtla import NixtlaClient
+
+
+def get_seasonality(freq: str) -> int:
+    return _get_seasonality(freq, seasonalities={"D": 7})
+
+
+def maybe_convert_col_to_datetime(df: pd.DataFrame, col_name: str) -> pd.DataFrame:
+    if not pd.api.types.is_datetime64_any_dtype(df[col_name]):
+        df = df.copy()
+        df[col_name] = pd.to_datetime(df[col_name])
+    return df
+
+
+def zero_pad_time_series(df, freq, min_length=36):
+    """If time_series length is less than min_length, front pad it with zeros."""
+    # 1. Calculate required padding for each unique_id
+    value_counts = df["unique_id"].value_counts()
+    to_pad = value_counts[value_counts < min_length].index
+
+    # 2. Create a new DataFrame to hold padded data
+    padded_data = []
+
+    for unique_id in to_pad:
+        # 2a. Filter data for the specific unique_id
+        subset = df[df["unique_id"] == unique_id]
+        if len(subset) > min_length:
+            padded_data.append(subset)
+        else:
+            # 2b. Determine earliest date and calculate padding dates
+            start_date = subset["ds"].min()
+            padding_dates = pd.date_range(
+                end=start_date,
+                periods=min_length - len(subset) + 1,
+                freq=freq,  # 'MS' for month start
+            )[
+                :-1
+            ]  # Exclude the start_date itself
+
+            # 2c. Create padding data
+            padding_df = pd.DataFrame(
+                {"ds": padding_dates, "unique_id": unique_id, "y": 0}  # Zero padding
+            )
+
+            # 2d. Combine original and padding data, and append to the list
+            padded_data.append(pd.concat([padding_df, subset]).sort_values("ds"))
+
+    # 3. Combine all padded data and original data (unchanged)
+    result_df = pd.concat(padded_data + [df[~df["unique_id"].isin(to_pad)]])
+    return result_df
+
+
+class Forecaster:
+    """Borrowed from
+    https://github.com/Nixtla/nixtla/tree/main/experiments/foundation-time-series-arena/xiuhmolpilli/models.
+    """
+
+    def forecast(
+        self,
+        df: pd.DataFrame,
+        h: int,
+        freq: str,
+    ) -> pd.DataFrame:
+        raise NotImplementedError
+
+    def cross_validation(
+        self,
+        df: pd.DataFrame,
+        h: int,
+        freq: str,
+        n_windows: int = 1,
+        step_size: int | None = None,
+    ) -> pd.DataFrame:
+        df = maybe_convert_col_to_datetime(df, "ds")
+        # mlforecast cv code
+        results = []
+        sort_idxs = maybe_compute_sort_indices(df, "unique_id", "ds")
+        if sort_idxs is not None:
+            df = take_rows(df, sort_idxs)
+        splits = backtest_splits(
+            df,
+            n_windows=n_windows,
+            h=h,
+            id_col="unique_id",
+            time_col="ds",
+            freq=pd.tseries.frequencies.to_offset(freq),
+            step_size=h if step_size is None else step_size,
+        )
+        for _, (cutoffs, train, valid) in tqdm(enumerate(splits)):
+            if len(valid.columns) > 3:
+                raise NotImplementedError(
+                    "Cross validation with exogenous variables is not yet supported."
+                )
+            y_pred = self.forecast(
+                df=train,
+                h=h,
+                freq=freq,
+            )
+            y_pred = join(y_pred, cutoffs, on="unique_id", how="left")
+            result = join(
+                valid[["unique_id", "ds", "y"]],
+                y_pred,
+                on=["unique_id", "ds"],
+            )
+            if result.shape[0] < valid.shape[0]:
+                raise ValueError(
+                    "Cross validation result produced less results than expected. "
+                    "Please verify that the frequency parameter (freq) matches your series' "
+                    "and that there aren't any missing periods."
+                )
+            results.append(result)
+        out = vertical_concat(results)
+        out = drop_index_if_pandas(out)
+        first_out_cols = ["unique_id", "ds", "cutoff", "y"]
+        remaining_cols = [c for c in out.columns if c not in first_out_cols]
+        fcst_cv_df = out[first_out_cols + remaining_cols]
+        return fcst_cv_df
+
+
+class TimeGPT(Forecaster):
+    """Borrowed from
+    https://github.com/Nixtla/nixtla/tree/main/experiments/foundation-time-series-arena/xiuhmolpilli/models.
+    We modify the class to take care of edge cases.
+    """
+
+    def __init__(
+        self,
+        api_key: str | None = None,
+        base_url: Optional[str] = None,
+        max_retries: int = 1,
+        model: str = "timegpt-1",
+        alias: str = "TimeGPT",
+    ):
+        self.api_key = api_key
+        self.base_url = base_url
+        self.max_retries = max_retries
+        self.model = model
+        self.alias = alias
+
+    def _get_client(self) -> NixtlaClient:
+        if self.api_key is None:
+            api_key = os.environ["NIXTLA_API_KEY"]
+        else:
+            api_key = self.api_key
+        return NixtlaClient(
+            api_key=api_key,
+            base_url=self.base_url,
+            max_retries=self.max_retries,
+        )
+
+    def forecast(
+        self,
+        df: pd.DataFrame,
+        h: int,
+        freq: str,
+        level: List = [90.0],
+        chunk_size: Optional[int] = None,
+    ) -> pd.DataFrame:
+        client = self._get_client()
+        fcst_df = None
+        if chunk_size is None:
+            fcst_df = client.forecast(
+                df=df,
+                h=h,
+                freq=freq,
+                level=level,
+                model=self.model,
+            )
+        else:
+            all_unique_ids = df["unique_id"].unique()
+            all_fcst_df = []
+            for i in range(0, len(all_unique_ids), chunk_size):
+                chunk_ids = all_unique_ids[i : i + chunk_size]
+                chunk_df = df[df["unique_id"].isin(chunk_ids)]
+                fct_chunk_df = client.forecast(
+                    df=chunk_df,
+                    h=h,
+                    freq=freq,
+                    level=level,
+                )
+                all_fcst_df.append(fct_chunk_df)
+            fcst_df = pd.concat(all_fcst_df)
+        fcst_df["ds"] = pd.to_datetime(fcst_df["ds"])
+        replace_dict = {}
+        for col in fcst_df.columns:
+            if col.startswith("TimeGPT"):
+                replace_dict[col] = col.replace("TimeGPT", self.alias)
+        fcst_df = fcst_df.rename(columns=replace_dict)
+        return fcst_df
+
+
+def run_timegpt(
+    train_df: pd.DataFrame,
+    horizon: int,
+    freq: str,
+    seasonality: int,
+    level: List[int],
+    dataset: str,
+    model: str = "timegpt-1",
+) -> Tuple[pd.DataFrame, float, str]:
+    os.environ["NIXTLA_ID_AS_COL"] = "true"
+    model = TimeGPT(model="timegpt-1", alias=model)
+    padded_train_df = zero_pad_time_series(train_df, freq)
+    init_time = time()
+    # For these datasets the API fails if we do not chunk.
+    if dataset in ["m5", "m4_quarterly"]:
+        chunk_size = 5000
+    else:
+        chunk_size = None
+    fcsts_df = model.forecast(
+        df=padded_train_df, h=horizon, level=level, freq=freq, chunk_size=chunk_size
+    )
+    total_time = time() - init_time
+    # In case levels are not returned we replace the levels with the mean predictions.
+    # Note that this does not affect the results table as we only compare on point
+    # forecastign metrics.
+    for lvl in level:
+        if f"{model.alias}-lo-{lvl}" not in fcsts_df.columns:
+            fcsts_df[f"{model.alias}-lo-{lvl}"] = fcsts_df[model.alias]
+        if f"{model.alias}-hi-{lvl}" not in fcsts_df.columns:
+            fcsts_df[f"{model.alias}-hi-{lvl}"] = fcsts_df[model.alias]
+    return fcsts_df, total_time, model.alias

experiments/environment.yml

@@ -20,5 +20,8 @@ dependencies:
     - git+https://github.com/amazon-science/chronos-forecasting.git
     - praxis
     - paxml
-    - jax[cuda12]==0.4.26
+    - jax[cuda12]
     - einshape
+    - python-dotenv
+    - nixtla>=0.5.1
+    - rich

experiments/environment_cpu.yml

@@ -20,5 +20,8 @@ dependencies:
     - git+https://github.com/amazon-science/chronos-forecasting.git
     - praxis
     - paxml
-    - jax[cpu]==0.4.26
+    - jax[cpu]
     - einshape
+    - python-dotenv
+    - nixtla>=0.5.1
+    - rich

experiments/extended_benchmarks/README.md

@@ -2,7 +2,6 @@
 
 The benchmark setting has been borrowed from Nixtla's original [benchmarking](https://github.com/AzulGarza/nixtla/tree/main/experiments/amazon-chronos) of time-series foundation models against a strong statistical ensemble. Later more datasets were added by the Chronos team in this [pull request](https://github.com/shchur/nixtla/tree/chronos-full-eval/experiments/amazon-chronos). We compare on all the datasets in this extended benchmarks.
 
-All experiments were performed on a [g2-standard-32](https://cloud.google.com/compute/docs/gpus). 
 
 ## Running TimesFM on the benchmark
 
@@ -19,7 +18,11 @@ Note: In the current version of TimesFM we focus on point forecasts and therefor
 
 ## Benchmark Results
 
-![Benchmark Results Table](./tfm_results.png)
+![Benchmark Results Table](./tfm_extended_new.png)
+
+__Update:__ We have added TimeGPT-1 to the benchmark results. We had to remove the Dominick dataset as we were not able to run TimeGPT-1 on this benchmark. Note that the previous results including Dominick remain available at `./tfm_results.png`. In order to reproduce the results for TimeGPT-1, please run `run_timegpt.py`.
+
+_Remark:_ All baselines except the ones involving TimeGPT were run performed on a [g2-standard-32](https://cloud.google.com/compute/docs/gpus). Since TimeGPT-1 can only be accessed by an API, the time column might not reflect the true speed of the model as it also includes the communication cost. Moreover, we are not sure about the exact backend hardware for TimeGPT.
 
 We can see that TimesFM performs the best in terms of both mase and smape. More importantly it is much faster than the other methods, in particular it is more than 600x faster than StatisticalEnsemble and 80x faster than Chronos (Large).
 

experiments/extended_benchmarks/run_timegpt.py

@@ -0,0 +1,108 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Evaluation script for timegpt."""
+
+import os
+import sys
+import time
+
+from absl import flags
+import numpy as np
+import pandas as pd
+from ..baselines.timegpt_pipeline import run_timegpt
+
+from .utils import ExperimentHandler
+
+
+dataset_names = [
+    "m1_monthly",
+    "m1_quarterly",
+    "m1_yearly",
+    "m3_monthly",
+    "m3_other",
+    "m3_quarterly",
+    "m3_yearly",
+    "m4_quarterly",
+    "m4_yearly",
+    "tourism_monthly",
+    "tourism_quarterly",
+    "tourism_yearly",
+    "nn5_daily_without_missing",
+    "m5",
+    "nn5_weekly",
+    "traffic",
+    "weather",
+    "australian_electricity_demand",
+    "car_parts_without_missing",
+    "cif_2016",
+    "covid_deaths",
+    "ercot",
+    "ett_small_15min",
+    "ett_small_1h",
+    "exchange_rate",
+    "fred_md",
+    "hospital",
+]
+
+_MODEL_NAME = flags.DEFINE_string(
+    "model_name",
+    "timegpt-1-long-horizon",
+    "Path to model, can also be set to timegpt-1",
+)
+_SAVE_DIR = flags.DEFINE_string("save_dir", "./results", "Save directory")
+
+
+QUANTILES = list(np.arange(1, 10) / 10.0)
+
+
+def main():
+    results_list = []
+    run_id = np.random.randint(100000)
+    model_name = _MODEL_NAME.value
+    for dataset in dataset_names:
+        print(f"Evaluating model {model_name} on dataset {dataset}", flush=True)
+        exp = ExperimentHandler(dataset, quantiles=QUANTILES)
+        train_df = exp.train_df
+        horizon = exp.horizon
+        seasonality = exp.seasonality
+        freq = exp.freq
+        level = exp.level
+        fcsts_df, total_time, model_name = run_timegpt(
+            train_df=train_df,
+            horizon=exp.horizon,
+            model=model_name,
+            seasonality=seasonality,
+            freq=freq,
+            dataset=dataset,
+            level=level,
+        )
+        time_df = pd.DataFrame({"time": [total_time], "model": model_name})
+        fcsts_df = exp.fcst_from_level_to_quantiles(fcsts_df, model_name)
+        results = exp.evaluate_from_predictions(
+            models=[model_name], fcsts_df=fcsts_df, times_df=time_df
+        )
+        print(results, flush=True)
+        results_list.append(results)
+        results_full = pd.concat(results_list)
+        save_path = os.path.join(_SAVE_DIR.value, str(run_id))
+        print(f"Saving results to {save_path}", flush=True)
+        os.makedirs(save_path, exist_ok=True)
+        results_full.to_csv(f"{save_path}/results.csv")
+
+
+if __name__ == "__main__":
+    FLAGS = flags.FLAGS
+    FLAGS(sys.argv)
+    main()

experiments/extended_benchmarks/run_timesfm.py

@@ -45,7 +45,6 @@ dataset_names = [
     "nn5_weekly",
     "traffic",
     "weather",
-    "dominick",
     "australian_electricity_demand",
     "car_parts_without_missing",
     "cif_2016",

pyproject.toml

@@ -8,7 +8,7 @@ dependencies = [
     "einshape>=1.0.0",
     "paxml>=1.4.0",
     "praxis>=1.4.0",
-    "jax==0.4.26",
+    "jax>=0.4.26",
     "numpy>=1.26.4",
     "pandas>=2.1.4",
 ]