autogluon.timeseries 1.0.1b20240304__py3-none-any.whl → 1.4.1b20251210__py3-none-any.whl

autogluon/timeseries/models/autogluon_tabular/mlforecast.py

@@ -1,16 +1,21 @@
+import copy
 import logging
 import math
-import os
 import time
-from typing import Any, Dict, List, Optional, Tuple
+import warnings
+from typing import Any, Callable, Collection, Type
 
 import numpy as np
 import pandas as pd
 from sklearn.base import BaseEstimator
 
 import autogluon.core as ag
-from autogluon.tabular import TabularPredictor
-from autogluon.timeseries.dataset.ts_dataframe import ITEMID, TIMESTAMP, TimeSeriesDataFrame
+from autogluon.core.models import AbstractModel as AbstractTabularModel
+from autogluon.features import AutoMLPipelineFeatureGenerator
+from autogluon.tabular.registry import ag_model_registry
+from autogluon.timeseries.dataset import TimeSeriesDataFrame
+from autogluon.timeseries.metrics.abstract import TimeSeriesScorer
+from autogluon.timeseries.metrics.utils import in_sample_squared_seasonal_error
 from autogluon.timeseries.models.abstract import AbstractTimeSeriesModel
 from autogluon.timeseries.models.local import SeasonalNaiveModel
 from autogluon.timeseries.utils.datetime import (
@@ -18,52 +23,53 @@ from autogluon.timeseries.utils.datetime import (
     get_seasonality,
     get_time_features_for_frequency,
 )
-from autogluon.timeseries.utils.forecast import get_forecast_horizon_index_ts_dataframe
-from autogluon.timeseries.utils.warning_filters import warning_filter
+from autogluon.timeseries.utils.warning_filters import set_loggers_level, warning_filter
 
-logger = logging.getLogger(__name__)
-
-MLF_TARGET = "y"
-MLF_ITEMID = "unique_id"
-MLF_TIMESTAMP = "ds"
+from .utils import MLF_ITEMID, MLF_TARGET, MLF_TIMESTAMP
 
+logger = logging.getLogger(__name__)
 
-class TabularEstimator(BaseEstimator):
-    """Scikit-learn compatible interface for TabularPredictor."""
 
-    def __init__(self, predictor_init_kwargs: Optional[dict] = None, predictor_fit_kwargs: Optional[dict] = None):
-        self.predictor_init_kwargs = predictor_init_kwargs if predictor_init_kwargs is not None else {}
-        self.predictor_fit_kwargs = predictor_fit_kwargs if predictor_fit_kwargs is not None else {}
+class TabularModel(BaseEstimator):
+    """A scikit-learn compatible wrapper for arbitrary autogluon.tabular models"""
 
-    def get_params(self, deep: bool = True) -> dict:
-        return {
-            "predictor_init_kwargs": self.predictor_init_kwargs,
-            "predictor_fit_kwargs": self.predictor_fit_kwargs,
-        }
+    def __init__(self, model_class: Type[AbstractTabularModel], model_kwargs: dict | None = None):
+        self.model_class = model_class
+        self.model_kwargs = {} if model_kwargs is None else model_kwargs
+        self.feature_pipeline = AutoMLPipelineFeatureGenerator(verbosity=0)
 
-    def fit(self, X: pd.DataFrame, y: pd.Series) -> "TabularEstimator":
-        assert isinstance(X, pd.DataFrame) and isinstance(y, pd.Series)
-        df = pd.concat([X, y.rename(MLF_TARGET).to_frame()], axis=1)
-        self.predictor = TabularPredictor(**self.predictor_init_kwargs)
-        with warning_filter():
-            self.predictor.fit(df, **self.predictor_fit_kwargs)
+    def fit(self, X: pd.DataFrame, y: pd.Series, X_val: pd.DataFrame, y_val: pd.Series, **kwargs):
+        self.model = self.model_class(**self.model_kwargs)
+        X = self.feature_pipeline.fit_transform(X=X)
+        X_val = self.feature_pipeline.transform(X=X_val)
+        self.model.fit(X=X, y=y, X_val=X_val, y_val=y_val, **kwargs)
         return self
 
-    def predict(self, X: pd.DataFrame) -> np.ndarray:
-        assert isinstance(X, pd.DataFrame)
-        return self.predictor.predict(X).values
+    def predict(self, X: pd.DataFrame, **kwargs):
+        X = self.feature_pipeline.transform(X=X)
+        return self.model.predict(X=X, **kwargs)
+
+    def get_params(self, deep=True):
+        params = {"model_class": self.model_class, "model_kwargs": self.model_kwargs}
+        if deep:
+            return copy.deepcopy(params)
+        else:
+            return params
 
 
 class AbstractMLForecastModel(AbstractTimeSeriesModel):
+    _supports_known_covariates = True
+    _supports_static_features = True
+
     def __init__(
         self,
-        freq: Optional[str] = None,
+        freq: str | None = None,
         prediction_length: int = 1,
-        path: Optional[str] = None,
-        name: Optional[str] = None,
-        eval_metric: str = None,
-        hyperparameters: Dict[str, Any] = None,
-        **kwargs,  # noqa
+        path: str | None = None,
+        name: str | None = None,
+        eval_metric: str | TimeSeriesScorer | None = None,
+        hyperparameters: dict[str, Any] | None = None,
+        **kwargs,
     ):
         super().__init__(
             path=path,
@@ -78,44 +84,90 @@ class AbstractMLForecastModel(AbstractTimeSeriesModel):
         from mlforecast.target_transforms import BaseTargetTransform
 
         self._sum_of_differences: int = 0  # number of time steps removed from each series by differencing
-        self._max_ts_length: Optional[int] = None
-        self._target_lags: Optional[List[int]] = None
-        self._date_features: Optional[List[str]] = None
-        self._mlf: Optional[MLForecast] = None
-        self._scaler: Optional[BaseTargetTransform] = None
-        self._residuals_std_per_item: Optional[pd.Series] = None
-        self._avg_residuals_std: Optional[float] = None
-
-    def _get_extra_tabular_init_kwargs(self) -> dict:
-        raise NotImplementedError
+        self._max_ts_length: int | None = None
+        self._target_lags: np.ndarray
+        self._date_features: list[Callable]
+        self._mlf: MLForecast
+        self._scaler: BaseTargetTransform | None = None
+        self._residuals_std_per_item: pd.Series
+        self._train_target_median: float | None = None
+        self._non_boolean_real_covariates: list[str] = []
+
+    def _initialize_transforms_and_regressor(self):
+        super()._initialize_transforms_and_regressor()
+        # Do not create a scaler in the model, scaler will be passed to MLForecast
+        self.target_scaler = None
 
-    def _get_model_params(self) -> dict:
-        model_params = super()._get_model_params().copy()
-        model_params.setdefault("max_num_items", 20_000)
-        model_params.setdefault("max_num_samples", 1_000_000)
-        model_params.setdefault("tabular_hyperparameters", {"GBM": {}})
-        model_params.setdefault("tabular_fit_kwargs", {})
-        return model_params
+    @property
+    def allowed_hyperparameters(self) -> list[str]:
+        return super().allowed_hyperparameters + [
+            "lags",
+            "date_features",
+            "differences",
+            "model_name",
+            "model_hyperparameters",
+            "max_num_items",
+            "max_num_samples",
+            "lag_transforms",
+        ]
+
+    def preprocess(
+        self,
+        data: TimeSeriesDataFrame,
+        known_covariates: TimeSeriesDataFrame | None = None,
+        is_train: bool = False,
+        **kwargs,
+    ) -> tuple[TimeSeriesDataFrame, TimeSeriesDataFrame | None]:
+        if is_train:
+            # All-NaN series are removed; partially-NaN series in train_data are handled inside _generate_train_val_dfs
+            all_nan_items = data.item_ids[
+                data[self.target].isna().groupby(TimeSeriesDataFrame.ITEMID, sort=False).all()
+            ]
+            if len(all_nan_items):
+                data = data.query("item_id not in @all_nan_items")
+        else:
+            data = data.fill_missing_values()
+            # Fill time series consisting of all NaNs with the median of target in train_data
+            if data.isna().any(axis=None):
+                data[self.target] = data[self.target].fillna(value=self._train_target_median)
+        return data, known_covariates
 
-    def _get_mlforecast_init_args(self, train_data: TimeSeriesDataFrame, model_params: dict) -> dict:
-        # TODO: Support lag generation for all pandas frequencies
-        # TODO: Support date_feature generation for all pandas frequencies
+    def _get_default_hyperparameters(self) -> dict[str, Any]:
+        return {
+            "max_num_items": 20_000,
+            "max_num_samples": 1_000_000,
+            "model_name": "GBM",
+            "model_hyperparameters": {},
+        }
+
+    def _create_tabular_model(self, model_name: str, model_hyperparameters: dict[str, Any]) -> TabularModel:
+        raise NotImplementedError
+
+    def _get_mlforecast_init_args(
+        self, train_data: TimeSeriesDataFrame, model_params: dict[str, Any]
+    ) -> dict[str, Any]:
         from mlforecast.target_transforms import Differences
 
-        from .utils import MeanAbsScaler, StandardScaler
+        from .transforms import MLForecastScaler
 
         lags = model_params.get("lags")
         if lags is None:
+            assert self.freq is not None
             lags = get_lags_for_frequency(self.freq)
         self._target_lags = np.array(sorted(set(lags)), dtype=np.int64)
 
         date_features = model_params.get("date_features")
         if date_features is None:
             date_features = get_time_features_for_frequency(self.freq)
-        self._date_features = date_features
+        known_covariates = self.covariate_metadata.known_covariates
+        conflicting = [f.__name__ for f in date_features if f.__name__ in known_covariates]
+        if conflicting:
+            logger.info(f"\tRemoved automatic date_features {conflicting} since they clash with known_covariates")
+        self._date_features = [f for f in date_features if f.__name__ not in known_covariates]
 
         target_transforms = []
         differences = model_params.get("differences")
+        assert isinstance(differences, Collection)
 
         ts_lengths = train_data.num_timesteps_per_item()
         required_ts_length = sum(differences) + 1
@@ -132,25 +184,22 @@ class AbstractMLForecastModel(AbstractTimeSeriesModel):
             target_transforms.append(Differences(differences))
             self._sum_of_differences = sum(differences)
 
-        scaler_name = model_params.get("scaler")
-        if scaler_name is None:
-            pass
-        elif scaler_name == "standard":
-            self._scaler = StandardScaler()
-        elif scaler_name == "mean_abs":
-            self._scaler = MeanAbsScaler()
-        else:
-            logger.warning(
-                f"Unrecognized `scaler` {scaler_name} (supported options: ['standard', 'mean_abs', None]). Scaling disabled."
+        if "target_scaler" in model_params and "scaler" in model_params:
+            warnings.warn(
+                f"Both 'target_scaler' and 'scaler' hyperparameters are provided to {self.__class__.__name__}. "
+                "Please only set the 'target_scaler' parameter."
             )
-
-        if self._scaler is not None:
+        # Support "scaler" for backward compatibility
+        scaler_type = model_params.get("target_scaler", model_params.get("scaler"))
+        if scaler_type is not None:
+            self._scaler = MLForecastScaler(scaler_type=scaler_type)
             target_transforms.append(self._scaler)
 
         return {
             "lags": self._target_lags.tolist(),
             "date_features": self._date_features,
             "target_transforms": target_transforms,
+            "lag_transforms": model_params.get("lag_transforms"),
         }
 
     def _mask_df(self, df: pd.DataFrame) -> pd.DataFrame:
@@ -162,13 +211,13 @@ class AbstractMLForecastModel(AbstractTimeSeriesModel):
         return df
 
     @staticmethod
-    def _shorten_all_series(mlforecast_df: pd.DataFrame, max_length: int):
+    def _shorten_all_series(mlforecast_df: pd.DataFrame, max_length: int) -> pd.DataFrame:
         logger.debug(f"Shortening all series to at most {max_length}")
         return mlforecast_df.groupby(MLF_ITEMID, as_index=False, sort=False).tail(max_length)
 
     def _generate_train_val_dfs(
-        self, data: TimeSeriesDataFrame, max_num_items: Optional[int] = None, max_num_samples: Optional[int] = None
-    ) -> Tuple[pd.DataFrame, pd.DataFrame]:
+        self, data: TimeSeriesDataFrame, max_num_items: int | None = None, max_num_samples: int | None = None
+    ) -> tuple[pd.DataFrame, pd.DataFrame]:
         # Exclude items that are too short for chosen differences - otherwise exception will be raised
         if self._sum_of_differences > 0:
             ts_lengths = data.num_timesteps_per_item()
@@ -181,6 +230,10 @@ class AbstractMLForecastModel(AbstractTimeSeriesModel):
             items_to_keep = data.item_ids.to_series().sample(n=int(max_num_items))  # noqa: F841
             data = data.query("item_id in @items_to_keep")
 
+        # MLForecast.preprocess does not support missing values, but we will exclude them later from the training set
+        missing_entries = data.index[data[self.target].isna()]
+        data = data.fill_missing_values()
+
         num_items = data.num_items
         mlforecast_df = self._to_mlforecast_df(data, data.static_features)
 
@@ -193,10 +246,14 @@ class AbstractMLForecastModel(AbstractTimeSeriesModel):
         # Unless we set static_features=[], MLForecast interprets all known covariates as static features
         df = self._mlf.preprocess(mlforecast_df, dropna=False, static_features=[])
         # df.query results in 2x memory saving compared to df.dropna(subset="y")
-        df = df.query("y.notnull()")
+        df = df.query("y.notnull()")  # type: ignore
 
         df = self._mask_df(df)
 
+        # We remove originally missing values filled via imputation from the training set
+        if len(missing_entries):
+            df = df.set_index(["unique_id", "ds"]).drop(missing_entries, errors="ignore").reset_index()
+
         if max_num_samples is not None and len(df) > max_num_samples:
             df = df.sample(n=max_num_samples)
 
@@ -208,12 +265,12 @@ class AbstractMLForecastModel(AbstractTimeSeriesModel):
         val_df = grouped_df.tail(val_rows_per_item)
         logger.debug(f"train_df shape: {train_df.shape}, val_df shape: {val_df.shape}")
 
-        return train_df.drop(columns=[MLF_TIMESTAMP]), val_df.drop(columns=[MLF_TIMESTAMP])
+        return train_df.drop(columns=[MLF_TIMESTAMP]), val_df.drop(columns=[MLF_TIMESTAMP])  # type: ignore
 
     def _to_mlforecast_df(
         self,
         data: TimeSeriesDataFrame,
-        static_features: pd.DataFrame,
+        static_features: pd.DataFrame | None,
         include_target: bool = True,
     ) -> pd.DataFrame:
         """Convert TimeSeriesDataFrame to a format expected by MLForecast methods `predict` and `preprocess`.
@@ -221,15 +278,33 @@ class AbstractMLForecastModel(AbstractTimeSeriesModel):
         Each row contains unique_id, ds, y, and (optionally) known covariates & static features.
         """
         # TODO: Add support for past_covariates
-        selected_columns = self.metadata.known_covariates_real.copy()
-        column_name_mapping = {ITEMID: MLF_ITEMID, TIMESTAMP: MLF_TIMESTAMP}
+        selected_columns = self.covariate_metadata.known_covariates.copy()
+        column_name_mapping = {TimeSeriesDataFrame.ITEMID: MLF_ITEMID, TimeSeriesDataFrame.TIMESTAMP: MLF_TIMESTAMP}
         if include_target:
             selected_columns += [self.target]
             column_name_mapping[self.target] = MLF_TARGET
 
         df = pd.DataFrame(data)[selected_columns].reset_index()
         if static_features is not None:
-            df = pd.merge(df, static_features, how="left", on=ITEMID, suffixes=(None, "_static_feat"))
+            df = pd.merge(
+                df, static_features, how="left", on=TimeSeriesDataFrame.ITEMID, suffixes=(None, "_static_feat")
+            )
+
+        for col in self._non_boolean_real_covariates:
+            # Normalize non-boolean features using mean_abs scaling
+            df[f"__scaled_{col}"] = (
+                df[col]
+                / df[col]
+                .abs()
+                .groupby(df[TimeSeriesDataFrame.ITEMID])
+                .mean()
+                .reindex(df[TimeSeriesDataFrame.ITEMID])
+                .values
+            )
+
+        # Convert float64 to float32 to reduce memory usage
+        float64_cols = list(df.select_dtypes(include="float64"))
+        df[float64_cols] = df[float64_cols].astype("float32")
 
         # We assume that df is sorted by 'unique_id' inside `TimeSeriesPredictor._check_and_prepare_data_frame`
         return df.rename(columns=column_name_mapping)
@@ -237,19 +312,26 @@ class AbstractMLForecastModel(AbstractTimeSeriesModel):
     def _fit(
         self,
         train_data: TimeSeriesDataFrame,
-        val_data: Optional[TimeSeriesDataFrame] = None,
-        time_limit: Optional[int] = None,
+        val_data: TimeSeriesDataFrame | None = None,
+        time_limit: float | None = None,
+        num_cpus: int | None = None,
+        num_gpus: int | None = None,
         verbosity: int = 2,
         **kwargs,
     ) -> None:
         from mlforecast import MLForecast
 
         self._check_fit_params()
+        self._log_unused_hyperparameters()
         fit_start_time = time.time()
-        # TabularEstimator is passed to MLForecast later to include tuning_data
-        model_params = self._get_model_params()
+        self._train_target_median = train_data[self.target].median()
+        for col in self.covariate_metadata.known_covariates_real:
+            if not set(train_data[col].unique()) == set([0, 1]):
+                self._non_boolean_real_covariates.append(col)
+        model_params = self.get_hyperparameters()
 
         mlforecast_init_args = self._get_mlforecast_init_args(train_data, model_params)
+        assert self.freq is not None
         self._mlf = MLForecast(models={}, freq=self.freq, **mlforecast_init_args)
 
         # We generate train/val splits from train_data and ignore val_data to avoid overfitting
@@ -259,54 +341,65 @@ class AbstractMLForecastModel(AbstractTimeSeriesModel):
             max_num_samples=model_params["max_num_samples"],
         )
 
-        estimator = TabularEstimator(
-            predictor_init_kwargs={
-                "path": os.path.join(self.path, "tabular_predictor"),
-                "verbosity": verbosity - 2,
-                "label": MLF_TARGET,
-                **self._get_extra_tabular_init_kwargs(),
-            },
-            predictor_fit_kwargs={
-                "tuning_data": val_df.drop(columns=[MLF_ITEMID]),
-                "time_limit": None if time_limit is None else time_limit - (time.time() - fit_start_time),
-                "hyperparameters": model_params["tabular_hyperparameters"],
-                **model_params["tabular_fit_kwargs"],
-            },
-        )
-        self._mlf.models = {"mean": estimator}
+        with set_loggers_level(regex=r"^autogluon\.(tabular|features).*", level=logging.ERROR):
+            tabular_model = self._create_tabular_model(
+                model_name=model_params["model_name"], model_hyperparameters=model_params["model_hyperparameters"]
+            )
+            tabular_model.fit(
+                X=train_df.drop(columns=[MLF_TARGET, MLF_ITEMID]),
+                y=train_df[MLF_TARGET],
+                X_val=val_df.drop(columns=[MLF_TARGET, MLF_ITEMID]),
+                y_val=val_df[MLF_TARGET],
+                time_limit=(None if time_limit is None else time_limit - (time.time() - fit_start_time)),
+                verbosity=verbosity - 1,
+            )
 
-        with warning_filter():
-            self._mlf.fit_models(X=train_df.drop(columns=[MLF_TARGET, MLF_ITEMID]), y=train_df[MLF_TARGET])
+        # We directly insert the trained model into models_ since calling _mlf.fit_models does not support X_val, y_val
+        self._mlf.models_ = {"mean": tabular_model}
 
         self._save_residuals_std(val_df)
 
+    def get_tabular_model(self) -> TabularModel:
+        """Get the underlying tabular regression model."""
+        assert "mean" in self._mlf.models_, "Call `fit` before calling `get_tabular_model`"
+        mean_estimator = self._mlf.models_["mean"]
+        assert isinstance(mean_estimator, TabularModel)
+        return mean_estimator
+
     def _save_residuals_std(self, val_df: pd.DataFrame) -> None:
         """Compute standard deviation of residuals for each item using the validation set.
 
-        Saves per-item residuals to `self.residuals_std_per_item` and average std to `self._avg_residuals_std`.
+        Saves per-item residuals to `self.residuals_std_per_item`.
         """
-        residuals = val_df[MLF_TARGET] - self._mlf.models_["mean"].predict(val_df)
-        self._residuals_std_per_item = residuals.pow(2.0).groupby(val_df[MLF_ITEMID], sort=False).mean().pow(0.5)
-        self._avg_residuals_std = np.sqrt(residuals.pow(2.0).mean())
+        residuals_df = val_df[[MLF_ITEMID, MLF_TARGET]]
+        mean_estimator = self.get_tabular_model()
 
-    def _get_scale_per_item(self, item_ids: pd.Index) -> pd.Series:
-        """Extract the '_scale' values from the scaler object, if available."""
-        raise NotImplementedError
+        residuals_df = residuals_df.assign(y_pred=mean_estimator.predict(val_df))
+        if self._scaler is not None:
+            # Scaler expects to find column MLF_TIMESTAMP even though it's not used - fill with dummy
+            residuals_df = residuals_df.assign(**{MLF_TIMESTAMP: np.datetime64("2010-01-01")})
+            residuals_df = self._scaler.inverse_transform(residuals_df)
+
+        assert isinstance(residuals_df, pd.DataFrame)
+        residuals = residuals_df[MLF_TARGET] - residuals_df["y_pred"]
+        self._residuals_std_per_item = (
+            residuals.pow(2.0).groupby(val_df[MLF_ITEMID].values, sort=False).mean().pow(0.5)  # type: ignore
+        )
 
     def _remove_short_ts_and_generate_fallback_forecast(
         self,
         data: TimeSeriesDataFrame,
-        known_covariates: Optional[TimeSeriesDataFrame] = None,
-    ) -> Tuple[TimeSeriesDataFrame, Optional[TimeSeriesDataFrame], Optional[TimeSeriesDataFrame]]:
+        known_covariates: TimeSeriesDataFrame | None = None,
+    ) -> tuple[TimeSeriesDataFrame, TimeSeriesDataFrame, TimeSeriesDataFrame | None]:
         """Remove series that are too short for chosen differencing from data and generate naive forecast for them.
 
         Returns
         -------
-        data_long : TimeSeriesDataFrame
+        data_long
             Data containing only time series that are long enough for the model to predict.
-        known_covariates_long : TimeSeriesDataFrame or None
+        known_covariates_long
             Future known covariates containing only time series that are long enough for the model to predict.
-        forecast_for_short_series : TimeSeriesDataFrame or None
+        forecast_for_short_series
             Seasonal naive forecast for short series, if there are any in the dataset.
         """
         ts_lengths = data.num_timesteps_per_item()
