File size: 4,280 Bytes
45e60de |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 |
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),
]
|