autogluon.timeseries 1.0.1b20240405__tar.gz → 1.0.1b20240407__tar.gz

{autogluon.timeseries-1.0.1b20240405 → autogluon.timeseries-1.0.1b20240407}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: autogluon.timeseries
-Version: 1.0.1b20240405
+Version: 1.0.1b20240407
 Summary: AutoML for Image, Text, and Tabular Data
 Home-page: https://github.com/autogluon/autogluon
 Author: AutoGluon Community

{autogluon.timeseries-1.0.1b20240405 → autogluon.timeseries-1.0.1b20240407}/src/autogluon/timeseries/learner.py

@@ -198,7 +198,7 @@ class TimeSeriesLearner(AbstractLearner):
 
     def evaluate(
         self,
-        data: Union[TimeSeriesDataFrame, pd.DataFrame, str],
+        data: TimeSeriesDataFrame,
         model: Optional[str] = None,
         metrics: Optional[Union[str, TimeSeriesScorer, List[Union[str, TimeSeriesScorer]]]] = None,
         use_cache: bool = True,
@@ -206,6 +206,75 @@ class TimeSeriesLearner(AbstractLearner):
         data = self.feature_generator.transform(data)
         return self.load_trainer().evaluate(data=data, model=model, metrics=metrics, use_cache=use_cache)
 
+    def get_feature_importance(
+        self,
+        data: Optional[TimeSeriesDataFrame] = None,
+        model: Optional[str] = None,
+        metric: Optional[Union[str, TimeSeriesScorer]] = None,
+        features: Optional[List[str]] = None,
+        time_limit: Optional[float] = None,
+        method: Literal["naive", "permutation"] = "permutation",
+        subsample_size: int = 50,
+        num_iterations: int = 1,
+        random_seed: Optional[int] = None,
+        relative_scores: bool = False,
+        include_confidence_band: bool = True,
+        confidence_level: float = 0.99,
+    ) -> pd.DataFrame:
+        trainer = self.load_trainer()
+        if data is None:
+            data = trainer.load_val_data() or trainer.load_train_data()
+
+        # if features are provided in the dataframe, check that they are valid features in the covariate metadata
+        provided_static_columns = [] if data.static_features is None else data.static_features.columns
+        unused_features = [
+            f
+            for f in set(provided_static_columns).union(set(data.columns) - {self.target})
+            if f not in self.feature_generator.covariate_metadata.all_features
+        ]
+
+        if features is None:
+            features = self.feature_generator.covariate_metadata.all_features
+        else:
+            if len(features) == 0:
+                raise ValueError(
+                    "No features provided to compute feature importance. At least some valid features should be provided."
+                )
+            for fn in features:
+                if fn not in self.feature_generator.covariate_metadata.all_features and fn not in unused_features:
+                    raise ValueError(f"Feature {fn} not found in covariate metadata or the dataset.")
+
+        if len(set(features)) < len(features):
+            logger.warning(
+                "Duplicate feature names provided to compute feature importance. This will lead to unexpected behavior. "
+                "Please provide unique feature names across both static features and covariates."
+            )
+
+        data = self.feature_generator.transform(data)
+
+        importance_df = trainer.get_feature_importance(
+            data=data,
+            features=features,
+            model=model,
+            metric=metric,
+            time_limit=time_limit,
+            method=method,
+            subsample_size=subsample_size,
+            num_iterations=num_iterations,
+            random_seed=random_seed,
+            relative_scores=relative_scores,
+            include_confidence_band=include_confidence_band,
+            confidence_level=confidence_level,
+        )
+
+        for feature in set(features).union(unused_features):
+            if feature not in importance_df.index:
+                importance_df.loc[feature] = (
+                    [0, 0, 0] if not include_confidence_band else [0, 0, 0, float("nan"), float("nan")]
+                )
+
+        return importance_df
+
     def leaderboard(self, data: Optional[TimeSeriesDataFrame] = None, use_cache: bool = True) -> pd.DataFrame:
         if data is not None:
             data = self.feature_generator.transform(data)

{autogluon.timeseries-1.0.1b20240405 → autogluon.timeseries-1.0.1b20240407}/src/autogluon/timeseries/models/abstract/abstract_timeseries_model.py

@@ -3,7 +3,7 @@ import os
 import re
 import time
 from contextlib import nullcontext
-from typing import Any, Dict, List, Optional, Union
+from typing import Dict, List, Optional, Union
 
 from autogluon.common import space
 from autogluon.common.loaders import load_pkl
@@ -74,6 +74,10 @@ class AbstractTimeSeriesModel(AbstractModel):
     _preprocess_nonadaptive = None
     _preprocess_set_features = None
 
+    supports_known_covariates: bool = False
+    supports_past_covariates: bool = False
+    supports_static_features: bool = False
+
     def __init__(
         self,
         freq: Optional[str] = None,
@@ -296,6 +300,7 @@ class AbstractTimeSeriesModel(AbstractModel):
             of input items.
         """
         data = self.preprocess(data, is_train=False)
+        known_covariates = self.preprocess_known_covariates(known_covariates)
         predictions = self._predict(data=data, known_covariates=known_covariates, **kwargs)
         logger.debug(f"Predicting with model {self.name}")
         # "0.5" might be missing from the quantiles if self is a wrapper (MultiWindowBacktestingModel or ensemble)
@@ -358,7 +363,7 @@ class AbstractTimeSeriesModel(AbstractModel):
             time steps of each time series.
         """
         past_data, known_covariates = data.get_model_inputs_for_scoring(
-            prediction_length=self.prediction_length, known_covariates_names=self.metadata.known_covariates_real
+            prediction_length=self.prediction_length, known_covariates_names=self.metadata.known_covariates
         )
         predictions = self.predict(past_data, known_covariates=known_covariates)
         return self._score_with_predictions(data=data, predictions=predictions, metric=metric)
@@ -371,7 +376,7 @@ class AbstractTimeSeriesModel(AbstractModel):
     ) -> None:
         """Compute val_score, predict_time and cache out-of-fold (OOF) predictions."""
         past_data, known_covariates = val_data.get_model_inputs_for_scoring(
-            prediction_length=self.prediction_length, known_covariates_names=self.metadata.known_covariates_real
+            prediction_length=self.prediction_length, known_covariates_names=self.metadata.known_covariates
         )
         predict_start_time = time.time()
         oof_predictions = self.predict(past_data, known_covariates=known_covariates)
@@ -494,9 +499,14 @@ class AbstractTimeSeriesModel(AbstractModel):
 
         return hpo_models, analysis
 
-    def preprocess(self, data: TimeSeriesDataFrame, is_train: bool = False, **kwargs) -> Any:
+    def preprocess(self, data: TimeSeriesDataFrame, is_train: bool = False, **kwargs) -> TimeSeriesDataFrame:
         return data
 
+    def preprocess_known_covariates(
+        self, known_covariates: Optional[TimeSeriesDataFrame]
+    ) -> Optional[TimeSeriesDataFrame]:
+        return known_covariates
+
     def get_memory_size(self, **kwargs) -> Optional[int]:
         return None
 

{autogluon.timeseries-1.0.1b20240405 → autogluon.timeseries-1.0.1b20240407}/src/autogluon/timeseries/models/autogluon_tabular/mlforecast.py

@@ -242,7 +242,7 @@ 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()
+        selected_columns = self.metadata.known_covariates.copy()
         column_name_mapping = {ITEMID: MLF_ITEMID, TIMESTAMP: MLF_TIMESTAMP}
         if include_target:
             selected_columns += [self.target]
@@ -425,6 +425,9 @@ class DirectTabularModel(AbstractMLForecastModel):
         end of each time series).
     """
 
+    supports_known_covariates = True
+    supports_static_features = True
+
     @property
     def is_quantile_model(self) -> bool:
         return self.eval_metric.needs_quantile
@@ -576,6 +579,9 @@ class RecursiveTabularModel(AbstractMLForecastModel):
         end of each time series).
     """
 