@@ -318,7 +411,12 @@ class AbstractMLForecastModel(AbstractTimeSeriesModel):
                 "Fallback model SeasonalNaive is used for these time series."
             )
             data_short = data.query("item_id in @short_series")
-            seasonal_naive = SeasonalNaiveModel(freq=self.freq, prediction_length=self.prediction_length)
+            seasonal_naive = SeasonalNaiveModel(
+                freq=self.freq,
+                prediction_length=self.prediction_length,
+                target=self.target,
+                quantile_levels=self.quantile_levels,
+            )
             seasonal_naive.fit(train_data=data_short)
             forecast_for_short_series = seasonal_naive.predict(data_short)
 
@@ -333,35 +431,41 @@ class AbstractMLForecastModel(AbstractTimeSeriesModel):
             forecast_for_short_series = None
         return data_long, known_covariates_long, forecast_for_short_series
 
-    def _add_gaussian_quantiles(self, predictions: pd.DataFrame, repeated_item_ids: pd.Series):
+    def _add_gaussian_quantiles(
+        self, predictions: pd.DataFrame, repeated_item_ids: pd.Series, past_target: pd.Series
+    ) -> pd.DataFrame:
         """
         Add quantile levels assuming that residuals follow normal distribution
         """
         from scipy.stats import norm
 
