Catboost native multi-output with RMSE

CHANGELOG.md

@@ -17,6 +17,7 @@ but cannot always guarantee backwards compatibility. Changes that may **break co
 - Added more resampling methods to `TimeSeries.resample()`. This allows to aggregate values when down-sampling and to fill or keep the holes when up-sampling. [#2654](https://github.com/unit8co/darts/pull/2654) by [Jonas Blanc](https://github.com/jonasblanc)
 - Added general function `darts.slice_intersect()` to intersect a sequence of `TimeSeries` along the time index. [#2592](https://github.com/unit8co/darts/pull/2592) by [Yoav Matzkevich](https://github.com/ymatzkevich).
 - Added new time aggregated metric `wmape()` (Weighted Mean Absolute Percentage Error). [#2544](https://github.com/unit8co/darts/pull/2648) by [He Weilin](https://github.com/cnhwl).
+- Added support for native `CatBoostModel` multioutput (by converting loss function from `RMSE` to `MultiRMSE` when required) instead of relying on Darts `MultiOutputRegressor`. [#2659](https://github.com/unit8co/darts/pull/2659) by [Jonas Blanc](https://github.com/jonasblanc)
 - Improvements to `ForecastingModel`:
   - Added parameter `clean: bool` to `ForecastingModel.save()` to store a cleaned version of the model (removes training data from global models, and Lightning Trainer-related parameters from torch models). [#2649](https://github.com/unit8co/darts/pull/2649) by [Jonas Blanc](https://github.com/jonasblanc).
 - Added parameter `pl_trainer_kwargs` to `TorchForecastingModel.load()` to setup a new Lightning Trainer used to configure the model for downstream tasks (e.g. prediction). [#2649](https://github.com/unit8co/darts/pull/2649) by [Jonas Blanc](https://github.com/jonasblanc).

darts/models/forecasting/catboost_model.py

@@ -16,6 +16,7 @@
 from darts.logging import get_logger
 from darts.models.forecasting.regression_model import RegressionModel, _LikelihoodMixin
 from darts.timeseries import TimeSeries
+from darts.utils.multioutput import MultiOutputRegressor
 
 logger = get_logger(__name__)
 
@@ -208,6 +209,10 @@ def encode_year(idx):
             use_static_covariates=use_static_covariates,
         )
 
+    def _native_support_multioutput(self):
+        # CatBoostRegressor supports multioutput natively, but only with the "MultiRMSE" loss function
+        return self.kwargs["loss_function"] == "MultiRMSE"
+
     def fit(
         self,
         series: Union[TimeSeries, Sequence[TimeSeries]],
@@ -249,9 +254,9 @@ def fit(
             creation) to know their sizes, which might be expensive on big datasets.
             If some series turn out to have a length that would allow more than `max_samples_per_ts`, only the
             most recent `max_samples_per_ts` samples will be considered.
-        n_jobs_multioutput_wrapper
-            Number of jobs of the MultiOutputRegressor wrapper to run in parallel. Only used if the model doesn't
-            support multi-output regression natively.
+        # n_jobs_multioutput_wrapper
+        #     Number of jobs of the MultiOutputRegressor wrapper to run in parallel. Only used if the model doesn't
+        #     support multi-output regression natively.
         sample_weight
             Optionally, some sample weights to apply to the target `series` labels. They are applied per observation,
             per label (each step in `output_chunk_length`), and per component.
@@ -294,6 +299,23 @@ def fit(
                 self._model_container[quantile] = self.model
             return self
 
+        # If multioutput, and notprobabilistic, use MultiRMSE loss for CatBoost native multioutput support
+        require_multioutput = not series[0].is_univariate or (
+            self.output_chunk_length > 1
+            and self.multi_models
+            and not isinstance(self.model, MultiOutputRegressor)
+        )
+
+        if require_multioutput and (
+            self.kwargs.get("loss_function") is None
+            or self.kwargs["loss_function"] == "RMSE"
+        ):
+            self.kwargs["loss_function"] = "MultiRMSE"
+            self.model = CatBoostRegressor(**self.kwargs)
+            logger.warning(
+                "Changed loss function to 'MultiRMSE' for multioutput support"
+            )
+
         super().fit(
             series=series,
             past_covariates=past_covariates,

darts/models/forecasting/regression_model.py

@@ -736,6 +736,16 @@ def _fit_model(
             )
         )
 
+    def _native_support_multioutput(self) -> bool:
+        """
+        Returns True if the model supports multi-output regression natively.
+        """
+        return (
+            callable(getattr(self.model, "_get_tags", None))
+            and isinstance(self.model._get_tags(), dict)
+            and self.model._get_tags().get("multioutput")
+        )
+
     def fit(
         self,
         series: Union[TimeSeries, Sequence[TimeSeries]],
@@ -846,36 +856,17 @@ def fit(
             "future": future_covariates[0].width if future_covariates else None,
         }
 
-        # if multi-output regression
-        use_mor = False
-        if not series[0].is_univariate or (
+        # Check if multi-output regression is required
+        require_multioutput = not series[0].is_univariate or (
             self.output_chunk_length > 1
             and self.multi_models
             and not isinstance(self.model, MultiOutputRegressor)
-        ):
-            if sample_weight is not None:
-                # we have 2D sample (and time) weights, only supported in Darts
-                use_mor = True
-            elif not (
-                callable(getattr(self.model, "_get_tags", None))
-                and isinstance(self.model._get_tags(), dict)
-                and self.model._get_tags().get("multioutput")
-            ):
-                # model does not support multi-output regression natively
-                use_mor = True
-            elif (
-                self.model.__class__.__name__ == "CatBoostRegressor"
-                and self.model.get_params()["loss_function"] == "RMSEWithUncertainty"
-            ):
-                use_mor = True
-            elif (
-                self.model.__class__.__name__ == "XGBRegressor"
-                and self.likelihood is not None
-            ):
-                # since xgboost==2.1.0, likelihoods do not support native multi output regression
-                use_mor = True
+        )
 
-        if use_mor:
+        # If multi-output required and model doesn't support it natively, wrap it in a MultiOutputRegressor
+        if require_multioutput and (
+            not self._native_support_multioutput() or sample_weight is not None
+        ):
             val_set_name, val_weight_name = self.val_set_params
             mor_kwargs = {
                 "eval_set_name": val_set_name,
@@ -884,7 +875,6 @@ def fit(
             }
             self.model = MultiOutputRegressor(self.model, **mor_kwargs)
 
-        # warn if n_jobs_multioutput_wrapper was provided but not used
         if (
             not isinstance(self.model, MultiOutputRegressor)
             and n_jobs_multioutput_wrapper is not None

darts/models/forecasting/xgboost.py

@@ -225,6 +225,10 @@ def encode_year(idx):
             use_static_covariates=use_static_covariates,
         )
 
+    def _native_support_multioutput(self):
+        # since xgboost==2.1.0, likelihoods do not support native multi output regression
+        return super()._native_support_multioutput() and self.likelihood is None
+
     def fit(
         self,
         series: Union[TimeSeries, Sequence[TimeSeries]],

darts/tests/explainability/test_shap_explainer.py

@@ -205,8 +205,9 @@ def test_creation(self):
 
         # Good type of explainers
         shap_explain = ShapExplainer(m)
-        if isinstance(m, XGBModel):
+        if m._native_support_multioutput():
             # since xgboost > 2.1.0, model supports native multi output regression
+            # CatBoostModel supports multioutput for certain loss functions
             assert isinstance(shap_explain.explainers.explainers, shap.explainers.Tree)
         else:
             assert isinstance(
@@ -270,7 +271,7 @@ def test_creation(self):
             future_covariates=self.fut_cov_ts,
         )
         shap_explain = ShapExplainer(m)
-        if isinstance(m, XGBModel):
+        if m._native_support_multioutput():
             assert isinstance(shap_explain.explainers.explainers, shap.explainers.Tree)
         else:
             assert isinstance(

darts/tests/models/forecasting/test_regression_models.py

@@ -182,11 +182,7 @@ class NewCls(cls):
 class TestRegressionModels:
     np.random.seed(42)
     # default regression models
-    models = [
-        RandomForest,
-        LinearRegressionModel,
-        RegressionModel,
-    ]
+    models = [RandomForest, LinearRegressionModel, RegressionModel]
 
     # register likelihood regression models
     QuantileLinearRegressionModel = partialclass(
@@ -1362,7 +1358,9 @@ def test_multioutput_wrapper(self, config):
     if lgbm_available:
         model_configs += [(LightGBMModel, lgbm_test_params)]
     if cb_available:
-        model_configs += [(CatBoostModel, cb_test_params)]
+        model_configs += [
+            (CatBoostModel, dict({"likelihood": "poisson"}, **cb_test_params))
+        ]
 
     @pytest.mark.parametrize("config", product(model_configs, [1, 2], [True, False]))
     def test_multioutput_validation(self, config):