Upload 10 files

src/__init__.py

@@ -0,0 +1,38 @@
+from .data import load_dataset, SEASONALITY_MAP
+from .fit_model import fit_predict_with_model, MODEL_NAME_TO_CLASS
+from .score import score_predictions
+
+
+AVAILABLE_MODELS = list(MODEL_NAME_TO_CLASS.keys())
+
+AVAILABLE_DATASETS = [
+    "car_parts_without_missing",
+    "cif_2016",
+    "covid_deaths",
+    "electricity_hourly",
+    "electricity_weekly",
+    "fred_md",
+    "hospital",
+    "kaggle_web_traffic_weekly",
+    "kdd_cup_2018_without_missing",
+    "m1_monthly",
+    "m1_quarterly",
+    "m1_yearly",
+    "m3_monthly",
+    "m3_other",
+    "m3_quarterly",
+    "m3_yearly",
+    "m4_daily",
+    "m4_hourly",
+    "m4_weekly",
+    "m4_yearly",
+    "m4_monthly",
+    "m4_quarterly",
+    "nn5_daily_without_missing",
+    "nn5_weekly",
+    "pedestrian_counts",
+    "tourism_monthly",
+    "tourism_quarterly",
+    "tourism_yearly",
+    "uber_tlc_without_missing",
+]

src/data.py

@@ -0,0 +1,31 @@
+import pandas as pd
+import copy
+from typing import Tuple
+from gluonts.dataset.common import TrainDatasets
+from gluonts.dataset.repository.datasets import get_dataset
+
+SEASONALITY_MAP = {
+    "Y": 1,
+    "Q": 4,
+    "M": 12,
+    "W": 1,
+    "D": 7,
+    "H": 24,
+}
+
+
+def fix_m3_other_start(ts: dict):
+    new_ts = copy.copy(ts)
+    new_ts["start"] = pd.Period("1750", freq="Y")
+    return new_ts
+
+
+def load_dataset(dataset_name) -> TrainDatasets:
+    data = get_dataset(dataset_name)
+    # m3_other provided by GluonTS has incorrect freq Q that should be replaced by Y
+    if dataset_name == "m3_other":
+        fixed_train = [fix_m3_other_start(ts) for ts in data.train]
+        fixed_test = [fix_m3_other_start(ts) for ts in data.test]
+        data = TrainDatasets(metadata=data.metadata, train=fixed_train, test=fixed_test)
+        data.metadata.freq = "Y"
+    return data

src/fit_model.py

@@ -0,0 +1,44 @@
+from gluonts.dataset.common import Dataset
+
+from .models import (
+    AbstractPredictor,
+    AutoGluonPredictor,
+    AutoPyTorchPredictor,
+    DeepARPredictor,
+    TFTPredictor,
+    AutoARIMAPredictor,
+    AutoETSPredictor,
+    AutoThetaPredictor,
+    StatsEnsemblePredictor,
+)
+
+MODEL_NAME_TO_CLASS = {
+    "autogluon": AutoGluonPredictor,
+    "autopytorch": AutoPyTorchPredictor,
+    "deepar": DeepARPredictor,
+    "tft": TFTPredictor,
+    "autoarima": AutoARIMAPredictor,
+    "autoets": AutoETSPredictor,
+    "autotheta": AutoThetaPredictor,
+    "statsensemble": StatsEnsemblePredictor,
+}
+
+
+def fit_predict_with_model(
+    model_name: str,
+    dataset: Dataset,
+    prediction_length: int,
+    freq: str,
+    seasonality: int,
+    **model_kwargs,
+):
+    model_class = MODEL_NAME_TO_CLASS[model_name.lower()]
+    model: AbstractPredictor = model_class(
+        prediction_length=prediction_length,
+        freq=freq,
+        seasonality=seasonality,
+        **model_kwargs,
+    )
+    predictions = model.fit_predict(dataset)
+    info = {"run_time": model.get_runtime()}
+    return predictions, info

src/models/__init__.py

@@ -0,0 +1,10 @@
+from .abstract import AbstractPredictor
+from .autogluon import AutoGluonPredictor
+from .autopytorch import AutoPyTorchPredictor
+from .deep import DeepARPredictor, TFTPredictor
+from .statsforecast import (
+    AutoARIMAPredictor,
+    AutoETSPredictor,
+    AutoThetaPredictor,
+    StatsEnsemblePredictor,
+)

src/models/abstract.py

