autogluon.timeseries 1.2.1b20250422__py3-none-any.whl → 1.2.1b20250424__py3-none-any.whl

autogluon/timeseries/models/ensemble/{greedy_ensemble.py → greedy.py}

@@ -9,9 +9,10 @@ import autogluon.core as ag
 from autogluon.core.models.greedy_ensemble.ensemble_selection import EnsembleSelection
 from autogluon.timeseries import TimeSeriesDataFrame
 from autogluon.timeseries.metrics import TimeSeriesScorer
-from autogluon.timeseries.models.ensemble import AbstractTimeSeriesEnsembleModel
 from autogluon.timeseries.utils.datetime import get_seasonality
 
+from .abstract import AbstractWeightedTimeSeriesEnsembleModel
+
 logger = logging.getLogger(__name__)
 
 
@@ -24,16 +25,15 @@ class TimeSeriesEnsembleSelection(EnsembleSelection):
         sorted_initialization: bool = False,
         bagging: bool = False,
         tie_breaker: str = "random",
-        random_state: np.random.RandomState = None,
+        random_state: Optional[np.random.RandomState] = None,
         prediction_length: int = 1,
         target: str = "target",
-        eval_metric_seasonal_period: Optional[int] = None,
-        horizon_weight: Optional[np.ndarray] = None,
+        eval_metric_seasonal_period: int = 1,
         **kwargs,
     ):
         super().__init__(
             ensemble_size=ensemble_size,
-            metric=metric,
+            metric=metric,  # type: ignore
             problem_type=problem_type,
             sorted_initialization=sorted_initialization,
             bagging=bagging,
@@ -44,14 +44,33 @@ class TimeSeriesEnsembleSelection(EnsembleSelection):
         self.prediction_length = prediction_length
         self.target = target
         self.eval_metric_seasonal_period = eval_metric_seasonal_period
-        self.horizon_weight = horizon_weight
+        self.metric: TimeSeriesScorer
 
-    def _fit(
+        self.dummy_pred_per_window = []
+        self.scorer_per_window = []
+
+        self.dummy_pred_per_window: Optional[List[TimeSeriesDataFrame]]
+        self.scorer_per_window: Optional[List[TimeSeriesScorer]]
+        self.data_future_per_window: Optional[List[TimeSeriesDataFrame]]
+
+    def fit(  # type: ignore
         self,
-        predictions: List[List[TimeSeriesDataFrame]],  # first dim: model, second dim: val window index
+        predictions: List[List[TimeSeriesDataFrame]],
         labels: List[TimeSeriesDataFrame],
-        time_limit: Optional[int] = None,
-        sample_weight=None,
+        time_limit: Optional[float] = None,
+    ):
+        return super().fit(
+            predictions=predictions,  # type: ignore
+            labels=labels,  # type: ignore
+            time_limit=time_limit,
+        )
+
+    def _fit(  # type: ignore
+        self,
+        predictions: List[List[TimeSeriesDataFrame]],
+        labels: List[TimeSeriesDataFrame],
+        time_limit: Optional[float] = None,
+        sample_weight: Optional[List[float]] = None,
     ):
         # Stack predictions for each model into a 3d tensor of shape [num_val_windows, num_rows, num_cols]
         stacked_predictions = [np.stack(preds) for preds in predictions]
@@ -77,26 +96,33 @@ class TimeSeriesEnsembleSelection(EnsembleSelection):
 
         super()._fit(
             predictions=stacked_predictions,
-            labels=data_future,
+            labels=data_future,  # type: ignore
             time_limit=time_limit,
         )
         self.dummy_pred_per_window = None
         self.evaluator_per_window = None
         self.data_future_per_window = None
 
-    def _calculate_regret(self, y_true, y_pred_proba, metric=None, sample_weight=None):  # noqa
+    def _calculate_regret(  # type: ignore
+        self,
+        y_true,
+        y_pred_proba,
+        metric: TimeSeriesScorer,
+        sample_weight=None,
+    ):
         # Compute average score across all validation windows
         total_score = 0.0
+
+        assert self.data_future_per_window is not None
+        assert self.dummy_pred_per_window is not None
+        assert self.scorer_per_window is not None
+
         for window_idx, data_future in enumerate(self.data_future_per_window):
             dummy_pred = self.dummy_pred_per_window[window_idx]
             dummy_pred[list(dummy_pred.columns)] = y_pred_proba[window_idx]
             # We use scorer.compute_metric instead of scorer.score to avoid repeated calls to scorer.save_past_metrics
             metric_value = self.scorer_per_window[window_idx].compute_metric(
-                data_future,
-                dummy_pred,
-                target=self.target,
-                prediction_length=self.prediction_length,
-                horizon_weight=self.horizon_weight,
+                data_future, dummy_pred, target=self.target
             )
             total_score += metric.sign * metric_value
         avg_score = total_score / len(self.data_future_per_window)
@@ -104,32 +130,46 @@ class TimeSeriesEnsembleSelection(EnsembleSelection):
         return -avg_score
 
 
-class TimeSeriesGreedyEnsemble(AbstractTimeSeriesEnsembleModel):
-    """Constructs a weighted ensemble using the greedy Ensemble Selection algorithm."""
+class GreedyEnsemble(AbstractWeightedTimeSeriesEnsembleModel):
+    """Constructs a weighted ensemble using the greedy Ensemble Selection algorithm by
+    Caruana et al. [Car2004]
+
+    Other Parameters
+    ----------------
+    ensemble_size: int, default = 100
+        Number of models (with replacement) to include in the ensemble.
+
+    References
+    ----------
+    .. [Car2024] Caruana, Rich, et al. "Ensemble selection from libraries of models."
+        Proceedings of the twenty-first international conference on Machine learning. 2004.
+    """
 
-    def __init__(self, name: Optional[str] = None, ensemble_size: int = 100, **kwargs):
+    def __init__(self, name: Optional[str] = None, **kwargs):
         if name is None:
+            # FIXME: the name here is kept for backward compatibility. it will be called
+            # GreedyEnsemble in v1.4 once ensemble choices are exposed
             name = "WeightedEnsemble"
         super().__init__(name=name, **kwargs)
-        self.ensemble_size = ensemble_size
-        self.model_to_weight: Dict[str, float] = {}
 
-    def _fit_ensemble(
+    def _get_default_hyperparameters(self) -> Dict:
+        return {"ensemble_size": 100}
+
+    def _fit(
         self,
         predictions_per_window: Dict[str, List[TimeSeriesDataFrame]],
         data_per_window: List[TimeSeriesDataFrame],
-        time_limit: Optional[int] = None,
-        **kwargs,
+        model_scores: Optional[Dict[str, float]] = None,
+        time_limit: Optional[float] = None,
     ):
         if self.eval_metric_seasonal_period is None:
             self.eval_metric_seasonal_period = get_seasonality(self.freq)
         ensemble_selection = TimeSeriesEnsembleSelection(
-            ensemble_size=self.ensemble_size,
+            ensemble_size=self.get_hyperparameters()["ensemble_size"],
             metric=self.eval_metric,
             prediction_length=self.prediction_length,
             target=self.target,
             eval_metric_seasonal_period=self.eval_metric_seasonal_period,
-            horizon_weight=self.horizon_weight,
         )
         ensemble_selection.fit(
             predictions=list(predictions_per_window.values()),
@@ -143,37 +183,3 @@ class TimeSeriesGreedyEnsemble(AbstractTimeSeriesEnsembleModel):
 
         weights_for_printing = {model: round(weight, 2) for model, weight in self.model_to_weight.items()}
         logger.info(f"\tEnsemble weights: {pprint.pformat(weights_for_printing, width=200)}")
-
-    @property
-    def model_names(self) -> List[str]:
-        return list(self.model_to_weight.keys())
-
-    @property
-    def model_weights(self) -> np.ndarray:
-        return np.array(list(self.model_to_weight.values()), dtype=np.float64)
-
-    def predict(self, data: Dict[str, Optional[TimeSeriesDataFrame]], **kwargs) -> TimeSeriesDataFrame:
-        if not set(self.model_names).issubset(set(data.keys())):
-            raise ValueError(
-                f"Set of models given for prediction in {self.name} differ from those provided during initialization."
-            )
-        for model_name, model_pred in data.items():
-            if model_pred is None:
-                raise RuntimeError(f"{self.name} cannot predict because base model {model_name} failed.")
-
-        # Make sure that all predictions have same shape
-        assert len(set(pred.shape for pred in data.values())) == 1
-
-        return sum(data[model_name] * weight for model_name, weight in self.model_to_weight.items())
-
-    def get_info(self) -> dict:
-        info = super().get_info()
-        info["model_weights"] = self.model_to_weight
-        return info
-
-    def remap_base_models(self, model_refit_map: Dict[str, str]) -> None:
-        updated_weights = {}
-        for model, weight in self.model_to_weight.items():
-            model_full_name = model_refit_map.get(model, model)
-            updated_weights[model_full_name] = weight
-        self.model_to_weight = updated_weights

autogluon/timeseries/models/presets.py

@@ -4,8 +4,6 @@ import re
 from collections import defaultdict
 from typing import Any, Dict, List, Optional, Type, Union
 
-import numpy as np
-
 from autogluon.common import space
 from autogluon.core import constants
 from autogluon.timeseries.metrics import TimeSeriesScorer
@@ -186,7 +184,6 @@ def get_preset_models(
     path: str,
     eval_metric: Union[str, TimeSeriesScorer],
     eval_metric_seasonal_period: Optional[int],
-    horizon_weight: Optional[np.ndarray],
     hyperparameters: Union[str, Dict, None],
     hyperparameter_tune: bool,
     covariate_metadata: CovariateMetadata,
@@ -265,7 +262,6 @@ def get_preset_models(
                 eval_metric=eval_metric,
                 eval_metric_seasonal_period=eval_metric_seasonal_period,
                 covariate_metadata=covariate_metadata,
-                horizon_weight=horizon_weight,
                 hyperparameters=model_hps,
                 **kwargs,
             )

autogluon/timeseries/predictor.py

@@ -24,7 +24,7 @@ from autogluon.timeseries import __version__ as current_ag_version
 from autogluon.timeseries.configs import TIMESERIES_PRESETS_CONFIGS
 from autogluon.timeseries.dataset.ts_dataframe import ITEMID, TimeSeriesDataFrame
 from autogluon.timeseries.learner import TimeSeriesLearner
-from autogluon.timeseries.metrics import TimeSeriesScorer, check_get_evaluation_metric, check_get_horizon_weight
+from autogluon.timeseries.metrics import TimeSeriesScorer, check_get_evaluation_metric
 from autogluon.timeseries.splitter import ExpandingWindowSplitter
 from autogluon.timeseries.trainer import TimeSeriesTrainer
 from autogluon.timeseries.utils.forecast import make_future_data_frame
@@ -93,14 +93,6 @@ class TimeSeriesPredictor:
     eval_metric_seasonal_period : int, optional
         Seasonal period used to compute some evaluation metrics such as mean absolute scaled error (MASE). Defaults to
         ``None``, in which case the seasonal period is computed based on the data frequency.
-    horizon_weight : List[float], optional
-        Weight assigned to each time step in the forecast horizon when computing the `eval_metric`. If provided, this
-        must be a list with `prediction_length` non-negative values, where at least some values are greater than zero.
-        AutoGluon will automatically normalize the weights so that they sum up to `prediction_length`. By default, all
-        time steps in the forecast horizon have the same weight, which is equivalent to setting `horizon_weight = [1] * prediction_length`.
-
-        This parameter only affects model selection and ensemble construction; it has no effect on the loss function of
-        the individual forecasting models.
     known_covariates_names: List[str], optional
         Names of the covariates that are known in advance for all time steps in the forecast horizon. These are also
         known as dynamic features, exogenous variables, additional regressors or related time series. Examples of such
@@ -152,7 +144,6 @@ class TimeSeriesPredictor:
         freq: Optional[str] = None,
         eval_metric: Union[str, TimeSeriesScorer, None] = None,
         eval_metric_seasonal_period: Optional[int] = None,
-        horizon_weight: list[float] | None = None,
         path: Optional[Union[str, Path]] = None,
         verbosity: int = 2,
         log_to_file: bool = True,
@@ -198,7 +189,6 @@ class TimeSeriesPredictor:
             self.freq = std_freq
         self.eval_metric = check_get_evaluation_metric(eval_metric)
         self.eval_metric_seasonal_period = eval_metric_seasonal_period
-        self.horizon_weight = check_get_horizon_weight(horizon_weight, prediction_length=self.prediction_length)
         if quantile_levels is None:
             quantile_levels = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
         self.quantile_levels = sorted(quantile_levels)
@@ -206,7 +196,6 @@ class TimeSeriesPredictor:
             path_context=self.path,
             eval_metric=eval_metric,
             eval_metric_seasonal_period=eval_metric_seasonal_period,
-            horizon_weight=self.horizon_weight,
             target=self.target,
             known_covariates_names=self.known_covariates_names,
             prediction_length=self.prediction_length,
@@ -316,15 +305,32 @@ class TimeSeriesPredictor:
                 df = df.convert_frequency(freq=self.freq)
         return df
 
-    def _check_data_for_evaluation(self, data: TimeSeriesDataFrame, name: str = "data") -> None:
-        """Make sure that provided evaluation data includes both historical and future time series values."""
-        if data.num_timesteps_per_item().min() <= self.prediction_length:
+    def _check_and_prepare_data_frame_for_evaluation(
+        self, data: TimeSeriesDataFrame, cutoff: Optional[int] = None, name: str = "data"
+    ) -> TimeSeriesDataFrame:
+        """
+        Make sure that provided evaluation data includes both historical and future time series values.
+        Slices the dataframe based on cutoff, if needed.
+        """
+        cutoff = -1 * self.prediction_length if cutoff is None else cutoff
+        if not (isinstance(cutoff, int) and cutoff <= -self.prediction_length):
+            raise ValueError(f"`cutoff` should be a negative integer <= -prediction_length, got: {cutoff=}")
+
+        expected_length = -cutoff
+
+        if data.num_timesteps_per_item().min() <= expected_length:
+            var_name = "-cutoff" if expected_length > self.prediction_length else "prediction_length"
             raise ValueError(
-                f"Cannot reserve last prediction_length={self.prediction_length} time steps for evaluation in some "
+                f"Cannot reserve last {expected_length} time steps for evaluation in some "
                 f"time series in {name}. Please make sure that {name} includes both historical and future data, and that"
-                f"all time series have length > prediction_length (at least {self.prediction_length + 1})"
+                f"all time series have length > {var_name} (at least {expected_length + 1})"
             )
 
+        if cutoff < -self.prediction_length:
+            data = data.slice_by_timestep(None, cutoff + self.prediction_length)
+
+        return data
+
     def _get_dataset_stats(self, data: TimeSeriesDataFrame) -> str:
         ts_lengths = data.num_timesteps_per_item()
         median_length = ts_lengths.median()
@@ -712,7 +718,7 @@ class TimeSeriesPredictor:
 
         if tuning_data is not None:
             tuning_data = self._check_and_prepare_data_frame(tuning_data, name="tuning_data")
-            self._check_data_for_evaluation(tuning_data, name="tuning_data")
+            tuning_data = self._check_and_prepare_data_frame_for_evaluation(tuning_data, name="tuning_data")
             logger.info(f"Provided tuning_data has {self._get_dataset_stats(tuning_data)}")
             # TODO: Use num_val_windows to perform multi-window backtests on tuning_data
             if num_val_windows > 0:
@@ -857,6 +863,7 @@ class TimeSeriesPredictor:
         data: Union[TimeSeriesDataFrame, pd.DataFrame, Path, str],
         model: Optional[str] = None,
         metrics: Optional[Union[str, TimeSeriesScorer, List[Union[str, TimeSeriesScorer]]]] = None,
+        cutoff: Optional[int] = None,
         display: bool = False,
         use_cache: bool = True,
     ) -> Dict[str, float]:
@@ -874,11 +881,13 @@ class TimeSeriesPredictor:
         Parameters
         ----------
         data : Union[TimeSeriesDataFrame, pd.DataFrame, Path, str]
-            The data to evaluate the best model on. The last ``prediction_length`` time steps of each time series in
-            ``data`` will be held out for prediction and forecast accuracy will be calculated on these time steps.
+            The data to evaluate the best model on. If a ``cutoff`` is not provided, the last ``prediction_length``
+            time steps of each time series in ``data`` will be held out for prediction and forecast accuracy will
+            be calculated on these time steps. When a ``cutoff`` is provided, the ``-cutoff``-th to the
+            ``-cutoff + prediction_length``-th time steps of each time series are used for evaluation.
 
             Must include both historical and future data (i.e., length of all time series in ``data`` must be at least
-            ``prediction_length + 1``).
+            ``prediction_length + 1``, if ``cutoff`` is not provided, ``-cutoff + 1`` otherwise).
 
             The names and dtypes of columns and static features in ``data`` must match the ``train_data`` used to train
             the predictor.
@@ -891,6 +900,11 @@ class TimeSeriesPredictor:
         metrics : str, TimeSeriesScorer or List[Union[str, TimeSeriesScorer]], optional
             Metric or a list of metrics to compute scores with. Defaults to ``self.eval_metric``. Supports both
             metric names as strings and custom metrics based on TimeSeriesScorer.
+        cutoff : int, optional
+            A *negative* integer less than or equal to ``-1 * prediction_length`` denoting the time step in ``data``
+            where the forecast evaluation starts, i.e., time series are evaluated from the ``-cutoff``-th to the
+            ``-cutoff + prediction_length``-th time step. Defaults to ``-1 * prediction_length``, using the last
+            ``prediction_length`` time steps of each time series for evaluation.
         display : bool, default = False
             If True, the scores will be printed.
         use_cache : bool, default = True
@@ -904,8 +918,10 @@ class TimeSeriesPredictor:
             will have their signs flipped to obey this convention. For example, negative MAPE values will be reported.
             To get the ``eval_metric`` score, do ``output[predictor.eval_metric.name]``.
         """
+
         data = self._check_and_prepare_data_frame(data)
-        self._check_data_for_evaluation(data)
+        data = self._check_and_prepare_data_frame_for_evaluation(data, cutoff=cutoff)
+
         scores_dict = self._learner.evaluate(data, model=model, metrics=metrics, use_cache=use_cache)
         if display:
             logger.info("Evaluations on test data:")
@@ -1021,7 +1037,7 @@ class TimeSeriesPredictor:
         """
         if data is not None:
             data = self._check_and_prepare_data_frame(data)
-            self._check_data_for_evaluation(data)
+            data = self._check_and_prepare_data_frame_for_evaluation(data)
 
         fi_df = self._learner.get_feature_importance(
             data=data,
@@ -1199,6 +1215,7 @@ class TimeSeriesPredictor:
     def leaderboard(
         self,
         data: Optional[Union[TimeSeriesDataFrame, pd.DataFrame, Path, str]] = None,
+        cutoff: Optional[int] = None,
         extra_info: bool = False,
         extra_metrics: Optional[List[Union[str, TimeSeriesScorer]]] = None,
         display: bool = False,
@@ -1227,13 +1244,19 @@ class TimeSeriesPredictor:
         ----------
         data : Union[TimeSeriesDataFrame, pd.DataFrame, Path, str], optional
             dataset used for additional evaluation. Must include both historical and future data (i.e., length of all
-            time series in ``data`` must be at least ``prediction_length + 1``).
+            time series in ``data`` must be at least ``prediction_length + 1``, if ``cutoff`` is not provided,
+            ``-cutoff + 1`` otherwise).
 
             The names and dtypes of columns and static features in ``data`` must match the ``train_data`` used to train
             the predictor.
 
             If provided data is a `pandas.DataFrame`, AutoGluon will attempt to convert it to a `TimeSeriesDataFrame`.
             If a `str` or a `Path` is provided, AutoGluon will attempt to load this file.
+        cutoff : int, optional
+            A *negative* integer less than or equal to ``-1 * prediction_length`` denoting the time step in ``data``
+            where the forecast evaluation starts, i.e., time series are evaluated from the ``-cutoff``-th to the
+            ``-cutoff + prediction_length``-th time step. Defaults to ``-1 * prediction_length``, using the last
+            ``prediction_length`` time steps of each time series for evaluation.
         extra_info : bool, default = False
             If True, the leaderboard will contain an additional column `hyperparameters` with the hyperparameters used
             by each model during training. An empty dictionary `{}` means that the model was trained with default
@@ -1269,10 +1292,12 @@ class TimeSeriesPredictor:
                 raise TypeError(f"TimeSeriesPredictor.leaderboard() got an unexpected keyword argument '{key}'")
         if data is None and extra_metrics is not None:
             raise ValueError("`extra_metrics` is only valid when `data` is specified.")
+        if data is None and cutoff is not None:
+            raise ValueError("`cutoff` is only valid when `data` is specified.")
 
         if data is not None:
             data = self._check_and_prepare_data_frame(data)
-            self._check_data_for_evaluation(data)
+            data = self._check_and_prepare_data_frame_for_evaluation(data, cutoff=cutoff)
 
         leaderboard = self._learner.leaderboard(
             data, extra_info=extra_info, extra_metrics=extra_metrics, use_cache=use_cache
@@ -1441,7 +1466,7 @@ class TimeSeriesPredictor:
             return cast(TimeSeriesDataFrame, ts_df[[self.target]])
 
         test_data = self._check_and_prepare_data_frame(test_data)
-        self._check_data_for_evaluation(test_data, name="test_data")
+        test_data = self._check_and_prepare_data_frame_for_evaluation(test_data, name="test_data")
         test_data = self._learner.feature_generator.transform(test_data)
 
         trainer = self._trainer

autogluon/timeseries/trainer.py

@@ -19,8 +19,8 @@ from autogluon.core.utils.loaders import load_pkl
 from autogluon.core.utils.savers import save_pkl
 from autogluon.timeseries import TimeSeriesDataFrame
 from autogluon.timeseries.metrics import TimeSeriesScorer, check_get_evaluation_metric
-from autogluon.timeseries.models.abstract import AbstractTimeSeriesModel
-from autogluon.timeseries.models.ensemble import AbstractTimeSeriesEnsembleModel, TimeSeriesGreedyEnsemble
+from autogluon.timeseries.models.abstract import AbstractTimeSeriesModel, TimeSeriesModelBase
+from autogluon.timeseries.models.ensemble import AbstractTimeSeriesEnsembleModel, GreedyEnsemble
 from autogluon.timeseries.models.multi_window import MultiWindowBacktestingModel
 from autogluon.timeseries.models.presets import contains_searchspace, get_preset_models
 from autogluon.timeseries.splitter import AbstractWindowSplitter, ExpandingWindowSplitter
@@ -34,7 +34,7 @@ from autogluon.timeseries.utils.warning_filters import disable_tqdm, warning_fil
 logger = logging.getLogger("autogluon.timeseries.trainer")
 
 
-class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
+class TimeSeriesTrainer(AbstractTrainer[TimeSeriesModelBase]):
     _cached_predictions_filename = "cached_predictions.pkl"
 
     max_rel_importance_score: float = 1e5
@@ -47,7 +47,6 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
         prediction_length: int = 1,
         eval_metric: Union[str, TimeSeriesScorer, None] = None,
         eval_metric_seasonal_period: Optional[int] = None,
-        horizon_weight: Optional[np.ndarray] = None,
         save_data: bool = True,
         skip_model_selection: bool = False,
         enable_ensemble: bool = True,
@@ -74,12 +73,12 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
         # Ensemble cannot be fit if val_scores are not computed
         self.enable_ensemble = enable_ensemble and not skip_model_selection
         if ensemble_model_type is None:
-            ensemble_model_type = TimeSeriesGreedyEnsemble
+            ensemble_model_type = GreedyEnsemble
         else:
             logger.warning(
                 "Using a custom `ensemble_model_type` is experimental functionality that may break in future versions."
             )
-        self.ensemble_model_type = ensemble_model_type
+        self.ensemble_model_type: Type[AbstractTimeSeriesEnsembleModel] = ensemble_model_type
 
         self.verbosity = verbosity
 
@@ -89,7 +88,6 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
 
         self.eval_metric: TimeSeriesScorer = check_get_evaluation_metric(eval_metric)
         self.eval_metric_seasonal_period = eval_metric_seasonal_period
-        self.horizon_weight = horizon_weight
         if val_splitter is None:
             val_splitter = ExpandingWindowSplitter(prediction_length=self.prediction_length)
         assert isinstance(val_splitter, AbstractWindowSplitter), "val_splitter must be of type AbstractWindowSplitter"
@@ -147,7 +145,7 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
 
     def _add_model(
         self,
-        model: AbstractTimeSeriesModel,
+        model: TimeSeriesModelBase,
         base_models: Optional[List[str]] = None,
     ):
         """Add a model to the model graph of the trainer. If the model is an ensemble, also add
@@ -155,7 +153,7 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
 
         Parameters
         ----------
-        model : AbstractTimeSeriesModel
+        model : TimeSeriesModelBase
             The model to be added to the model graph.
         base_models : List[str], optional, default None
             If the model is an ensemble, the list of base model names that are included in the ensemble.
@@ -444,6 +442,8 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
         num_base_models = len(models)
         model_names_trained = []
         for i, model in enumerate(models):
+            assert isinstance(model, AbstractTimeSeriesModel)
+
             if time_limit is None:
                 time_left = None
                 time_left_for_model = None
@@ -560,20 +560,24 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
         return ensemble_name
 
     def fit_ensemble(
-        self, data_per_window: List[TimeSeriesDataFrame], model_names: List[str], time_limit: Optional[float] = None
+        self,
+        data_per_window: List[TimeSeriesDataFrame],
+        model_names: List[str],
+        time_limit: Optional[float] = None,
     ) -> str:
         logger.info("Fitting simple weighted ensemble.")
 
-        model_preds: Dict[str, List[TimeSeriesDataFrame]] = {}
+        predictions_per_window: Dict[str, List[TimeSeriesDataFrame]] = {}
+        base_model_scores = self.get_models_attribute_dict(attribute="val_score", models=self.get_model_names(0))
+
         for model_name in model_names:
-            model_preds[model_name] = self._get_model_oof_predictions(model_name=model_name)
+            predictions_per_window[model_name] = self._get_model_oof_predictions(model_name=model_name)
 
         time_start = time.time()
         ensemble = self.ensemble_model_type(
             name=self._get_ensemble_model_name(),
             eval_metric=self.eval_metric,
             eval_metric_seasonal_period=self.eval_metric_seasonal_period,
-            horizon_weight=self.horizon_weight,
             target=self.target,
             prediction_length=self.prediction_length,
             path=self.path,
@@ -582,7 +586,12 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
             covariate_metadata=self.covariate_metadata,
         )
         with warning_filter():
-            ensemble.fit_ensemble(model_preds, data_per_window=data_per_window, time_limit=time_limit)
+            ensemble.fit(
+                predictions_per_window=predictions_per_window,
+                data_per_window=data_per_window,
+                model_scores=base_model_scores,
+                time_limit=time_limit,
+            )
         ensemble.fit_time = time.time() - time_start
 
         predict_time = 0
@@ -592,7 +601,7 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
 
         score_per_fold = []
         for window_idx, data in enumerate(data_per_window):
-            predictions = ensemble.predict({n: model_preds[n][window_idx] for n in ensemble.model_names})
+            predictions = ensemble.predict({n: predictions_per_window[n][window_idx] for n in ensemble.model_names})
             score_per_fold.append(self._score_with_predictions(data, predictions))
         ensemble.val_score = float(np.mean(score_per_fold, dtype=np.float64))
 
@@ -737,7 +746,7 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
         return unpersisted_models
 
     def _get_model_for_prediction(
-        self, model: Optional[Union[str, AbstractTimeSeriesModel]] = None, verbose: bool = True
+        self, model: Optional[Union[str, TimeSeriesModelBase]] = None, verbose: bool = True
     ) -> str:
         """Given an optional identifier or model object, return the name of the model with which to predict.
 
@@ -754,7 +763,7 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
                 )
             return self.model_best
         else:
-            if isinstance(model, AbstractTimeSeriesModel):
+            if isinstance(model, TimeSeriesModelBase):
                 return model.name
             else:
                 if model not in self.get_model_names():
@@ -765,7 +774,7 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
         self,
         data: TimeSeriesDataFrame,
         known_covariates: Optional[TimeSeriesDataFrame] = None,
-        model: Optional[Union[str, AbstractTimeSeriesModel]] = None,
+        model: Optional[Union[str, TimeSeriesModelBase]] = None,
         use_cache: bool = True,
         random_seed: Optional[int] = None,
     ) -> TimeSeriesDataFrame:
@@ -796,13 +805,12 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
             prediction_length=self.prediction_length,
             target=self.target,
             seasonal_period=self.eval_metric_seasonal_period,
-            horizon_weight=self.horizon_weight,
         )
 
     def score(
         self,
         data: TimeSeriesDataFrame,
-        model: Optional[Union[str, AbstractTimeSeriesModel]] = None,
+        model: Optional[Union[str, TimeSeriesModelBase]] = None,
         metric: Union[str, TimeSeriesScorer, None] = None,
         use_cache: bool = True,
     ) -> float:
@@ -813,7 +821,7 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
     def evaluate(
         self,
         data: TimeSeriesDataFrame,
-        model: Optional[Union[str, AbstractTimeSeriesModel]] = None,
+        model: Optional[Union[str, TimeSeriesModelBase]] = None,
         metrics: Optional[Union[str, TimeSeriesScorer, List[Union[str, TimeSeriesScorer]]]] = None,
         use_cache: bool = True,
     ) -> Dict[str, float]:
@@ -835,7 +843,7 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
         self,
         data: TimeSeriesDataFrame,
         features: List[str],
-        model: Optional[Union[str, AbstractTimeSeriesModel]] = None,
+        model: Optional[Union[str, TimeSeriesModelBase]] = None,
         metric: Optional[Union[str, TimeSeriesScorer]] = None,
         time_limit: Optional[float] = None,
         method: Literal["naive", "permutation"] = "permutation",
@@ -938,7 +946,7 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
 
         return importance_df
 
-    def _model_uses_feature(self, model: Union[str, AbstractTimeSeriesModel], feature: str) -> bool:
+    def _model_uses_feature(self, model: Union[str, TimeSeriesModelBase], feature: str) -> bool:
         """Check if the given model uses the given feature."""
         models_with_ancestors = set(self.get_minimum_model_set(model))
 
@@ -980,7 +988,7 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
 
     def _predict_model(
         self,
-        model: Union[str, AbstractTimeSeriesModel],
+        model: Union[str, TimeSeriesModelBase],
         data: TimeSeriesDataFrame,
         model_pred_dict: Dict[str, Optional[TimeSeriesDataFrame]],
         known_covariates: Optional[TimeSeriesDataFrame] = None,
@@ -996,7 +1004,7 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
 
     def _get_inputs_to_model(
         self,
-        model: Union[str, AbstractTimeSeriesModel],
+        model: Union[str, TimeSeriesModelBase],
         data: TimeSeriesDataFrame,
         model_pred_dict: Dict[str, Optional[TimeSeriesDataFrame]],
     ) -> Union[TimeSeriesDataFrame, Dict[str, Optional[TimeSeriesDataFrame]]]:
@@ -1188,6 +1196,7 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
             model_name = model.name
             if model._get_tags()["can_refit_full"]:
                 model_full = model.convert_to_refit_full_template()
+                assert isinstance(model_full, AbstractTimeSeriesModel)
                 logger.info(f"Fitting model: {model_full.name}")
                 models_trained = self._train_and_save(
                     train_data=refit_full_data,
@@ -1253,12 +1262,11 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
         freq: Optional[str] = None,
         excluded_model_types: Optional[List[str]] = None,
         hyperparameter_tune: bool = False,
-    ) -> List[AbstractTimeSeriesModel]:
+    ) -> List[TimeSeriesModelBase]:
         return get_preset_models(
             path=self.path,
             eval_metric=self.eval_metric,
             eval_metric_seasonal_period=self.eval_metric_seasonal_period,
-            horizon_weight=self.horizon_weight,
             prediction_length=self.prediction_length,
             freq=freq,
             hyperparameters=hyperparameters,