autogluon.timeseries 1.4.1b20250907__py3-none-any.whl → 1.5.1b20260122__py3-none-any.whl

autogluon/timeseries/models/ensemble/__init__.py

@@ -1,3 +1,37 @@
 from .abstract import AbstractTimeSeriesEnsembleModel
-from .basic import PerformanceWeightedEnsemble, SimpleAverageEnsemble
-from .greedy import GreedyEnsemble
+from .array_based import LinearStackerEnsemble, MedianEnsemble, PerQuantileTabularEnsemble, TabularEnsemble
+from .per_item_greedy import PerItemGreedyEnsemble
+from .weighted import GreedyEnsemble, PerformanceWeightedEnsemble, SimpleAverageEnsemble
+
+
+def get_ensemble_class(name: str):
+    mapping = {
+        "Greedy": GreedyEnsemble,
+        "PerItemGreedy": PerItemGreedyEnsemble,
+        "PerformanceWeighted": PerformanceWeightedEnsemble,
+        "SimpleAverage": SimpleAverageEnsemble,
+        "Weighted": GreedyEnsemble,  # old alias for this model
+        "Median": MedianEnsemble,
+        "Tabular": TabularEnsemble,
+        "PerQuantileTabular": PerQuantileTabularEnsemble,
+        "LinearStacker": LinearStackerEnsemble,
+    }
+
+    name_clean = name.removesuffix("Ensemble")
+    if name_clean not in mapping:
+        raise ValueError(f"Unknown ensemble type: {name}. Available: {list(mapping.keys())}")
+    return mapping[name_clean]
+
+
+__all__ = [
+    "AbstractTimeSeriesEnsembleModel",
+    "GreedyEnsemble",
+    "LinearStackerEnsemble",
+    "MedianEnsemble",
+    "PerformanceWeightedEnsemble",
+    "PerItemGreedyEnsemble",
+    "PerQuantileTabularEnsemble",
+    "SimpleAverageEnsemble",
+    "TabularEnsemble",
+    "get_ensemble_class",
+]

autogluon/timeseries/models/ensemble/abstract.py

@@ -1,9 +1,6 @@
-import functools
 import logging
 from abc import ABC, abstractmethod
-from typing import Optional
 
-import numpy as np
 from typing_extensions import final
 
 from autogluon.core.utils.exceptions import TimeLimitExceeded
@@ -14,7 +11,12 @@ logger = logging.getLogger(__name__)
 
 
 class AbstractTimeSeriesEnsembleModel(TimeSeriesModelBase, ABC):
-    """Abstract class for time series ensemble models."""
+    """Abstract base class for time series ensemble models that combine predictions from multiple base models.
+
+    Ensemble training process operates on validation predictions from base models rather than raw time series
+    data. This allows the ensemble to learn optimal combination strategies based on each model's performance
+    across different validation windows and time series patterns.
+    """
 
     @property
     @abstractmethod
@@ -27,8 +29,8 @@ class AbstractTimeSeriesEnsembleModel(TimeSeriesModelBase, ABC):
         self,
         predictions_per_window: dict[str, list[TimeSeriesDataFrame]],
         data_per_window: list[TimeSeriesDataFrame],