-        scale_per_item = self._get_scale_per_item(repeated_item_ids.unique())
         num_items = int(len(predictions) / self.prediction_length)
         sqrt_h = np.sqrt(np.arange(1, self.prediction_length + 1))
         # Series where normal_scale_per_timestep.loc[item_id].loc[N] = sqrt(1 + N) for N in range(prediction_length)
         normal_scale_per_timestep = pd.Series(np.tile(sqrt_h, num_items), index=repeated_item_ids)
 
         residuals_std_per_timestep = self._residuals_std_per_item.reindex(repeated_item_ids)
-        # Use avg_residuals_std in case unseen item received for prediction
-        if residuals_std_per_timestep.isna().any():
-            residuals_std_per_timestep = residuals_std_per_timestep.fillna(value=self._avg_residuals_std)
-        std_per_timestep = residuals_std_per_timestep * scale_per_item * normal_scale_per_timestep
+        # Use in-sample seasonal error in for items not seen during fit
+        items_not_seen_during_fit = residuals_std_per_timestep.index[residuals_std_per_timestep.isna()].unique()
+        if len(items_not_seen_during_fit) > 0:
+            scale_for_new_items: pd.Series = in_sample_squared_seasonal_error(
+                y_past=past_target.loc[items_not_seen_during_fit]
+            ).pow(0.5)
+            residuals_std_per_timestep = residuals_std_per_timestep.fillna(scale_for_new_items)
+
+        std_per_timestep = residuals_std_per_timestep * normal_scale_per_timestep
         for q in self.quantile_levels:
             predictions[str(q)] = predictions["mean"] + norm.ppf(q) * std_per_timestep.to_numpy()
         return predictions
 
