Implementation of gamma forecasting. (#77)

- Implement Gamma regression with two-parameter prediction.
- Add special functions.

Author: trivialfis

docs/source/api/metrics.rst

@@ -1,5 +1,5 @@
 Metrics
-====================
+=======
 
 .. autoclass:: legateboost.BaseMetric
     :members:
@@ -16,6 +16,9 @@ Metrics
 .. autoclass:: legateboost.GammaDevianceMetric
     :members:
 
+.. autoclass:: legateboost.GammaLLMetric
+    :members:
+
 .. autoclass:: legateboost.QuantileMetric
     :members:
 

docs/source/api/objectives.rst

@@ -13,6 +13,9 @@ Objectives
 .. autoclass:: legateboost.GammaDevianceObjective
     :members:
 
+.. autoclass:: legateboost.GammaObjective
+    :members:
+
 .. autoclass:: legateboost.QuantileObjective
     :members:
 

examples/probabalistic_regression/README.md

@@ -8,4 +8,4 @@ Running the example produces a gif animation. This animation shows the boosting
 
 The shaded area on the left hand figure shows the 95% confidence interval. The right hand example shows different quantile values.
 
-Notice that the normal distribution is symmetric about the mean, while the data is somewhat skewed. To better fit the data we can also use quantile regression, which is able to model the skewed distribution of the data.
+Notice that the normal distribution is symmetric about the mean, while the data is somewhat skewed. To better fit the data we can also use quantile regression, which is able to model the skewed distribution of the data. A Gamma distribution is another possibility that can fit well for skewed, strictly positive data.

examples/probabalistic_regression/probabilistic_regression.gif

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import os
 from pathlib import Path
 
@@ -6,9 +8,11 @@
 import pandas as pd
 import seaborn as sns
 from matplotlib import animation
-from scipy.stats import norm
+from matplotlib.figure import Figure
+from scipy.stats import gamma, norm
 from sklearn.datasets import fetch_california_housing
 
+import cunumeric as cn
 import legateboost as lb
 
 sns.set()
@@ -26,23 +30,38 @@
 n_frames = 2 if os.environ.get("CI") else 40
 
 
-def fit_normal_distribution():
+def fit_normal_distribution() -> tuple[lb.LBRegressor, list[cn.ndarray]]:
+    obj = lb.NormalObjective()
     model = lb.LBRegressor(
         verbose=True,
         init="average",
         base_models=(lb.models.Tree(max_depth=2),),
         n_estimators=n_estimators,
         learning_rate=0.1,
         random_state=rs,
-        objective="normal",
+        objective=obj,
+    )
+    return model, [model.partial_fit(X, y).predict(X_test) for _ in range(n_frames)]
+
+
+def fit_gamma_distribution() -> tuple[lb.LBRegressor, list[cn.ndarray]]:
+    obj = lb.GammaObjective()
+    model = lb.LBRegressor(
+        verbose=True,
+        init="average",
+        base_models=(lb.models.Tree(max_depth=2),),
+        n_estimators=n_estimators,
+        learning_rate=0.1,
+        random_state=rs,
+        objective=obj,
     )
     return model, [model.partial_fit(X, y).predict(X_test) for _ in range(n_frames)]
 
 
 quantiles = np.array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9])
 
 