@@ -0,0 +1,41 @@
+from typing import List, Optional
+
+from gluonts.dataset.common import Dataset
+from gluonts.model.forecast import Forecast
+
+
+class AbstractPredictor:
+    def __init__(
+        self,
+        prediction_length: int,
+        freq: str,
+        seasonality: int,
+        quantile_levels: Optional[List[float]] = None,
+    ):
+        self.prediction_length = prediction_length
+        self.freq = freq
+        self.seasonality = seasonality
+        self.quantile_levels = quantile_levels or [
+            0.1,
+            0.2,
+            0.3,
+            0.4,
+            0.5,
+            0.6,
+            0.7,
+            0.8,
+            0.9,
+        ]
+        self._runtime = None
+
+    def fit_predict(
+        self,
+        dataset: Dataset
+    ) -> List[Forecast]:
+        raise NotImplementedError
+
+    def save_runtime(self, time: float) -> None:
+        self._runtime = time
+
+    def get_runtime(self) -> float:
+        return self._runtime

src/models/autogluon.py

@@ -0,0 +1,76 @@
+import time
+from typing import List, Optional
+import pandas as pd
+
+from gluonts.dataset.common import Dataset
+from gluonts.model.forecast import Forecast, QuantileForecast
+
+from .abstract import AbstractPredictor
+
+
+class AutoGluonPredictor(AbstractPredictor):
+    def __init__(
+        self,
+        prediction_length: int,
+        freq: str,
+        seasonality: int,
+        time_limit: Optional[int] = None,
+        presets: str = "high_quality",
+        eval_metric: str = "MASE",
+        seed: int = 1,
+        enable_ensemble: bool = True,
+        hyperparameters: Optional[dict] = None,
+        **kwargs
+    ):
+        super().__init__(prediction_length, freq, seasonality)
+        self.presets = presets
+        self.eval_metric = eval_metric
+        self.time_limit = time_limit
+        self.seed = seed
+        self.enable_ensemble = enable_ensemble
+        self.hyperparameters = hyperparameters
+
+    def fit_predict(self, dataset: Dataset) -> List[Forecast]:
+        from autogluon.timeseries import TimeSeriesDataFrame, TimeSeriesPredictor
+
+        train_data = TimeSeriesDataFrame(dataset)
+        predictor = TimeSeriesPredictor(
+            prediction_length=self.prediction_length,
+            eval_metric=self.eval_metric,
+            eval_metric_seasonal_period=self.seasonality,
+            quantile_levels=self.quantile_levels,
+        )
+        start_time = time.time()
+        predictor.fit(
+            train_data,
+            time_limit=self.time_limit,
+            presets=self.presets,
+            random_seed=self.seed,
+            enable_ensemble=self.enable_ensemble,
+            hyperparameters=self.hyperparameters,
+        )
+        predictions = predictor.predict(train_data)
+        self.save_runtime(time.time() - start_time)
+        return self._predictions_df_to_gluonts_forecast(
+            predictions_df=predictions.drop("mean", axis=1), dataset=dataset
+        )
+
+    def _predictions_df_to_gluonts_forecast(
+        self, predictions_df, dataset: Dataset
+    ) -> List[Forecast]:
+        agts_forecasts = [
+            f.droplevel("item_id")
+            for _, f in predictions_df.groupby(level="item_id", sort=False)
+        ]
+        forecast_list = []
+        for ts, f in zip(dataset, agts_forecasts):
+            item_id = ts["item_id"]
+            forecast_list.append(
+                QuantileForecast(
+                    forecast_arrays=f.values.T,
+                    forecast_keys=f.columns,
+                    start_date=pd.Period(f.index[0], freq=self.freq),
+                    item_id=item_id,
+                )
+            )
+        return forecast_list

src/models/autopytorch.py

