Refactor / unify likelihood models (#2742)

* make base likelihood model

* move likelihoods to dedicated module

* refactor regression model likelihoods

* fix failing tests

* remove remaining old likelihood references

* rename likelihood.likelihood_models to likelihood_models.torch

* update changelog

* fix imports

* address some missed lines

* make torch likelihoods backwards compatible

* add sklearn likelihood tests

* remane likelihood.py to base.py

* apply suggestions from PR review

* update gaussian sklearn likelihood

* rename torch base likelihood

* update docs

Author: dennisbader

CHANGELOG.md

@@ -11,17 +11,28 @@ but cannot always guarantee backwards compatibility. Changes that may **break co
 
 **Improved**
 
-**Removed**
+**Removed / moved**
+
+- 🔴 Removed model `AutoARIMA`. To support `numpy>=2.0.0`, we unfortunately had to remove the `pmdarima` dependency. Use `StatsForecastAutoARIMA` instead. [#2734](https://github.com/unit8co/darts/pull/2734) by [Dennis Bader](https://github.com/dennisbader).
 - 🔴 Removed deprecated method `TimeSeries.pd_dataframe()`. Use `TimeSeries.to_dataframe()` instead. [#2733](https://github.com/unit8co/darts/pull/2733) by [Dennis Bader](https://github.com/dennisbader).
 - 🔴 Removed deprecated method `TimeSeries.pd_serise()`. Use `TimeSeries.to_series()` instead. [#2733](https://github.com/unit8co/darts/pull/2733) by [Dennis Bader](https://github.com/dennisbader).
-- 🔴 Removed model `AutoARIMA`. To support `numpy>=2.0.0`, we unfortunately had to remove the `pmdarima` dependency. Use `StatsForecastAutoARIMA` instead. [#2734](https://github.com/unit8co/darts/pull/2734) by [Dennis Bader](https://github.com/dennisbader).
 
 **Fixed**
 
+- Fixed a bug in `CatBoostModel` with `likelihood="gaussian"`, where predicting with `predict_likelihood_parameters=True` resulted in wrong ordering of the predicted parameters. [#2742](https://github.com/unit8co/darts/pull/2742) by [Dennis Bader](https://github.com/dennisbader).
+
 **Dependencies**
 
 ### For developers of the library:
 
+- Refactored likelihoods: [#2742](https://github.com/unit8co/darts/pull/2742) by [Dennis Bader](https://github.com/dennisbader).
+  - Moved all likelihood related files into `darts/utils/likelihood_models/`.
+  - Added `BaseLikelihood` as a base class for all likelihood models to `darts.utils.likelihood_models.base`.
+  - Added `RegressionModel` likelihoods to `darts.utils.likelihood_models.sklearn`.
+  - Moved `TorchForecastingModel` likelihoods to `darts.utils.likelihood_models.torch`. They can still be imported through `darts.utils.likelihood_models`.
+  - Removed `darts.models.forecasting.regression_model._LikelihoodMixin`. Use dedicated Likelihood models instead.
+
+
 ## [0.34.0](https://github.com/unit8co/darts/tree/0.34.0) (2025-03-09)
 
 ### For users of the library:

darts/metrics/metrics.py

@@ -17,13 +17,15 @@
 from darts import TimeSeries
 from darts.dataprocessing import dtw
 from darts.logging import get_logger, raise_log
+from darts.utils.likelihood_models.base import (
+    likelihood_component_names,
+    quantile_names,
+)
 from darts.utils.ts_utils import SeriesType, get_series_seq_type, series2seq
 from darts.utils.utils import (
     _build_tqdm_iterator,
     _parallel_apply,
-    likelihood_component_names,
     n_steps_between,
-    quantile_names,
 )
 
 logger = get_logger(__name__)

darts/models/forecasting/baselines.py

@@ -414,12 +414,8 @@ def _target_average(self, prediction: TimeSeries, series: TimeSeries) -> TimeSer
 
     def _params_average(self, prediction: TimeSeries, series: TimeSeries) -> TimeSeries:
         """Average across the components after grouping by likelihood parameter, rename components"""
-        # str or torch Likelihood
-        likelihood = getattr(self.forecasting_models[0], "likelihood")
-        if isinstance(likelihood, str):
-            likelihood_n_params = self.forecasting_models[0].num_parameters
-        else:  # Likelihood
-            likelihood_n_params = likelihood.num_parameters
+        likelihood = self.forecasting_models[0].likelihood
+        likelihood_n_params = likelihood.num_parameters
         n_forecasting_models = len(self.forecasting_models)
         n_components = series.n_components
         # aggregate across predictions [model1_param0, model1_param1, ..., modeln_param0, modeln_param1]

darts/models/forecasting/block_rnn_model.py

@@ -270,7 +270,7 @@ def __init__(
             This parameter will be ignored for probabilistic models if the ``likelihood`` parameter is specified.
             Default: ``torch.nn.MSELoss()``.
         likelihood
-            One of Darts' :meth:`Likelihood <darts.utils.likelihood_models.Likelihood>` models to be used for
+            One of Darts' :meth:`Likelihood <darts.utils.likelihood_models.torch.TorchLikelihood>` models to be used for
             probabilistic forecasts. Default: ``None``.
         torch_metrics
             A torch metric or a ``MetricCollection`` used for evaluation. A full list of available metrics can be found

darts/models/forecasting/catboost_model.py

@@ -10,17 +10,24 @@
 from collections.abc import Sequence
 from typing import Optional, Union
 
-import numpy as np
 from catboost import CatBoostRegressor, Pool
 
 from darts.logging import get_logger
-from darts.models.forecasting.regression_model import RegressionModel, _LikelihoodMixin
+from darts.models.forecasting.regression_model import (
+    RegressionModel,
+    _QuantileModelContainer,
+)
 from darts.timeseries import TimeSeries
+from darts.utils.likelihood_models.sklearn import (
+    QuantileRegression,
+    _check_likelihood,
+    _get_likelihood,
+)
 
 logger = get_logger(__name__)
 
 
-class CatBoostModel(RegressionModel, _LikelihoodMixin):
+class CatBoostModel(RegressionModel):
     def __init__(
         self,
         lags: Union[int, list] = None,
@@ -29,7 +36,7 @@ def __init__(
         output_chunk_length: int = 1,
         output_chunk_shift: int = 0,
         add_encoders: Optional[dict] = None,
-        likelihood: str = None,
+        likelihood: Optional[str] = None,
         quantiles: list = None,
         random_state: Optional[int] = None,
         multi_models: Optional[bool] = True,
@@ -167,34 +174,33 @@ def encode_year(idx):
         """
         kwargs["random_state"] = random_state  # seed for tree learner
         self.kwargs = kwargs
-        self._median_idx = None
         self._model_container = None
-        self._rng = None
-        self._likelihood = likelihood
-        self.quantiles = None
-
-        self._output_chunk_length = output_chunk_length
-
-        likelihood_map = {
-            "quantile": None,
-            "poisson": "Poisson",
-            "gaussian": "RMSEWithUncertainty",
-            "RMSEWithUncertainty": "RMSEWithUncertainty",
-        }
-
-        available_likelihoods = list(likelihood_map.keys())
 
+        # parse likelihood
         if likelihood is not None:
-            self._check_likelihood(likelihood, available_likelihoods)
-            self._rng = np.random.default_rng(seed=random_state)  # seed for sampling
+            likelihood_map = {
+                "quantile": None,
+                "poisson": "Poisson",
+                "gaussian": "RMSEWithUncertainty",
+                "RMSEWithUncertainty": "RMSEWithUncertainty",
+            }
+            _check_likelihood(likelihood, list(likelihood_map.keys()))
+            if likelihood == "RMSEWithUncertainty":
+                # RMSEWithUncertainty returns mean and variance which is equivalent to gaussian
+                likelihood = "gaussian"
 
             if likelihood == "quantile":
-                self.quantiles, self._median_idx = self._prepare_quantiles(quantiles)
-                self._model_container = self._get_model_container()
-
+                self._model_container = _QuantileModelContainer()
             else:
                 self.kwargs["loss_function"] = likelihood_map[likelihood]
 
+        self._likelihood = _get_likelihood(
+            likelihood=likelihood,
+            n_outputs=output_chunk_length if multi_models else 1,
+            random_state=random_state,
+            quantiles=quantiles,
+        )
+
         # suppress writing catboost info files when user does not specifically ask to
         if "allow_writing_files" not in kwargs:
             kwargs["allow_writing_files"] = False
@@ -275,10 +281,11 @@ def fit(
         **kwargs
             Additional kwargs passed to `catboost.CatboostRegressor.fit()`
         """
-        if self.likelihood == "quantile":
+        likelihood = self.likelihood
+        if isinstance(likelihood, QuantileRegression):
             # empty model container in case of multiple calls to fit, e.g. when backtesting
             self._model_container.clear()
-            for quantile in self.quantiles:
+            for quantile in likelihood.quantiles:
                 this_quantile = str(quantile)
                 # translating to catboost argument
                 self.kwargs["loss_function"] = f"Quantile:alpha={this_quantile}"
@@ -316,42 +323,6 @@ def fit(
         )
         return self
 
-    def _predict_and_sample(
-        self,
-        x: np.ndarray,
-        num_samples: int,
-        predict_likelihood_parameters: bool,
-        **kwargs,
-    ) -> np.ndarray:
-        """Override of RegressionModel's method to allow for the probabilistic case"""
-        if self.likelihood in ["gaussian", "RMSEWithUncertainty"]:
-            return self._predict_and_sample_likelihood(
-                x, num_samples, "normal", predict_likelihood_parameters, **kwargs
-            )
-        elif self.likelihood is not None:
-            return self._predict_and_sample_likelihood(
-                x, num_samples, self.likelihood, predict_likelihood_parameters, **kwargs
-            )
-        else:
-            return super()._predict_and_sample(
-                x, num_samples, predict_likelihood_parameters, **kwargs
-            )
-
-    def _likelihood_components_names(
-        self, input_series: TimeSeries
-    ) -> Optional[list[str]]:
-        """Override of RegressionModel's method to support the gaussian/normal likelihood"""
-        if self.likelihood == "quantile":
-            return self._quantiles_generate_components_names(input_series)
-        elif self.likelihood == "poisson":
-            return self._likelihood_generate_components_names(input_series, ["lamba"])
-        elif self.likelihood in ["gaussian", "RMSEWithUncertainty"]:
-            return self._likelihood_generate_components_names(
-                input_series, ["mu", "sigma"]
-            )
-        else:
-            return None
-
     def _add_val_set_to_kwargs(
         self,
         kwargs: dict,
@@ -386,10 +357,6 @@ def _add_val_set_to_kwargs(
         kwargs[val_set_name] = val_pools
         return kwargs
 
-    @property
-    def supports_probabilistic_prediction(self) -> bool:
-        return self.likelihood is not None
-
     @property
     def supports_val_set(self) -> bool:
         return True

darts/models/forecasting/conformal_models.py

@@ -13,6 +13,12 @@
 from collections.abc import Sequence
 from typing import Any, BinaryIO, Callable, Optional, Union
 
+from darts.utils.likelihood_models.base import (
+    Likelihood,
+    LikelihoodType,
+    quantile_names,
+)
+
 try:
     from typing import Literal
 except ImportError:
@@ -46,9 +52,7 @@
 from darts.utils.utils import (
     _check_quantiles,
     generate_index,
-    likelihood_component_names,
     n_steps_between,
-    quantile_names,
     random_method,
     sample_from_quantiles,
 )
@@ -180,7 +184,10 @@ def __init__(
         self.cal_num_samples = (
             cal_num_samples if model.supports_probabilistic_prediction else 1
         )
-        self._likelihood = "quantile"
+        self._likelihood = Likelihood(
+            likelihood_type=LikelihoodType.Quantile,
+            parameter_names=quantile_names(quantiles),
+        )
         self._fit_called = True
 
     def fit(
@@ -1251,7 +1258,7 @@ def conformal_predict(idx_, pred_vals_):
                 inner_iterator = enumerate(s_hfcs[first_fc_idx:last_fc_idx:rel_stride])
 
             comp_names_out = (
-                self._cp_component_names(series_)
+                self.likelihood.component_names(series_)
                 if predict_likelihood_parameters
                 else None
             )
@@ -1426,12 +1433,6 @@ def _residuals_metric(self) -> tuple[METRIC_TYPE, Optional[dict]]:
         """Gives the "per time step" metric and optional metric kwargs used to compute residuals /
         non-conformity scores."""
 
-    def _cp_component_names(self, input_series) -> list[str]:
-        """Gives the component names for generated forecasts."""
-        return likelihood_component_names(
-            input_series.components, quantile_names(self.quantiles)
-        )
-
     def _historical_forecasts_sanity_checks(self, *args: Any, **kwargs: Any) -> None:
         super()._historical_forecasts_sanity_checks(*args, **kwargs, is_conformal=True)
 
@@ -1516,7 +1517,7 @@ def considers_static_covariates(self) -> bool:
         return self.model.considers_static_covariates
 
     @property
-    def likelihood(self) -> str:
+    def likelihood(self) -> Likelihood:
         return self._likelihood
 
 

darts/models/forecasting/dlinear.py

@@ -294,7 +294,7 @@ def __init__(
             This parameter will be ignored for probabilistic models if the ``likelihood`` parameter is specified.
             Default: ``torch.nn.MSELoss()``.
         likelihood
-            One of Darts' :meth:`Likelihood <darts.utils.likelihood_models.Likelihood>` models to be used for
+            One of Darts' :meth:`Likelihood <darts.utils.likelihood_models.torch.TorchLikelihood>` models to be used for
             probabilistic forecasts. Default: ``None``.
         torch_metrics
             A torch metric or a ``MetricCollection`` used for evaluation. A full list of available metrics can be found

darts/models/forecasting/ensemble_model.py

@@ -8,6 +8,8 @@
 from collections.abc import Sequence
 from typing import BinaryIO, Optional, Union
 
+from darts.utils.likelihood_models.base import LikelihoodType
+
 if sys.version_info >= (3, 11):
     from typing import Self
 else:
@@ -563,19 +565,13 @@ def _models_same_likelihood(self) -> bool:
         lkl_same_params = True
         tmp_quantiles = None
         for m in self.forecasting_models:
-            # regression model likelihood is a string, torch-based model likelihoods is an object
-            likelihood = getattr(m, "likelihood")
-            is_obj_lkl = not isinstance(likelihood, str)
-            lkl_simplified_name = (
-                likelihood.simplified_name() if is_obj_lkl else likelihood
-            )
-            models_likelihood.add(lkl_simplified_name)
+            likelihood = m.likelihood
+            lkl_type = likelihood.type
+            models_likelihood.add(lkl_type)
 
             # check the quantiles
-            if lkl_simplified_name == "quantile":
-                quantiles: list[str] = (
-                    likelihood.quantiles if is_obj_lkl else m.quantiles
-                )
+            if lkl_type is LikelihoodType.Quantile:
+                quantiles: list[str] = likelihood.quantiles
                 if tmp_quantiles is None:
                     tmp_quantiles = quantiles
                 elif tmp_quantiles != quantiles:

darts/models/forecasting/forecasting_model.py

@@ -27,6 +27,13 @@
 from random import sample
 from typing import Any, BinaryIO, Callable, Literal, Optional, Union
 
+from darts.utils.likelihood_models.base import (
+    Likelihood,
+    likelihood_component_names,
+    quantile_interval_names,
+    quantile_names,
+)
+
 if sys.version_info >= (3, 11):
     from typing import Self
 else:
@@ -68,9 +75,6 @@
 )
 from darts.utils.utils import (
     generate_index,
-    likelihood_component_names,
-    quantile_interval_names,
-    quantile_names,
 )
 
 logger = get_logger(__name__)
@@ -206,6 +210,11 @@ def _supports_range_index(self) -> bool:
         """
         return True
 
+    @property
+    def likelihood(self) -> Optional[Likelihood]:
+        """Returns the likelihood (if any) that the model uses for probabilistic forecasts."""
+        return None
+
     @property
     def supports_probabilistic_prediction(self) -> bool:
         """
@@ -267,7 +276,7 @@ def supports_likelihood_parameter_prediction(self) -> bool:
         """
         Whether model instance supports direct prediction of likelihood parameters
         """
-        return getattr(self, "likelihood", None) is not None
+        return self.likelihood is not None
 
     @property
     @abstractmethod