-    def _more_tags(self) -> dict:
-        return {"can_refit_full": True}
+    def _more_tags(self) -> dict[str, Any]:
+        return {"allow_nan": True, "can_refit_full": True}
 
 
 class DirectTabularModel(AbstractMLForecastModel):
-    """Predict all future time series values simultaneously using TabularPredictor from AutoGluon-Tabular.
+    """Predict all future time series values simultaneously using a regression model from AutoGluon-Tabular.
 
-    A single TabularPredictor is used to forecast all future time series values using the following features:
+    A single tabular model is used to forecast all future time series values using the following features:
 
     - lag features (observed time series values) based on ``freq`` of the data
     - time features (e.g., day of the week) based on the timestamp of the measurement
@@ -370,8 +474,8 @@ class DirectTabularModel(AbstractMLForecastModel):
 
     Features not known during the forecast horizon (e.g., future target values) are replaced by NaNs.
 
-    If ``eval_metric.needs_quantile``, the TabularPredictor will be trained with ``"quantile"`` problem type.
-    Otherwise, TabularPredictor will be trained with ``"regression"`` problem type, and dummy quantiles will be
+    If ``eval_metric.needs_quantile``, the tabular regression model will be trained with ``"quantile"`` problem type.
+    Otherwise, the model will be trained with ``"regression"`` problem type, and dummy quantiles will be
     obtained by assuming that the residuals follow zero-mean normal distribution.
 
     Based on the `mlforecast <https://github.com/Nixtla/mlforecast>`_ library.
