autogluon.timeseries 1.4.1b20250820__py3-none-any.whl → 1.4.1b20250902__py3-none-any.whl

autogluon/timeseries/models/gluonts/abstract.py

@@ -3,7 +3,7 @@ import os
 import shutil
 from datetime import timedelta
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Type, Union, cast, overload
+from typing import TYPE_CHECKING, Any, Callable, Optional, Type, Union, cast, overload
 
 import gluonts
 import gluonts.core.settings
@@ -42,20 +42,20 @@ class AbstractGluonTSModel(AbstractTimeSeriesModel):
 
     Parameters
     ----------
-    path: str
+    path
         directory to store model artifacts.
-    freq: str
+    freq
         string representation (compatible with GluonTS frequency strings) for the data provided.
         For example, "1D" for daily data, "1H" for hourly data, etc.
-    prediction_length: int
+    prediction_length
         Number of time steps ahead (length of the forecast horizon) the model will be optimized
         to predict. At inference time, this will be the number of time steps the model will
         predict.
-    name: str
+    name
         Name of the model. Also, name of subdirectory inside path where model will be saved.
-    eval_metric: str
+    eval_metric
         objective function the model intends to optimize, will use WQL by default.
-    hyperparameters:
+    hyperparameters
         various hyperparameters that will be used by model (can be search spaces instead of
         fixed values). See *Other Parameters* in each inheriting model's documentation for
         possible values.
@@ -77,7 +77,7 @@ class AbstractGluonTSModel(AbstractTimeSeriesModel):
         path: Optional[str] = None,
         name: Optional[str] = None,
         eval_metric: Optional[str] = None,
-        hyperparameters: Optional[Dict[str, Any]] = None,
+        hyperparameters: Optional[dict[str, Any]] = None,
         **kwargs,  # noqa
     ):
         super().__init__(
@@ -100,9 +100,9 @@ class AbstractGluonTSModel(AbstractTimeSeriesModel):
         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.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: Optional[str] = None, verbose: bool = True) -> str:
@@ -234,7 +234,7 @@ class AbstractGluonTSModel(AbstractTimeSeriesModel):
 
         return self._get_default_hyperparameters() | init_args
 
-    def _get_estimator_init_args(self) -> Dict[str, Any]:
+    def _get_estimator_init_args(self) -> dict[str, Any]:
         """Get GluonTS specific constructor arguments for estimator objects, an alias to `self.get_hyperparameters`
         for better readability."""
         return self.get_hyperparameters()
@@ -277,8 +277,8 @@ class AbstractGluonTSModel(AbstractTimeSeriesModel):
 
         return torch.cuda.is_available()
 
-    def get_minimum_resources(self, is_gpu_available: bool = False) -> Dict[str, Union[int, float]]:
-        minimum_resources: Dict[str, Union[int, float]] = {"num_cpus": 1}
+    def get_minimum_resources(self, is_gpu_available: bool = False) -> dict[str, Union[int, float]]:
+        minimum_resources: dict[str, Union[int, float]] = {"num_cpus": 1}
         # if GPU is available, we train with 1 GPU per trial
         if is_gpu_available:
             minimum_resources["num_gpus"] = 1
@@ -440,7 +440,7 @@ class AbstractGluonTSModel(AbstractTimeSeriesModel):
         self,
         time_limit: Optional[float],
         early_stopping_patience: Optional[int] = None,
-    ) -> List[Callable]:
+    ) -> list[Callable]:
         """Retrieve a list of callback objects for the GluonTS trainer"""
         from lightning.pytorch.callbacks import EarlyStopping, Timer
 
@@ -473,7 +473,7 @@ class AbstractGluonTSModel(AbstractTimeSeriesModel):
         data: TimeSeriesDataFrame,
         known_covariates: Optional[TimeSeriesDataFrame] = None,
         num_samples: Optional[int] = None,
-    ) -> List[Forecast]:
+    ) -> list[Forecast]:
         assert self.gts_predictor is not None, "GluonTS models must be fit before predicting."
         gts_data = self._to_gluonts_dataset(data, known_covariates=known_covariates)
         return list(
@@ -483,7 +483,7 @@ class AbstractGluonTSModel(AbstractTimeSeriesModel):
             )
         )
 