-        model_scores: Optional[dict[str, float]] = None,
-        time_limit: Optional[float] = None,
+        model_scores: dict[str, float] | None = None,
+        time_limit: float | None = None,
     ):
         """Fit ensemble model given predictions of candidate base models and the true data.
 
@@ -52,7 +54,7 @@ class AbstractTimeSeriesEnsembleModel(TimeSeriesModelBase, ABC):
             )
             raise TimeLimitExceeded
         if isinstance(data_per_window, TimeSeriesDataFrame):
-            raise ValueError("When fitting ensemble, `data` should contain ground truth for each validation window")
+            raise ValueError("When fitting ensemble, ``data`` should contain ground truth for each validation window")
         num_val_windows = len(data_per_window)
         for model, preds in predictions_per_window.items():
             if len(preds) != num_val_windows:
@@ -69,11 +71,11 @@ class AbstractTimeSeriesEnsembleModel(TimeSeriesModelBase, ABC):
         self,
         predictions_per_window: dict[str, list[TimeSeriesDataFrame]],
         data_per_window: list[TimeSeriesDataFrame],
-        model_scores: Optional[dict[str, float]] = None,
-        time_limit: Optional[float] = None,
-    ):
-        """Private method for `fit`. See `fit` for documentation of arguments. Apart from the model
-        training logic, `fit` additionally implements other logic such as keeping track of the time limit.
+        model_scores: dict[str, float] | None = None,
+        time_limit: float | None = None,
+    ) -> None:
+        """Private method for ``fit``. See ``fit`` for documentation of arguments. Apart from the model
+        training logic, ``fit`` additionally implements other logic such as keeping track of the time limit.
         """
         raise NotImplementedError
 