@@ -379,45 +483,55 @@ class DirectTabularModel(AbstractMLForecastModel):
 
     Other Parameters
     ----------------
-    lags : List[int], default = None
+    lags : list[int], default = None
         Lags of the target that will be used as features for predictions. If None, will be determined automatically
         based on the frequency of the data.
-    date_features : List[Union[str, Callable]], default = None
+    date_features : list[str | Callable], default = None
         Features computed from the dates. Can be pandas date attributes or functions that will take the dates as input.
         If None, will be determined automatically based on the frequency of the data.
-    differences : List[int], default = []
+    differences : list[int], default = []
         Differences to take of the target before computing the features. These are restored at the forecasting step.
-        If None, will be set to ``[seasonal_period]``, where seasonal_period is determined based on the data frequency.
         Defaults to no differencing.
-    scaler : {"standard", "mean_abs", None}, default = "mean_abs"
-        Scaling applied to each time series.
-    tabular_hyperparameters : Dict[Dict[str, Any]], optional
-        Hyperparameters dictionary passed to ``TabularPredictor.fit``. Contains the names of models that should be fit.
-        Defaults to ``{"GBM": {}}``.
-    tabular_fit_kwargs : Dict[str, Any], optional
-        Additional keyword arguments passed to ``TabularPredictor.fit``. Defaults to an empty dict.
+    target_scaler : {"standard", "mean_abs", "min_max", "robust", None}, default = "mean_abs"
+        Scaling applied to each time series. Scaling is applied after differencing.
+    model_name : str, default = "GBM"
+        Name of the tabular regression model. See ``autogluon.tabular.registry.ag_model_registry`` or
+        `the documentation <https://auto.gluon.ai/stable/api/autogluon.tabular.models.html>`_ for the list of available
+        tabular models.
+    model_hyperparameters : dict[str, Any], optional
+        Hyperparameters passed to the tabular regression model.
     max_num_items : int or None, default = 20_000
         If not None, the model will randomly select this many time series for training and validation.
     max_num_samples : int or None, default = 1_000_000
