A demo for multi-output regression

The demo is adopted from scikit-learn:

https://scikit-learn.org/stable/auto_examples/ensemble/plot_random_forest_regression_multioutput.html#sphx-glr-auto-examples-ensemble-plot-random-forest-regression-multioutput-py

See Multiple Outputs for more information.

Note

The feature is experimental. For the multi_output_tree strategy, many features are missing.

See also

A demo for multi-output regression using reduced gradient

import argparse
from typing import Dict, List, Optional, Tuple

import matplotlib
import numpy as np
from matplotlib import pyplot as plt

import xgboost as xgb


def plot_predt(
    y: np.ndarray, y_predt: np.ndarray, name: str, ax: matplotlib.axes.Axes
) -> None:
    s = 25
    ax.scatter(y[:, 0], y[:, 1], c="navy", s=s, edgecolor="black", label=name)
    ax.scatter(y_predt[:, 0], y_predt[:, 1], c="cornflowerblue", s=s, edgecolor="black")
    ax.legend()


def gen_circle() -> Tuple[np.ndarray, np.ndarray]:
    "Generate a sample dataset that y is a 2 dim circle."
    rng = np.random.RandomState(1994)
    X = np.sort(200 * rng.rand(100, 1) - 100, axis=0)
    y = np.array([np.pi * np.sin(X).ravel(), np.pi * np.cos(X).ravel()]).T
    y[::5, :] += 0.5 - rng.rand(20, 2)
    y = y - y.min()
    y = y / y.max()
    return X, y


def rmse_model(strategy: str, ax: Optional[matplotlib.axes.Axes]) -> None:
    """Draw a circle with 2-dim coordinate as target variables."""
    X, y = gen_circle()
    # Train a regressor on it
    reg = xgb.XGBRegressor(
        tree_method="hist",
        n_estimators=128,
        n_jobs=16,
        max_depth=8,
        multi_strategy=strategy,
        subsample=0.6,
    )
    reg.fit(X, y, eval_set=[(X, y)])

    y_predt = reg.predict(X)
    if ax:
        plot_predt(y, y_predt, f"RMSE-{strategy}", ax)


def custom_rmse_model(strategy: str, ax: Optional[matplotlib.axes.Axes]) -> None:
    """Train using Python implementation of Squared Error."""

    def gradient(predt: np.ndarray, dtrain: xgb.DMatrix) -> np.ndarray:
        """Compute the gradient squared error."""
        y = dtrain.get_label().reshape(predt.shape)
        return predt - y

    def hessian(predt: np.ndarray, dtrain: xgb.DMatrix) -> np.ndarray:
        """Compute the hessian for squared error."""
        return np.ones(predt.shape)

    def squared_log(
        predt: np.ndarray, dtrain: xgb.DMatrix
    ) -> Tuple[np.ndarray, np.ndarray]:
        grad = gradient(predt, dtrain)
        hess = hessian(predt, dtrain)
        # both numpy.ndarray and cupy.ndarray works.
        return grad, hess

    def rmse(predt: np.ndarray, dtrain: xgb.DMatrix) -> Tuple[str, float]:
        y = dtrain.get_label().reshape(predt.shape)
        v = np.sqrt(np.mean(np.power(y - predt, 2)))
        return "PyRMSE", v

    X, y = gen_circle()
    Xy = xgb.DMatrix(X, y)
    results: Dict[str, Dict[str, List[float]]] = {}
    # Make sure the `num_target` is passed to XGBoost when custom objective is used.
    # When builtin objective is used, XGBoost can figure out the number of targets
    # automatically.
    booster = xgb.train(
        {
            "tree_method": "hist",
            "num_target": y.shape[1],
            "multi_strategy": strategy,
        },
        dtrain=Xy,
        num_boost_round=128,
        obj=squared_log,
        evals=[(Xy, "Train")],
        evals_result=results,
        custom_metric=rmse,
    )

    y_predt = booster.inplace_predict(X)
    if ax:
        plot_predt(y, y_predt, f"PyRMSE-{strategy}", ax)

    np.testing.assert_allclose(
        results["Train"]["rmse"], results["Train"]["PyRMSE"], rtol=1e-2
    )


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--plot", choices=[0, 1], type=int, default=1)
    args = parser.parse_args()
    if args.plot == 1:
        _, axs = plt.subplots(2, 2)
    else:
        axs = np.full(shape=(2, 2), fill_value=None)
    assert isinstance(axs, np.ndarray)

    # Train with builtin RMSE objective
    # - One model per output.
    rmse_model("one_output_per_tree", axs[0, 0])
    # - One model for all outputs, this is still working in progress, many features are
    # missing.
    rmse_model("multi_output_tree", axs[0, 1])

    # Train with custom objective.
    # - One model per output.
    custom_rmse_model("one_output_per_tree", axs[1, 0])
    # - One model for all outputs, this is still working in progress, many features are
    # missing.
    custom_rmse_model("multi_output_tree", axs[1, 1])
    if args.plot == 1:
        plt.show()