diff --git a/darts/models/forecasting/catboost_model.py b/darts/models/forecasting/catboost_model.py index fd52053f14..486dfd4b14 100644 --- a/darts/models/forecasting/catboost_model.py +++ b/darts/models/forecasting/catboost_model.py @@ -4,13 +4,12 @@ This is a wrapper that enables using the CatBoost regressor as model """ - from typing import List, Optional, Sequence, Tuple, Union import numpy as np from catboost import CatBoostRegressor -from darts.logging import get_logger, raise_if +from darts.logging import get_logger from darts.models.forecasting.regression_model import RegressionModel, _LikelihoodMixin from darts.timeseries import TimeSeries @@ -50,8 +49,10 @@ def __init__( horizon `n` used in `predict()`. However, setting `output_chunk_length` equal to the forecast horizon may be useful if the covariates don't extend far enough into the future. likelihood - Can be set to `quantile` or 'poisson'. If set, the model will be probabilistic, allowing sampling at - prediction time. + Can be set to 'quantile', 'poisson' or 'gaussian'. If set, the model will be probabilistic, + allowing sampling at prediction time. When set to 'gaussian', the model will use CatBoost's + 'RMSEWithUncertainty' loss function. When using this loss function, CatBoost returns a mean + and variance couple, which capture data (aleatoric) uncertainty. quantiles Fit the model to these quantiles if the `likelihood` is set to `quantile`. random_state @@ -68,21 +69,13 @@ def __init__( self.likelihood = likelihood self.quantiles = None - if "loss_function" in kwargs.keys(): - raise_if( - kwargs["loss_function"] == "RMSEWithUncertainty", - "The loss function RMSEWithUncertainty is not supported by darts.", - ) - elif "objective" in kwargs.keys(): - raise_if( - kwargs["objective"] == "RMSEWithUncertainty", - "The objective RMSEWithUncertainty is not supported by darts.", - ) + self.output_chunk_length = output_chunk_length - # to be extended to RMSEWithUncertainty likelihood_map = { "quantile": None, "poisson": "Poisson", + "gaussian": "RMSEWithUncertainty", + "RMSEWithUncertainty": "RMSEWithUncertainty", } available_likelihoods = list(likelihood_map.keys()) @@ -193,10 +186,15 @@ def fit( def _predict_and_sample( self, x: np.ndarray, num_samples: int, **kwargs ) -> np.ndarray: + """Override of RegressionModel's predict method, + to allow for the probabilistic case + """ if self.likelihood == "quantile": return self._predict_quantiles(x, num_samples, **kwargs) elif self.likelihood == "poisson": return self._predict_poisson(x, num_samples, **kwargs) + elif self.likelihood in ["gaussian", "RMSEWithUncertainty"]: + return self._predict_normal(x, num_samples, **kwargs) else: return super()._predict_and_sample(x, num_samples, **kwargs) diff --git a/darts/models/forecasting/regression_model.py b/darts/models/forecasting/regression_model.py index 9b5de591c1..f5ada37a2b 100644 --- a/darts/models/forecasting/regression_model.py +++ b/darts/models/forecasting/regression_model.py @@ -29,6 +29,7 @@ import numpy as np import pandas as pd +from catboost import CatBoostRegressor from sklearn.linear_model import LinearRegression from sklearn.multioutput import MultiOutputRegressor @@ -417,6 +418,15 @@ def fit( self.model = MultiOutputRegressor( self.model, n_jobs=n_jobs_multioutput_wrapper ) + elif isinstance(self.model, CatBoostRegressor): + if ( + self.model.get_params()["loss_function"] + == "RMSEWithUncertainty" + ): + self.model = MultiOutputRegressor( + self.model, n_jobs=n_jobs_multioutput_wrapper + ) + # warn if n_jobs_multioutput_wrapper was provided but not used if ( not isinstance(self.model, MultiOutputRegressor) @@ -635,7 +645,7 @@ def __str__(self): class _LikelihoodMixin: """ - A class containing functions supporting quantile and poisson regression, to be used as a mixin for some + A class containing functions supporting quantile, poisson and gaussian regression, to be used as a mixin for some `RegressionModel` subclasses. """ @@ -700,6 +710,67 @@ def _predict_quantiles( return sampled + def _predict_normal(self, x: np.ndarray, num_samples: int, **kwargs) -> np.ndarray: + """Method intended for CatBoost's RMSEWithUncertainty loss. Returns samples + computed from double-valued inputs [mean, variance]. + X is of shape (n_series * n_samples, n_regression_features) + """ + k = x.shape[0] + + # model_output shape: + # if univariate & output_chunk_length = 1: (num_samples, 2) + # else: (2, num_samples, n_components * output_chunk_length) + # where the axis with 2 dims is mu, sigma + model_output = self.model.predict(x, **kwargs) + output_dim = len(model_output.shape) + + # deterministic case: we return the mean only + if num_samples == 1: + # univariate & single-chunk output + if output_dim <= 2: + output_slice = model_output[:, 0] + else: + output_slice = model_output[0, :, :] + + return output_slice.reshape(k, self.output_chunk_length, -1) + + # probabilistic case + # univariate & single-chunk output + if output_dim <= 2: + # embedding well shaped 2D output into 3D + model_output = np.expand_dims(model_output, axis=0) + + else: + # we transpose to get mu, sigma couples on last axis + # shape becomes: (n_components * output_chunk_length, num_samples, 2) + model_output = model_output.transpose() + + return self._normal_sampling(model_output, num_samples) + + def _normal_sampling(self, model_output: np.ndarray, n_samples: int) -> np.ndarray: + """Sampling method for CatBoost's [mean, variance] output. + model_output is of shape (n_components * output_chunk_length, n_samples, 2), + where the last 2 dimensions are mu and sigma. + """ + shape = model_output.shape + chunk_len = self.output_chunk_length + + # treating each component separately + mu_sigma_list = [model_output[i, :, :] for i in range(shape[0])] + + list_of_samples = [ + self._rng.normal( + mu_sigma[:, 0], # mean vector + mu_sigma[:, 1], # diagonal covariance matrix + ) + for mu_sigma in mu_sigma_list + ] + + samples_transposed = np.array(list_of_samples).transpose() + samples_reshaped = samples_transposed.reshape(n_samples, chunk_len, -1) + + return samples_reshaped + def _predict_poisson(self, x: np.ndarray, num_samples: int, **kwargs) -> np.ndarray: """ X is of shape (n_series * n_samples, n_regression_features) diff --git a/darts/tests/models/forecasting/test_regression_models.py b/darts/tests/models/forecasting/test_regression_models.py index 0cc65d1335..b7b1f7a734 100644 --- a/darts/tests/models/forecasting/test_regression_models.py +++ b/darts/tests/models/forecasting/test_regression_models.py @@ -180,6 +180,11 @@ class RegressionModelsTestCase(DartsBaseTestClass): likelihood="poisson", random_state=42, ) + NormalCatBoostModel = partialclass( + CatBoostModel, + likelihood="gaussian", + random_state=42, + ) QuantileLightGBMModel = partialclass( LightGBMModel, likelihood="quantile", @@ -207,6 +212,7 @@ class RegressionModelsTestCase(DartsBaseTestClass): PoissonLightGBMModel, PoissonLinearRegressionModel, PoissonCatBoostModel, + NormalCatBoostModel, ] ) @@ -222,6 +228,7 @@ class RegressionModelsTestCase(DartsBaseTestClass): 0.4, # PoissonLightGBMModel 0.4, # PoissonLinearRegressionModel 1e-01, # PoissonCatBoostModel + 1e-05, # NormalCatBoostModel ] multivariate_accuracies = [ 0.3, @@ -235,6 +242,7 @@ class RegressionModelsTestCase(DartsBaseTestClass): 0.4, 0.4, 0.15, + 1e-05, ] multivariate_multiseries_accuracies = [ 0.05, @@ -248,6 +256,7 @@ class RegressionModelsTestCase(DartsBaseTestClass): 0.4, 0.4, 1e-01, + 1e-03, ] # dummy feature and target TimeSeries instances @@ -991,6 +1000,23 @@ def test_gradient_boosted_model_with_eval_set(self, lgb_fit_patch): assert lgb_fit_patch.call_args[1]["eval_set"] is not None assert lgb_fit_patch.call_args[1]["early_stopping_rounds"] == 2 + @patch.object(darts.models.forecasting.catboost_model.CatBoostRegressor, "fit") + def test_catboost_model_with_eval_set(self, lgb_fit_patch): + """Test whether these evaluation set parameters are passed to CatBoostRegressor""" + model = CatBoostModel(lags=4, lags_past_covariates=2) + model.fit( + series=self.sine_univariate1, + past_covariates=self.sine_multivariate1, + val_series=self.sine_univariate1, + val_past_covariates=self.sine_multivariate1, + early_stopping_rounds=2, + ) + + lgb_fit_patch.assert_called_once() + + assert lgb_fit_patch.call_args[1]["eval_set"] is not None + assert lgb_fit_patch.call_args[1]["early_stopping_rounds"] == 2 + class ProbabilisticRegressionModelsTestCase(DartsBaseTestClass): models_cls_kwargs_errs = [ ( @@ -1033,6 +1059,11 @@ class ProbabilisticRegressionModelsTestCase(DartsBaseTestClass): {"lags": 2, "likelihood": "poisson", "random_state": 42}, 0.6, ), + ( + CatBoostModel, + {"lags": 2, "likelihood": "gaussian", "random_state": 42}, + 0.05, + ), ( LinearRegressionModel, {"lags": 2, "likelihood": "quantile", "random_state": 42},