-def fit_quantile_regression():
+def fit_quantile_regression() -> tuple[lb.LBRegressor, list]:
     model = lb.LBRegressor(
         verbose=True,
         base_models=(lb.models.Tree(max_depth=2),),
@@ -55,12 +74,15 @@ def fit_quantile_regression():
 
 
 normal_model, normal_preds = fit_normal_distribution()
+gamma_mode, gamma_preds = fit_gamma_distribution()
 quantile_model, quantile_preds = fit_quantile_regression()
 
-fig, ax = plt.subplots(1, 2, figsize=(12, 6))
+fig, ax = plt.subplots(1, 3, figsize=(12, 6))
+
 
+def animate(i: int) -> tuple[Figure]:
+    lower, upper = -0.5, 6.5
 
-def animate(i):
     fig.suptitle(
         "Distribution of House Values: Boosting iterations {}".format(
             (i + 1) * n_estimators
@@ -70,12 +92,13 @@ def animate(i):
     # Plot the normal distribution
     ax[0].cla()
     ax[0].set_title("Normal Distribution - 95% Confidence Interval")
+    norm_obj = lb.NormalObjective()
     data = pd.DataFrame(
         {
             feature_name: X_test[:, 0],
             "y": y_test,
-            "Predicted house value": normal_preds[i][:, 0],
-            "sigma": np.exp(normal_preds[i][:, 1]),
+            "Predicted house value": norm_obj.mean(normal_preds[i]),
+            "sigma": norm_obj.var(normal_preds[i]),
         }
     ).sort_values(by=feature_name)
     sns.lineplot(
@@ -89,15 +112,42 @@ def animate(i):
         0.95, loc=data["Predicted house value"], scale=data["sigma"]
     )
     ax[0].fill_between(data[feature_name], interval[0], interval[1], alpha=0.2)
-    ax[0].set_ylim(-0.5, 5.5)
+    ax[0].set_ylim(lower, upper)
 
     sns.scatterplot(
         x=feature_name, y="y", data=data, ax=ax[0], s=15, color=".2", alpha=0.2
     )
 
-    # Plot the quantile regression
+    # Plot the gamma distribution
     ax[1].cla()
-    ax[1].set_title("Quantile Regression")
+    ax[1].set_title("Gamma Distribution - 95% Confidence Interval")
+    gamma_obj = lb.GammaObjective()
+    data = pd.DataFrame(
+        {
+            feature_name: X_test[:, 0],
+            "y": y_test,
+            "Predicted house value": gamma_obj.mean(gamma_preds[i]),
+            "shape": gamma_obj.shape(gamma_preds[i]),
+            "scale": gamma_obj.scale(gamma_preds[i]),
+        }
+    ).sort_values(by=feature_name)
+    sns.lineplot(
+        x=feature_name,
+        y="Predicted house value",
+        data=data[[feature_name, "Predicted house value"]],
+        ax=ax[1],
+        errorbar=("sd", 0),
+    )
+    interval = gamma.interval(0.95, data["shape"], scale=data["scale"])
+    ax[1].fill_between(data[feature_name], interval[0], interval[1], alpha=0.2)
+    ax[1].set_ylim(lower, upper)
+    sns.scatterplot(
+        x=feature_name, y="y", data=data, ax=ax[1], s=15, color=".2", alpha=0.2
+    )
+
+    # Plot the quantile regression
+    ax[2].cla()
+    ax[2].set_title("Quantile Regression")
 
     data = {
         feature_name: X_test[:, 0],
@@ -115,14 +165,14 @@ def animate(i):
         y="Predicted house value",
         data=lines,
         style="quantile",
-        ax=ax[1],
+        ax=ax[2],
         dashes=dashes,
         errorbar=("sd", 0),
     )
-    ax[1].set_ylim(-0.5, 5.5)
+    ax[2].set_ylim(lower, upper)
 
     sns.scatterplot(
-        x=feature_name, y="y", data=data, ax=ax[1], s=15, color=".2", alpha=0.2
+        x=feature_name, y="y", data=data, ax=ax[2], s=15, color=".2", alpha=0.2
     )
 
     plt.tight_layout()

examples/probabalistic_regression/probabilistic_regression.py

@@ -1,25 +1,23 @@
-from .legateboost import LBRegressor, LBClassifier
+from .legateboost import LBClassifier, LBRegressor
 from .metrics import (
+    BaseMetric,
+    ExponentialMetric,
+    GammaDevianceMetric,
+    LogLossMetric,
     MSEMetric,
-    NormalLLMetric,
     NormalCRPSMetric,
-    GammaDevianceMetric,
+    NormalLLMetric,
+    GammaLLMetric,
     QuantileMetric,
-    LogLossMetric,
-    ExponentialMetric,
-    BaseMetric,
 )
 from .objectives import (
+    BaseObjective,
+    ExponentialObjective,
+    GammaDevianceObjective,
+    GammaObjective,
     LogLossObjective,
-    SquaredErrorObjective,
     NormalObjective,
-    GammaDevianceObjective,
     QuantileObjective,
-    ExponentialObjective,
-    BaseObjective,
-)
-from .utils import (
-    pick_col_by_idx,
-    set_col_by_idx,
-    mod_col_by_idx,
+    SquaredErrorObjective,
 )
+from .utils import mod_col_by_idx, pick_col_by_idx, set_col_by_idx

legateboost/__init__.py

@@ -166,7 +166,7 @@ def _get_weighted_gradient(
         point summation.
         """
         # check input dimensions are consistent
-        assert y.ndim == pred.ndim == 2
+        assert y.ndim == pred.ndim == 2, (y.shape, pred.shape)
         g, h = self._objective_instance.gradient(
             y, self._objective_instance.transform(pred)
         )

legateboost/legateboost.py

@@ -1,11 +1,11 @@
 from abc import ABC, abstractmethod
-from typing import Tuple
+from typing import Dict, Tuple, Type
 
-from typing_extensions import Self
+from typing_extensions import Self, override
 
 import cunumeric as cn
 
-from .special import erf
+from .special import erf, loggamma
 from .utils import pick_col_by_idx, sample_average, set_col_by_idx
 
 
@@ -75,7 +75,7 @@ def name(self) -> str:
         return "mse"
 
 
-def check_normal(y: cn.ndarray, pred: cn.ndarray) -> Tuple[cn.ndarray, cn.ndarray]:
+def check_dist_param(y: cn.ndarray, pred: cn.ndarray) -> Tuple[cn.ndarray, cn.ndarray]:
     """Checks for normal distribution inputs."""
     if y.size * 2 != pred.size:
         raise ValueError("Expected pred to contain mean and sd for each y_i")
@@ -98,7 +98,7 @@ class NormalLLMetric(BaseMetric):
     """  # noqa: E501
 
     def metric(self, y: cn.ndarray, pred: cn.ndarray, w: cn.ndarray) -> float:
-        y, pred = check_normal(y, pred)
+        y, pred = check_dist_param(y, pred)
         w_sum = w.sum()
         if w_sum == 0:
             return 0
@@ -119,6 +119,29 @@ def name(self) -> str:
         return "normal_neg_ll"
 
 
+class GammaLLMetric(BaseMetric):
+    """The mean negative log likelihood of the labels, given parameters
+    predicted by the model."""
+
+    @override
+    def metric(self, y: cn.ndarray, pred: cn.ndarray, w: cn.ndarray) -> float:
+        y, pred = check_dist_param(y, pred)
+
+        w_sum = w.sum()
+        if w_sum == 0:
+            return 0
+
+        k = pred[:, :, 0]
+        b = pred[:, :, 1]
+        error = -(k - 1) * cn.log(y) + y / (b + 1e-6) + k * cn.log(b) + loggamma(k)
+
+        return float(sample_average(error, w))
+
+    @override
+    def name(self) -> str:
+        return "gamma_neg_ll"
+
+
 def norm_cdf(x: cn.ndarray) -> cn.ndarray:
     """CDF function for standard normal distribution."""
     return 0.5 * (1.0 + erf(x / cn.sqrt(2.0)))
@@ -140,7 +163,7 @@ class NormalCRPSMetric(BaseMetric):
     """
 
     def metric(self, y: cn.ndarray, pred: cn.ndarray, w: cn.ndarray) -> float:
-        y, pred = check_normal(y, pred)
+        y, pred = check_dist_param(y, pred)
         loc = pred[:, :, 0]
         # `NormalObjective` outputs variance instead of scale.
         scale = cn.sqrt(pred[:, :, 1])
@@ -287,11 +310,12 @@ def name(self) -> str:
         return "exp"
 
 
-metrics = {
+metrics: Dict[str, Type[BaseMetric]] = {
     "log_loss": LogLossMetric,
     "mse": MSEMetric,
     "exp": ExponentialMetric,
     "normal_neg_ll": NormalLLMetric,
+    "gamma_neg_ll": GammaLLMetric,
     "normal_crps": NormalCRPSMetric,
     "deviance_gamma": GammaDevianceMetric,
 }