Fix/sklearn model quantile reg (#2838)

* improve likelihood handling for sklearn models

* update tests

* update changelog

Author: dennisbader

CHANGELOG.md

@@ -13,6 +13,8 @@ but cannot always guarantee backwards compatibility. Changes that may **break co
 
 **Fixed**
 
+- Fixed a bug in `SKLearnModel.get_estimator()` for univariate quantile models that use `multi_models=False` , where using `quantile` did not return the correct fitted quantile model / estimator. [#2838](https://github.com/unit8co/darts/pull/2838) by [Dennis Bader](https://github.com/dennisbader).
+
 **Dependencies**
 
 ### For developers of the library:

darts/models/forecasting/catboost_model.py

@@ -192,7 +192,6 @@ def encode_year(idx):
         """
         kwargs["random_state"] = random_state  # seed for tree learner
         self.kwargs = kwargs
-        self._model_container = None
 
         # parse likelihood
         if likelihood is not None:
@@ -207,12 +206,10 @@ def encode_year(idx):
                 # RMSEWithUncertainty returns mean and variance which is equivalent to gaussian
                 likelihood = "gaussian"
 
-            if likelihood == "quantile":
-                self._model_container = _QuantileModelContainer()
-            else:
+            if likelihood != "quantile":
                 self.kwargs["loss_function"] = likelihood_map[likelihood]
 
-        self._likelihood = _get_likelihood(
+        likelihood = _get_likelihood(
             likelihood=likelihood,
             n_outputs=output_chunk_length if multi_models else 1,
             quantiles=quantiles,
@@ -241,6 +238,10 @@ def encode_year(idx):
         # if no loss provided, get the default loss from the model
         self.kwargs["loss_function"] = self.model.get_params().get("loss_function")
 
+        self._likelihood = likelihood
+        if isinstance(likelihood, QuantileRegression):
+            self._model_container = _QuantileModelContainer()
+
     def fit(
         self,
         series: Union[TimeSeries, Sequence[TimeSeries]],
@@ -325,6 +326,7 @@ def fit(
                     verbose=verbose,
                     **kwargs,
                 )
+                # store the trained model in the container as it might have been wrapped by MultiOutputRegressor
                 self._model_container[quantile] = self.model
             return self
 

darts/models/forecasting/lgbm.py

@@ -190,16 +190,13 @@ def encode_year(idx):
         """
         kwargs["random_state"] = random_state  # seed for tree learner
         self.kwargs = kwargs
-        self._model_container = None
 
         # parse likelihood
         if likelihood is not None:
             _check_likelihood(likelihood, ["quantile", "poisson"])
             self.kwargs["objective"] = likelihood
-            if likelihood == "quantile":
-                self._model_container = _QuantileModelContainer()
 
-        self._likelihood = _get_likelihood(
+        likelihood = _get_likelihood(
             likelihood=likelihood,
             n_outputs=output_chunk_length if multi_models else 1,
             quantiles=quantiles,
@@ -221,6 +218,10 @@ def encode_year(idx):
             random_state=random_state,
         )
 
+        self._likelihood = likelihood
+        if isinstance(likelihood, QuantileRegression):
+            self._model_container = _QuantileModelContainer()
+
     def fit(
         self,
         series: Union[TimeSeries, Sequence[TimeSeries]],
@@ -299,6 +300,7 @@ def fit(
                     val_sample_weight=val_sample_weight,
                     **kwargs,
                 )
+                # store the trained model in the container as it might have been wrapped by MultiOutputRegressor
                 self._model_container[quantile] = self.model
             return self
 

darts/models/forecasting/linear_regression_model.py

@@ -173,7 +173,6 @@ def encode_year(idx):
                [1005.81830675]])
         """
         self.kwargs = kwargs
-        self._model_container = None
 
         # parse likelihood
         if likelihood is not None:
@@ -182,11 +181,10 @@ def encode_year(idx):
                 model = PoissonRegressor(**kwargs)
             if likelihood == "quantile":
                 model = QuantileRegressor(**kwargs)
-                self._model_container = _QuantileModelContainer()
         else:
             model = LinearRegression(**kwargs)
 
-        self._likelihood = _get_likelihood(
+        likelihood = _get_likelihood(
             likelihood=likelihood,
             n_outputs=output_chunk_length if multi_models else 1,
             quantiles=quantiles,
@@ -205,6 +203,10 @@ def encode_year(idx):
             random_state=random_state,
         )
 
+        self._likelihood = likelihood
+        if isinstance(likelihood, QuantileRegression):
+            self._model_container = _QuantileModelContainer()
+
     def fit(
         self,
         series: Union[TimeSeries, Sequence[TimeSeries]],
@@ -249,7 +251,7 @@ def fit(
                     sample_weight=sample_weight,
                     **kwargs,
                 )
-
+                # store the trained model in the container as it might have been wrapped by MultiOutputRegressor
                 self._model_container[quantile] = self.model
 
             # replace the last trained QuantileRegressor with the dictionary of Regressors.

darts/models/forecasting/sklearn_model.py

@@ -227,6 +227,11 @@ def encode_year(idx):
         self._lagged_feature_names: Optional[list[str]] = None
         self._lagged_label_names: Optional[list[str]] = None
 
+        # optionally, the model can be wrapped in a likelihood model
+        self._likelihood: Optional[SKLearnLikelihood] = None
+        # for quantile likelihood models, the model container is a dict of quantile -> model
+        self._model_container: Optional[dict[float, Any]] = None
+
         # check and set output_chunk_length
         raise_if_not(
             isinstance(output_chunk_length, int) and output_chunk_length > 0,
@@ -525,9 +530,11 @@ def get_estimator(
     ):
         """Returns the estimator that forecasts the `horizon`th step of the `target_dim`th target component.
 
-        For probabilistic models fitting quantiles, it is possible to also specify the quantile.
+        For probabilistic models fitting quantiles, a desired `quantile` can also be passed. If not passed, it will
+        return the model predicting the median (quantile=0.5).
 
-        The model is returned directly if it supports multi-output natively.
+        If the (quantile) model supports multi-output natively, it will return the model that can predict the entire
+        horizon and all target components jointly.
 
         Note: Internally, estimators are grouped by `output_chunk_length` position, then by component. For probabilistic
         models fitting quantiles, there is an additional abstraction layer, grouping the estimators by `quantile`.
@@ -539,13 +546,38 @@ def get_estimator(
         target_dim
             The index of the target component.
         quantile
-            Optionally, for probabilistic model with `likelihood="quantile"`, a quantile value.
+            Optionally, for probabilistic model with `likelihood="quantile"`, the desired quantile value. If `None` and
+            `likelihood="quantile"`, returns the model predicting the median (quantile=0.5).
         """
-        if not isinstance(self.model, MultiOutputRegressor):
+        likelihood = self.likelihood
+        if isinstance(likelihood, QuantileRegression):
+            # for quantile-models, the estimators are grouped by quantiles
+            if quantile is None:
+                quantile = likelihood.quantiles[likelihood._median_idx]
+            elif quantile not in self._model_container:
+                raise_log(
+                    ValueError(
+                        f"Invalid `quantile={quantile}`. Must be one of the fitted quantiles "
+                        f"`{list(self._model_container.keys())}`."
+                    ),
+                    logger,
+                )
+            model = self._model_container[quantile]
+        elif quantile is not None:
+            raise_log(
+                ValueError(
+                    "`quantile` is only supported for probabilistic models that use `likelihood='quantile'`."
+                ),
+                logger=logger,
+            )
+        else:
+            model = self.model
+
+        if not isinstance(model, MultiOutputRegressor):
             logger.warning(
                 "Model supports multi-output; a single estimator forecasts all the horizons and components."
             )
-            return self.model
+            return model
 
         if not 0 <= horizon < self.output_chunk_length:
             raise_log(
@@ -566,27 +598,7 @@ def get_estimator(
         idx_estimator = (
             self.multi_models * self.input_dim["target"] * horizon + target_dim
         )
-        if quantile is None:
-            return self.model.estimators_[idx_estimator]
-
-        # for quantile-models, the estimators are also grouped by quantiles
-        if not isinstance(self.likelihood, QuantileRegression):
-            raise_log(
-                ValueError(
-                    "`quantile` is only supported for probabilistic models that "
-                    "use `likelihood='quantile'`."
-                ),
-                logger,
-            )
-        if quantile not in self._model_container:
-            raise_log(
-                ValueError(
-                    f"Invalid `quantile={quantile}`. Must be one of the fitted quantiles "
-                    f"`{list(self._model_container.keys())}`."
-                ),
-                logger,
-            )
-        return self._model_container[quantile].estimators_[idx_estimator]
+        return model.estimators_[idx_estimator]
 
     def _add_val_set_to_kwargs(
         self,

darts/models/forecasting/xgboost.py

@@ -191,7 +191,6 @@ def encode_year(idx):
         """
         kwargs["random_state"] = random_state  # seed for tree learner
         self.kwargs = kwargs
-        self._model_container = None
 
         # parse likelihood
         if likelihood is not None:
@@ -200,9 +199,8 @@ def encode_year(idx):
                 self.kwargs["objective"] = f"count:{likelihood}"
             elif likelihood == "quantile":
                 self.kwargs["objective"] = "reg:quantileerror"
-                self._model_container = _QuantileModelContainer()
 
-        self._likelihood = _get_likelihood(
+        likelihood = _get_likelihood(
             likelihood=likelihood,
             n_outputs=output_chunk_length if multi_models else 1,
             quantiles=quantiles,
@@ -221,6 +219,10 @@ def encode_year(idx):
             random_state=random_state,
         )
 
+        self._likelihood = likelihood
+        if isinstance(likelihood, QuantileRegression):
+            self._model_container = _QuantileModelContainer()
+
     def fit(
         self,
         series: Union[TimeSeries, Sequence[TimeSeries]],
@@ -301,6 +303,7 @@ def fit(
                     val_sample_weight=val_sample_weight,
                     **kwargs,
                 )
+                # store the trained model in the container as it might have been wrapped by MultiOutputRegressor
                 self._model_container[quantile] = self.model
             return self