emhass 0.11.4__py3-none-any.whl → 0.15.5__py3-none-any.whl

emhass/machine_learning_forecaster.py

@@ -1,22 +1,30 @@
-#!/usr/bin/env python3
-# -*- coding: utf-8 -*-
-
+import asyncio
 import logging
-import copy
 import time
-from typing import Optional, Tuple
-import pandas as pd
-import numpy as np
-
-from sklearn.linear_model import LinearRegression
-from sklearn.linear_model import ElasticNet
-from sklearn.neighbors import KNeighborsRegressor
-from sklearn.metrics import r2_score
+import warnings
 
+import numpy as np
+import pandas as pd
+from skforecast.model_selection import (
+    TimeSeriesFold,
+    backtesting_forecaster,
+    bayesian_search_forecaster,
+)
 from skforecast.recursive import ForecasterRecursive
-from skforecast.model_selection import bayesian_search_forecaster, backtesting_forecaster, TimeSeriesFold
+from sklearn.ensemble import (
+    AdaBoostRegressor,
+    ExtraTreesRegressor,
+    GradientBoostingRegressor,
+    RandomForestRegressor,
+)
+from sklearn.linear_model import ElasticNet, Lasso, LinearRegression, Ridge
+from sklearn.metrics import r2_score
+from sklearn.neighbors import KNeighborsRegressor
+from sklearn.neural_network import MLPRegressor
+from sklearn.svm import SVR
+from sklearn.tree import DecisionTreeRegressor
 
-import warnings
+from emhass import utils
 
 warnings.filterwarnings("ignore", category=DeprecationWarning)
 