-    def _stack_quantile_forecasts(self, forecasts: List[QuantileForecast], item_ids: pd.Index) -> pd.DataFrame:
+    def _stack_quantile_forecasts(self, forecasts: list[QuantileForecast], item_ids: pd.Index) -> pd.DataFrame:
         # GluonTS always saves item_id as a string
         item_id_to_forecast = {str(f.item_id): f for f in forecasts}
         result_dfs = []
@@ -496,7 +496,7 @@ class AbstractGluonTSModel(AbstractTimeSeriesModel):
         columns_order = ["mean"] + [str(q) for q in self.quantile_levels]
         return forecast_df[columns_order]
 
-    def _stack_sample_forecasts(self, forecasts: List[SampleForecast], item_ids: pd.Index) -> pd.DataFrame:
+    def _stack_sample_forecasts(self, forecasts: list[SampleForecast], item_ids: pd.Index) -> pd.DataFrame:
         item_id_to_forecast = {str(f.item_id): f for f in forecasts}
         samples_per_item = []
         for item_id in item_ids:
@@ -509,7 +509,7 @@ class AbstractGluonTSModel(AbstractTimeSeriesModel):
         return pd.DataFrame(forecast_array, columns=["mean"] + [str(q) for q in self.quantile_levels])
 
     def _stack_distribution_forecasts(
-        self, forecasts: List["DistributionForecast"], item_ids: pd.Index
+        self, forecasts: list["DistributionForecast"], item_ids: pd.Index
     ) -> pd.DataFrame:
         import torch
         from gluonts.torch.distributions import AffineTransformed
@@ -523,7 +523,7 @@ class AbstractGluonTSModel(AbstractTimeSeriesModel):
             "Expected forecast.distribution to be an instance of AffineTransformed"
         )
 
-        def stack_distributions(distributions: List[Distribution]) -> Distribution:
+        def stack_distributions(distributions: list[Distribution]) -> Distribution:
             """Stack multiple torch.Distribution objects into a single distribution"""
             last_dist: Distribution = distributions[-1]
 
@@ -561,18 +561,18 @@ class AbstractGluonTSModel(AbstractTimeSeriesModel):
 
     def _gluonts_forecasts_to_data_frame(
         self,
-        forecasts: List[Forecast],
+        forecasts: list[Forecast],
         forecast_index: pd.MultiIndex,
     ) -> TimeSeriesDataFrame:
         from gluonts.torch.model.forecast import DistributionForecast
 
         item_ids = forecast_index.unique(level=ITEMID)
         if isinstance(forecasts[0], SampleForecast):
-            forecast_df = self._stack_sample_forecasts(cast(List[SampleForecast], forecasts), item_ids)
+            forecast_df = self._stack_sample_forecasts(cast(list[SampleForecast], forecasts), item_ids)
         elif isinstance(forecasts[0], QuantileForecast):
-            forecast_df = self._stack_quantile_forecasts(cast(List[QuantileForecast], forecasts), item_ids)
+            forecast_df = self._stack_quantile_forecasts(cast(list[QuantileForecast], forecasts), item_ids)
         elif isinstance(forecasts[0], DistributionForecast):
-            forecast_df = self._stack_distribution_forecasts(cast(List[DistributionForecast], forecasts), item_ids)
+            forecast_df = self._stack_distribution_forecasts(cast(list[DistributionForecast], forecasts), item_ids)
         else:
             raise ValueError(f"Unrecognized forecast type {type(forecasts[0])}")
 

autogluon/timeseries/models/gluonts/dataset.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, Iterator, Optional, Type
+from typing import Any, Iterator, Optional, Type
 
 import numpy as np
 import pandas as pd
@@ -76,7 +76,7 @@ class SimpleGluonTSDataset(GluonTSDataset):
     def __len__(self):
         return len(self.indptr) - 1  # noqa
 
-    def __iter__(self) -> Iterator[Dict[str, Any]]:
+    def __iter__(self) -> Iterator[dict[str, Any]]:
         for j in range(len(self.indptr) - 1):
             start_idx = self.indptr[j]
             end_idx = self.indptr[j + 1]

