strategies_direct.py

from etna.ensembles import DirectEnsemble
from etna.pipeline import assemble_pipelines


def get_metrics_by_horizon_step(
    ts: TSDataset, horizon_step: int, horizon: int, n_folds: int
):
    if horizon % horizon_step != 0:
        raise ValueError("Horizon step should devide horizon!")

    num_models = horizon // horizon_step
    horizons = [horizon_step * (i + 1) for i in range(num_models)]
    lags = [
        LagTransform(
            in_column="target",
            lags=list(range(horizon, MAX_LAG)),
            out_column="lag_target",
        )
        for horizon in horizons
    ]
    transforms = [lags, SegmentEncoderTransform()]

    model = CatBoostMultiSegmentModel()
    pipelines = assemble_pipelines(
        models=model, transforms=transforms, horizons=horizons
    )
    pipeline = DirectEnsemble(pipelines=pipelines)

    start_time = time.perf_counter()
    metrics_df, forecast_df, fold_info_df = 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}
  
  
horizon_step_values = [12, 6, 4, 3, 2, 1]
horizon_step_results = []

for horizon_step in horizon_step_values:
    cur_res = get_metrics_by_horizon_step(
        ts=ts, horizon_step=horizon_step, horizon=HORIZON, n_folds=N_FOLDS
    )
    horizon_step_results.append({"horizon_step": horizon_step, **cur_res})