@@ -25,17 +33,17 @@ class MLForecaster:
     r"""
     A forecaster class using machine learning models with auto-regressive approach and features\
     based on timestamp information (hour, day, week, etc).
-    
+
     This class uses the `skforecast` module and the machine learning models are from `scikit-learn`.
-    
+
     It exposes three main methods:
-    
+
     - `fit`: to train a model with the passed data.
-    
+
     - `predict`: to obtain a forecast from a pre-trained model.
-    
-    - `tune`: to optimize the models hyperparameters using bayesian optimization. 
-    
+
+    - `tune`: to optimize the models hyperparameters using bayesian optimization.
+
     """
 
     def __init__(
@@ -78,36 +86,57 @@ class MLForecaster:
         self.emhass_conf = emhass_conf
         self.logger = logger
         self.is_tuned = False
+        self.forecaster: ForecasterRecursive | None = None
+        self.optimize_results: pd.DataFrame | None = None
+        self.optimize_results_object = None
+
         # A quick data preparation
+        self._prepare_data()
+
+    def _prepare_data(self):
+        """Prepare the input data by cleaning and sorting."""
         self.data.index = pd.to_datetime(self.data.index)
         self.data.sort_index(inplace=True)
         self.data = self.data[~self.data.index.duplicated(keep="first")]
 
-    @staticmethod
-    def add_date_features(data: pd.DataFrame) -> pd.DataFrame:
-        """Add date features from the input DataFrame timestamp
-
-        :param data: The input DataFrame
-        :type data: pd.DataFrame
-        :return: The DataFrame with the added features
-        :rtype: pd.DataFrame
-        """
-        df = copy.deepcopy(data)
-        df["year"] = [i.year for i in df.index]
-        df["month"] = [i.month for i in df.index]
-        df["day_of_week"] = [i.dayofweek for i in df.index]
-        df["day_of_year"] = [i.dayofyear for i in df.index]
-        df["day"] = [i.day for i in df.index]
-        df["hour"] = [i.hour for i in df.index]
-        return df
-
     @staticmethod
     def neg_r2_score(y_true, y_pred):
         """The negative of the r2 score."""
         return -r2_score(y_true, y_pred)
 
     @staticmethod
-    def generate_exog(data_last_window, periods, var_name):
+    async def interpolate_async(data: pd.DataFrame) -> pd.DataFrame:
+        """Interpolate missing values asynchronously."""
+        return await asyncio.to_thread(data.interpolate, method="linear", axis=0, limit=None)
+
+    @staticmethod
+    def get_lags_list_from_frequency(freq: pd.Timedelta) -> list[int]:
+        """Calculate appropriate lag values based on data frequency.
+
+        The lags represent different time horizons (6h, 12h, 1d, 1.5d, 2d, 2.5d, 3d).
+        This method scales these horizons according to the actual data frequency.
+
+        :param freq: The frequency of the data as a pandas Timedelta
+        :type freq: pd.Timedelta
+        :return: A list of lag values appropriate for the data frequency
+        :rtype: list[int]
+        """
+        # Define target time horizons in hours
+        target_horizons_hours = [6, 12, 24, 36, 48, 60, 72]
+
+        # Calculate frequency in hours
+        freq_hours = freq.total_seconds() / 3600
+
+        # Calculate lags for each horizon
+        lags = [int(round(horizon / freq_hours)) for horizon in target_horizons_hours]
+
+        # Remove duplicates and ensure minimum value of 1
+        lags = sorted({max(1, lag) for lag in lags})
+
+        return lags
+
+    @staticmethod
+    async def generate_exog(data_last_window, periods, var_name):
         """Generate the exogenous data for future timestamps."""
         forecast_dates = pd.date_range(
             start=data_last_window.index[-1] + data_last_window.index.freq,
@@ -115,14 +144,46 @@ class MLForecaster:
             freq=data_last_window.index.freq,
         )
         exog = pd.DataFrame({var_name: [np.nan] * periods}, index=forecast_dates)
-        exog = MLForecaster.add_date_features(exog)
+        exog = utils.add_date_features(exog)
         return exog
 
-    def fit(
+    def _get_sklearn_model(self, model_name: str):
+        """Get the sklearn model instance based on the model name."""
+        seed = 42
+        models = {
+            "LinearRegression": LinearRegression(),
+            "RidgeRegression": Ridge(),
+            "LassoRegression": Lasso(random_state=seed),
+            "ElasticNet": ElasticNet(alpha=1.0, l1_ratio=0.5, random_state=seed),
+            "KNeighborsRegressor": KNeighborsRegressor(),
+            "DecisionTreeRegressor": DecisionTreeRegressor(ccp_alpha=0.0, random_state=seed),
+            "SVR": SVR(),
+            "RandomForestRegressor": RandomForestRegressor(
+                min_samples_leaf=1, max_features=1.0, random_state=seed
+            ),
+            "ExtraTreesRegressor": ExtraTreesRegressor(
+                min_samples_leaf=1, max_features=1.0, random_state=seed
+            ),
+            "GradientBoostingRegressor": GradientBoostingRegressor(
+                learning_rate=0.1, random_state=seed
+            ),
+            "AdaBoostRegressor": AdaBoostRegressor(learning_rate=1.0, random_state=seed),
+            "MLPRegressor": MLPRegressor(hidden_layer_sizes=(100,), random_state=seed),
+        }
+
+        if model_name not in models:
+            self.logger.error(
+                f"Passed sklearn model {model_name} is not valid. Defaulting to KNeighborsRegressor"
+            )
+            return KNeighborsRegressor()
+
+        return models[model_name]
+
+    async def fit(
         self,
-        split_date_delta: Optional[str] = "48h",
-        perform_backtest: Optional[bool] = False,
-    ) -> Tuple[pd.DataFrame, pd.DataFrame]:
+        split_date_delta: str | None = "48h",
+        perform_backtest: bool | None = False,
+    ) -> tuple[pd.DataFrame, pd.DataFrame]:
         r"""The fit method to train the ML model.
 
         :param split_date_delta: The delta from now to `split_date_delta` that will be used \
@@ -134,100 +195,133 @@ class MLForecaster:
         :return: The DataFrame containing the forecast data results without and with backtest
         :rtype: Tuple[pd.DataFrame, pd.DataFrame]
         """
-        self.logger.info("Performing a forecast model fit for " + self.model_type)
-        # Preparing the data: adding exogenous features
-        self.data_exo = pd.DataFrame(index=self.data.index)
-        self.data_exo = MLForecaster.add_date_features(self.data_exo)
-        self.data_exo[self.var_model] = self.data[self.var_model]
-        self.data_exo = self.data_exo.interpolate(method="linear", axis=0, limit=None)
-        # train/test split
-        self.date_train = (
-            self.data_exo.index[-1] - pd.Timedelta("5days") + self.data_exo.index.freq
-        )  # The last 5 days
-        self.date_split = (
-            self.data_exo.index[-1]
-            - pd.Timedelta(split_date_delta)
-            + self.data_exo.index.freq
-        )  # The last 48h
-        self.data_train = self.data_exo.loc[
-            : self.date_split - self.data_exo.index.freq, :
-        ]
-        self.data_test = self.data_exo.loc[self.date_split :, :]
-        self.steps = len(self.data_test)
-        # Pick correct sklearn model
-        if self.sklearn_model == "LinearRegression":
-            base_model = LinearRegression()
-        elif self.sklearn_model == "ElasticNet":
-            base_model = ElasticNet()
-        elif self.sklearn_model == "KNeighborsRegressor":
-            base_model = KNeighborsRegressor()
-        else:
-            self.logger.error(
-                "Passed sklearn model "
-                + self.sklearn_model
-                + " is not valid. Defaulting to KNeighborsRegressor"
-            )
-            base_model = KNeighborsRegressor()
-        # Define the forecaster object
-        self.forecaster = ForecasterRecursive(
-            regressor = base_model,
-            lags      = self.num_lags
-            )
-        # Fit and time it
-        self.logger.info("Training a " + self.sklearn_model + " model")
-        start_time = time.time()
-        self.forecaster.fit(
-            y=self.data_train[self.var_model],
-            exog=self.data_train.drop(self.var_model, axis=1),
-        )
-        self.logger.info(f"Elapsed time for model fit: {time.time() - start_time}")
-        # Make a prediction to print metrics
-        predictions = self.forecaster.predict(
-            steps=self.steps, exog=self.data_test.drop(self.var_model, axis=1)
-        )
-        pred_metric = r2_score(self.data_test[self.var_model], predictions)
-        self.logger.info(
-            f"Prediction R2 score of fitted model on test data: {pred_metric}"
-        )
-        # Packing results in a DataFrame
-        df_pred = pd.DataFrame(
-            index=self.data_exo.index, columns=["train", "test", "pred"]
-        )
-        df_pred["train"] = self.data_train[self.var_model]
-        df_pred["test"] = self.data_test[self.var_model]
-        df_pred["pred"] = predictions
-        df_pred_backtest = None
-        if perform_backtest is True:
-            # Using backtesting tool to evaluate the model
-            self.logger.info("Performing simple backtesting of fitted model")
+        try:
+            self.logger.info("Performing a forecast model fit for " + self.model_type)
+
+            # Check if variable exists in data
+            if self.var_model not in self.data.columns:
+                raise KeyError(
+                    f"Variable '{self.var_model}' not found in data columns: {list(self.data.columns)}"
+                )
+
+            # Preparing the data: adding exogenous features
+            self.data_exo = pd.DataFrame(index=self.data.index)
+            self.data_exo = utils.add_date_features(self.data_exo)
+            self.data_exo[self.var_model] = self.data[self.var_model]
+
+            self.data_exo = await self.interpolate_async(self.data_exo)
+
+            # train/test split
+            self.date_train = (
+                self.data_exo.index[-1] - pd.Timedelta("5days") + self.data_exo.index.freq
+            )  # The last 5 days
+            self.date_split = (
+                self.data_exo.index[-1] - pd.Timedelta(split_date_delta) + self.data_exo.index.freq
+            )  # The last 48h
+            self.data_train = self.data_exo.loc[: self.date_split - self.data_exo.index.freq, :]
+            self.data_test = self.data_exo.loc[self.date_split :, :]
+            self.steps = len(self.data_test)
+
+            # Pick correct sklearn model
+            base_model = self._get_sklearn_model(self.sklearn_model)
+
+            # Define the forecaster object
+            self.forecaster = ForecasterRecursive(estimator=base_model, lags=self.num_lags)
+
+            # Fit and time it
+            self.logger.info("Training a " + self.sklearn_model + " model")
             start_time = time.time()
-            cv = TimeSeriesFold(
-                steps                 = self.num_lags,
-                initial_train_size    = None,
-                fixed_train_size      = False,
-                gap                   = 0,
-                allow_incomplete_fold = True,
-                refit                 = False
+
+            await asyncio.to_thread(
+                self.forecaster.fit,
+                y=self.data_train[self.var_model],
+                exog=self.data_train.drop(self.var_model, axis=1),
+                store_in_sample_residuals=True,
             )
-            metric, predictions_backtest = backtesting_forecaster(
-                forecaster         = self.forecaster,
-                y                  = self.data_train[self.var_model],
-                exog               = self.data_train.drop(self.var_model, axis=1),
-                cv                 = cv,
-                metric             = MLForecaster.neg_r2_score,
-                verbose            = False,
-                show_progress      = True
+
+            fit_time = time.time() - start_time
+            self.logger.info(f"Elapsed time for model fit: {fit_time}")
+
+            # Make a prediction to print metrics
+            predictions = await asyncio.to_thread(
+                self.forecaster.predict,
+                steps=self.steps,
+                exog=self.data_test.drop(self.var_model, axis=1),
             )
-            self.logger.info(f"Elapsed backtesting time: {time.time() - start_time}")
-            self.logger.info(f"Backtest R2 score: {-metric}")
-            df_pred_backtest = pd.DataFrame(
-                index=self.data_exo.index, columns=["train", "pred"]
+            pred_metric = await asyncio.to_thread(
+                r2_score, self.data_test[self.var_model], predictions
             )
-            df_pred_backtest["train"] = self.data_exo[self.var_model]
-            df_pred_backtest["pred"] = predictions_backtest
-        return df_pred, df_pred_backtest
+            self.logger.info(f"Prediction R2 score of fitted model on test data: {pred_metric}")
+
+            # Packing results in a DataFrame
+            df_pred = pd.DataFrame(index=self.data_exo.index, columns=["train", "test", "pred"])
+
+            df_pred["train"] = self.data_train[self.var_model]
+            df_pred["test"] = self.data_test[self.var_model]
+            df_pred["pred"] = predictions
+
+            df_pred_backtest = None
+
+            if perform_backtest is True:
+                # Using backtesting tool to evaluate the model
+                self.logger.info("Performing simple backtesting of fitted model")
+                start_time = time.time()
+                cv = TimeSeriesFold(
+                    steps=self.num_lags,
+                    initial_train_size=None,
+                    fixed_train_size=False,
+                    gap=0,
+                    allow_incomplete_fold=True,
+                    refit=False,
+                )
+
+                metric, predictions_backtest = await asyncio.to_thread(
+                    backtesting_forecaster,
+                    forecaster=self.forecaster,
+                    y=self.data_train[self.var_model],
+                    exog=self.data_train.drop(self.var_model, axis=1),
+                    cv=cv,
+                    metric=MLForecaster.neg_r2_score,
+                    verbose=False,
+                    show_progress=True,
+                )
+
+                backtest_time = time.time() - start_time
+                backtest_r2 = -metric
+                self.logger.info(f"Elapsed backtesting time: {backtest_time}")
+                self.logger.info(f"Backtest R2 score: {backtest_r2}")
+                df_pred_backtest = pd.DataFrame(
+                    index=self.data_exo.index, columns=["train", "pred"]
+                )
+                df_pred_backtest["train"] = self.data_exo[self.var_model]
+                # Handle skforecast 0.18.0+ DataFrame output with fold column
+                if isinstance(predictions_backtest, pd.DataFrame):
+                    # Extract the 'pred' column from the DataFrame
+                    pred_values = (
+                        predictions_backtest["pred"]
+                        if "pred" in predictions_backtest.columns
+                        else predictions_backtest.iloc[:, -1]
+                    )
+                else:
+                    # If it's a Series, use it directly
+                    pred_values = predictions_backtest
+
+                # Use loc to align indices properly - only assign where indices match
+                df_pred_backtest.loc[pred_values.index, "pred"] = pred_values
+
+            return df_pred, df_pred_backtest
 
-    def predict(self, data_last_window: Optional[pd.DataFrame] = None) -> pd.Series:
+        except asyncio.CancelledError:
+            self.logger.info("Model training was cancelled")
+            raise
+        except Exception as e:
+            self.logger.error(f"Error during model fitting: {e}")
+            raise
+
+    async def predict(
+        self,
+        data_last_window: pd.DataFrame | None = None,
+    ) -> pd.Series:
         """The predict method to generate forecasts from a previously fitted ML model.
 
         :param data_last_window: The data that will be used to generate the new forecast, this \
@@ -238,168 +332,327 @@ class MLForecaster:
         :return: A pandas series containing the generated forecasts.
         :rtype: pd.Series
         """
-        if data_last_window is None:
-            predictions = self.forecaster.predict(
-                steps=self.num_lags, exog=self.data_test.drop(self.var_model, axis=1)
-            )
-        else:
-            data_last_window = data_last_window.interpolate(
-                method="linear", axis=0, limit=None
-            )
-            if self.is_tuned:
-                exog = MLForecaster.generate_exog(
-                    data_last_window, self.lags_opt, self.var_model
-                )
-                predictions = self.forecaster.predict(
-                    steps=self.lags_opt,
-                    last_window=data_last_window[self.var_model],
-                    exog=exog.drop(self.var_model, axis=1),
-                )
-            else:
-                exog = MLForecaster.generate_exog(
-                    data_last_window, self.num_lags, self.var_model
-                )
-                predictions = self.forecaster.predict(
+        try:
+            if self.forecaster is None:
+                raise ValueError("Model has not been fitted yet. Call fit() first.")
+
+            if data_last_window is None:
+                predictions = await asyncio.to_thread(
+                    self.forecaster.predict,
                     steps=self.num_lags,
-                    last_window=data_last_window[self.var_model],
-                    exog=exog.drop(self.var_model, axis=1),
+                    exog=self.data_test.drop(self.var_model, axis=1),
                 )
-        return predictions
+            else:
+                data_last_window = await self.interpolate_async(data_last_window)
+
+                if self.is_tuned:
+                    exog = await self.generate_exog(data_last_window, self.lags_opt, self.var_model)
+
+                    predictions = await asyncio.to_thread(
+                        self.forecaster.predict,
+                        steps=self.lags_opt,
+                        last_window=data_last_window[self.var_model],
+                        exog=exog.drop(self.var_model, axis=1),
+                    )
+                else:
+                    exog = await self.generate_exog(data_last_window, self.num_lags, self.var_model)
+
+                    predictions = await asyncio.to_thread(
+                        self.forecaster.predict,
+                        steps=self.num_lags,
+                        last_window=data_last_window[self.var_model],
+                        exog=exog.drop(self.var_model, axis=1),
+                    )
+
+            return predictions
+
+        except asyncio.CancelledError:
+            self.logger.info("Prediction was cancelled")
+            raise
+        except Exception as e:
+            self.logger.error(f"Error during prediction: {e}")
+            raise
 
-    def tune(self, debug: Optional[bool] = False) -> pd.DataFrame:
+    def _get_search_space(self, debug: bool, lags_list: list[int] | None = None):
+        """Get the hyperparameter search space for the given model.
+
+        :param debug: If True, use simplified search space for faster testing
+        :type debug: bool
+        :param lags_list: List of lag values to use. If None, uses default values
+        :type lags_list: list[int] | None
+        """
+        if lags_list is None:
+            lags_list = [6, 12, 24, 36, 48, 60, 72]
+
+        debug_lags = [3]
+
+        def get_lags(trial):
+            return trial.suggest_categorical("lags", debug_lags if debug else lags_list)
+
+        def svr_search_space(trial):
+            # Base SVR parameters
+            search = {
+                "C": trial.suggest_float("C", 0.1, 1.0)
+                if debug
+                else trial.suggest_float("C", 1e-2, 100.0, log=True),
+                "epsilon": trial.suggest_float("epsilon", 0.01, 1.0),
+                "kernel": trial.suggest_categorical("kernel", ["linear", "rbf"]),
+                "gamma": trial.suggest_categorical(
+                    "gamma", ["scale", "auto", 0.01, 0.1, 1.0, 10.0]
+                ),
+                "lags": get_lags(trial),
+            }
+            return search
+
+        # Registry of search space generators
+        search_spaces = {
+            "LinearRegression": lambda trial: {
+                "fit_intercept": trial.suggest_categorical(
+                    "fit_intercept", [True] if debug else [True, False]
+                ),
+                "lags": get_lags(trial),
+            },
+            "RidgeRegression": lambda trial: {
+                "alpha": trial.suggest_float("alpha", 0.1, 1.0)
+                if debug
+                else trial.suggest_float("alpha", 1e-4, 100.0, log=True),
+                "lags": get_lags(trial),
+            },
+            "LassoRegression": lambda trial: {
+                "alpha": trial.suggest_float("alpha", 0.1, 1.0)
+                if debug
+                else trial.suggest_float("alpha", 1e-4, 100.0, log=True),
+                "lags": get_lags(trial),
+            },
+            "ElasticNet": lambda trial: {
+                "alpha": trial.suggest_float("alpha", 0.0, 2.0),
+                "l1_ratio": trial.suggest_float("l1_ratio", 0.0, 1.0),
+                "selection": trial.suggest_categorical(
+                    "selection", ["random"] if debug else ["cyclic", "random"]
+                ),
+                "lags": get_lags(trial),
+            },
+            "KNeighborsRegressor": lambda trial: {
+                "n_neighbors": trial.suggest_int("n_neighbors", 2, 2)
+                if debug
+                else trial.suggest_int("n_neighbors", 2, 20),
+                "leaf_size": trial.suggest_int("leaf_size", 20, 20)
+                if debug
+                else trial.suggest_int("leaf_size", 20, 40),
+                "weights": trial.suggest_categorical(
+                    "weights", ["uniform"] if debug else ["uniform", "distance"]
+                ),
+                "lags": get_lags(trial),
+            },
+            "DecisionTreeRegressor": lambda trial: {
+                "max_depth": trial.suggest_int("max_depth", 2, 5)
+                if debug
+                else trial.suggest_int("max_depth", 2, 20),
+                "min_samples_split": trial.suggest_int("min_samples_split", 2, 20),
+                "min_samples_leaf": trial.suggest_int("min_samples_leaf", 1, 10),
+                "lags": get_lags(trial),
+            },
+            "SVR": svr_search_space,
+            "RandomForestRegressor": lambda trial: {
+                "n_estimators": trial.suggest_int("n_estimators", 10, 20)
+                if debug
+                else trial.suggest_int("n_estimators", 50, 300),
+                "max_depth": trial.suggest_int("max_depth", 3, 20),
+                "min_samples_split": trial.suggest_int("min_samples_split", 2, 10),
+                "max_features": trial.suggest_categorical("max_features", ["sqrt", "log2", None]),
+                "lags": get_lags(trial),
+            },
+            "ExtraTreesRegressor": lambda trial: {
+                "n_estimators": trial.suggest_int("n_estimators", 10, 20)
+                if debug
+                else trial.suggest_int("n_estimators", 50, 300),
+                "max_depth": trial.suggest_int("max_depth", 3, 20),
+                "min_samples_split": trial.suggest_int("min_samples_split", 2, 10),
+                "max_features": trial.suggest_categorical("max_features", ["sqrt", "log2", None]),
+                "lags": get_lags(trial),
+            },
+            "GradientBoostingRegressor": lambda trial: {
+                "n_estimators": trial.suggest_int("n_estimators", 10, 20)
+                if debug
+                else trial.suggest_int("n_estimators", 50, 300),
+                "learning_rate": trial.suggest_float("learning_rate", 0.01, 0.5),
+                "max_depth": trial.suggest_int("max_depth", 3, 10),
+                "lags": get_lags(trial),
+            },
+            "AdaBoostRegressor": lambda trial: {
+                "n_estimators": trial.suggest_int("n_estimators", 10, 20)
+                if debug
+                else trial.suggest_int("n_estimators", 50, 300),
+                "learning_rate": trial.suggest_float("learning_rate", 0.01, 1.0),
+                "lags": get_lags(trial),
+            },
+            "MLPRegressor": lambda trial: {
+                "learning_rate_init": trial.suggest_float("learning_rate_init", 0.001, 0.01),
+                "hidden_layer_sizes": trial.suggest_categorical(
+                    "hidden_layer_sizes", [(50,), (100,), (50, 50)]
+                ),
+                "activation": trial.suggest_categorical("activation", ["relu", "tanh"]),
+                "alpha": trial.suggest_float("alpha", 1e-5, 1e-1, log=True),
+                "lags": get_lags(trial),
+            },
+        }
+
+        if self.sklearn_model not in search_spaces:
+            raise ValueError(f"Unsupported model for tuning: {self.sklearn_model}")
+
+        return search_spaces[self.sklearn_model]
+
+    async def tune(
+        self,
+        split_date_delta: str | None = "48h",
+        n_trials: int = 10,
+        debug: bool | None = False,
+    ) -> pd.DataFrame:
         """Tuning a previously fitted model using bayesian optimization.
 
+        :param split_date_delta: The delta from now to `split_date_delta` that will be used \
+            as the test period to evaluate the model, defaults to '48h'.\
+            This define the training/validation split for the tuning process.
+        :type split_date_delta: Optional[str], optional
         :param debug: Set to True for testing and faster optimizations, defaults to False
         :type debug: Optional[bool], optional
+        :param n_trials: Number of trials for bayesian optimization, defaults to 10
+        :type n_trials: Optional[int], optional
         :return: The DataFrame with the forecasts using the optimized model.
         :rtype: pd.DataFrame
         """
-        # Regressor hyperparameters search space
-        if self.sklearn_model == "LinearRegression":
-            if debug:
+        try:
+            if self.forecaster is None:
+                raise ValueError("Model has not been fitted yet. Call fit() first.")
 
-                def search_space(trial):
-                    search_space = {
-                        "fit_intercept": trial.suggest_categorical(
-                            "fit_intercept", [True]
-                        ),
-                        "lags": trial.suggest_categorical("lags", [3]),
-                    }
-                    return search_space
-            else:
+            # Calculate appropriate lags based on data frequency
+            freq_timedelta = pd.Timedelta(self.data_exo.index.freq)
+            lags_list = MLForecaster.get_lags_list_from_frequency(freq_timedelta)
+            self.logger.info(
+                f"Using lags list based on data frequency ({self.data_exo.index.freq}): {lags_list}"
+            )
 
-                def search_space(trial):
-                    search_space = {
-                        "fit_intercept": trial.suggest_categorical(
-                            "fit_intercept", [True, False]
-                        ),
-                        "lags": trial.suggest_categorical(
-                            "lags", [6, 12, 24, 36, 48, 60, 72]
-                        ),
-                    }
-                    return search_space
-        elif self.sklearn_model == "ElasticNet":
-            if debug:
+            # Get the search space for this model
+            search_space = self._get_search_space(debug, lags_list)
 
-                def search_space(trial):
-                    search_space = {
-                        "selection": trial.suggest_categorical("selection", ["random"]),
-                        "lags": trial.suggest_categorical("lags", [3]),
-                    }
-                    return search_space
+            # Bayesian search hyperparameter and lags with skforecast/optuna
+            if debug:
+                refit = False
+                num_lags = 3
             else:
+                refit = True
+                num_lags = self.num_lags
+            # The optimization routine call
+            self.logger.info("Bayesian hyperparameter optimization with backtesting")
+            start_time = time.time()
+
+            # Use the 'y' data that will be passed to the optimizer
+            data_to_tune = self.data_train[self.var_model]
 
-                def search_space(trial):
-                    search_space = {
-                        "alpha": trial.suggest_float("alpha", 0.0, 2.0),
-                        "l1_ratio": trial.suggest_float("l1_ratio", 0.0, 1.0),
-                        "selection": trial.suggest_categorical(
-                            "selection", ["cyclic", "random"]
-                        ),
-                        "lags": trial.suggest_categorical(
-                            "lags", [6, 12, 24, 36, 48, 60, 72]
-                        ),
-                    }
-                    return search_space
-        elif self.sklearn_model == "KNeighborsRegressor":
+            # Calculate the new split date and initial_train_size based on the passed split_date_delta
+            try:
+                date_split = (
+                    data_to_tune.index[-1]
+                    - pd.Timedelta(split_date_delta)
+                    + data_to_tune.index.freq
+                )
+                initial_train_size = len(data_to_tune.loc[: date_split - data_to_tune.index.freq])
+            except (ValueError, TypeError):
+                self.logger.warning(
+                    f"Invalid split_date_delta: {split_date_delta}. Falling back to 5 days."
+                )
+                date_split = (
+                    data_to_tune.index[-1] - pd.Timedelta("5days") + data_to_tune.index.freq
+                )
+                initial_train_size = len(data_to_tune.loc[: date_split - data_to_tune.index.freq])
+
+            # Check if the calculated initial_train_size is valid
+            window_size = num_lags  # This is what skforecast will use as window_size
             if debug:
+                window_size = 3  # Match debug lags
 
-                def search_space(trial):
-                    search_space = {
-                        "weights": trial.suggest_categorical("weights", ["uniform"]),
-                        "lags": trial.suggest_categorical("lags", [3]),
-                    }
-                    return search_space
-            else:
+            if initial_train_size <= window_size:
+                self.logger.warning(
+                    f"Calculated initial_train_size ({initial_train_size}) is <= window_size ({window_size})."
+                )
+                self.logger.warning(
+                    "This is likely because split_date_delta is too large for the dataset."
+                )
+                MIN_SAMPLES_FOR_KNN = 6
+                new_train_size = window_size + MIN_SAMPLES_FOR_KNN
+                self.logger.warning(
+                    f"Adjusting initial_train_size to {new_train_size} to attempt recovery."
+                )
+                initial_train_size = new_train_size
 
-                def search_space(trial):
-                    search_space = {
-                        "n_neighbors": trial.suggest_int("n_neighbors", 2, 20),
-                        "leaf_size": trial.suggest_int("leaf_size", 20, 40),
-                        "weights": trial.suggest_categorical(
-                            "weights", ["uniform", "distance"]
-                        ),
-                        "lags": trial.suggest_categorical(
-                            "lags", [6, 12, 24, 36, 48, 60, 72]
-                        ),
-                    }
-                    return search_space
-
-        # Bayesian search hyperparameter and lags with skforecast/optuna
-        # Lags used as predictors
-        if debug:
-            refit = False
-            num_lags = 3
-        else:
-            refit = True
-            num_lags = self.num_lags
-        # The optimization routine call
-        self.logger.info("Bayesian hyperparameter optimization with backtesting")
-        start_time = time.time()
-        cv = TimeSeriesFold(
-            steps                 = num_lags,
-            initial_train_size    = len(self.data_exo.loc[:self.date_train]),
-            fixed_train_size      = True,
-            gap                   = 0,
-            skip_folds            = None,
-            allow_incomplete_fold = True,
-            refit                 = refit
-        )
-        self.optimize_results, self.optimize_results_object = bayesian_search_forecaster(
-            forecaster         = self.forecaster,
-            y                  = self.data_train[self.var_model],
-            exog               = self.data_train.drop(self.var_model, axis=1),
-            cv                 = cv,
-            search_space       = search_space,
-            metric             = MLForecaster.neg_r2_score,
-            n_trials           = 10,
-            random_state       = 123,
-            return_best        = True
-        )
-        self.logger.info(f"Elapsed time: {time.time() - start_time}")
-        self.is_tuned = True
-        predictions_opt = self.forecaster.predict(
-            steps=self.num_lags, exog=self.data_test.drop(self.var_model, axis=1)
-        )
-        freq_hours = self.data_exo.index.freq.delta.seconds / 3600
-        self.lags_opt = int(np.round(len(self.optimize_results.iloc[0]["lags"])))
-        self.days_needed = int(np.round(self.lags_opt * freq_hours / 24))
-        df_pred_opt = pd.DataFrame(
-            index=self.data_exo.index, columns=["train", "test", "pred_optim"]
-        )
-        df_pred_opt["train"] = self.data_train[self.var_model]
-        df_pred_opt["test"] = self.data_test[self.var_model]
-        df_pred_opt["pred_optim"] = predictions_opt
-        pred_optim_metric_train = -self.optimize_results.iloc[0]["neg_r2_score"]
-        self.logger.info(
-            f"R2 score for optimized prediction in train period: {pred_optim_metric_train}"
-        )
-        pred_optim_metric_test = r2_score(
-            df_pred_opt.loc[predictions_opt.index, "test"],
-            df_pred_opt.loc[predictions_opt.index, "pred_optim"],
-        )
-        self.logger.info(
-            f"R2 score for optimized prediction in test period: {pred_optim_metric_test}"
-        )
-        self.logger.info("Number of optimal lags obtained: " + str(self.lags_opt))
-        return df_pred_opt
+            cv = TimeSeriesFold(
+                steps=num_lags,
+                initial_train_size=initial_train_size,
+                fixed_train_size=True,
+                gap=0,
+                skip_folds=None,
+                allow_incomplete_fold=True,
+                refit=refit,
+            )
+
+            (
+                self.optimize_results,
+                self.optimize_results_object,
+            ) = await asyncio.to_thread(
+                bayesian_search_forecaster,
+                forecaster=self.forecaster,
+                y=self.data_train[self.var_model],
+                exog=self.data_train.drop(self.var_model, axis=1),
+                cv=cv,
+                search_space=search_space,
+                metric=MLForecaster.neg_r2_score,
+                n_trials=n_trials,
+                random_state=123,
+                return_best=True,
+            )
+
+            optimization_time = time.time() - start_time
+            self.logger.info(f"Elapsed time: {optimization_time}")
+
+            self.is_tuned = True
+
+            predictions_opt = await asyncio.to_thread(
+                self.forecaster.predict,
+                steps=self.num_lags,
+                exog=self.data_test.drop(self.var_model, axis=1),
+            )
+
+            freq_hours = self.data_exo.index.freq.delta.seconds / 3600
+            self.lags_opt = int(np.round(len(self.optimize_results.iloc[0]["lags"])))
+            self.days_needed = int(np.round(self.lags_opt * freq_hours / 24))
+
+            df_pred_opt = pd.DataFrame(
+                index=self.data_exo.index, columns=["train", "test", "pred_optim"]
+            )
+            df_pred_opt["train"] = self.data_train[self.var_model]
+            df_pred_opt["test"] = self.data_test[self.var_model]
+            df_pred_opt["pred_optim"] = predictions_opt
+
+            pred_optim_metric_train = -self.optimize_results.iloc[0]["neg_r2_score"]
+            self.logger.info(
+                f"R2 score for optimized prediction in train period: {pred_optim_metric_train}"
+            )
+
+            pred_optim_metric_test = await asyncio.to_thread(
+                r2_score,
+                df_pred_opt.loc[predictions_opt.index, "test"],
+                df_pred_opt.loc[predictions_opt.index, "pred_optim"],
+            )
+            self.logger.info(
+                f"R2 score for optimized prediction in test period: {pred_optim_metric_test}"
+            )
+            self.logger.info("Number of optimal lags obtained: " + str(self.lags_opt))
+
+            return df_pred_opt
+
+        except asyncio.CancelledError:
+            self.logger.info("Model tuning was cancelled")
+            raise
+        except Exception as e:
+            self.logger.error(f"Error during model tuning: {e}")
+            raise