Source code for esmvaltool.diag_scripts.mlr.models.gbr_sklearn

"""Gradient Boosting Regression model (using :mod:`sklearn`).

Use ``mlr_model_type: gbr_sklearn`` to use this MLR model in the recipe.

"""

import logging
import os

import numpy as np
from sklearn.ensemble import GradientBoostingRegressor
from sklearn.metrics import mean_squared_error

from esmvaltool.diag_scripts.mlr.models import MLRModel
from esmvaltool.diag_scripts.mlr.models.gbr_base import GBRModel

logger = logging.getLogger(os.path.basename(__file__))




@MLRModel.register_mlr_model("gbr_sklearn")
class SklearnGBRModel(GBRModel):
    """Gradient Boosting Regression model (:mod:`sklearn` implementation)."""

    _CLF_TYPE = GradientBoostingRegressor



    def plot_training_progress(self, filename=None):
        """Plot training progress for training and (if possible) test data.

        Parameters
        ----------
        filename : str, optional (default: 'training_progress')
            Name of the plot file.

        """
        clf = self._clf.steps[-1][1].regressor_
        train_score = clf.train_score_
        test_score = None
        if "test" in self.data:
            test_score = np.zeros((len(clf.train_score_),), dtype=np.float64)
            x_test = self._clf.transform_only(self.data["test"].x)
            y_test = self._clf.transform_target_only(self.get_y_array("test"))
            sample_weights = self._get_sample_weights("test")
            for idx, y_pred in enumerate(clf.staged_predict(x_test)):
                test_score[idx] = np.sqrt(
                    mean_squared_error(
                        y_test,
                        y_pred,
                        sample_weight=sample_weights,
                    )
                )
        self._plot_training_progress(
            train_score, test_score=test_score, filename=filename
        )