+    supports_known_covariates = True
+    supports_static_features = True
+
     def _get_model_params(self) -> dict:
         model_params = super()._get_model_params()
         model_params.setdefault("scaler", "standard")

{autogluon.timeseries-1.0.1b20240405 → autogluon.timeseries-1.0.1b20240407}/src/autogluon/timeseries/models/chronos/model.py

@@ -181,7 +181,8 @@ class ChronosModel(AbstractTimeSeriesModel):
             )
             self.context_length = self.maximum_context_length
 
-        model_path_safe = str.replace(model_path_input, "/", "__")
+        # we truncate the name to avoid long path errors on Windows
+        model_path_safe = str(model_path_input).replace("/", "__").replace(os.path.sep, "__")[-50:]
         name = (name if name is not None else "Chronos") + f"[{model_path_safe}]"
 
         super().__init__(

{autogluon.timeseries-1.0.1b20240405 → autogluon.timeseries-1.0.1b20240407}/src/autogluon/timeseries/models/gluonts/abstract_gluonts.py

@@ -3,7 +3,7 @@ import os
 import shutil
 from datetime import timedelta
 from pathlib import Path
-from typing import Any, Callable, Dict, Iterator, List, Optional, Type, Union
+from typing import Any, Callable, Dict, Iterator, List, Literal, Optional, Type, Union
 
 import gluonts
 import gluonts.core.settings
@@ -16,9 +16,14 @@ from gluonts.model.estimator import Estimator as GluonTSEstimator
 from gluonts.model.forecast import Forecast, QuantileForecast, SampleForecast
 from gluonts.model.predictor import Predictor as GluonTSPredictor
 from pandas.tseries.frequencies import to_offset
+from sklearn.compose import ColumnTransformer
+from sklearn.preprocessing import QuantileTransformer, StandardScaler
 
 from autogluon.common.loaders import load_pkl
 from autogluon.core.hpo.constants import RAY_BACKEND
+from autogluon.tabular.models.tabular_nn.utils.categorical_encoders import (
+    OneHotMergeRaresHandleUnknownEncoder as OneHotEncoder,
+)
 from autogluon.timeseries.dataset.ts_dataframe import ITEMID, TIMESTAMP, TimeSeriesDataFrame
 from autogluon.timeseries.models.abstract import AbstractTimeSeriesModel
 from autogluon.timeseries.utils.datetime import norm_freq_str
@@ -42,21 +47,25 @@ class SimpleGluonTSDataset(GluonTSDataset):
         self,
         target_df: TimeSeriesDataFrame,
         target_column: str = "target",
-        feat_static_cat: Optional[pd.DataFrame] = None,
-        feat_static_real: Optional[pd.DataFrame] = None,
-        feat_dynamic_real: Optional[pd.DataFrame] = None,
-        past_feat_dynamic_real: Optional[pd.DataFrame] = None,
+        feat_static_cat: Optional[np.ndarray] = None,
+        feat_static_real: Optional[np.ndarray] = None,
+        feat_dynamic_cat: Optional[np.ndarray] = None,
+        feat_dynamic_real: Optional[np.ndarray] = None,
+        past_feat_dynamic_cat: Optional[np.ndarray] = None,
+        past_feat_dynamic_real: Optional[np.ndarray] = None,
         includes_future: bool = False,
         prediction_length: int = None,
     ):
         assert target_df is not None
         assert target_df.freq, "Initializing GluonTS data sets without freq is not allowed"
         # Convert TimeSeriesDataFrame to pd.Series for faster processing