-        If not None, training dataset passed to TabularPredictor will contain at most this many rows (starting from the
-        end of each time series).
+        If not None, training dataset passed to the tabular regression model will contain at most this many rows
+        (starting from the end of each time series).
     """
 
+    ag_priority = 85
+
     @property
     def is_quantile_model(self) -> bool:
         return self.eval_metric.needs_quantile
 
-    def _get_model_params(self) -> dict:
-        model_params = super()._get_model_params()
-        model_params.setdefault("scaler", "mean_abs")
-        model_params.setdefault("differences", [])
+    def get_hyperparameters(self) -> dict[str, Any]:
+        model_params = super().get_hyperparameters()
+        # We don't set 'target_scaler' if user already provided 'scaler' to avoid overriding the user-provided value
+        if "scaler" not in model_params:
+            model_params.setdefault("target_scaler", "mean_abs")
+        if "differences" not in model_params or model_params["differences"] is None:
+            model_params["differences"] = []
+        if "lag_transforms" in model_params:
+            model_params.pop("lag_transforms")
+            logger.warning(f"{self.name} does not support the 'lag_transforms' hyperparameter.")
         return model_params
 
     def _mask_df(self, df: pd.DataFrame) -> pd.DataFrame:
         """Apply a mask that mimics the situation at prediction time when target/covariates are unknown during the
         forecast horizon.
         """
-        num_hidden = np.random.randint(0, self.prediction_length, size=len(df))
+        # Fix seed to make the model deterministic
+        rng = np.random.default_rng(seed=123)
+        num_hidden = rng.integers(0, self.prediction_length, size=len(df))
         lag_cols = [f"lag{lag}" for lag in self._target_lags]
         mask = num_hidden[:, None] < self._target_lags[None]  # shape [len(num_hidden), len(_target_lags)]
         # use df.loc[:, lag_cols] instead of df[lag_cols] to avoid SettingWithCopyWarning
@@ -428,41 +542,46 @@ class DirectTabularModel(AbstractMLForecastModel):
         if self.is_quantile_model:
             # Quantile model does not require residuals to produce prediction intervals
             self._residuals_std_per_item = pd.Series(1.0, index=val_df[MLF_ITEMID].unique())
-            self._avg_residuals_std = 1.0
         else:
             super()._save_residuals_std(val_df=val_df)
 
     def _predict(
         self,
         data: TimeSeriesDataFrame,
-        known_covariates: Optional[TimeSeriesDataFrame] = None,
+        known_covariates: TimeSeriesDataFrame | None = None,
         **kwargs,
     ) -> TimeSeriesDataFrame:
+        from .transforms import apply_inverse_transform
+
         original_item_id_order = data.item_ids
         data, known_covariates, forecast_for_short_series = self._remove_short_ts_and_generate_fallback_forecast(
             data=data, known_covariates=known_covariates
         )
         if len(data) == 0:
             # All time series are too short for chosen differences
+            assert forecast_for_short_series is not None
             return forecast_for_short_series
 
         if known_covariates is not None:
             data_future = known_covariates.copy()
         else:
-            future_index = get_forecast_horizon_index_ts_dataframe(data, self.prediction_length)
+            future_index = self.get_forecast_horizon_index(data)
             data_future = pd.DataFrame(columns=[self.target], index=future_index, dtype="float32")
         # MLForecast raises exception of target contains NaN. We use inf as placeholder, replace them by NaN afterwards
         data_future[self.target] = float("inf")
         data_extended = pd.concat([data, data_future])
-        mlforecast_df = self._to_mlforecast_df(data_extended, data.static_features)
+        mlforecast_df = self._to_mlforecast_df(data_extended, data.static_features)  # type: ignore
         if self._max_ts_length is not None:
             # We appended `prediction_length` time steps to each series, so increase length
             mlforecast_df = self._shorten_all_series(mlforecast_df, self._max_ts_length + self.prediction_length)
         df = self._mlf.preprocess(mlforecast_df, dropna=False, static_features=[])
+        assert isinstance(df, pd.DataFrame)
+
         df = df.groupby(MLF_ITEMID, sort=False).tail(self.prediction_length)
         df = df.replace(float("inf"), float("nan"))
 
-        raw_predictions = self._mlf.models_["mean"].predict(df)
+        mean_estimator = self.get_tabular_model()
+        raw_predictions = mean_estimator.predict(df)
         predictions = self._postprocess_predictions(raw_predictions, repeated_item_ids=df[MLF_ITEMID])
         # Paste columns one by one to preserve dtypes
         predictions[MLF_ITEMID] = df[MLF_ITEMID].values
@@ -473,57 +592,72 @@ class DirectTabularModel(AbstractMLForecastModel):
             mlforecast_df_past = self._to_mlforecast_df(data, None)
             if self._max_ts_length is not None:
                 mlforecast_df_past = self._shorten_all_series(mlforecast_df_past, self._max_ts_length)
-            self._mlf.preprocess(mlforecast_df_past, static_features=[])
+            self._mlf.preprocess(mlforecast_df_past, static_features=[], dropna=False)
+            assert self._mlf.ts.target_transforms is not None
             for tfm in self._mlf.ts.target_transforms[::-1]:
-                predictions = tfm.inverse_transform(predictions)
-        predictions = TimeSeriesDataFrame(predictions.rename(columns={MLF_ITEMID: ITEMID, MLF_TIMESTAMP: TIMESTAMP}))
+                predictions = apply_inverse_transform(predictions, transform=tfm)
+
+        if not self.is_quantile_model:
+            predictions = self._add_gaussian_quantiles(
+                predictions, repeated_item_ids=predictions[MLF_ITEMID], past_target=data[self.target]
+            )
+        predictions_tsdf: TimeSeriesDataFrame = TimeSeriesDataFrame(
+            predictions.rename(
+                columns={MLF_ITEMID: TimeSeriesDataFrame.ITEMID, MLF_TIMESTAMP: TimeSeriesDataFrame.TIMESTAMP}
+            )
+        )
 
         if forecast_for_short_series is not None:
-            predictions = pd.concat([predictions, forecast_for_short_series])
-            predictions = predictions.reindex(original_item_id_order, level=ITEMID)
-        return predictions
+            predictions_tsdf = pd.concat([predictions_tsdf, forecast_for_short_series])  # type: ignore
+            predictions_tsdf = predictions_tsdf.reindex(original_item_id_order, level=TimeSeriesDataFrame.ITEMID)
+
+        return predictions_tsdf
 
-    def _postprocess_predictions(self, predictions: np.ndarray, repeated_item_ids: pd.Series) -> pd.DataFrame:
+    def _postprocess_predictions(
+        self, predictions: np.ndarray | pd.Series, repeated_item_ids: pd.Series
+    ) -> pd.DataFrame:
         if self.is_quantile_model:
-            predictions = pd.DataFrame(predictions, columns=[str(q) for q in self.quantile_levels])
-            predictions.values.sort(axis=1)
-            predictions["mean"] = predictions["0.5"]
+            predictions_df = pd.DataFrame(predictions, columns=[str(q) for q in self.quantile_levels])
+            predictions_df.values.sort(axis=1)
+            predictions_df["mean"] = predictions_df["0.5"]
         else:
-            predictions = pd.DataFrame(predictions, columns=["mean"])
-            predictions = self._add_gaussian_quantiles(predictions, repeated_item_ids=repeated_item_ids)
-
-        column_order = ["mean"] + [col for col in predictions.columns if col != "mean"]
-        return predictions[column_order]
+            predictions_df = pd.DataFrame(predictions, columns=["mean"])
 
-    def _get_scale_per_item(self, item_ids: pd.Index) -> pd.Series:
-        # Rescaling is applied in the inverse_transform step, no need to scale predictions
-        return pd.Series(1.0, index=item_ids)
+        column_order = ["mean"] + [col for col in predictions_df.columns if col != "mean"]
+        return predictions_df[column_order]
 
-    def _get_extra_tabular_init_kwargs(self) -> dict:
+    def _create_tabular_model(self, model_name: str, model_hyperparameters: dict[str, Any]) -> TabularModel:
+        model_class = ag_model_registry.key_to_cls(model_name)
         if self.is_quantile_model:
-            return {
-                "problem_type": ag.constants.QUANTILE,
-                "quantile_levels": self.quantile_levels,
-                "eval_metric": "pinball_loss",
-            }
+            problem_type = ag.constants.QUANTILE
+            eval_metric = "pinball_loss"
+            model_hyperparameters["ag.quantile_levels"] = self.quantile_levels
         else:
-            return {
-                "problem_type": ag.constants.REGRESSION,
-                "eval_metric": self.eval_metric.equivalent_tabular_regression_metric or "mean_absolute_error",
-            }
+            problem_type = ag.constants.REGRESSION
+            eval_metric = self.eval_metric.equivalent_tabular_regression_metric or "mean_absolute_error"
+        return TabularModel(
+            model_class=model_class,
+            model_kwargs={
+                "path": "",
+                "name": model_class.__name__,
+                "hyperparameters": model_hyperparameters,
+                "problem_type": problem_type,
+                "eval_metric": eval_metric,
+            },
+        )
 
 
 class RecursiveTabularModel(AbstractMLForecastModel):
-    """Predict future time series values one by one using TabularPredictor from AutoGluon-Tabular.
+    """Predict future time series values one by one using a regression model from AutoGluon-Tabular.
 
-    A single TabularPredictor is used to forecast the future time series values using the following features:
+    A single tabular regression model is used to forecast the future time series values using the following features:
 
     - lag features (observed time series values) based on ``freq`` of the data
     - time features (e.g., day of the week) based on the timestamp of the measurement
     - known covariates (if available)
     - static features of each item (if available)
 
-    TabularPredictor will always be trained with ``"regression"`` problem type, and dummy quantiles will be
+    The tabular model will always be trained with ``"regression"`` problem type, and dummy quantiles will be
     obtained by assuming that the residuals follow zero-mean normal distribution.
 
     Based on the `mlforecast <https://github.com/Nixtla/mlforecast>`_ library.
@@ -531,39 +665,48 @@ class RecursiveTabularModel(AbstractMLForecastModel):
 
     Other Parameters
     ----------------
-    lags : List[int], default = None
+    lags : list[int], default = None
         Lags of the target that will be used as features for predictions. If None, will be determined automatically
         based on the frequency of the data.
-    date_features : List[Union[str, Callable]], default = None
+    date_features : list[str | Callable], default = None
         Features computed from the dates. Can be pandas date attributes or functions that will take the dates as input.
         If None, will be determined automatically based on the frequency of the data.
-    differences : List[int], default = None
+    differences : list[int], default = None
         Differences to take of the target before computing the features. These are restored at the forecasting step.
         If None, will be set to ``[seasonal_period]``, where seasonal_period is determined based on the data frequency.
-    scaler : {"standard", "mean_abs", None}, default = "standard"
-        Scaling applied to each time series.
-    tabular_hyperparameters : Dict[Dict[str, Any]], optional
-        Hyperparameters dictionary passed to ``TabularPredictor.fit``. Contains the names of models that should be fit.
-        Defaults to ``{"GBM": {}}``.
-    tabular_fit_kwargs : Dict[str, Any], optional
-        Additional keyword arguments passed to ``TabularPredictor.fit``. Defaults to an empty dict.
+    target_scaler : {"standard", "mean_abs", "min_max", "robust", None}, default = "standard"
+        Scaling applied to each time series. Scaling is applied after differencing.
+    lag_transforms : dict[int, list[Callable]], default = None
+        Dictionary mapping lag periods to transformation functions applied to lagged target values (e.g., rolling mean).
+        See `MLForecast documentation <https://nixtlaverse.nixtla.io/mlforecast/lag_transforms.html>`_ for more details.
+    model_name : str, default = "GBM"
+        Name of the tabular regression model. See ``autogluon.tabular.registry.ag_model_registry`` or
+        `the documentation <https://auto.gluon.ai/stable/api/autogluon.tabular.models.html>`_ for the list of available
+        tabular models.
+    model_hyperparameters : dict[str, Any], optional
+        Hyperparameters passed to the tabular regression model.
     max_num_items : int or None, default = 20_000
         If not None, the model will randomly select this many time series for training and validation.
     max_num_samples : int or None, default = 1_000_000
-        If not None, training dataset passed to TabularPredictor will contain at most this many rows (starting from the
-        end of each time series).
+        If not None, training dataset passed to the tabular regression model will contain at most this many rows
+        (starting from the end of each time series).
     """
 