autogluon/timeseries/models/gluonts/models.py

@@ -3,7 +3,7 @@ Module including wrappers for PyTorch implementations of models in GluonTS
 """
 
 import logging
-from typing import Any, Dict, Type
+from typing import Any, Type
 
 from gluonts.model.estimator import Estimator as GluonTSEstimator
 
@@ -91,7 +91,7 @@ class DeepARModel(AbstractGluonTSModel):
 
         return DeepAREstimator
 
-    def _get_estimator_init_args(self) -> Dict[str, Any]:
+    def _get_estimator_init_args(self) -> dict[str, Any]:
         init_kwargs = super()._get_estimator_init_args()
         init_kwargs["num_feat_static_cat"] = self.num_feat_static_cat
         init_kwargs["num_feat_static_real"] = self.num_feat_static_real
@@ -113,7 +113,7 @@ class SimpleFeedForwardModel(AbstractGluonTSModel):
     ----------------
     context_length : int, default = max(10, 2 * prediction_length)
         Number of time units that condition the predictions
-    hidden_dimensions: List[int], default = [20, 20]
+    hidden_dimensions: list[int], default = [20, 20]
         Size of hidden layers in the feedforward network
     distr_output : gluonts.torch.distributions.Output, default = StudentTOutput()
         Distribution output object that defines how the model output is converted to a forecast, and how the loss is computed.
@@ -221,7 +221,7 @@ class TemporalFusionTransformerModel(AbstractGluonTSModel):
             "context_length": min(512, max(64, 2 * self.prediction_length)),
         }
 
-    def _get_estimator_init_args(self) -> Dict[str, Any]:
+    def _get_estimator_init_args(self) -> dict[str, Any]:
         init_kwargs = super()._get_estimator_init_args()
         if self.num_feat_dynamic_real > 0:
             init_kwargs["dynamic_dims"] = [self.num_feat_dynamic_real]
@@ -361,7 +361,7 @@ class PatchTSTModel(AbstractGluonTSModel):
     def _get_default_hyperparameters(self):
         return super()._get_default_hyperparameters() | {"context_length": 96, "patch_len": 16}
 
-    def _get_estimator_init_args(self) -> Dict[str, Any]:
+    def _get_estimator_init_args(self) -> dict[str, Any]:
         init_kwargs = super()._get_estimator_init_args()
         init_kwargs["num_feat_dynamic_real"] = self.num_feat_dynamic_real
         return init_kwargs
@@ -438,7 +438,7 @@ class WaveNetModel(AbstractGluonTSModel):
 
         return WaveNetEstimator
 
-    def _get_estimator_init_args(self) -> Dict[str, Any]:
+    def _get_estimator_init_args(self) -> dict[str, Any]:
         init_kwargs = super()._get_estimator_init_args()
         init_kwargs["num_feat_static_cat"] = self.num_feat_static_cat
         init_kwargs["num_feat_static_real"] = self.num_feat_static_real
@@ -547,7 +547,7 @@ class TiDEModel(AbstractGluonTSModel):
             "batch_size": 256,
         }
 
-    def _get_estimator_init_args(self) -> Dict[str, Any]:
+    def _get_estimator_init_args(self) -> dict[str, Any]:
         init_kwargs = super()._get_estimator_init_args()
         init_kwargs["num_feat_static_cat"] = self.num_feat_static_cat
         init_kwargs["num_feat_static_real"] = self.num_feat_static_real

autogluon/timeseries/models/local/abstract_local_model.py

@@ -1,7 +1,7 @@
 import logging
 import time
 from multiprocessing import TimeoutError
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable, Optional, Union
 
 import numpy as np
 import pandas as pd
@@ -29,17 +29,17 @@ class AbstractLocalModel(AbstractTimeSeriesModel):
 
     Attributes
     ----------
-    allowed_local_model_args : List[str]
+    allowed_local_model_args
         Argument that can be passed to the underlying local model.
-    default_max_ts_length : Optional[int]
+    default_max_ts_length
         If not None, only the last ``max_ts_length`` time steps of each time series will be used to train the model.
         This significantly speeds up fitting and usually leads to no change in accuracy.
-    init_time_in_seconds : int
+    init_time_in_seconds
         Time that it takes to initialize the model in seconds (e.g., because of JIT compilation by Numba).
         If time_limit is below this number, model won't be trained.
     """
 