-        self.target_array = self._to_array(target_df[target_column], dtype=np.float32)
-        self.feat_static_cat = self._to_array(feat_static_cat, dtype=np.int64)
-        self.feat_static_real = self._to_array(feat_static_real, dtype=np.float32)
-        self.feat_dynamic_real = self._to_array(feat_dynamic_real, dtype=np.float32)
-        self.past_feat_dynamic_real = self._to_array(past_feat_dynamic_real, dtype=np.float32)
+        self.target_array = target_df[target_column].to_numpy(np.float32)
+        self.feat_static_cat = self._astype(feat_static_cat, dtype=np.int64)
+        self.feat_static_real = self._astype(feat_static_real, dtype=np.float32)
+        self.feat_dynamic_cat = self._astype(feat_dynamic_cat, dtype=np.int64)
+        self.feat_dynamic_real = self._astype(feat_dynamic_real, dtype=np.float32)
+        self.past_feat_dynamic_cat = self._astype(past_feat_dynamic_cat, dtype=np.int64)
+        self.past_feat_dynamic_real = self._astype(past_feat_dynamic_real, dtype=np.float32)
         self.freq = self._to_gluonts_freq(target_df.freq)
 
         # Necessary to compute indptr for known_covariates at prediction time
@@ -73,11 +82,11 @@ class SimpleGluonTSDataset(GluonTSDataset):
         assert len(self.item_ids) == len(self.start_timestamps)
 
     @staticmethod
