option to re-use cv splits during tuning

mlforecast/auto.py

@@ -244,6 +244,8 @@ class AutoMLForecast:
         fit_config (callable, optional): Function that takes an optuna trial and produces a configuration passed to the MLForecast fit method.
             Defaults to None.
         num_threads (int): Number of threads to use when computing the features. Defaults to 1.
+        reuse_cv_splits (bool): Creates splits for cv once and re-uses them for tuning instead of generating the splits in each tuning round.
+            Default is set to False.
     """
 
     def __init__(
@@ -254,6 +256,7 @@ def __init__(
         init_config: Optional[_TrialToConfig] = None,
         fit_config: Optional[_TrialToConfig] = None,
         num_threads: int = 1,
+        reuse_cv_splits: bool = False,
     ):
         self.freq = freq
         if season_length is None and init_config is None:
@@ -279,6 +282,7 @@ def __init__(
         else:
             models_with_names = models
         self.models = models_with_names
+        self.reuse_cv_splits = reuse_cv_splits
 
     def __repr__(self):
         return f"AutoMLForecast(models={self.models})"
@@ -501,11 +505,23 @@ def loss(df, train_df):  # noqa: ARG001
             study_kwargs["sampler"] = optuna.samplers.TPESampler(seed=0)
         if optimize_kwargs is None:
             optimize_kwargs = {}
-
         self.results_ = {}
         self.models_ = {}
+        cv_splits = None
+        if self.reuse_cv_splits:
+            cv_splits = list(
+                ufp.backtest_splits(
+                    df,
+                    n_windows=n_windows,
+                    h=h,
+                    id_col=id_col,
+                    time_col=time_col,
+                    freq=self.freq,
+                    step_size=step_size,
+                    input_size=input_size,
+                )
+            )
         for name, auto_model in self.models.items():
-
             def config_fn(trial: optuna.Trial) -> Dict[str, Any]:
                 return {
                     "model_params": auto_model.config(trial),
@@ -529,7 +545,8 @@ def config_fn(trial: optuna.Trial) -> Dict[str, Any]:
                 id_col=id_col,
                 time_col=time_col,
                 target_col=target_col,
-                weight_col=weight_col
+                weight_col=weight_col,
+                cv_splits=cv_splits,
             )
             study = optuna.create_study(direction="minimize", **study_kwargs)
             study.optimize(objective, n_trials=num_samples, **optimize_kwargs)

mlforecast/optimization.py

@@ -2,7 +2,7 @@
 
 
 import copy
-from typing import Any, Callable, Dict, Optional, Union
+from typing import Any, Callable, Dict, Optional, Union, List, Tuple
 
 import numpy as np
 import pandas as pd
@@ -16,6 +16,7 @@
 from .core import Freq
 
 _TrialToConfig = Callable[[optuna.Trial], Dict[str, Any]]
+CVSplit = Tuple[DataFrame, DataFrame, DataFrame]
 
 
 def mlforecast_objective(
@@ -32,7 +33,8 @@ def mlforecast_objective(
     id_col: str = "unique_id",
     time_col: str = "ds",
     target_col: str = "y",
-    weight_col: Optional[str] = None
+    weight_col: Optional[str] = None,
+    cv_splits: Optional[List[CVSplit]] = None
 ) -> Callable[[optuna.Trial], float]:
     """optuna objective function for the MLForecast class
 
@@ -58,45 +60,41 @@ def mlforecast_objective(
         time_col (str): Column that identifies each timestep, its values can be timestamps or integers. Defaults to 'ds'.
         target_col (str): Column that contains the target. Defaults to 'y'.
         weight_col (str): Column that contains sample weights. Defaults to None.
+        cv_splits (List[Tuple[DataFrame, DataFrame, DataFrame]] | None): Optional cached CV splits (cutoffs, train, valid) to 
+            reuse across trials. If None, backtest splits are generated on each trial.
 
     Returns:
         (Callable[[optuna.Trial], float]): optuna objective function
     """
-
     def objective(trial: optuna.Trial) -> float:
         config = config_fn(trial)
         trial.set_user_attr("config", copy.deepcopy(config))
-        if all(
-            config["mlf_init_params"].get(k, None) is None
-            for k in ["lags", "lag_transforms", "date_features"]
-        ):
-            # no features
-            return np.inf
-        splits = ufp.backtest_splits(
-            df,
-            n_windows=n_windows,
-            h=h,
-            id_col=id_col,
-            time_col=time_col,
-            freq=freq,
-            step_size=step_size,
-            input_size=input_size,
-        )
+
         model_copy = clone(model)
         model_params = config["model_params"]
-        if config["mlf_fit_params"].get("static_features", []) and isinstance(
-            model, CatBoostRegressor
-        ):
-            # catboost needs the categorical features in the init signature
-            # we assume all statics are categoricals
+        if config["mlf_fit_params"].get("static_features", []) and isinstance(model, CatBoostRegressor):
             model_params["cat_features"] = config["mlf_fit_params"]["static_features"]
-        model_copy.set_params(**config["model_params"])
-        metrics = []
+        model_copy.set_params(**model_params)
         mlf = MLForecast(
-            models={"model": model_copy},
+            models={"model": model_copy}, 
             freq=freq,
             **config["mlf_init_params"],
         )
+        splits = cv_splits
+        if splits is None:
+            splits = ufp.backtest_splits(
+                df,
+                n_windows=n_windows,
+                h=h,
+                id_col=id_col,
+                time_col=time_col,
+                freq=freq,
+                step_size=step_size,
+                input_size=input_size,
+            )
+        elif not isinstance(splits, list):
+            splits = list(splits)
+        metrics = []
         for i, (_, train, valid) in enumerate(splits):
             should_fit = i == 0 or (refit > 0 and i % refit == 0)
             if should_fit:

tests/test_auto.py

@@ -269,4 +269,48 @@ def test_input_size_speedup(weekly_data):
     model.fit(df=train, input_size=50, **fit_kwargs)
     time_with_limit = time.perf_counter() - start
 
-    assert time_with_limit < time_no_limit
+    assert time_with_limit < time_no_limit
+
+
+def test_reuse_cv_splits_same_predictions(weekly_data):
+    train, valid, info = weekly_data
+    h = info.horizon
+    n_windows = 2
+    num_samples = 5 
+
+    def ridge_config(trial):  
+        return {"alpha": 1.0, "fit_intercept": True, "solver": "svd"}
+    def fit_config(trial):  
+        return {"dropna": True}
+    def init_config(trial):  
+        return {
+            "lags": [1, 2, 3],
+        }
+    ridge_auto = AutoModel(model=Ridge(), config=ridge_config)
+
+    common_kwargs = dict(
+        models={"ridge": ridge_auto},
+        freq=1,
+        fit_config=fit_config,
+        init_config=init_config,
+    )
+
+    automl_a = AutoMLForecast(**common_kwargs, reuse_cv_splits=False).fit(
+        train,
+        n_windows=n_windows,
+        h=h,
+        num_samples=num_samples,
+    )
+
+    automl_b = AutoMLForecast(**common_kwargs, reuse_cv_splits=True).fit(
+        train,
+        n_windows=n_windows,
+        h=h,
+        num_samples=num_samples,
+    )
+
+    preds_a = automl_a.predict(h=h)
+    preds_b = automl_b.predict(h=h)
+
+    assert preds_a.columns.tolist() == preds_b.columns.tolist()
+    assert (preds_a["ridge"].to_numpy() == preds_b["ridge"].to_numpy()).all()