-    allowed_local_model_args: List[str] = []
+    allowed_local_model_args: list[str] = []
     default_max_ts_length: Optional[int] = 2500
     default_max_time_limit_ratio = 1.0
     init_time_in_seconds: int = 0
@@ -51,7 +51,7 @@ class AbstractLocalModel(AbstractTimeSeriesModel):
         path: Optional[str] = None,
         name: Optional[str] = None,
         eval_metric: Union[str, TimeSeriesScorer, None] = None,
-        hyperparameters: Optional[Dict[str, Any]] = None,
+        hyperparameters: Optional[dict[str, Any]] = None,
         **kwargs,  # noqa
     ):
         super().__init__(
@@ -64,12 +64,12 @@ class AbstractLocalModel(AbstractTimeSeriesModel):
             **kwargs,
         )
 
-        self._local_model_args: Dict[str, Any]
+        self._local_model_args: dict[str, Any]
         self._seasonal_period: int
         self._dummy_forecast: pd.DataFrame
 
     @property
-    def allowed_hyperparameters(self) -> List[str]:
+    def allowed_hyperparameters(self) -> list[str]:
         return (
             super().allowed_hyperparameters
             + ["use_fallback_model", "max_ts_length", "n_jobs"]
@@ -82,7 +82,7 @@ class AbstractLocalModel(AbstractTimeSeriesModel):
         known_covariates: Optional[TimeSeriesDataFrame] = None,
         is_train: bool = False,
         **kwargs,
-    ) -> Tuple[TimeSeriesDataFrame, Optional[TimeSeriesDataFrame]]:
+    ) -> tuple[TimeSeriesDataFrame, Optional[TimeSeriesDataFrame]]:
         if not self._get_tags()["allow_nan"]:
             data = data.fill_missing_values()
         return data, known_covariates
@@ -134,7 +134,7 @@ class AbstractLocalModel(AbstractTimeSeriesModel):
         stats_repeated = np.tile(stats_marginal.values, [self.prediction_length, 1])
         return pd.DataFrame(stats_repeated, columns=stats_marginal.index)
 
-    def _update_local_model_args(self, local_model_args: Dict[str, Any]) -> Dict[str, Any]:
+    def _update_local_model_args(self, local_model_args: dict[str, Any]) -> dict[str, Any]:
         return local_model_args
 
     def _predict(self, data: TimeSeriesDataFrame, **kwargs) -> TimeSeriesDataFrame:
@@ -185,7 +185,7 @@ class AbstractLocalModel(AbstractTimeSeriesModel):
         time_series: pd.Series,
         use_fallback_model: bool,
         end_time: Optional[float] = None,
-    ) -> Tuple[pd.DataFrame, bool]:
+    ) -> tuple[pd.DataFrame, bool]:
         if end_time is not None and time.time() >= end_time:
             raise TimeLimitExceeded
 
@@ -222,7 +222,7 @@ class AbstractLocalModel(AbstractTimeSeriesModel):
 
 
 def seasonal_naive_forecast(
-    target: np.ndarray, prediction_length: int, quantile_levels: List[float], seasonal_period: int
+    target: np.ndarray, prediction_length: int, quantile_levels: list[float], seasonal_period: int
 ) -> pd.DataFrame:
     """Generate seasonal naive forecast, predicting the last observed value from the same period."""
 

autogluon/timeseries/models/local/statsforecast.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Any, Dict, Optional, Type
+from typing import Any, Optional, Type
 
 import numpy as np
 import pandas as pd
@@ -14,7 +14,7 @@ class AbstractStatsForecastModel(AbstractLocalModel):
 
     init_time_in_seconds = 15  # numba compilation for the first run
 
-    def _update_local_model_args(self, local_model_args: Dict[str, Any]) -> Dict[str, Any]:
+    def _update_local_model_args(self, local_model_args: dict[str, Any]) -> dict[str, Any]:
         seasonal_period = local_model_args.pop("seasonal_period")
         local_model_args["season_length"] = seasonal_period
         return local_model_args