-    def _to_array(df: Optional[pd.DataFrame], dtype: np.dtype) -> Optional[np.ndarray]:
-        if df is None:
+    def _astype(array: Optional[np.ndarray], dtype: np.dtype) -> Optional[np.ndarray]:
+        if array is None:
             return None
         else:
-            return df.to_numpy(dtype=dtype)
+            return array.astype(dtype)
 
     @staticmethod
     def _to_gluonts_freq(freq: str) -> str:
@@ -111,12 +120,18 @@ class SimpleGluonTSDataset(GluonTSDataset):
                 ts[FieldName.FEAT_STATIC_CAT] = self.feat_static_cat[j]
             if self.feat_static_real is not None:
                 ts[FieldName.FEAT_STATIC_REAL] = self.feat_static_real[j]
+            if self.past_feat_dynamic_cat is not None:
+                ts[FieldName.PAST_FEAT_DYNAMIC_CAT] = self.past_feat_dynamic_cat[start_idx:end_idx].T
             if self.past_feat_dynamic_real is not None:
                 ts[FieldName.PAST_FEAT_DYNAMIC_REAL] = self.past_feat_dynamic_real[start_idx:end_idx].T
+
+            # Dynamic features that may extend into the future
+            if self.includes_future:
+                start_idx = start_idx + j * self.prediction_length
+                end_idx = end_idx + (j + 1) * self.prediction_length
+            if self.feat_dynamic_cat is not None:
+                ts[FieldName.FEAT_DYNAMIC_CAT] = self.feat_dynamic_cat[start_idx:end_idx].T
             if self.feat_dynamic_real is not None:
-                if self.includes_future:
-                    start_idx = start_idx + j * self.prediction_length
-                    end_idx = end_idx + (j + 1) * self.prediction_length
                 ts[FieldName.FEAT_DYNAMIC_REAL] = self.feat_dynamic_real[start_idx:end_idx].T
             yield ts
 
@@ -148,8 +163,7 @@ class AbstractGluonTSModel(AbstractTimeSeriesModel):
     gluonts_model_path = "gluon_ts"
     # default number of samples for prediction
     default_num_samples: int = 250
-    supports_known_covariates: bool = False
-    supports_past_covariates: bool = False
+    supports_cat_covariates: bool = False
 
     def __init__(
         self,
@@ -171,12 +185,20 @@ class AbstractGluonTSModel(AbstractTimeSeriesModel):
             **kwargs,
         )
         self.gts_predictor: Optional[GluonTSPredictor] = None
+        self._real_column_transformers: Dict[Literal["known", "past", "static"], ColumnTransformer] = {}
+        self._ohe_generator_known: Optional[OneHotEncoder] = None
+        self._ohe_generator_past: Optional[OneHotEncoder] = None
         self.callbacks = []
+        # Following attributes may be overridden during fit() based on train_data & model parameters
         self.num_feat_static_cat = 0
         self.num_feat_static_real = 0
+        self.num_feat_dynamic_cat = 0
         self.num_feat_dynamic_real = 0
+        self.num_past_feat_dynamic_cat = 0
         self.num_past_feat_dynamic_real = 0
         self.feat_static_cat_cardinality: List[int] = []
+        self.feat_dynamic_cat_cardinality: List[int] = []
+        self.past_feat_dynamic_cat_cardinality: List[int] = []
         self.negative_data = True
 
     def save(self, path: str = None, verbose: bool = True) -> str:
@@ -210,42 +232,136 @@ class AbstractGluonTSModel(AbstractTimeSeriesModel):
     def _get_hpo_backend(self):
         return RAY_BACKEND
 
