Fix/add fit/predict_kwargs argument to historical_forecasts

darts/models/forecasting/forecasting_model.py

@@ -31,10 +31,12 @@
 
 import numpy as np
 import pandas as pd
+from sklearn.multioutput import MultiOutputRegressor
 
 from darts import metrics
 from darts.dataprocessing.encoders import SequentialEncoder
 from darts.logging import get_logger, raise_if, raise_if_not, raise_log
+from darts.models.utils import _check_kwargs_keys
 from darts.timeseries import TimeSeries
 from darts.utils import _build_tqdm_iterator, _parallel_apply, _with_sanity_checks
 from darts.utils.historical_forecasts.utils import (
@@ -316,6 +318,7 @@ def _fit_wrapper(
         series: TimeSeries,
         past_covariates: Optional[TimeSeries],
         future_covariates: Optional[TimeSeries],
+        **kwargs,
     ):
         self.fit(series)
 
@@ -328,10 +331,21 @@ def _predict_wrapper(
         num_samples: int,
         verbose: bool = False,
         predict_likelihood_parameters: bool = False,
+        num_loader_workers: int = 0,
+        batch_size: Optional[int] = None,
+        n_jobs: int = 1,
+        roll_size: Optional[int] = None,
+        mc_dropout: bool = False,
     ) -> TimeSeries:
         kwargs = dict()
         if self.supports_likelihood_parameter_prediction:
             kwargs["predict_likelihood_parameters"] = predict_likelihood_parameters
+        if getattr(self, "trainer_params", False):
+            kwargs["num_loader_workers"] = num_loader_workers
+            kwargs["batch_size"] = batch_size
+            kwargs["n_jobs"] = n_jobs
+            kwargs["roll_size"] = roll_size
+            kwargs["mc_dropout"] = mc_dropout
         return self.predict(n, num_samples=num_samples, verbose=verbose, **kwargs)
 
     @property
@@ -586,6 +600,8 @@ def historical_forecasts(
         show_warnings: bool = True,
         predict_likelihood_parameters: bool = False,
         enable_optimization: bool = True,
+        fit_kwargs: Optional[Dict[str, Any]] = None,
+        predict_kwargs: Optional[Dict[str, Any]] = None,
     ) -> Union[
         TimeSeries, List[TimeSeries], Sequence[TimeSeries], Sequence[List[TimeSeries]]
     ]:
@@ -692,6 +708,12 @@ def historical_forecasts(
             Default: ``False``
         enable_optimization
             Whether to use the optimized version of historical_forecasts when supported and available.
+        fit_kwargs
+            Additional arguments passed to the model `fit()` method, for example `max_samples_per_ts`,
+            `n_jobs_multiouput_wrapper` or `num_loader_workers`.
+        predict_kwargs
+            Additional arguments passed to the model `predict()` method, for example `num_samples`,
+            `predict_likelihood_parameters` or `num_loader_workers`.
 
         Returns
         -------
@@ -802,6 +824,51 @@ def retrain_func(
                 logger,
             )
 
+        if fit_kwargs is None:
+            fit_kwargs = dict()
+        if predict_kwargs is None:
+            predict_kwargs = dict()
+
+        # sanity checks of the arguments directly exposed by historical_forecasts
+        if "predict_likelihood_parameters" not in predict_kwargs:
+            predict_kwargs[
+                "predict_likelihood_parameters"
+            ] = predict_likelihood_parameters
+        elif (
+            predict_kwargs["predict_likelihood_parameters"]
+            != predict_likelihood_parameters
+        ):
+            logger.warning(
+                "`predict_likelihood_parameters` was provided with contradictory values, "
+                "retaining the value passed with `predict_kwargs`."
+            )
+        if "num_samples" not in predict_kwargs:
+            predict_kwargs["num_samples"] = num_samples
+        elif predict_kwargs["num_samples"] != num_samples:
+            logger.warning(
+                "`num_samples` was provided with contradictory values, "
+                "retaining the value passed with `predict_kwargs`."
+            )
+
+        # fit/predict_kwargs cannot be used to pass arguments used by historical_forecast logic
+        forbiden_args = ["series", "past_covariates", "future_covariates"]
+        fit_invalid_args = forbiden_args + [
+            "val_series",
+            "val_past_covariates",
+            "val_future_covariates",
+        ]
+        _check_kwargs_keys(
+            param_name="fit_kwargs",
+            kwargs_dict=fit_kwargs,
+            invalid_keys=fit_invalid_args,
+        )
+        predict_invalid_args = forbiden_args + ["n", "trainer"]
+        _check_kwargs_keys(
+            param_name="predict_kwargs",
+            kwargs_dict=predict_kwargs,
+            invalid_keys=predict_invalid_args,
+        )
+
         series = series2seq(series)
         past_covariates = series2seq(past_covariates)
         future_covariates = series2seq(future_covariates)
@@ -819,7 +886,6 @@ def retrain_func(
                 series=series,
                 past_covariates=past_covariates,
                 future_covariates=future_covariates,
-                num_samples=num_samples,
                 start=start,
                 start_format=start_format,
                 forecast_horizon=forecast_horizon,
@@ -828,7 +894,7 @@ def retrain_func(
                 last_points_only=last_points_only,
                 verbose=verbose,
                 show_warnings=show_warnings,
-                predict_likelihood_parameters=predict_likelihood_parameters,
+                predict_kwargs=predict_kwargs,
             )
 
         if len(series) == 1:
@@ -969,6 +1035,7 @@ def retrain_func(
                             series=train_series,
                             past_covariates=past_covariates_,
                             future_covariates=future_covariates_,
+                            **fit_kwargs,
                         )
                     else:
                         # untrained model was not trained on the first trainable timestamp
@@ -1016,9 +1083,8 @@ def retrain_func(
                     series=train_series,
                     past_covariates=past_covariates_,
                     future_covariates=future_covariates_,
-                    num_samples=num_samples,
                     verbose=verbose,
-                    predict_likelihood_parameters=predict_likelihood_parameters,
+                    **predict_kwargs,
                 )
                 if forecast_components is None:
                     forecast_components = forecast.columns
@@ -1076,6 +1142,8 @@ def backtest(
         reduction: Union[Callable[[np.ndarray], float], None] = np.mean,
         verbose: bool = False,
         show_warnings: bool = True,
+        fit_kwargs: Optional[Dict[str, Any]] = None,
+        predict_kwargs: Optional[Dict[str, Any]] = None,
     ) -> Union[float, List[float], Sequence[float], List[Sequence[float]]]:
         """Compute error values that the model would have produced when
         used on (potentially multiple) `series`.
@@ -1185,6 +1253,12 @@ def backtest(
             Whether to print progress.
         show_warnings
             Whether to show warnings related to parameters `start`, and `train_length`.
+        fit_kwargs
+            Additional arguments passed to the model `fit()` method, for example `max_samples_per_ts`,
+            `n_jobs_multiouput_wrapper` or `num_loader_workers`.
+        predict_kwargs
+            Additional arguments passed to the model `predict()` method, for example `num_samples`,
+            `predict_likelihood_parameters` or `num_loader_workers`.
 
         Returns
         -------
@@ -1208,6 +1282,8 @@ def backtest(
                 last_points_only=last_points_only,
                 verbose=verbose,
                 show_warnings=show_warnings,
+                fit_kwargs=fit_kwargs,
+                predict_kwargs=predict_kwargs,
             )
         else:
             forecasts = historical_forecasts
@@ -1261,6 +1337,8 @@ def gridsearch(
         verbose=False,
         n_jobs: int = 1,
         n_random_samples: Optional[Union[int, float]] = None,
+        fit_kwargs: Optional[Dict[str, Any]] = None,
+        predict_kwargs: Optional[Dict[str, Any]] = None,
     ) -> Tuple["ForecastingModel", Dict[str, Any], float]:
         """
         Find the best hyper-parameters among a given set using a grid search.
@@ -1374,6 +1452,12 @@ def gridsearch(
             must be between `0` and the total number of parameter combinations.
             If a float, `n_random_samples` is the ratio of parameter combinations selected from the full grid and must
             be between `0` and `1`. Defaults to `None`, for which random selection will be ignored.
+        fit_kwargs
+            Additional arguments passed to the model `fit()` method, for example `max_samples_per_ts`,
+            `n_jobs_multiouput_wrapper` or `num_loader_workers`.
+        predict_kwargs
+            Additional arguments passed to the model `predict()` method, for example `predict_likelihood_parameters` or
+            `num_loader_workers`.
 
         Returns
         -------
@@ -1406,10 +1490,28 @@ def gridsearch(
                 logger,
             )
 
-        # TODO: here too I'd say we can leave these checks to the models
-        # if covariates is not None:
-        #     raise_if_not(series.has_same_time_as(covariates), 'The provided series and covariates must have the '
-        #                                                       'same time axes.')
+        if fit_kwargs is None:
+            fit_kwargs = dict()
+        if predict_kwargs is None:
+            predict_kwargs = dict()
+
+        forbiden_args = ["series", "past_covariates", "future_covariates"]
+        fit_invalid_args = forbiden_args + [
+            "val_series",
+            "val_past_covariates",
+            "val_future_covariates",
+        ]
+        _check_kwargs_keys(
+            param_name="fit_kwargs",
+            kwargs_dict=fit_kwargs,
+            invalid_keys=fit_invalid_args,
+        )
+        predict_invalid_args = forbiden_args + ["n", "trainer", "num_samples"]
+        _check_kwargs_keys(
+            param_name="predict_kwargs",
+            kwargs_dict=predict_kwargs,
+            invalid_keys=predict_invalid_args,
+        )
 
         # compute all hyperparameter combinations from selection
         params_cross_product = list(product(*parameters.values()))
@@ -1437,7 +1539,12 @@ def _evaluate_combination(param_combination) -> float:
 
             model = model_class(**param_combination_dict)
             if use_fitted_values:  # fitted value mode
-                model._fit_wrapper(series, past_covariates, future_covariates)
+                model._fit_wrapper(
+                    series,
+                    past_covariates,
+                    future_covariates,
+                    **fit_kwargs,
+                )
                 fitted_values = TimeSeries.from_times_and_values(
                     series.time_index, model.fitted_values
                 )
@@ -1457,16 +1564,20 @@ def _evaluate_combination(param_combination) -> float:
                     last_points_only=last_points_only,
                     verbose=verbose,
                     show_warnings=show_warnings,
+                    predict_kwargs=predict_kwargs,
                 )
             else:  # split mode
-                model._fit_wrapper(series, past_covariates, future_covariates)
+                model._fit_wrapper(
+                    series, past_covariates, future_covariates, **fit_kwargs
+                )
                 pred = model._predict_wrapper(
                     len(val_series),
                     series,
                     past_covariates,
                     future_covariates,
                     num_samples=1,
                     verbose=verbose,
+                    **predict_kwargs,
                 )
                 error = metric(val_series, pred)
 
@@ -2220,10 +2331,21 @@ def _predict_wrapper(
         num_samples: int,
         verbose: bool = False,
         predict_likelihood_parameters: bool = False,
+        num_loader_workers: int = 0,
+        batch_size: Optional[int] = None,
+        n_jobs: int = 1,
+        roll_size: Optional[int] = None,
+        mc_dropout: bool = False,
     ) -> Union[TimeSeries, Sequence[TimeSeries]]:
         kwargs = dict()
         if self.supports_likelihood_parameter_prediction:
             kwargs["predict_likelihood_parameters"] = predict_likelihood_parameters
+        if getattr(self, "trainer_params", False):
+            kwargs["num_loader_workers"] = num_loader_workers
+            kwargs["batch_size"] = batch_size
+            kwargs["n_jobs"] = n_jobs
+            kwargs["roll_size"] = roll_size
+            kwargs["mc_dropout"] = mc_dropout
         return self.predict(
             n,
             series,
@@ -2239,13 +2361,30 @@ def _fit_wrapper(
         series: Union[TimeSeries, Sequence[TimeSeries]],
         past_covariates: Optional[Union[TimeSeries, Sequence[TimeSeries]]],
         future_covariates: Optional[Union[TimeSeries, Sequence[TimeSeries]]],
+        max_samples_per_ts: Optional[int] = None,
+        n_jobs_multioutput_wrapper: Optional[int] = None,
+        trainer=None,
+        verbose: Optional[bool] = None,
+        epochs: int = 0,
+        num_loader_workers: int = 0,
     ):
+        kwargs = dict()
+        if getattr(self, "trainer_params", False):
+            kwargs["trainer"] = trainer
+            kwargs["epochs"] = epochs
+            kwargs["verbose"] = verbose
+            kwargs["num_loader_workers"] = num_loader_workers
+            kwargs["max_samples_per_ts"] = max_samples_per_ts
+        elif isinstance(self, MultiOutputRegressor):
+            kwargs["n_jobs_multioutput_wrapper"] = n_jobs_multioutput_wrapper
+            kwargs["max_samples_per_ts"] = max_samples_per_ts
         self.fit(
             series=series,
             past_covariates=past_covariates if self.supports_past_covariates else None,
             future_covariates=future_covariates
             if self.supports_future_covariates
             else None,
+            **kwargs,
         )
 
     @property
@@ -2453,6 +2592,7 @@ def _fit_wrapper(
         series: TimeSeries,
         past_covariates: Optional[TimeSeries],
         future_covariates: Optional[TimeSeries],
+        **kwargs,
     ):
         self.fit(series, future_covariates=future_covariates)
 

darts/models/forecasting/regression_model.py

@@ -1094,7 +1094,6 @@ def _optimized_historical_forecasts(
         series: Optional[Sequence[TimeSeries]],
         past_covariates: Optional[Sequence[TimeSeries]] = None,
         future_covariates: Optional[Sequence[TimeSeries]] = None,
-        num_samples: int = 1,
         start: Optional[Union[pd.Timestamp, float, int]] = None,
         start_format: Literal["position", "value"] = "value",
         forecast_horizon: int = 1,
@@ -1103,7 +1102,7 @@ def _optimized_historical_forecasts(
         last_points_only: bool = True,
         verbose: bool = False,
         show_warnings: bool = True,
-        predict_likelihood_parameters: bool = False,
+        predict_kwargs: Optional[Dict[str, Any]] = None,
     ) -> Union[
         TimeSeries, List[TimeSeries], Sequence[TimeSeries], Sequence[List[TimeSeries]]
     ]:
@@ -1124,36 +1123,37 @@ def _optimized_historical_forecasts(
             allow_autoregression=False,
         )
 
+        if predict_kwargs is None:
+            predict_kwargs = dict()
+
         # TODO: move the loop here instead of duplicated code in each sub-routine?
         if last_points_only:
             return _optimized_historical_forecasts_last_points_only(
                 model=self,
                 series=series,
                 past_covariates=past_covariates,
                 future_covariates=future_covariates,
-                num_samples=num_samples,
                 start=start,
                 start_format=start_format,
                 forecast_horizon=forecast_horizon,
                 stride=stride,
                 overlap_end=overlap_end,
                 show_warnings=show_warnings,
-                predict_likelihood_parameters=predict_likelihood_parameters,
+                **predict_kwargs,
             )
         else:
             return _optimized_historical_forecasts_all_points(
                 model=self,
                 series=series,
                 past_covariates=past_covariates,
                 future_covariates=future_covariates,
-                num_samples=num_samples,
                 start=start,
                 start_format=start_format,
                 forecast_horizon=forecast_horizon,
                 stride=stride,
                 overlap_end=overlap_end,
                 show_warnings=show_warnings,
-                predict_likelihood_parameters=predict_likelihood_parameters,
+                **predict_kwargs,
             )