-    def _get_model_params(self) -> dict:
-        model_params = super()._get_model_params()
-        model_params.setdefault("scaler", "standard")
-        model_params.setdefault("differences", [get_seasonality(self.freq)])
+    ag_priority = 90
+
+    def get_hyperparameters(self) -> dict[str, Any]:
+        model_params = super().get_hyperparameters()
+        # We don't set 'target_scaler' if user already provided 'scaler' to avoid overriding the user-provided value
+        if "scaler" not in model_params:
+            model_params.setdefault("target_scaler", "standard")
+        if "differences" not in model_params or model_params["differences"] is None:
+            model_params["differences"] = [get_seasonality(self.freq)]
         return model_params
 
     def _predict(
         self,
         data: TimeSeriesDataFrame,
-        known_covariates: Optional[TimeSeriesDataFrame] = None,
+        known_covariates: TimeSeriesDataFrame | None = None,
         **kwargs,
     ) -> TimeSeriesDataFrame:
         original_item_id_order = data.item_ids
@@ -572,14 +715,17 @@ class RecursiveTabularModel(AbstractMLForecastModel):
         )
         if len(data) == 0:
             # All time series are too short for chosen differences
+            assert forecast_for_short_series is not None
             return forecast_for_short_series
 
         new_df = self._to_mlforecast_df(data, data.static_features)
         if self._max_ts_length is not None:
             new_df = self._shorten_all_series(new_df, self._max_ts_length)
         if known_covariates is None:
