Repository: unit8co/darts

Issue:
[BUG] Can't access underlying models in multi-quantile regression 
**Describe the bug**

I'm using a multi-quantile forecaster on multivariate target data. E.g, a `CatBoostModel(likelihood='quantile', quantile=[0.01, 0.05, 0.1, 0.25, 0.5, 0.75, 0.9, 0.95, 0.99], ...)`. 

Darts fits a separate GBT model for each quantile level and each target component. However, these aren't accessible to me. 

Suppose my target series has two components and my `CatBoostModel` is called `model`. Then `model.model.estimators_` returns only 2 models (corresponding to quantile 0.99 for each component). 

This means that `model.get_multioutput_estimator` and `model.get_estimator` are incapable of returning estimators for any quantile other than 0.99.

Trying to access the models using `model.model` or `model._model_container` or `model._model_container[0.5]` all give a similar error: 

```python
{
	"name": "RuntimeError",
	"message": "scikit-learn estimators should always specify their parameters in the signature of their __init__ (no varargs). <class 'darts.utils.multioutput.MultiOutputRegressor'> with constructor (self, *args, eval_set_name: Optional[str] = None, eval_weight_name: Optional[str] = None, **kwargs) doesn't  follow this convention.",
	"stack": "---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
File ~/projects/equilibrium/helios/.venv/lib/python3.11/site-packages/IPython/core/formatters.py:347, in BaseFormatter.__call__(self, obj)
    345     method = get_real_method(obj, self.print_method)
    346     if method is not None:
--> 347         return method()
    348     return None
    349 else:

File ~/projects/equilibrium/helios/.venv/lib/python3.11/site-packages/sklearn/base.py:693, in BaseEstimator._repr_html_inner(self)
    688 def _repr_html_inner(self):
    689     \"\"\"This function is returned by the @property `_repr_html_` to make
    690     `hasattr(estimator, \"_repr_html_\") return `True` or `False` depending
    691     on `get_config()[\"display\"]`.
    692     \"\"\"
--> 693     return estimator_html_repr(self)

File ~/projects/equilibrium/helios/.venv/lib/python3.11/site-packages/sklearn/utils/_estimator_html_repr.py:363, in estimator_html_repr(estimator)
    361 style_template = Template(_CSS_STYLE)
    362 style_with_id = style_template.substitute(id=container_id)
--> 363 estimator_str = str(estimator)
    365 # The fallback message is shown by default and loading the CSS sets
    366 # div.sk-text-repr-fallback to display: none to hide the fallback message.
    367 #
   (...)
    372 # The reverse logic applies to HTML repr div.sk-container.
    373 # div.sk-container is hidden by default and the loading the CSS displays it.
    374 fallback_msg = (
    375     \"In a Jupyter environment, please rerun this cell to show the HTML\"
    376     \" representation or trust the notebook. <br />On GitHub, the\"
    377     \" HTML representation is unable to render, please try loading this page\"
    378     \" with nbviewer.org.\"
    379 )

File ~/projects/equilibrium/helios/.venv/lib/python3.11/site-packages/sklearn/base.py:315, in BaseEstimator.__repr__(self, N_CHAR_MAX)
    307 # use ellipsis for sequences with a lot of elements
    308 pp = _EstimatorPrettyPrinter(
    309     compact=True,
    310     indent=1,
    311     indent_at_name=True,
    312     n_max_elements_to_show=N_MAX_ELEMENTS_TO_SHOW,
    313 )
--> 315 repr_ = pp.pformat(self)
    317 # Use bruteforce ellipsis when there are a lot of non-blank characters
    318 n_nonblank = len(\"\".join(repr_.split()))

File ~/.pyenv/versions/3.11.10/lib/python3.11/pprint.py:161, in PrettyPrinter.pformat(self, object)
    159 def pformat(self, object):
    160     sio = _StringIO()
--> 161     self._format(object, sio, 0, 0, {}, 0)
    162     return sio.getvalue()

File ~/.pyenv/versions/3.11.10/lib/python3.11/pprint.py:178, in PrettyPrinter._format(self, object, stream, indent, allowance, context, level)
    176     self._readable = False
    177     return
--> 178 rep = self._repr(object, context, level)
    179 max_width = self._width - indent - allowance
    180 if len(rep) > max_width:

File ~/.pyenv/versions/3.11.10/lib/python3.11/pprint.py:458, in PrettyPrinter._repr(self, object, context, level)
    457 def _repr(self, object, context, level):
--> 458     repr, readable, recursive = self.format(object, context.copy(),
    459                                             self._depth, level)
    460     if not readable:
    461         self._readable = False

File ~/projects/equilibrium/helios/.venv/lib/python3.11/site-packages/sklearn/utils/_pprint.py:189, in _EstimatorPrettyPrinter.format(self, object, context, maxlevels, level)
    188 def format(self, object, context, maxlevels, level):
--> 189     return _safe_repr(
    190         object, context, maxlevels, level, changed_only=self._changed_only
    191     )

File ~/projects/equilibrium/helios/.venv/lib/python3.11/site-packages/sklearn/utils/_pprint.py:440, in _safe_repr(object, context, maxlevels, level, changed_only)
    438 recursive = False
    439 if changed_only:
--> 440     params = _changed_params(object)
    441 else:
    442     params = object.get_params(deep=False)

File ~/projects/equilibrium/helios/.venv/lib/python3.11/site-packages/sklearn/utils/_pprint.py:93, in _changed_params(estimator)
     89 def _changed_params(estimator):
     90     \"\"\"Return dict (param_name: value) of parameters that were given to
     91     estimator with non-default values.\"\"\"
---> 93     params = estimator.get_params(deep=False)
     94     init_func = getattr(estimator.__init__, \"deprecated_original\", estimator.__init__)
     95     init_params = inspect.signature(init_func).parameters

File ~/projects/equilibrium/helios/.venv/lib/python3.11/site-packages/sklearn/base.py:243, in BaseEstimator.get_params(self, deep)
    228 \"\"\"
    229 Get parameters for this estimator.
    230 
   (...)
    240     Parameter names mapped to their values.
    241 \"\"\"
    242 out = dict()
--> 243 for key in self._get_param_names():
    244     value = getattr(self, key)
    245     if deep and hasattr(value, \"get_params\") and not isinstance(value, type):

File ~/projects/equilibrium/helios/.venv/lib/python3.11/site-packages/sklearn/base.py:217, in BaseEstimator._get_param_names(cls)
    215 for p in parameters:
    216     if p.kind == p.VAR_POSITIONAL:
--> 217         raise RuntimeError(
    218             \"scikit-learn estimators should always \"
    219             \"specify their parameters in the signature\"
    220             \" of their __init__ (no varargs).\"
    221             \" %s with constructor %s doesn't \"
    222             \" follow this convention.\" % (cls, init_signature)
    223         )
    224 # Extract and sort argument names excluding 'self'
    225 return sorted([p.name for p in parameters])

RuntimeError: scikit-learn estimators should always specify their parameters in the signature of their __init__ (no varargs). <class 'darts.utils.multioutput.MultiOutputRegressor'> with constructor (self, *args, eval_set_name: Optional[str] = None, eval_weight_name: Optional[str] = None, **kwargs) doesn't  follow this convention."
}
```

**To Reproduce**

```python
import numpy as np
import pandas as pd
from darts import TimeSeries
from darts.models import CatBoostModel
from darts.utils.timeseries_generation import linear_timeseries

# Generate a synthetic multivariate time series
np.random.seed(42)
series1 = linear_timeseries(length=100, start_value=0, end_value=10)
series2 = linear_timeseries(length=100, start_value=10, end_value=0)
multivariate_series = series1.stack(series2)

# Define future covariates (optional, here just using a simple linear trend)
future_covariates = linear_timeseries(length=100, start_value=0, end_value=5)

# Initialize the CatBoostModel with quantile regression
model = CatBoostModel(
    lags=12,
    lags_future_covariates=[0],
    likelihood='quantile',
    quantiles=[0.01, 0.05, 0.1, 0.25, 0.5, 0.75, 0.9, 0.95, 0.99],
    random_state=42
)

# Fit the model
model.fit(multivariate_series, future_covariates=future_covariates)

# This is only 2 instead of 2 * 9
len(model.model.estimators_)

# Both of these give the above error
model.model
model._model_container
```

**Expected behavior**

Ability to access all the underlying estimators. 

**System (please complete the following information):**

 - Python version: 3.11.10
 - darts version: 0.31.0



Discussion:
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/unit8co/darts/pull/2716?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=unit8co) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 94.09%. Comparing base [(`24cec52`)](https://app.codecov.io/gh/unit8co/darts/commit/24cec52d2f9f89d1e436c5686c955baccf457c90?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=unit8co) to head [(`41ca608`)](https://app.codecov.io/gh/unit8co/darts/commit/41ca60891430c8da6fe51755f71f288052abf725?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=unit8co).

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #2716      +/-   ##
==========================================
- Coverage   94.16%   94.09%   -0.08%     
==========================================
  Files         141      141              
  Lines       15601    15610       +9     
==========================================
- Hits        14691    14688       -3     
- Misses        910      922      +12     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/unit8co/darts/pull/2716?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=unit8co).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=unit8co).


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index c133fca5..7a2eba14 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -17,6 +17,11 @@ but cannot always guarantee backwards compatibility. Changes that may **break co
 - Made method `ForecastingModel.untrained_model()` public. Use this method to get a new (untrained) model instance created with the same parameters. [#2684](https://github.com/unit8co/darts/pull/2684) by [Timon Erhart](https://github.com/turbotimon)
 - `TimeSeries.plot()` now supports setting the color for each component in the series. Simply pass a list / sequence of colors with length matching the number of components as parameters "c" or "colors". [#2680](https://github.com/unit8co/darts/pull/2680) by [Jules Authier](https://github.com/authierj)
 - Made it possible to run the quickstart notebook `00-quickstart.ipynb` locally. [#2691](https://github.com/unit8co/darts/pull/2691) by [Jules Authier](https://github.com/authierj)
+- Added `quantile` parameter to `RegressionModel.get_estimator()` to get the specific quantile estimator for probabilistic regression models using the `quantile` likelihood. [#2716](https://github.com/unit8co/darts/pull/2716) by [Antoine Madrona](https://github.com/madtoinou)
+
+**Removed**
+
+- 🔴 Removed method `RegressionModel.get_multioutput_estimator()`. Use `RegressionModel.get_estimator()` instead. [#2716](https://github.com/unit8co/darts/pull/2716) by [Antoine Madrona](https://github.com/madtoinou)
 
 **Fixed**
 
diff --git a/darts/explainability/shap_explainer.py b/darts/explainability/shap_explainer.py
index b931bbdc..d80091ac 100644
--- a/darts/explainability/shap_explainer.py
+++ b/darts/explainability/shap_explainer.py
@@ -611,7 +611,7 @@ class _RegressionShapExplainers:
                 self.explainers[i] = {}
                 for j in range(self.target_dim):
                     self.explainers[i][j] = self._build_explainer_sklearn(
-                        self.model.get_multioutput_estimator(horizon=i, target_dim=j),
+                        self.model.get_estimator(horizon=i, target_dim=j),
                         self.background_X,
                         self.shap_method,
                         **kwargs,
diff --git a/darts/models/forecasting/regression_model.py b/darts/models/forecasting/regression_model.py
index 23258c12..2c8444a8 100644
--- a/darts/models/forecasting/regression_model.py
+++ b/darts/models/forecasting/regression_model.py
@@ -505,10 +505,17 @@ class RegressionModel(GlobalForecastingModel):
     def output_chunk_shift(self) -> int:
         return self._output_chunk_shift
 
-    def get_multioutput_estimator(self, horizon: int, target_dim: int):
+    def get_estimator(
+        self, horizon: int, target_dim: int, quantile: Optional[float] = None
+    ):
         """Returns the estimator that forecasts the `horizon`th step of the `target_dim`th target component.
 
-        Internally, estimators are grouped by `output_chunk_length` position, then by component.
+        For probabilistic models fitting quantiles, it is possible to also specify the quantile.
+
+        The model is returned directly if it supports multi-output natively.
+
+        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`.
 
         Parameters
         ----------
@@ -516,51 +523,55 @@ class RegressionModel(GlobalForecastingModel):
             The index of the forecasting point within `output_chunk_length`.
         target_dim
             The index of the target component.
+        quantile
+            Optionally, for probabilistic model with `likelihood="quantile"`, a quantile value.
         """
-        raise_if_not(
-            isinstance(self.model, MultiOutputRegressor),
-            "The sklearn model is not a MultiOutputRegressor object.",
-            logger,
-        )
-        raise_if_not(
-            0 <= horizon < self.output_chunk_length,
-            f"`horizon` must be `>= 0` and `< output_chunk_length={self.output_chunk_length}`.",
-            logger,
-        )
-        raise_if_not(
-            0 <= target_dim < self.input_dim["target"],
-            f"`target_dim` must be `>= 0`, and `< n_target_components={self.input_dim['target']}`.",
-            logger,
-        )
+        if not isinstance(self.model, MultiOutputRegressor):
+            logger.warning(
+                "Model supports multi-output; a single estimator forecasts all the horizons and components."
+            )
+            return self.model
+
+        if not 0 <= horizon < self.output_chunk_length:
+            raise_log(
+                ValueError(
+                    f"`horizon` must be `>= 0` and `< output_chunk_length={self.output_chunk_length}`."
+                ),
+                logger,
+            )
+        if not 0 <= target_dim < self.input_dim["target"]:
+            raise_log(
+                ValueError(
+                    f"`target_dim` must be `>= 0`, and `< n_target_components={self.input_dim['target']}`."
+                ),
+                logger,
+            )
 
         # when multi_models=True, one model per horizon and target component
         idx_estimator = (
             self.multi_models * self.input_dim["target"] * horizon + target_dim
         )
-        return self.model.estimators_[idx_estimator]
+        if quantile is None:
+            return self.model.estimators_[idx_estimator]
 
-    def get_estimator(self, horizon: int, target_dim: int):
-        """Returns the estimator that forecasts the `horizon`th step of the `target_dim`th target component.
-
-        The model is returned directly if it supports multi-output natively.
-
-        Parameters
-        ----------
-        horizon
-            The index of the forecasting point within `output_chunk_length`.
-        target_dim
-            The index of the target component.
-        """
-
-        if isinstance(self.model, MultiOutputRegressor):
-            return self.get_multioutput_estimator(
-                horizon=horizon, target_dim=target_dim
+        # for quantile-models, the estimators are also grouped by quantiles
+        if self.likelihood != "quantile":
+            raise_log(
+                ValueError(
+                    "`quantile` is only supported for probabilistic models that "
+                    "use `likelihood='quantile'`."
+                ),
+                logger,
             )
-        else:
-            logger.info(
-                "Model supports multi-output; a single estimator forecasts all the horizons and components."
+        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
+        return self._model_container[quantile].estimators_[idx_estimator]
 
     def _add_val_set_to_kwargs(
         self,


Test patch:
diff --git a/darts/tests/models/forecasting/test_regression_models.py b/darts/tests/models/forecasting/test_regression_models.py
index 195dd460..26ed894c 100644
--- a/darts/tests/models/forecasting/test_regression_models.py
+++ b/darts/tests/models/forecasting/test_regression_models.py
@@ -1,6 +1,7 @@
 import functools
 import importlib
 import inspect
+import logging
 import math
 from copy import deepcopy
 from itertools import product
@@ -1333,7 +1334,7 @@ class TestRegressionModels:
         ],
     )
     def test_multioutput_wrapper(self, config):
-        """Check that with input_chunk_length=1, wrapping in MultiOutputRegressor is not happening"""
+        """Check that with input_chunk_length=1, wrapping in MultiOutputRegressor occurs only when necessary"""
         model, supports_multioutput_natively = config
         model.fit(series=self.sine_multivariate1)
         if supports_multioutput_natively:
@@ -1396,7 +1397,7 @@ class TestRegressionModels:
         else:
             assert not isinstance(model.model, MultiOutputRegressor)
 
-    def test_get_multioutput_estimator_multi_models(self):
+    def test_get_estimator_multi_models(self):
         """Craft training data so that estimator_[i].predict(X) == i + 1"""
 
         def helper_check_overfitted_estimators(ts: TimeSeries, ocl: int):
@@ -1417,7 +1418,7 @@ class TestRegressionModels:
             estimator_counter = 0
             for i in range(ocl):
                 for j in range(ts.width):
-                    sub_model = m.get_multioutput_estimator(horizon=i, target_dim=j)
+                    sub_model = m.get_estimator(horizon=i, target_dim=j)
                     pred = sub_model.predict(dummy_feats)[0]
                     # sub-model is overfitted on the training series
                     assert np.abs(estimator_counter - pred) < 1e-2
@@ -1455,7 +1456,7 @@ class TestRegressionModels:
         # estimators_[3] labels : [3]
         helper_check_overfitted_estimators(ts, ocl)
 
-    def test_get_multioutput_estimator_single_model(self):
+    def test_get_estimator_single_model(self):
         """Check estimator getter when multi_models=False"""
         # multivariate, one sub-model per component
         ocl = 2
@@ -1485,11 +1486,144 @@ class TestRegressionModels:
         dummy_feats = np.array([[0, 0, 0] * ts.width])
         for i in range(ocl):
             for j in range(ts.width):
-                sub_model = m.get_multioutput_estimator(horizon=i, target_dim=j)
+                sub_model = m.get_estimator(horizon=i, target_dim=j)
                 pred = sub_model.predict(dummy_feats)[0]
                 # sub-model forecast only depend on the target_dim
                 assert np.abs(j + 1 - pred) < 1e-2
 
+    @pytest.mark.parametrize("multi_models", [True, False])
+    def test_get_estimator_quantile(self, multi_models):
+        """Check estimator getter when using quantile value"""
+        ocl = 3
+        lags = 3
+        quantiles = [0.01, 0.5, 0.99]
+        ts = tg.sine_timeseries(length=100, column_name="sine").stack(
+            tg.linear_timeseries(length=100, column_name="linear"),
+        )
+
+        m = XGBModel(
+            lags=lags,
+            output_chunk_length=ocl,
+            multi_models=multi_models,
+            likelihood="quantile",
+            quantiles=quantiles,
+            random_state=1,
+        )
+        m.fit(ts)
+
+        assert len(m._model_container) == len(quantiles)
+        assert sorted(list(m._model_container.keys())) == sorted(quantiles)
+        for quantile_container in m._model_container.values():
+            # one sub-model per quantile, per component, per horizon
+            if multi_models:
+                assert len(quantile_container.estimators_) == ocl * ts.width
+            # one sub-model per quantile, per component
+            else:
+                assert len(quantile_container.estimators_) == ts.width
+
+        # check that retrieve sub-models prediction match the "wrapper" model predictions
+        pred_input = ts[-lags:] if multi_models else ts[-lags - ocl + 1 :]
+        pred = m.predict(
+            n=ocl,
+            series=pred_input,
+            num_samples=1,
+            predict_likelihood_parameters=True,
+        )
+        for j in range(ts.width):
+            for i in range(ocl):
+                if multi_models:
+                    dummy_feats = pred_input.values()[:lags]
+                else:
+                    dummy_feats = pred_input.values()[i : +i + lags]
+                dummy_feats = np.expand_dims(dummy_feats.flatten(), 0)
+                for q in quantiles:
+                    sub_model = m.get_estimator(horizon=i, target_dim=j, quantile=q)
+                    pred_sub_model = sub_model.predict(dummy_feats)[0]
+                    assert (
+                        pred_sub_model
+                        == pred[f"{ts.components[j]}_q{q:.2f}"].values()[i][0]
+                    )
+
+    def test_get_estimator_exceptions(self, caplog):
+        """Check that all the corner-cases are properly covered by the method"""
+        ts = TimeSeries.from_values(
+            values=np.array([
+                [0, 0, 0, 0, 1],
+                [0, 0, 0, 0, 2],
+            ]).T,
+            columns=["a", "b"],
+        )
+        m = LinearRegressionModel(
+            lags=2,
+            output_chunk_length=2,
+            random_state=1,
+        )
+        m.fit(ts["a"])
+        # not wrapped in MultiOutputRegressor because of native multi-output support
+        with caplog.at_level(logging.WARNING):
+            m.get_estimator(horizon=0, target_dim=0)
+        assert len(caplog.records) == 1
+        assert caplog.records[0].levelname == "WARNING"
+        assert caplog.records[0].message == (
+            "Model supports multi-output; a single estimator "
+            "forecasts all the horizons and components."
+        )
+
+        # univariate, deterministic, ocl > 2
+        m = RegressionModel(
+            model=HistGradientBoostingRegressor(),
+            lags=2,
+            output_chunk_length=2,
+        )
+        m.fit(ts["a"])
+        # horizon > ocl
+        with pytest.raises(ValueError) as err:
+            m.get_estimator(horizon=3, target_dim=0)
+        assert str(err.value).startswith(
+            "`horizon` must be `>= 0` and `< output_chunk_length"
+        )
+        # target dim > training series width
+        with pytest.raises(ValueError) as err:
+            m.get_estimator(horizon=0, target_dim=1)
+        assert str(err.value).startswith(
+            "`target_dim` must be `>= 0`, and `< n_target_components="
+        )
+
+        # univariate, probabilistic
+        # using the quantiles argument to force wrapping in MultiOutputRegressor
+        m = XGBModel(
+            lags=2,
+            output_chunk_length=2,
+            random_state=1,
+            likelihood="poisson",
+            quantiles=[0.5],
+        )
+        m.fit(ts["a"])
+        # incorrect likelihood
+        with pytest.raises(ValueError) as err:
+            m.get_estimator(horizon=0, target_dim=0, quantile=0.1)
+        assert str(err.value).startswith(
+            "`quantile` is only supported for probabilistic models that "
+            "use `likelihood='quantile'`."
+        )
+
+        # univariate, probabilistic
+        m = XGBModel(
+            lags=2,
+            output_chunk_length=2,
+            random_state=1,
+            likelihood="quantile",
+            quantiles=[0.01, 0.5, 0.99],
+        )
+        m.fit(ts["a"])
+        # retrieving a non-defined quantile
+        with pytest.raises(ValueError) as err:
+            m.get_estimator(horizon=0, target_dim=0, quantile=0.1)
+        assert str(err.value).startswith(
+            "Invalid `quantile=0.1`. Must be one of the fitted quantiles "
+            "`[0.01, 0.5, 0.99]`."
+        )
+
     @pytest.mark.parametrize("mode", [True, False])
     def test_regression_model(self, mode):
         lags = 4