-    def _deferred_init_params_aux(self, **kwargs) -> None:
-        """Update GluonTS specific parameters with information available
-        only at training time.
-        """
-        if "dataset" in kwargs:
-            ds = kwargs.get("dataset")
-            self.freq = ds.freq or self.freq
-            if not self.freq:
-                raise ValueError(
-                    "Dataset frequency not provided in the dataset, fit arguments or "
-                    "during initialization. Please provide a `freq` string to `fit`."
-                )
+    def _deferred_init_params_aux(self, dataset: TimeSeriesDataFrame) -> None:
+        """Update GluonTS specific parameters with information available only at training time."""
+        self.freq = dataset.freq or self.freq
+        if not self.freq:
+            raise ValueError(
+                "Dataset frequency not provided in the dataset, fit arguments or "
+                "during initialization. Please provide a `freq` string to `fit`."
+            )
 
-            model_params = self._get_model_params()
-            disable_static_features = model_params.get("disable_static_features", False)
-            if not disable_static_features:
-                self.num_feat_static_cat = len(self.metadata.static_features_cat)
-                self.num_feat_static_real = len(self.metadata.static_features_real)
-                if self.num_feat_static_cat > 0:
-                    feat_static_cat = ds.static_features[self.metadata.static_features_cat]
-                    self.feat_static_cat_cardinality = feat_static_cat.nunique().tolist()
-            disable_known_covariates = model_params.get("disable_known_covariates", False)
-            if not disable_known_covariates and self.supports_known_covariates:
-                self.num_feat_dynamic_real = len(self.metadata.known_covariates_real)
-            disable_past_covariates = model_params.get("disable_past_covariates", False)
-            if not disable_past_covariates and self.supports_past_covariates:
-                self.num_past_feat_dynamic_real = len(self.metadata.past_covariates_real)
-            self.negative_data = (ds[self.target] < 0).any()
-
-        if "callbacks" in kwargs:
-            self.callbacks += kwargs["callbacks"]
+        model_params = self._get_model_params()
+        disable_static_features = model_params.get("disable_static_features", False)
+        if not disable_static_features:
+            self.num_feat_static_cat = len(self.metadata.static_features_cat)
+            self.num_feat_static_real = len(self.metadata.static_features_real)
+            if self.num_feat_static_cat > 0:
+                feat_static_cat = dataset.static_features[self.metadata.static_features_cat]
+                self.feat_static_cat_cardinality = feat_static_cat.nunique().tolist()
+
+        disable_known_covariates = model_params.get("disable_known_covariates", False)
+        if not disable_known_covariates and self.supports_known_covariates:
+            self.num_feat_dynamic_cat = len(self.metadata.known_covariates_cat)
+            self.num_feat_dynamic_real = len(self.metadata.known_covariates_real)
+            if self.num_feat_dynamic_cat > 0:
+                feat_dynamic_cat = dataset[self.metadata.known_covariates_cat]
+                if self.supports_cat_covariates:
+                    self.feat_dynamic_cat_cardinality = feat_dynamic_cat.nunique().tolist()
+                else:
+                    # If model doesn't support categorical covariates, convert them to real via one hot encoding
+                    self._ohe_generator_known = OneHotEncoder(
+                        max_levels=model_params.get("max_cat_cardinality", 100),
+                        sparse=False,
+                        dtype="float32",
+                    )
+                    feat_dynamic_cat_ohe = self._ohe_generator_known.fit_transform(pd.DataFrame(feat_dynamic_cat))
+                    self.num_feat_dynamic_cat = 0
+                    self.num_feat_dynamic_real += feat_dynamic_cat_ohe.shape[1]
+
+        disable_past_covariates = model_params.get("disable_past_covariates", False)
+        if not disable_past_covariates and self.supports_past_covariates:
+            self.num_past_feat_dynamic_cat = len(self.metadata.past_covariates_cat)
+            self.num_past_feat_dynamic_real = len(self.metadata.past_covariates_real)
+            if self.num_past_feat_dynamic_cat > 0:
+                past_feat_dynamic_cat = dataset[self.metadata.past_covariates_cat]
+                if self.supports_cat_covariates:
+                    self.past_feat_dynamic_cat_cardinality = past_feat_dynamic_cat.nunique().tolist()
+                else:
+                    # If model doesn't support categorical covariates, convert them to real via one hot encoding
+                    self._ohe_generator_past = OneHotEncoder(
+                        max_levels=model_params.get("max_cat_cardinality", 100),
+                        sparse=False,
+                        dtype="float32",
+                    )
+                    past_feat_dynamic_cat_ohe = self._ohe_generator_past.fit_transform(
+                        pd.DataFrame(past_feat_dynamic_cat)
+                    )
+                    self.num_past_feat_dynamic_cat = 0
+                    self.num_past_feat_dynamic_real += past_feat_dynamic_cat_ohe.shape[1]
+
+        self.negative_data = (dataset[self.target] < 0).any()
 
     @property
     def default_context_length(self) -> int:
         return min(512, max(10, 2 * self.prediction_length))
 