@@ -0,0 +1,76 @@
+import copy
+import multiprocessing as mp
+import time
+from typing import List
+
+from gluonts.dataset.common import Dataset
+from gluonts.model.forecast import Forecast, SampleForecast
+
+from .abstract import AbstractPredictor
+
+
+class AutoPyTorchPredictor(AbstractPredictor):
+    def __init__(
+        self,
+        prediction_length: int,
+        freq: str,
+        seasonality: int,
+        time_limit: int = 6 * 60 * 60,
+        optimize_metric: str = "mean_MASE_forecasting",
+        seed: int = 1,
+        **kwargs
+    ):
+        super().__init__(prediction_length, freq, seasonality)
+        self.optimize_metric = optimize_metric
+        self.run_time = time_limit
+        self.seed = seed
+
+    def fit_predict(self, dataset: Dataset) -> List[Forecast]:
+        from autoPyTorch.api.time_series_forecasting import TimeSeriesForecastingTask
+        from autoPyTorch.datasets.resampling_strategy import HoldoutValTypes
+
+        y_train = [item["target"] for item in dataset]
+        start_times = [item["start"].to_timestamp(how="S") for item in dataset]
+
+        api = TimeSeriesForecastingTask(
+            seed=self.seed,
+            ensemble_size=20,
+            resampling_strategy=HoldoutValTypes.time_series_hold_out_validation,
+            resampling_strategy_args=None,
+        )
+        api.set_pipeline_options(early_stopping=20, torch_num_threads=mp.cpu_count())
+
+        fit_start_time = time.time()
+        api.search(
+            X_train=None,
+            y_train=copy.deepcopy(y_train),
+            optimize_metric=self.optimize_metric,
+            n_prediction_steps=self.prediction_length,
+            memory_limit=16 * 1024,
+            freq="1" + self.freq,
+            start_times=start_times,
+            normalize_y=False,
+            total_walltime_limit=self.run_time,
+            min_num_test_instances=1000,
+            budget_type="epochs",
+            max_budget=50,
+            min_budget=5,
+        )
+        # # Skip refitting as this raises exceptions for all models as of v0.2.1
+        # refit_dataset = api.dataset.create_refit_set()
+        # api.refit(refit_dataset, 0)
+
+        # Predict for the test set
+        test_sets = api.dataset.generate_test_seqs()
+        predictions = api.predict(test_sets)
+        self.save_runtime(time.time() - fit_start_time)
+        forecast_list = []
+        for ts, pred in zip(dataset, predictions):
+            forecast_list.append(
+                SampleForecast(
+                    samples=pred[None],
+                    start_date=ts["start"] + len(ts["target"]),
+                    item_id=ts["item_id"],
+                )
+            )
+        return forecast_list

src/models/deep.py

@@ -0,0 +1,64 @@
+import time
+from typing import List, Optional
+from datetime import timedelta
+
+from gluonts.dataset.split import split
+from gluonts.dataset.common import Dataset
+from gluonts.model.forecast import Forecast
+from gluonts.torch.model.estimator import Estimator
+
+
+from .abstract import AbstractPredictor
+
+
+class GluonTSPredictor(AbstractPredictor):
+    def __init__(
+        self,
+        prediction_length: int,
+        freq: str,
+        seasonality: int,
+        time_limit: Optional[int] = None,
+        **kwargs,
+    ):
+        super().__init__(prediction_length, freq, seasonality)
+        self.time_limit = time_limit
+
+    def fit_predict(self, dataset: Dataset) -> List[Forecast]:
+        estimator = self._get_estimator()
+        train_data, _ = split(dataset, offset=-self.prediction_length)
+        fit_start_time = time.time()
+        predictor = estimator.train(training_data=train_data, validation_data=dataset)
+        predictions = predictor.predict(dataset)
+        self.save_runtime(time.time() - fit_start_time)
+        return predictions
+
+    def _get_estimator(self) -> Estimator:
+        raise NotImplementedError
+
+    def _get_trainer_kwargs(self):
+        from pytorch_lightning.callbacks import Timer
+
+        # Train until time limit
+        return {"max_epochs": 100_000, "callbacks": [Timer(timedelta(seconds=self.time_limit))]}
+
+
+class DeepARPredictor(GluonTSPredictor):
+    def _get_estimator(self) -> Estimator:
+        from gluonts.torch.model.deepar import DeepAREstimator
+
+        return DeepAREstimator(
+            freq=self.freq,
+            prediction_length=self.prediction_length,
+            trainer_kwargs=self._get_trainer_kwargs(),
+        )
+
+
+class TFTPredictor(GluonTSPredictor):
+    def _get_estimator(self) -> Estimator:
+        from gluonts.torch.model.tft import TemporalFusionTransformerEstimator
+
+        return TemporalFusionTransformerEstimator(
+            freq=self.freq,
+            prediction_length=self.prediction_length,
+            trainer_kwargs=self._get_trainer_kwargs(),
+        )

src/models/statsforecast.py

@@ -0,0 +1,132 @@
+import time
+from typing import List
+import pandas as pd
+
+from gluonts.dataset.common import Dataset
+from gluonts.model.forecast import Forecast, QuantileForecast
+
+from .abstract import AbstractPredictor
+
+
+class StatsForecastPredictor(AbstractPredictor):
+    def __init__(self, prediction_length: int, freq: str, seasonality: int, **kwargs):
+        super().__init__(prediction_length, freq, seasonality)
+
+    def fit_predict(self, dataset: Dataset) -> List[Forecast]:
+        from statsforecast import StatsForecast
+        from statsforecast.models import SeasonalNaive
+
+        df = self._to_statsforecast_df(dataset)
+        models = self._get_models()
+        predictor = StatsForecast(
+            df=df,
+            freq=self.freq,
+            models=models,
+            fallback_model=SeasonalNaive(season_length=self.seasonality),
+            n_jobs=-1,
+        )
+        start_time = time.time()
+        predictions_df = predictor.forecast(
+            h=self.prediction_length, level=[0, 20, 40, 60, 80]
+        )
+        self.save_runtime(time.time() - start_time)
+        return self._predictions_df_to_gluonts_forecast(
+            predictions_df, dataset, model_names=[str(m) for m in models]
+        )
+
+    def _predictions_df_to_gluonts_forecast(
+        self,
+        predictions_df: pd.DataFrame,
+        dataset: Dataset,
+        model_names: List[str],
+    ) -> List[Forecast]:
+        def quantile_to_suffix(q: float) -> str:
+            if q < 0.5:
+                prefix = "-lo-"
+                level = 100 - 200 * q
+            else:
+                prefix = "-hi-"
+                level = 200 * q - 100
+            return prefix + str(int(level))
+
+        # Convert StatsForecast output -> DataFrame with quantile_levels as outputs
+        columns = {}
+        for q in self.quantile_levels:
+            suffix = quantile_to_suffix(q)
+            columns[str(q)] = predictions_df[[m + suffix for m in model_names]].median(
+                axis=1
+            )
+
+        # Convert quantiles DataFrame -> list of QuantileForecasts
+        forecast_df = pd.DataFrame(columns)
+        forecast_list = []
+        for ts in dataset:
+            item_id = ts["item_id"]
+            f = forecast_df.loc[item_id]
+            forecast_list.append(
+                QuantileForecast(
+                    forecast_arrays=f.values.T,
+                    forecast_keys=f.columns,
+                    start_date=pd.Period(
+                        predictions_df["ds"].loc[item_id].iloc[0], freq=self.freq
+                    ),
+                    item_id=item_id,
+                )
+            )
+        return forecast_list
+
+    def _to_statsforecast_df(self, dataset: Dataset) -> pd.DataFrame:
+        """Convert GluonTS Dataset to StatsForecast compatible DataFrame."""
+        dfs = []
+        for item in dataset:
+            target = item["target"]
+            timestamps = pd.date_range(
+                start=item["start"].to_timestamp(how="S"),
+                periods=len(target),
+                freq=self.freq,
+            )
+            df = pd.DataFrame(
+                {
+                    "unique_id": [item["item_id"]] * len(target),
+                    "ds": timestamps,
+                    "y": target,
+                }
+            )
+            dfs.append(df)
+        return pd.concat(dfs)
+
+
+class AutoARIMAPredictor(StatsForecastPredictor):
+    def _get_models(self):
+        from statsforecast.models import AutoARIMA
+
+        return [AutoARIMA(season_length=self.seasonality)]
+
+
+class AutoETSPredictor(StatsForecastPredictor):
+    def _get_models(self):
+        from statsforecast.models import AutoETS
+
+        return [AutoETS(season_length=self.seasonality)]
+
+
+class AutoThetaPredictor(StatsForecastPredictor):
+    def _get_models(self):
+        from statsforecast.models import AutoTheta
+
+        return [AutoTheta(season_length=self.seasonality)]
+
+
+class StatsEnsemblePredictor(StatsForecastPredictor):
+    def _get_models(self):
+        from statsforecast.models import (
+            AutoETS,
+            AutoARIMA,
+            AutoTheta,
+        )
+
+        return [
+            AutoETS(season_length=self.seasonality),
+            AutoTheta(season_length=self.seasonality),
+            AutoARIMA(season_length=self.seasonality),
+        ]

src/score.py

@@ -0,0 +1,22 @@
+from typing import List
+
+from gluonts.dataset.split import split
+from gluonts.dataset.common import Dataset
+from gluonts.model.forecast import Forecast
+from gluonts.evaluation.backtest import _to_dataframe, Evaluator
+
+
+def score_predictions(
+    dataset: Dataset,
+    predictions: List[Forecast],
+    prediction_length: int,
+    seasonality: int,
+):
+    _, test_template = split(dataset, offset=-prediction_length)
+    test_data = test_template.generate_instances(prediction_length)
+    ts_iterator = map(_to_dataframe, test_data)
+    evaluator = Evaluator(
+        quantiles=[0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9], seasonality=seasonality
+    )
+    metrics, _ = evaluator(ts_iterator=ts_iterator, fcst_iterator=predictions)
+    return metrics