@@ -22,7 +22,7 @@ class AbstractStatsForecastModel(AbstractLocalModel):
     def _get_model_type(self, variant: Optional[str] = None) -> Type:
         raise NotImplementedError
 
-    def _get_local_model(self, local_model_args: Dict):
+    def _get_local_model(self, local_model_args: dict):
         local_model_args = local_model_args.copy()
         variant = local_model_args.pop("variant", None)
         model_type = self._get_model_type(variant)
@@ -31,7 +31,7 @@ class AbstractStatsForecastModel(AbstractLocalModel):
     def _get_point_forecast(
         self,
         time_series: pd.Series,
-        local_model_args: Dict,
+        local_model_args: dict,
     ) -> np.ndarray:
         return self._get_local_model(local_model_args).forecast(
             h=self.prediction_length, y=time_series.values.ravel()
@@ -176,9 +176,9 @@ class ARIMAModel(AbstractProbabilisticStatsForecastModel):
 
     Other Parameters
     ----------------
-    order: Tuple[int, int, int], default = (1, 1, 1)
+    order: tuple[int, int, int], default = (1, 1, 1)
         The (p, d, q) order of the model for the number of AR parameters, differences, and MA parameters to use.
-    seasonal_order: Tuple[int, int, int], default = (0, 0, 0)
+    seasonal_order: tuple[int, int, int], default = (0, 0, 0)
         The (P, D, Q) parameters of the seasonal ARIMA model. Setting to (0, 0, 0) disables seasonality.
     include_mean : bool, default = True
         Should the ARIMA model include a mean term?
@@ -194,7 +194,7 @@ class ARIMAModel(AbstractProbabilisticStatsForecastModel):
     method : {"CSS-ML", "CSS", "ML"}, default = "CSS-ML"
         Fitting method: CSS (conditional sum of squares), ML (maximum likelihood), CSS-ML (initialize with CSS, then
         optimize with ML).
-    fixed : Dict[str, float], optional
+    fixed : dict[str, float], optional
         Dictionary containing fixed coefficients for the ARIMA model.
     seasonal_period : int or None, default = None
         Number of time steps in a complete seasonal cycle for seasonal models. For example, 7 for daily data with a
@@ -449,7 +449,7 @@ class AbstractConformalizedStatsForecastModel(AbstractStatsForecastModel):
     def _get_nonconformity_scores(
         self,
         time_series: pd.Series,
-        local_model_args: Dict,
+        local_model_args: dict,
     ) -> np.ndarray:
         h = self.prediction_length
         y = time_series.values.ravel()
@@ -556,7 +556,7 @@ class AutoCESModel(AbstractProbabilisticStatsForecastModel):
         local_model_args.setdefault("model", "Z")
         return local_model_args
 
-    def _get_point_forecast(self, time_series: pd.Series, local_model_args: Dict):
+    def _get_point_forecast(self, time_series: pd.Series, local_model_args: dict):
         # Disable seasonality if time series too short for chosen season_length or season_length == 1,
         # otherwise model will crash
         if len(time_series) < 5:
@@ -568,7 +568,7 @@ class AutoCESModel(AbstractProbabilisticStatsForecastModel):
 
 
 class AbstractStatsForecastIntermittentDemandModel(AbstractConformalizedStatsForecastModel):
-    def _update_local_model_args(self, local_model_args: Dict[str, Any]) -> Dict[str, Any]:
+    def _update_local_model_args(self, local_model_args: dict[str, Any]) -> dict[str, Any]:
         _ = local_model_args.pop("seasonal_period")
         return local_model_args
 
@@ -733,6 +733,6 @@ class ZeroModel(AbstractStatsForecastIntermittentDemandModel):
     def _get_point_forecast(
         self,
         time_series: pd.Series,
-        local_model_args: Dict,
+        local_model_args: dict,
     ):
         return np.zeros(self.prediction_length)

autogluon/timeseries/models/multi_window/multi_window_model.py

@@ -4,9 +4,10 @@ import logging
 import math
 import os
 import time
-from typing import Any, Dict, Optional, Type, Union
+from typing import Any, Optional, Type, Union
 
 import numpy as np
+from typing_extensions import Self
 
 import autogluon.core as ag
 from autogluon.timeseries.dataset.ts_dataframe import TimeSeriesDataFrame
@@ -25,10 +26,10 @@ class MultiWindowBacktestingModel(AbstractTimeSeriesModel):
 
     Parameters
     ----------
-    model_base : Union[AbstractTimeSeriesModel, Type[AbstractTimeSeriesModel]]
+    model_base
         The base model to repeatedly train. If a AbstractTimeSeriesModel class, then also provide model_base_kwargs
         which will be used to initialize the model via model_base(**model_base_kwargs).
-    model_base_kwargs : Optional[Dict[str, any]], default = None
+    model_base_kwargs
         kwargs used to initialize model_base if model_base is a class.
     """
 
@@ -38,7 +39,7 @@ class MultiWindowBacktestingModel(AbstractTimeSeriesModel):
     def __init__(
         self,
         model_base: Union[AbstractTimeSeriesModel, Type[AbstractTimeSeriesModel]],
-        model_base_kwargs: Optional[Dict[str, Any]] = None,
+        model_base_kwargs: Optional[dict[str, Any]] = None,
         **kwargs,
     ):
         if inspect.isclass(model_base) and issubclass(model_base, AbstractTimeSeriesModel):
@@ -73,10 +74,6 @@ class MultiWindowBacktestingModel(AbstractTimeSeriesModel):
     def supports_past_covariates(self) -> bool:
         return self.model_base.supports_past_covariates
 
-    @property
-    def supports_cat_covariates(self) -> bool:
-        return self.model_base.supports_cat_covariates
-
     def _get_model_base(self):
         return self.model_base
 
@@ -86,15 +83,18 @@ class MultiWindowBacktestingModel(AbstractTimeSeriesModel):
     def _is_gpu_available(self) -> bool:
         return self._get_model_base()._is_gpu_available()
 
-    def get_minimum_resources(self, is_gpu_available: bool = False) -> bool:
+    def get_minimum_resources(self, is_gpu_available: bool = False) -> dict[str, Union[int, float]]:
         return self._get_model_base().get_minimum_resources(is_gpu_available)
 
     def _fit(
         self,
         train_data: TimeSeriesDataFrame,
         val_data: Optional[TimeSeriesDataFrame] = None,
-        time_limit: Optional[int] = None,
-        val_splitter: AbstractWindowSplitter = None,
+        time_limit: Optional[float] = None,
+        num_cpus: Optional[int] = None,
+        num_gpus: Optional[int] = None,
+        verbosity: int = 2,
+        val_splitter: Optional[AbstractWindowSplitter] = None,
         refit_every_n_windows: Optional[int] = 1,
         **kwargs,
     ):
@@ -111,11 +111,15 @@ class MultiWindowBacktestingModel(AbstractTimeSeriesModel):
 
         oof_predictions_per_window = []
         global_fit_start_time = time.time()
+        model: Optional[AbstractTimeSeriesModel] = None
 
         for window_index, (train_fold, val_fold) in enumerate(val_splitter.split(train_data)):
             logger.debug(f"\tWindow {window_index}")
+
             # refit_this_window is always True for the 0th window
             refit_this_window = window_index % refit_every_n_windows == 0
+            assert window_index != 0 or refit_this_window
+
             if time_limit is None:
                 time_left_for_window = None
             else:
@@ -148,6 +152,7 @@ class MultiWindowBacktestingModel(AbstractTimeSeriesModel):
             else:
                 time_left_for_prediction = time_limit - (time.time() - global_fit_start_time)
 
+            assert model is not None
             model.score_and_cache_oof(
                 val_fold, store_val_score=True, store_predict_time=True, time_limit=time_left_for_prediction
             )
@@ -172,11 +177,13 @@ class MultiWindowBacktestingModel(AbstractTimeSeriesModel):
 
         # Only the model trained on most recent data is saved & used for prediction
         self.most_recent_model = model
-        self.most_recent_model_folder = most_recent_refit_window
+        assert self.most_recent_model is not None
+
+        self.most_recent_model_folder = most_recent_refit_window  # type: ignore
         self.predict_time = self.most_recent_model.predict_time
-        self.fit_time = time.time() - global_fit_start_time - self.predict_time
+        self.fit_time = time.time() - global_fit_start_time - self.predict_time  # type: ignore
         self._oof_predictions = oof_predictions_per_window
-        self.val_score = np.mean([info["val_score"] for info in self.info_per_val_window])
+        self.val_score = np.mean([info["val_score"] for info in self.info_per_val_window])  # type: ignore
 
     def get_info(self) -> dict:
         info = super().get_info()
@@ -227,7 +234,7 @@ class MultiWindowBacktestingModel(AbstractTimeSeriesModel):
         train_fn_kwargs["init_params"]["model_base_kwargs"] = self.get_params()
         return train_fn_kwargs
 
-    def save(self, path: str = None, verbose=True) -> str:
+    def save(self, path: Optional[str] = None, verbose: bool = True) -> str:
         most_recent_model = self.most_recent_model
         self.most_recent_model = None
         save_path = super().save(path, verbose)
@@ -238,32 +245,36 @@ class MultiWindowBacktestingModel(AbstractTimeSeriesModel):
             most_recent_model.save()
         return save_path
 
-    def persist(self):
+    def persist(self) -> Self:
         if self.most_recent_model is None:
             raise ValueError(f"{self.name} must be fit before persisting")
         self.most_recent_model.persist()
+        return self
 
     @classmethod
     def load(
         cls, path: str, reset_paths: bool = True, load_oof: bool = False, verbose: bool = True
     ) -> AbstractTimeSeriesModel:
         model = super().load(path=path, reset_paths=reset_paths, load_oof=load_oof, verbose=verbose)
-        most_recent_model_path = os.path.join(model.path, model.most_recent_model_folder)
-        model.most_recent_model = model.model_base_type.load(
-            most_recent_model_path,
-            reset_paths=reset_paths,
-            verbose=verbose,
-        )
+        if model.most_recent_model_folder is not None:
+            most_recent_model_path = os.path.join(model.path, model.most_recent_model_folder)
+            model.most_recent_model = model.model_base_type.load(
+                most_recent_model_path,
+                reset_paths=reset_paths,
+                verbose=verbose,
+            )
         return model
 
     def convert_to_refit_full_template(self) -> AbstractTimeSeriesModel:
         # refit_model is an instance of base model type, not MultiWindowBacktestingModel
+        assert self.most_recent_model is not None, "Most recent model is None. Model must be fit first."
         refit_model = self.most_recent_model.convert_to_refit_full_template()
         refit_model.rename(self.name + ag.constants.REFIT_FULL_SUFFIX)
         return refit_model
 
     def convert_to_refit_full_via_copy(self) -> AbstractTimeSeriesModel:
         # refit_model is an instance of base model type, not MultiWindowBacktestingModel
+        assert self.most_recent_model is not None, "Most recent model is None. Model must be fit first."
         refit_model = self.most_recent_model.convert_to_refit_full_via_copy()
         refit_model.rename(self.name + ag.constants.REFIT_FULL_SUFFIX)
         return refit_model

autogluon/timeseries/models/registry.py

@@ -1,7 +1,7 @@
 from abc import ABCMeta
 from dataclasses import dataclass
 from inspect import isabstract
-from typing import Dict, List, Union
+from typing import Union
 
 
 @dataclass
@@ -18,7 +18,7 @@ class ModelRegistry(ABCMeta):
     See, https://github.com/faif/python-patterns.
     """
 
-    REGISTRY: Dict[str, ModelRecord] = {}
+    REGISTRY: dict[str, ModelRecord] = {}
 
     def __new__(cls, name, bases, attrs):
         new_cls = super().__new__(cls, name, bases, attrs)
@@ -61,5 +61,5 @@ class ModelRegistry(ABCMeta):
         return cls._get_model_record(alias).ag_priority
 
     @classmethod
-    def available_aliases(cls) -> List[str]:
+    def available_aliases(cls) -> list[str]:
         return sorted(cls.REGISTRY.keys())