@@ -103,37 +105,3 @@ class AbstractTimeSeriesEnsembleModel(TimeSeriesModelBase, ABC):
         This method should be called after performing refit_full to point to the refitted base models, if necessary.
         """
         pass
-
-
-class AbstractWeightedTimeSeriesEnsembleModel(AbstractTimeSeriesEnsembleModel, ABC):
-    """Abstract class for weighted ensembles which assign one (global) weight per model."""
-
-    def __init__(self, name: Optional[str] = None, **kwargs):
-        if name is None:
-            name = "WeightedEnsemble"
-        super().__init__(name=name, **kwargs)
-        self.model_to_weight: dict[str, float] = {}
-
-    @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, TimeSeriesDataFrame], **kwargs) -> TimeSeriesDataFrame:
-        weighted_predictions = [data[model_name] * weight for model_name, weight in self.model_to_weight.items()]
-        return functools.reduce(lambda x, y: x + y, weighted_predictions)
-
-    def get_info(self) -> dict:
-        info = super().get_info()
-        info["model_weights"] = self.model_to_weight.copy()
-        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/ensemble/array_based/__init__.py

@@ -0,0 +1,3 @@
+from .models import LinearStackerEnsemble, MedianEnsemble, PerQuantileTabularEnsemble, TabularEnsemble
+
+__all__ = ["LinearStackerEnsemble", "MedianEnsemble", "PerQuantileTabularEnsemble", "TabularEnsemble"]

autogluon/timeseries/models/ensemble/array_based/abstract.py

@@ -0,0 +1,240 @@
+from abc import ABC, abstractmethod
+from typing import Any, Sequence
+
+import numpy as np
+
+from autogluon.timeseries.dataset import TimeSeriesDataFrame
+from autogluon.timeseries.metrics.abstract import TimeSeriesScorer
+from autogluon.timeseries.utils.features import CovariateMetadata
+
+from ..abstract import AbstractTimeSeriesEnsembleModel
+from .regressor import EnsembleRegressor
+
+
+class ArrayBasedTimeSeriesEnsembleModel(AbstractTimeSeriesEnsembleModel, ABC):
+    """Abstract base class for ensemble models that operate on multi-dimensional arrays of base model predictions.
+
+    Array-based ensembles convert time series predictions into structured numpy arrays for efficient processing
+    and enable sophisticated combination strategies beyond simple weighted averaging. Array-based ensembles also
+    support isotonization in quantile forecasts--ensuring quantile crossing does not occur. They also have built-in
+    failed model detection and filtering capabilities.
+
+    Other Parameters
+    ----------------
+    isotonization : str, default = "sort"
+        The isotonization method to use (i.e. the algorithm to prevent quantile non-crossing).
+        Currently only "sort" is supported.
+    detect_and_ignore_failures : bool, default = True
+        Whether to detect and ignore "failed models", defined as models which have a loss that is larger
+        than 10x the median loss of all the models. This can be very important for the regression-based
+        ensembles, as moving the weight from such a "failed model" to zero can require a long training
+        time.
+    """
+
+    def __init__(
+        self,
+        path: str | None = None,
+        name: str | None = None,
+        hyperparameters: dict[str, Any] | None = None,
+        freq: str | None = None,
+        prediction_length: int = 1,
+        covariate_metadata: CovariateMetadata | None = None,
+        target: str = "target",
+        quantile_levels: Sequence[float] = (0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9),
+        eval_metric: str | TimeSeriesScorer | None = None,
+    ):
+        super().__init__(
+            path=path,
+            name=name,
+            hyperparameters=hyperparameters,
+            freq=freq,
+            prediction_length=prediction_length,
+            covariate_metadata=covariate_metadata,
+            target=target,
+            quantile_levels=quantile_levels,
+            eval_metric=eval_metric,
+        )
+        self.ensemble_regressor: EnsembleRegressor | None = None
+        self._model_names: list[str] = []
+
+    def _get_default_hyperparameters(self) -> dict[str, Any]:
+        return {
+            "isotonization": "sort",
+            "detect_and_ignore_failures": True,
+        }
+
+    @staticmethod
+    def to_array(df: TimeSeriesDataFrame) -> np.ndarray:
+        """Given a TimeSeriesDataFrame object, return a single array composing the values contained
+        in the data frame.
+
+        Parameters
+        ----------
+        df
+            TimeSeriesDataFrame to convert to an array. Must contain exactly ``prediction_length``
+            values for each item. The columns of ``df`` can correspond to ground truth values
+            or predictions (in which case, these will be the mean or quantile forecasts).
+
+        Returns
+        -------
+        array
+            of shape (num_items, prediction_length, num_outputs).
+        """
+        assert df.index.is_monotonic_increasing
+        array = df.to_numpy()
+        num_items = df.num_items
+        shape = (
+            num_items,
+            df.shape[0] // num_items,  # timesteps per item
+            df.shape[1],  # num_outputs
+        )
+        return array.reshape(shape)
+
+    def _get_base_model_predictions(
+        self,
+        predictions_per_window: dict[str, list[TimeSeriesDataFrame]] | dict[str, TimeSeriesDataFrame],
+    ) -> tuple[np.ndarray, np.ndarray]:
+        """Given a mapping from model names to a list of data frames representing
+        their predictions per window, return a multidimensional array representation.
+
+        Parameters
+        ----------
+        predictions_per_window
+            A dictionary with list[TimeSeriesDataFrame] values, where each TimeSeriesDataFrame
+            contains predictions for the window in question. If the dictionary values are
+            TimeSeriesDataFrame, they will be treated like a single window.
+
+        Returns
+        -------
+        base_model_mean_predictions
+            Array of shape (num_windows, num_items, prediction_length, 1, num_models)
+        base_model_quantile_predictions
+            Array of shape (num_windows, num_items, prediction_length, num_quantiles, num_models)
+        """
+
+        if not predictions_per_window:
+            raise ValueError("No base model predictions are provided.")
+
+        first_prediction = list(predictions_per_window.values())[0]
+        if isinstance(first_prediction, TimeSeriesDataFrame):
+            predictions_per_window = {k: [v] for k, v in predictions_per_window.items()}  # type: ignore
+
+        predictions = {
+            model_name: [self.to_array(window) for window in windows]  # type: ignore
+            for model_name, windows in predictions_per_window.items()
+        }
+        base_model_predictions = np.stack([x for x in predictions.values()], axis=-1)
+
+        return base_model_predictions[:, :, :, :1, :], base_model_predictions[:, :, :, 1:, :]
+
+    def _isotonize(self, prediction_array: np.ndarray) -> np.ndarray:
+        """Apply isotonization to ensure quantile non-crossing.
+
+        Parameters
+        ----------
+        prediction_array
+            Array of shape (num_windows, num_items, prediction_length, num_quantiles)
+
+        Returns
+        -------
+        isotonized_array
+            Array with same shape but quantiles sorted along last dimension
+        """
+        isotonization = self.get_hyperparameter("isotonization")
+        if isotonization == "sort":
+            return np.sort(prediction_array, axis=-1)
+        return prediction_array
+
+    def _fit(
+        self,
+        predictions_per_window: dict[str, list[TimeSeriesDataFrame]],
+        data_per_window: list[TimeSeriesDataFrame],
+        model_scores: dict[str, float] | None = None,
+        time_limit: float | None = None,
+    ) -> None:
+        # process inputs
+        filtered_predictions = self._filter_failed_models(predictions_per_window, model_scores)
+        base_model_mean_predictions, base_model_quantile_predictions = self._get_base_model_predictions(
+            filtered_predictions
+        )
+
+        # process labels
+        ground_truth_per_window = [y.slice_by_timestep(-self.prediction_length, None) for y in data_per_window]
+        labels = np.stack(
+            [self.to_array(gt) for gt in ground_truth_per_window], axis=0
+        )  # (num_windows, num_items, prediction_length, 1)
+
+        self._model_names = list(filtered_predictions.keys())
+        self.ensemble_regressor = self._get_ensemble_regressor()
+        self.ensemble_regressor.fit(
+            base_model_mean_predictions=base_model_mean_predictions,
+            base_model_quantile_predictions=base_model_quantile_predictions,
+            labels=labels,
+            time_limit=time_limit,
+        )
+
+    @abstractmethod
+    def _get_ensemble_regressor(self) -> EnsembleRegressor:
+        pass
+
+    def _predict(self, data: dict[str, TimeSeriesDataFrame], **kwargs) -> TimeSeriesDataFrame:
+        if self.ensemble_regressor is None:
+            if not self._model_names:
+                raise ValueError("Ensemble model has not been fitted yet.")
+            # Try to recreate the regressor (for loaded models)
+            self.ensemble_regressor = self._get_ensemble_regressor()
+
+        input_data = {}
+        for m in self.model_names:
+            assert m in data, f"Predictions for model {m} not provided during ensemble prediction."
+            input_data[m] = data[m]
+
+        base_model_mean_predictions, base_model_quantile_predictions = self._get_base_model_predictions(input_data)
+
+        mean_predictions, quantile_predictions = self.ensemble_regressor.predict(
+            base_model_mean_predictions=base_model_mean_predictions,
+            base_model_quantile_predictions=base_model_quantile_predictions,
+        )
+
+        quantile_predictions = self._isotonize(quantile_predictions)
+        prediction_array = np.concatenate([mean_predictions, quantile_predictions], axis=-1)
+
+        output = list(input_data.values())[0].copy()
+        num_folds, num_items, num_timesteps, num_outputs = prediction_array.shape
+        assert (num_folds, num_timesteps) == (1, self.prediction_length)
+        assert len(output.columns) == num_outputs
+
+        output[output.columns] = prediction_array.reshape((num_items * num_timesteps, num_outputs))
+
+        return output
+
+    @property
+    def model_names(self) -> list[str]:
+        return self._model_names
+
+    def remap_base_models(self, model_refit_map: dict[str, str]) -> None:
+        """Update names of the base models based on the mapping in model_refit_map."""
+        self._model_names = [model_refit_map.get(name, name) for name in self._model_names]
+
+    def _filter_failed_models(
+        self,
+        predictions_per_window: dict[str, list[TimeSeriesDataFrame]],
+        model_scores: dict[str, float] | None,
+    ) -> dict[str, list[TimeSeriesDataFrame]]:
+        """Filter out failed models based on detect_and_ignore_failures setting."""
+        if not self.get_hyperparameter("detect_and_ignore_failures"):
+            return predictions_per_window
+
+        if model_scores is None or len(model_scores) == 0:
+            return predictions_per_window
+
+        valid_scores = {k: v for k, v in model_scores.items() if np.isfinite(v)}
+        if len(valid_scores) == 0:
+            raise ValueError("All models have NaN scores. At least one model must run successfully to fit an ensemble")
+
+        losses = {k: -v for k, v in valid_scores.items()}
+        median_loss = np.nanmedian(list(losses.values()))
+        threshold = 10 * median_loss
+        good_models = {k for k, loss in losses.items() if loss <= threshold}
+
+        return {k: v for k, v in predictions_per_window.items() if k in good_models}

autogluon/timeseries/models/ensemble/array_based/models.py

@@ -0,0 +1,185 @@
+from abc import ABC
+from typing import Any, Type
+
+from autogluon.timeseries.dataset import TimeSeriesDataFrame
+
+from .abstract import ArrayBasedTimeSeriesEnsembleModel
+from .regressor import (
+    EnsembleRegressor,
+    LinearStackerEnsembleRegressor,
+    MedianEnsembleRegressor,
+    PerQuantileTabularEnsembleRegressor,
+    TabularEnsembleRegressor,
+)
+
+
+class MedianEnsemble(ArrayBasedTimeSeriesEnsembleModel):
+    """Robust ensemble that computes predictions as the element-wise median of base model mean
+    and quantile forecasts, providing robustness to outlier predictions.
+
+    Other Parameters
+    ----------------
+    isotonization : str, default = "sort"
+        The isotonization method to use (i.e. the algorithm to prevent quantile non-crossing).
+        Currently only "sort" is supported.
+    detect_and_ignore_failures : bool, default = True
+        Whether to detect and ignore "failed models", defined as models which have a loss that is larger
+        than 10x the median loss of all the models. This can be very important for the regression-based
+        ensembles, as moving the weight from such a "failed model" to zero can require a long training
+        time.
+    """
+
+    def _get_ensemble_regressor(self) -> MedianEnsembleRegressor:
+        return MedianEnsembleRegressor()
+
+
+class BaseTabularEnsemble(ArrayBasedTimeSeriesEnsembleModel, ABC):
+    ensemble_regressor_type: Type[EnsembleRegressor]
+
+    def _get_default_hyperparameters(self) -> dict[str, Any]:
+        default_hps = super()._get_default_hyperparameters()
+        default_hps.update({"model_name": "CAT", "model_hyperparameters": {}})
+        return default_hps
+
+    def _get_ensemble_regressor(self):
+        hyperparameters = self.get_hyperparameters()
+        return self.ensemble_regressor_type(
+            quantile_levels=list(self.quantile_levels),
+            model_name=hyperparameters["model_name"],
+            model_hyperparameters=hyperparameters["model_hyperparameters"],
+        )
+
+
+class TabularEnsemble(BaseTabularEnsemble):
+    """Tabular ensemble that uses a single AutoGluon-Tabular model to learn ensemble combinations.
+
+    This ensemble trains a single tabular model (such as gradient boosting machines) to predict all
+    quantiles simultaneously from base model predictions. The tabular model learns complex non-linear
+    patterns in how base models should be combined, potentially capturing interactions and conditional
+    dependencies that simple weighted averages cannot represent.
+
+    Other Parameters
+    ----------------
+    model_name : str, default = "CAT"
+        Name of the AutoGluon-Tabular model to use for ensemble learning. Model name should be registered
+        in AutoGluon-Tabular model registry.
+    model_hyperparameters : dict, default = {}
+        Hyperparameters to pass to the underlying AutoGluon-Tabular model.
+    isotonization : str, default = "sort"
+        The isotonization method to use (i.e. the algorithm to prevent quantile non-crossing).
+        Currently only "sort" is supported.
+    detect_and_ignore_failures : bool, default = True
+        Whether to detect and ignore "failed models", defined as models which have a loss that is larger
+        than 10x the median loss of all the models. This can be very important for the regression-based
+        ensembles, as moving the weight from such a "failed model" to zero can require a long training
+        time.
+    """
+
+    ensemble_regressor_type = TabularEnsembleRegressor
+
+
+class PerQuantileTabularEnsemble(BaseTabularEnsemble):
+    """Tabular ensemble using separate AutoGluon-Tabular models for each quantile and mean forecast.
+
+    This ensemble trains dedicated tabular models for each quantile level plus a separate model
+    for the mean prediction. Each model specializes in learning optimal combinations for its
+    specific target, allowing for quantile-specific ensemble strategies that can capture different
+    model behaviors across the prediction distribution.
+
+    Other Parameters
+    ----------------
+    model_name : str, default = "GBM"
+        Name of the AutoGluon-Tabular model to use for ensemble learning. Model name should be registered
+        in AutoGluon-Tabular model registry.
+    model_hyperparameters : dict, default = {}
+        Hyperparameters to pass to the underlying AutoGluon-Tabular model.
+    isotonization : str, default = "sort"
+        The isotonization method to use (i.e. the algorithm to prevent quantile non-crossing).
+        Currently only "sort" is supported.
+    detect_and_ignore_failures : bool, default = True
+        Whether to detect and ignore "failed models", defined as models which have a loss that is larger
+        than 10x the median loss of all the models. This can be very important for the regression-based
+        ensembles, as moving the weight from such a "failed model" to zero can require a long training
+        time.
+    """
+
+    ensemble_regressor_type = PerQuantileTabularEnsembleRegressor
+
+
+class LinearStackerEnsemble(ArrayBasedTimeSeriesEnsembleModel):
+    """Linear stacking ensemble that learns optimal linear combination weights through gradient-based
+    optimization.
+
+    Weighted combinations can be per model or per model-quantile, model-horizon, model-quantile-horizon
+    combinations. These choices are controlled by the ``weights_per`` hyperparameter.
+
+    The optimization process uses gradient descent with configurable learning rates and convergence
+    criteria, allowing for flexible training dynamics. Weight pruning can be applied to remove
+    models with negligible contributions, resulting in sparse and interpretable ensembles.
+
+    Other Parameters
+    ----------------
+    weights_per : str, default = "m"
+        Granularity of weight learning.
+
+        - "m": single weight per model
+        - "mq": single weight for each model-quantile combination
+        - "mt": single weight for each model-time step where time steps run across the prediction horizon
+        - "mtq": single weight for each model-quantile-time step combination
+    lr : float, default = 0.1
+        Learning rate for PyTorch optimizer during weight training.
+    max_epochs : int, default = 10000
+        Maximum number of training epochs for weight optimization.
+    relative_tolerance : float, default = 1e-7
+        Relative tolerance for convergence detection during training.
+    prune_below : float, default = 0.0
+        Threshold below which weights are pruned to zero for sparsity. The weights are redistributed across
+        remaining models after pruning.
+    isotonization : str, default = "sort"
+        The isotonization method to use (i.e. the algorithm to prevent quantile non-crossing).
+        Currently only "sort" is supported.
+    detect_and_ignore_failures : bool, default = True
+        Whether to detect and ignore "failed models", defined as models which have a loss that is larger
+        than 10x the median loss of all the models. This can be very important for the regression-based
+        ensembles, as moving the weight from such a "failed model" to zero can require a long training
+        time.
+    """
+
+    def _get_default_hyperparameters(self) -> dict[str, Any]:
+        default_hps = super()._get_default_hyperparameters()
+        default_hps.update(
+            {
+                "weights_per": "m",
+                "lr": 0.1,
+                "max_epochs": 10000,
+                "relative_tolerance": 1e-7,
+                "prune_below": 0.0,
+            }
+        )
+        return default_hps
+
+    def _get_ensemble_regressor(self) -> LinearStackerEnsembleRegressor:
+        hps = self.get_hyperparameters()
+        return LinearStackerEnsembleRegressor(
+            quantile_levels=list(self.quantile_levels),
+            weights_per=hps["weights_per"],
+            lr=hps["lr"],
+            max_epochs=hps["max_epochs"],
+            relative_tolerance=hps["relative_tolerance"],
+            prune_below=hps["prune_below"],
+        )
+
+    def _fit(
+        self,
+        predictions_per_window: dict[str, list[TimeSeriesDataFrame]],
+        data_per_window: list[TimeSeriesDataFrame],
+        model_scores: dict[str, float] | None = None,
+        time_limit: float | None = None,
+    ) -> None:
+        super()._fit(predictions_per_window, data_per_window, model_scores, time_limit)
+
+        assert isinstance(self.ensemble_regressor, LinearStackerEnsembleRegressor)
+
+        if self.ensemble_regressor.kept_indices is not None:
+            original_names = self._model_names
+            self._model_names = [original_names[i] for i in self.ensemble_regressor.kept_indices]

autogluon/timeseries/models/ensemble/array_based/regressor/__init__.py

@@ -0,0 +1,12 @@
+from .abstract import EnsembleRegressor, MedianEnsembleRegressor
+from .linear_stacker import LinearStackerEnsembleRegressor
+from .per_quantile_tabular import PerQuantileTabularEnsembleRegressor
+from .tabular import TabularEnsembleRegressor
+
+__all__ = [
+    "EnsembleRegressor",
+    "LinearStackerEnsembleRegressor",
+    "MedianEnsembleRegressor",
+    "PerQuantileTabularEnsembleRegressor",
+    "TabularEnsembleRegressor",
+]

autogluon/timeseries/models/ensemble/array_based/regressor/abstract.py

@@ -0,0 +1,88 @@
+from abc import ABC, abstractmethod
+
+import numpy as np
+from typing_extensions import Self
+
+
+class EnsembleRegressor(ABC):
+    def __init__(self, *args, **kwargs):
+        pass
+
+    @abstractmethod
+    def fit(
+        self,
+        base_model_mean_predictions: np.ndarray,
+        base_model_quantile_predictions: np.ndarray,
+        labels: np.ndarray,
+        time_limit: float | None = None,
+    ) -> Self:
+        """
+        Parameters
+        ----------
+        base_model_mean_predictions
+            Mean (point) predictions of base models. Array of shape
+            (num_windows, num_items, prediction_length, 1, num_models)
+
+        base_model_quantile_predictions
+            Quantile predictions of base models. Array of shape
+            (num_windows, num_items, prediction_length, num_quantiles, num_models)
+
+        labels
+            Ground truth array of shape
+            (num_windows, num_items, prediction_length, 1)
+
+        time_limit
+            Approximately how long ``fit`` will run (wall-clock time in seconds). If
+            not specified, training time will not be limited.
+        """
+        pass
+
+    @abstractmethod
+    def predict(
+        self,
+        base_model_mean_predictions: np.ndarray,
+        base_model_quantile_predictions: np.ndarray,
+    ) -> tuple[np.ndarray, np.ndarray]:
+        """Predict with the fitted ensemble regressor for a single window.
+        The items do not have to refer to the same item indices used when fitting
+        the model.
+
+        Parameters
+        ----------
+        base_model_mean_predictions
+            Mean (point) predictions of base models. Array of shape
+            (1, num_items, prediction_length, 1, num_models)
+
+        base_model_quantile_predictions
+            Quantile predictions of base models. Array of shape
+            (1, num_items, prediction_length, num_quantiles, num_models)
+
+        Returns
+        -------
+        ensemble_mean_predictions
+            Array of shape (1, num_items, prediction_length, 1)
+        ensemble_quantile_predictions
+            Array of shape (1, num_items, prediction_length, num_quantiles)
+        """
+        pass
+
+
+class MedianEnsembleRegressor(EnsembleRegressor):
+    def fit(
+        self,
+        base_model_mean_predictions: np.ndarray,
+        base_model_quantile_predictions: np.ndarray,
+        labels: np.ndarray,
+        time_limit: float | None = None,
+    ) -> Self:
+        return self
+
+    def predict(
+        self,
+        base_model_mean_predictions: np.ndarray,
+        base_model_quantile_predictions: np.ndarray,
+    ) -> tuple[np.ndarray, np.ndarray]:
+        return (
+            np.nanmedian(base_model_mean_predictions, axis=-1),
+            np.nanmedian(base_model_quantile_predictions, axis=-1),
+        )