strategies_baseline.py

import time

from etna.metrics import SMAPE
from etna.models import CatBoostMultiSegmentModel
from etna.pipeline import Pipeline
from etna.transforms import LagTransform
from etna.transforms import SegmentEncoderTransform

HORIZON = 24
N_FOLDS = 3
MAX_LAG = 37


def get_baseline_metrics(ts: TSDataset, horizon: int, n_folds: int):
    transforms = [
        LagTransform(
            in_column="target",
            lags=list(range(horizon, MAX_LAG)),
            out_column="lag_target",
        ),
        SegmentEncoderTransform(),
    ]
    model = CatBoostMultiSegmentModel()
    simple_pipeline = Pipeline(model=model, transforms=transforms, horizon=horizon)
    start_time = time.perf_counter()
    metrics_df, forecast_df, fold_info_df = simple_pipeline.backtest(
        ts=ts, metrics=[SMAPE()], n_folds=n_folds
    )
    elapsed_time = time.perf_counter() - start_time
    smape_value = metrics_df.mean().to_dict()["SMAPE"]
    return {"SMAPE": smape_value, "time": elapsed_time}