-            future_index = get_forecast_horizon_index_ts_dataframe(data, self.prediction_length)
-            known_covariates = pd.DataFrame(columns=[self.target], index=future_index, dtype="float32")
+            future_index = self.get_forecast_horizon_index(data)
+            known_covariates = TimeSeriesDataFrame(
+                pd.DataFrame(columns=[self.target], index=future_index, dtype="float32")
+            )
         X_df = self._to_mlforecast_df(known_covariates, data.static_features, include_target=False)
         # If both covariates & static features are missing, set X_df = None to avoid exception from MLForecast
         if len(X_df.columns.difference([MLF_ITEMID, MLF_TIMESTAMP])) == 0:
@@ -590,23 +736,31 @@ class RecursiveTabularModel(AbstractMLForecastModel):
                 new_df=new_df,
                 X_df=X_df,
             )
-        predictions = raw_predictions.rename(columns={MLF_ITEMID: ITEMID, MLF_TIMESTAMP: TIMESTAMP})
-        predictions = TimeSeriesDataFrame(
-            self._add_gaussian_quantiles(predictions, repeated_item_ids=predictions[ITEMID])
+        assert isinstance(raw_predictions, pd.DataFrame)
+        raw_predictions = raw_predictions.rename(
+            columns={MLF_ITEMID: TimeSeriesDataFrame.ITEMID, MLF_TIMESTAMP: TimeSeriesDataFrame.TIMESTAMP}
         )
 
+        predictions: TimeSeriesDataFrame = TimeSeriesDataFrame(
+            self._add_gaussian_quantiles(
+                raw_predictions,
+                repeated_item_ids=raw_predictions[TimeSeriesDataFrame.ITEMID],
+                past_target=data[self.target],
+            )
+        )
         if forecast_for_short_series is not None:
-            predictions = pd.concat([predictions, forecast_for_short_series])
-        return predictions.reindex(original_item_id_order, level=ITEMID)
-
-    def _get_extra_tabular_init_kwargs(self) -> dict:
-        return {
-            "problem_type": ag.constants.REGRESSION,
-            "eval_metric": self.eval_metric.equivalent_tabular_regression_metric or "mean_absolute_error",
-        }
-
-    def _get_scale_per_item(self, item_ids: pd.Index) -> pd.Series:
-        if self._scaler is not None:
-            return self._scaler.stats_["_scale"].copy().reindex(item_ids)
-        else:
-            return pd.Series(1.0, index=item_ids)
+            predictions = pd.concat([predictions, forecast_for_short_series])  # type: ignore
+        return predictions.reindex(original_item_id_order, level=TimeSeriesDataFrame.ITEMID)
+
+    def _create_tabular_model(self, model_name: str, model_hyperparameters: dict[str, Any]) -> TabularModel:
+        model_class = ag_model_registry.key_to_cls(model_name)
+        return TabularModel(
+            model_class=model_class,
+            model_kwargs={
+                "path": "",
+                "name": model_class.__name__,
+                "hyperparameters": model_hyperparameters,
+                "problem_type": ag.constants.REGRESSION,
+                "eval_metric": self.eval_metric.equivalent_tabular_regression_metric or "mean_absolute_error",
+            },
+        )