+    def preprocess(self, data: TimeSeriesDataFrame, is_train: bool = False, **kwargs) -> TimeSeriesDataFrame:
+        # Copy data to avoid SettingWithCopyWarning from pandas
+        data = data.copy()
+        if self.supports_known_covariates and len(self.metadata.known_covariates_real) > 0:
+            columns = self.metadata.known_covariates_real
+            if is_train:
+                self._real_column_transformers["known"] = self._get_transformer_for_columns(data, columns=columns)
+            assert "known" in self._real_column_transformers, "Preprocessing pipeline must be fit first"
+            data[columns] = self._real_column_transformers["known"].transform(data[columns])
+
+        if self.supports_past_covariates and len(self.metadata.past_covariates_real) > 0:
+            columns = self.metadata.past_covariates_real
+            if is_train:
+                self._real_column_transformers["past"] = self._get_transformer_for_columns(data, columns=columns)
+            assert "past" in self._real_column_transformers, "Preprocessing pipeline must be fit first"
+            data[columns] = self._real_column_transformers["past"].transform(data[columns])
+
+        if self.supports_static_features and len(self.metadata.static_features_real) > 0:
+            columns = self.metadata.static_features_real
+            if is_train:
+                self._real_column_transformers["static"] = self._get_transformer_for_columns(
+                    data.static_features, columns=columns
+                )
+            assert "static" in self._real_column_transformers, "Preprocessing pipeline must be fit first"
+            data.static_features[columns] = self._real_column_transformers["static"].transform(
+                data.static_features[columns]
+            )
+        return data
+
+    def _get_transformer_for_columns(self, df: pd.DataFrame, columns: List[str]) -> Dict[str, str]:
+        """Passthrough bool features, use QuantileTransform for skewed features, and use StandardScaler for the rest.
+
+        The preprocessing logic is similar to the TORCH_NN model from Tabular.
+        """
+        skew_threshold = self._get_model_params().get("proc.skew_threshold", 0.99)
+        bool_features = []
+        skewed_features = []
+        continuous_features = []
+        for col in columns:
+            if df[col].isin([0, 1]).all():
+                bool_features.append(col)
+            elif np.abs(df[col].skew()) > skew_threshold:
+                skewed_features.append(col)
+            else:
+                continuous_features.append(col)
+        transformers = []
+        logger.debug(
+            f"\tbool_features: {bool_features}, continuous_features: {continuous_features}, skewed_features: {skewed_features}"
+        )
+        if continuous_features:
+            transformers.append(("scaler", StandardScaler(), continuous_features))
+        if skewed_features:
+            transformers.append(("skew", QuantileTransformer(output_distribution="normal"), skewed_features))
+        with warning_filter():
+            column_transformer = ColumnTransformer(transformers=transformers, remainder="passthrough").fit(df[columns])
+        return column_transformer
+
+    def preprocess_known_covariates(
+        self, known_covariates: Optional[TimeSeriesDataFrame]
+    ) -> Optional[TimeSeriesDataFrame]:
+        columns = self.metadata.known_covariates_real
+        if self.supports_known_covariates and len(columns) > 0:
+            assert "known" in self._real_column_transformers, "Preprocessing pipeline must be fit first"
+            known_covariates[columns] = self._real_column_transformers["known"].transform(known_covariates[columns])
+        return known_covariates
+
     def _get_model_params(self) -> dict:
         """Gets params that are passed to the inner model."""
         init_args = super()._get_model_params().copy()
@@ -322,42 +438,76 @@ class AbstractGluonTSModel(AbstractTimeSeriesModel):
         if time_series_df is not None:
             # TODO: Preprocess real-valued features with StdScaler?
             if self.num_feat_static_cat > 0:
-                feat_static_cat = time_series_df.static_features[self.metadata.static_features_cat]
+                feat_static_cat = time_series_df.static_features[self.metadata.static_features_cat].to_numpy()
             else:
                 feat_static_cat = None
 
             if self.num_feat_static_real > 0:
-                feat_static_real = time_series_df.static_features[self.metadata.static_features_real]
+                feat_static_real = time_series_df.static_features[self.metadata.static_features_real].to_numpy()
             else:
                 feat_static_real = None
 
+            expected_known_covariates_len = len(time_series_df) + self.prediction_length * time_series_df.num_items
+            # Convert TSDF -> DF to avoid overhead / input validation
+            df = pd.DataFrame(time_series_df)
+            if known_covariates is not None:
+                known_covariates = pd.DataFrame(known_covariates)
+            if self.num_feat_dynamic_cat > 0:
+                feat_dynamic_cat = df[self.metadata.known_covariates_cat].to_numpy()
+                if known_covariates is not None:
+                    feat_dynamic_cat = np.concatenate(
+                        [feat_dynamic_cat, known_covariates[self.metadata.known_covariates_cat].to_numpy()]
+                    )
+                    assert len(feat_dynamic_cat) == expected_known_covariates_len
+            else:
+                feat_dynamic_cat = None
+
             if self.num_feat_dynamic_real > 0:
-                # Convert TSDF -> DF to avoid overhead / input validation
-                feat_dynamic_real = pd.DataFrame(time_series_df[self.metadata.known_covariates_real])
+                feat_dynamic_real = df[self.metadata.known_covariates_real].to_numpy()
                 # Append future values of known covariates
                 if known_covariates is not None:
-                    feat_dynamic_real = pd.concat([feat_dynamic_real, known_covariates], axis=0)
-                    expected_length = len(time_series_df) + self.prediction_length * time_series_df.num_items
-                    if len(feat_dynamic_real) != expected_length:
-                        raise ValueError(
-                            f"known_covariates must contain values for the next prediction_length = "
-                            f"{self.prediction_length} time steps in each time series."
+                    feat_dynamic_real = np.concatenate(
+                        [feat_dynamic_real, known_covariates[self.metadata.known_covariates_real].to_numpy()]
+                    )
+                    assert len(feat_dynamic_real) == expected_known_covariates_len
+                # Categorical covariates are one-hot-encoded as real
+                if self._ohe_generator_known is not None:
+                    feat_dynamic_cat_ohe = self._ohe_generator_known.transform(df[self.metadata.known_covariates_cat])
+                    if known_covariates is not None:
+                        future_dynamic_cat_ohe = self._ohe_generator_known.transform(
+                            known_covariates[self.metadata.known_covariates_cat]
                         )
+                        feat_dynamic_cat_ohe = np.concatenate([feat_dynamic_cat_ohe, future_dynamic_cat_ohe])
+                        assert len(feat_dynamic_cat_ohe) == expected_known_covariates_len
+                    feat_dynamic_real = np.concatenate([feat_dynamic_real, feat_dynamic_cat_ohe], axis=1)
             else:
                 feat_dynamic_real = None
 
+            if self.num_past_feat_dynamic_cat > 0:
+                past_feat_dynamic_cat = df[self.metadata.past_covariates_cat].to_numpy()
+            else:
+                past_feat_dynamic_cat = None
+
             if self.num_past_feat_dynamic_real > 0:
-                # Convert TSDF -> DF to avoid overhead / input validation
-                past_feat_dynamic_real = pd.DataFrame(time_series_df[self.metadata.past_covariates_real])
+                past_feat_dynamic_real = df[self.metadata.past_covariates_real].to_numpy()
+                if self._ohe_generator_past is not None:
+                    past_feat_dynamic_cat_ohe = self._ohe_generator_past.transform(
+                        df[self.metadata.past_covariates_cat]
+                    )
+                    past_feat_dynamic_real = np.concatenate(
+                        [past_feat_dynamic_real, past_feat_dynamic_cat_ohe], axis=1
+                    )
             else:
                 past_feat_dynamic_real = None
 
             return SimpleGluonTSDataset(
-                target_df=time_series_df,
+                target_df=time_series_df[[self.target]],
                 target_column=self.target,
                 feat_static_cat=feat_static_cat,
                 feat_static_real=feat_static_real,
+                feat_dynamic_cat=feat_dynamic_cat,
                 feat_dynamic_real=feat_dynamic_real,
+                past_feat_dynamic_cat=past_feat_dynamic_cat,
                 past_feat_dynamic_real=past_feat_dynamic_real,
                 includes_future=known_covariates is not None,
                 prediction_length=self.prediction_length,
@@ -392,11 +542,11 @@ class AbstractGluonTSModel(AbstractTimeSeriesModel):
         # update auxiliary parameters
         init_args = self._get_estimator_init_args()
         keep_lightning_logs = init_args.pop("keep_lightning_logs", False)
-        callbacks = self._get_callbacks(
+        self.callbacks = self._get_callbacks(
             time_limit=time_limit,
             early_stopping_patience=None if val_data is None else init_args["early_stopping_patience"],
         )
-        self._deferred_init_params_aux(dataset=train_data, callbacks=callbacks)
+        self._deferred_init_params_aux(train_data)
 
         estimator = self._get_estimator()
         with warning_filter(), disable_root_logger(), gluonts.core.settings.let(gluonts.env.env, use_tqdm=False):

{autogluon.timeseries-1.0.1b20240405 → autogluon.timeseries-1.0.1b20240407}/src/autogluon/timeseries/models/gluonts/torch/models.py

@@ -61,6 +61,8 @@ class DeepARModel(AbstractGluonTSModel):
     embedding_dimension : int, optional
         Dimension of the embeddings for categorical features
         (if None, defaults to [min(50, (cat+1)//2) for cat in cardinality])
+    max_cat_cardinality : int, default = 100
+        Maximum number of dimensions to use when one-hot-encoding categorical known_covariates.
     distr_output : gluonts.torch.distributions.DistributionOutput, default = StudentTOutput()
         Distribution to use to evaluate observations and sample predictions
     scaling: bool, default = True
@@ -84,6 +86,7 @@ class DeepARModel(AbstractGluonTSModel):
     """
 
     supports_known_covariates = True
+    supports_static_features = True
 
     def _get_estimator_class(self) -> Type[GluonTSEstimator]:
         from gluonts.torch.model.deepar import DeepAREstimator
@@ -199,6 +202,8 @@ class TemporalFusionTransformerModel(AbstractGluonTSModel):
 
     supports_known_covariates = True
     supports_past_covariates = True
+    supports_cat_covariates = True
+    supports_static_features = True
 
     @property
     def default_context_length(self) -> int:
@@ -219,6 +224,11 @@ class TemporalFusionTransformerModel(AbstractGluonTSModel):
             init_kwargs["static_dims"] = [self.num_feat_static_real]
         if len(self.feat_static_cat_cardinality):
             init_kwargs["static_cardinalities"] = self.feat_static_cat_cardinality
+        if len(self.feat_dynamic_cat_cardinality):
+            init_kwargs["dynamic_cardinalities"] = self.feat_dynamic_cat_cardinality
+        if len(self.past_feat_dynamic_cat_cardinality):
+            init_kwargs["past_dynamic_cardinalities"] = self.past_feat_dynamic_cat_cardinality
+
         init_kwargs.setdefault("time_features", get_time_features_for_frequency(self.freq))
         return init_kwargs
 
@@ -372,6 +382,8 @@ class WaveNetModel(AbstractGluonTSModel):
         If True, logarithm of the scale of the past data will be used as an additional static feature.
     negative_data : bool, default = True
         Flag indicating whether the time series take negative values.
+    max_cat_cardinality : int, default = 100
+        Maximum number of dimensions to use when one-hot-encoding categorical known_covariates.
     max_epochs : int, default = 100
         Number of epochs the model will be trained for
     batch_size : int, default = 64
@@ -393,6 +405,7 @@ class WaveNetModel(AbstractGluonTSModel):
     """
 
     supports_known_covariates = True
+    supports_static_features = True
     default_num_samples: int = 100
 
     def _get_estimator_class(self) -> Type[GluonTSEstimator]: