Add NegativeBinomial support to XGBoostLSS

skpro/regression/tests/test_xgboostlss.py

@@ -1,8 +1,10 @@
 """Tests for the XGBoostLSS regressor."""
 
+import numpy as np
 import pandas as pd
 import pytest
 
+from skpro.distributions import NegativeBinomial
 from skpro.regression.xgboostlss import XGBoostLSS
 from skpro.tests.test_switch import run_test_for_class
 
@@ -68,3 +70,57 @@ def test_xgboostlss_param_handling(params, expected_xgb_params, should_error):
 
         for key, value in expected_xgb_params.items():
             assert reg.xgb_params_.get(key) == value
+
+
+@pytest.mark.skipif(
+    not run_test_for_class(XGBoostLSS),
+    reason="run test only if softdeps are present and incrementally (if requested)",
+)
+def test_xgboostlss_negative_binomial_param_conversion():
+    """Test conversion of xgboostlss NegativeBinomial parameters to skpro."""
+    import torch
+
+    reg = XGBoostLSS(dist="NegativeBinomial")
+
+    # mean and variance must match the torch distribution fitted by xgboostlss
+    params = pd.DataFrame(
+        {"total_count": [0.05, 0.8, 3.0, 25.0], "probs": [0.001, 0.2, 0.5, 0.999]}
+    )
+    skpro_dist = NegativeBinomial(**reg._get_skpro_val_dict("NegativeBinomial", params))
+    torch_dist = torch.distributions.NegativeBinomial(
+        total_count=torch.tensor(params["total_count"].to_numpy()),
+        probs=torch.tensor(params["probs"].to_numpy()),
+    )
+    np.testing.assert_allclose(skpro_dist.mean().values.ravel(), torch_dist.mean)
+    np.testing.assert_allclose(skpro_dist.var().values.ravel(), torch_dist.variance)
+
+    # saturated float32 probabilities must still give a finite mean
+    params = pd.DataFrame({"total_count": [2.0], "probs": [1.0]}, dtype="float32")
+    skpro_dist = NegativeBinomial(**reg._get_skpro_val_dict("NegativeBinomial", params))
+    assert np.isfinite(skpro_dist.mean().values).all()
+
+
+@pytest.mark.skipif(
+    not run_test_for_class(XGBoostLSS),
+    reason="run test only if softdeps are present and incrementally (if requested)",
+)
+def test_xgboostlss_negative_binomial():
+    """Test XGBoostLSS with NegativeBinomial distribution on count data."""
+    rng = np.random.default_rng(0)
+    X = pd.DataFrame(rng.normal(size=(100, 3)))
+    y = pd.DataFrame(rng.negative_binomial(3, 3 / (3 + np.exp(1.5 + X[0]))))
+
+    reg = XGBoostLSS(
+        dist="NegativeBinomial", response_fn="softplus", n_trials=0, n_cpu=1
+    )
+    reg.fit(X, y)
+
+    # response_fn applies to total_count, probs always uses sigmoid
+    response_fns = reg.xgblss_.dist.param_dict
+    assert response_fns["total_count"].__name__ == "softplus_fn"
+    assert response_fns["probs"].__name__ == "sigmoid_fn"
+
+    y_pred = reg.predict_proba(X)
+    assert isinstance(y_pred, NegativeBinomial)
+    quantiles = y_pred.ppf(0.84).values
+    np.testing.assert_array_equal(quantiles, np.round(quantiles))

skpro/regression/xgboostlss.py

@@ -2,6 +2,7 @@
 
 import warnings
 
+import numpy as np
 from skbase.utils.dependencies import _check_soft_dependencies
 
 from skpro.regression.base import BaseProbaRegressor
@@ -23,6 +24,7 @@ class XGBoostLSS(BaseProbaRegressor):
         * "Gamma": Gamma distribution.
         * "Laplace": Laplace distribution.
         * "LogNormal": LogNormal distribution.
+        * "NegativeBinomial": Negative binomial distribution for count data.
         * "TDistribution": Student's T distribution.
         * "Weibull": Weibull distribution.
         * "Beta": Beta distribution.
@@ -34,12 +36,15 @@ class XGBoostLSS(BaseProbaRegressor):
     response_fn: str, optional, default="exp"
         Response function for transforming the distributional parameters to the
         support of the distribution. Options are "exp" (exponential) or
-        "softplus" (softplus).
+        "softplus" (softplus). For ``dist="NegativeBinomial"``, this transforms
+        ``total_count`` and also accepts "relu" (rectified linear unit);
+        ``probs`` always uses "sigmoid".
 
     loss_fn: str, optional, default="nll"
         Loss function used in tuning and fitting.
         Options are "nll" (negative log-likelihood)
         or "crps" (continuous ranked probability score).
+        For ``dist="NegativeBinomial"``, only "nll" is supported.
 
     initialize: bool, optional, default=True
         Whether to initialize the distribution parameters.
@@ -95,6 +100,18 @@ class XGBoostLSS(BaseProbaRegressor):
         Keyword arguments of xgboost parameters to pass to the model.
         Used only if n_trials=0 (i.e., no hyperparameter optimization).
         See https://xgboost.readthedocs.io/en/stable/parameter.html for valid params.
+
+    Examples
+    --------
+    >>> from skpro.regression.xgboostlss import XGBoostLSS  # doctest: +SKIP
+    >>> from sklearn.datasets import load_diabetes
+    >>> from sklearn.model_selection import train_test_split
+    >>> X, y = load_diabetes(return_X_y=True, as_frame=True)
+    >>> X_train, X_test, y_train, y_test = train_test_split(X, y)
+    >>> reg_proba = XGBoostLSS(n_trials=0, num_boost_round=5)  # doctest: +SKIP
+    >>> reg_proba.fit(X_train, y_train)  # doctest: +SKIP
+    XGBoostLSS(...)
+    >>> y_pred = reg_proba.predict_proba(X_test)  # doctest: +SKIP
     """
 
     _tags = {
@@ -318,6 +335,16 @@ def _get_skpro_val_dict(self, distr, df):
         df : pd.DataFrame
             DataFrame of parameters as returned by predict, in xgboostlss.
         """
+        if distr == "NegativeBinomial":
+            total_count = df.loc[:, ["total_count"]].values
+            # Convert to float64 so the upper clipping bound stays below 1,
+            # including when xgboostlss returns float32 probabilities.
+            probs = df.loc[:, ["probs"]].to_numpy(dtype=float)
+            probs = np.clip(probs, 1e-8, 1 - 1e-8)
+            # PyTorch uses mean = total_count * probs / (1 - probs), whereas
+            # skpro uses mean = mu and variance = mu + mu**2 / alpha.
+            return {"mu": total_count * probs / (1 - probs), "alpha": total_count}
+
         name_map = {
             "Normal": {"mu": "loc", "sigma": "scale"},
             "Gamma": {"alpha": "concentration", "beta": "rate"},
@@ -391,6 +418,10 @@ def _fit(self, X, y):
             "initialize": self.initialize,
         }
 
+        if self.dist == "NegativeBinomial":
+            distr_params["response_fn_total_count"] = distr_params.pop("response_fn")
+            distr_params["response_fn_probs"] = "sigmoid"
+
         # initialize parameter was introduced in xgboostlss v0.6.1
         if _check_soft_dependencies("xgboostlss<0.6.1", severity="none"):
             distr_params.pop("initialize")
@@ -573,6 +604,13 @@ def get_test_params(cls, parameter_set="default"):
             "max_depth": 3,
         }
 
+        params9 = {
+            "dist": "NegativeBinomial",
+            "n_trials": 0,
+            "num_boost_round": 5,
+            "n_cpu": 1,
+        }
+
         return [
             params0,
             params1,
@@ -583,4 +621,5 @@ def get_test_params(cls, parameter_set="default"):
             params6,
             params7,
             params8,
+            params9,
         ]