sklearn/ensemble/_hist_gradient_boosting/binning.py - main - githubmirror/scikit-learn

sklearn/ensemble/_hist_gradient_boosting/binning.py
388 строк16 KB
Arthur Lacote
PRF: Optimize binning in HGBT (#34248)
24 июн 2026, 22:28
24 июн 2026, 22:28c8f2a90
"""
This module contains the BinMapper class.

BinMapper is used for mapping a real-valued dataset into integer-valued bins.
Bin thresholds are computed with the quantiles so that each bin contains
approximately the same number of samples.
"""

# Authors: The scikit-learn developers
# SPDX-License-Identifier: BSD-3-Clause

import numpy as np
from numpy.lib.stride_tricks import sliding_window_view

from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.ensemble._hist_gradient_boosting._binning import _map_to_bins
from sklearn.ensemble._hist_gradient_boosting.common import (
    ALMOST_INF,
    X_BINNED_DTYPE,
    X_BITSET_INNER_DTYPE,
    X_DTYPE,
)
from sklearn.utils import check_array, check_random_state
from sklearn.utils._bitset import set_bitset_memoryview
from sklearn.utils._openmp_helpers import _openmp_effective_n_threads
from sklearn.utils.parallel import Parallel, delayed
from sklearn.utils.stats import _weighted_percentile_1d_sorted
from sklearn.utils.validation import check_is_fitted


def _find_binning_thresholds(col_data, max_bins, sample_weight=None):
    """Extract quantiles from a continuous feature.

    Missing values are ignored for finding the thresholds.

    Parameters
    ----------
    col_data : array-like, shape (n_samples,)
        The continuous feature to bin.
    max_bins: int
        The maximum number of bins to use for non-missing values. If for a
        given feature the number of unique values is less than ``max_bins``,
        then those unique values will be used to compute the bin thresholds,
        instead of the quantiles

    Return
    ------
    binning_thresholds : ndarray of shape(min(max_bins, n_unique_values) - 1,)
        The increasing numeric values that can be used to separate the bins.
        A given value x will be mapped into bin value i iff
        bining_thresholds[i - 1] < x <= binning_thresholds[i]
    """
    # The data will be sorted anyway to find distinct values and again in percentile,
    # so we do it here, like once and for all. Sorting also returns a contiguous array.
    if sample_weight is None:
        col_data = np.sort(col_data)
        # ignore missing values when computing bin thresholds
        idx_nan = np.searchsorted(col_data, np.nan)
        col_data = col_data[:idx_nan]
    else:
        # First, remove missing values because argsort is much slower when missing
        # values are present (which is not the case for sort).
        missing_mask = np.isnan(col_data)
        if missing_mask.any():
            col_data = col_data[~missing_mask]
            sample_weight = sample_weight[~missing_mask]

        # If 0-weighted values exist, we need to remove those
        # before calculating the distinct points.
        nnz_sw = sample_weight != 0
        col_data = col_data[nnz_sw]
        sample_weight = sample_weight[nnz_sw]

        sort_idx = np.argsort(col_data)
        col_data = col_data[sort_idx]
        sample_weight = sample_weight[sort_idx]

    # fast way for n_distinct = len(np.unique(col_data))
    distinct_mask = np.empty(len(col_data), dtype=bool)
    distinct_mask[0] = True
    distinct_mask[1:] = col_data[1:] != col_data[:-1]
    n_distincts = distinct_mask.sum()

    if n_distincts == 1:
        return np.asarray([])
    elif n_distincts <= max_bins:
        # Calculate midpoints if distinct values <= max_bins
        distinct_values = col_data[distinct_mask]
        bin_thresholds = sliding_window_view(distinct_values, 2).mean(axis=1)
    elif sample_weight is None:
        # We compute bin edges using the output of np.percentile with
        # the "averaged_inverted_cdf" interpolation method that is consistent
        # with the code for the sample_weight != None case.
        percentiles = np.linspace(0, 100, num=max_bins + 1)
        percentiles = percentiles[1:-1]
        bin_thresholds = np.percentile(
            col_data, percentiles, method="averaged_inverted_cdf"
        )
        assert bin_thresholds.shape[0] == max_bins - 1
    else:
        percentiles = np.linspace(0, 100, num=max_bins + 1)
        percentiles = percentiles[1:-1]
        bin_thresholds = _weighted_percentile_1d_sorted(
            col_data, sample_weight, percentiles
        )
        assert bin_thresholds.shape[0] == max_bins - 1
    # Remove duplicated thresholds if they exist.
    bin_thresholds = np.unique(bin_thresholds)

    # We avoid having +inf thresholds: +inf thresholds are only allowed in
    # a "split on nan" situation.
    np.clip(bin_thresholds, a_min=None, a_max=ALMOST_INF, out=bin_thresholds)
    return bin_thresholds


class _BinMapper(TransformerMixin, BaseEstimator):
    """Transformer that maps a dataset into integer-valued bins.

    For continuous features, the bins are created in a feature-wise fashion,
    using quantiles so that each bins contains approximately the same number
    of samples. For large datasets, quantiles are computed on a subset of the
    data to speed-up the binning, but the quantiles should remain stable.

    For categorical features, the raw categorical values are expected to be
    in [0, 254] (this is not validated here though) and each category
    corresponds to a bin. All categorical values must be known at
    initialization: transform() doesn't know how to bin unknown categorical
    values. Note that transform() is only used on non-training data in the
    case of early stopping.

    Features with a small number of values may be binned into less than
    ``n_bins`` bins. The last bin (at index ``n_bins - 1``) is always reserved
    for missing values.

    Parameters
    ----------
    n_bins : int, default=256
        The maximum number of bins to use (including the bin for missing
        values). Should be in [3, 256]. Non-missing values are binned on
        ``max_bins = n_bins - 1`` bins. The last bin is always reserved for
        missing values. If for a given feature the number of unique values is
        less than ``max_bins``, then those unique values will be used to
        compute the bin thresholds, instead of the quantiles. For categorical
        features indicated by ``is_categorical``, the docstring for
        ``is_categorical`` details on this procedure.
    subsample : int or None, default=2e5
        If ``n_samples > subsample``, then ``sub_samples`` samples will be
        randomly chosen to compute the quantiles. If ``None``, the whole data
        is used.
    is_categorical : ndarray of bool of shape (n_features,), default=None
        Indicates categorical features. By default, all features are
        considered continuous.
    known_categories : list of {ndarray, None} of shape (n_features,), \
            default=none
        For each categorical feature, the array indicates the set of unique
        categorical values. These should be the possible values over all the
        data, not just the training data. For continuous features, the
        corresponding entry should be None.
    random_state: int, RandomState instance or None, default=None
        Pseudo-random number generator to control the random sub-sampling.
        Pass an int for reproducible output across multiple
        function calls.
        See :term:`Glossary <random_state>`.
    n_threads : int, default=None
        Number of OpenMP threads to use. `_openmp_effective_n_threads` is called
        to determine the effective number of threads use, which takes cgroups CPU
        quotes into account. See the docstring of `_openmp_effective_n_threads`
        for details.

    Attributes
    ----------
    bin_thresholds_ : list of ndarray
        For each feature, each array indicates how to map a feature into a
        binned feature. The semantic and size depends on the nature of the
        feature:
        - for real-valued features, the array corresponds to the real-valued
          bin thresholds (the upper bound of each bin). There are ``max_bins
          - 1`` thresholds, where ``max_bins = n_bins - 1`` is the number of
          bins used for non-missing values.
        - for categorical features, the array is a map from a binned category
          value to the raw category value. The size of the array is equal to
          ``min(max_bins, category_cardinality)`` where we ignore missing
          values in the cardinality.
    n_bins_non_missing_ : ndarray, dtype=np.uint32
        For each feature, gives the number of bins actually used for
        non-missing values. For features with a lot of unique values, this is
        equal to ``n_bins - 1``.
    is_categorical_ : ndarray of shape (n_features,), dtype=np.uint8
        Indicator for categorical features.
    missing_values_bin_idx_ : np.uint8
        The index of the bin where missing values are mapped. This is a
        constant across all features. This corresponds to the last bin, and
        it is always equal to ``n_bins - 1``. Note that if ``n_bins_non_missing_``
        is less than ``n_bins - 1`` for a given feature, then there are
        empty (and unused) bins.
    """

    def __init__(
        self,
        n_bins=256,
        subsample=int(2e5),
        is_categorical=None,
        known_categories=None,
        random_state=None,
        n_threads=None,
    ):
        self.n_bins = n_bins
        self.subsample = subsample
        self.is_categorical = is_categorical
        self.known_categories = known_categories
        self.random_state = random_state
        self.n_threads = n_threads

    def fit(self, X, y=None, sample_weight=None):
        """Fit data X by computing the binning thresholds.

        The last bin is reserved for missing values, whether missing values
        are present in the data or not.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            The data to bin.
        y: None
            Ignored.

        Returns
        -------
        self : object
        """
        if not (3 <= self.n_bins <= 256):
            # min is 3: at least 2 distinct bins and a missing values bin
            raise ValueError(
                "n_bins={} should be no smaller than 3 and no larger than 256.".format(
                    self.n_bins
                )
            )

        X = check_array(X, dtype=[X_DTYPE], ensure_all_finite=False)
        max_bins = self.n_bins - 1
        rng = check_random_state(self.random_state)
        if self.subsample is not None and X.shape[0] > self.subsample:
            subsampling_probabilities = None
            if sample_weight is not None:
                subsampling_probabilities = sample_weight / np.sum(sample_weight)
            # Sampling with replacement to implement frequency semantics
            # for sample weights. Note that we need `replace=True` even when
            # `sample_weight is None` to make sure that passing no weights is
            # statistically equivalent to passing unit weights.
            subset = rng.choice(
                X.shape[0], self.subsample, p=subsampling_probabilities, replace=True
            )
            X = X.take(subset, axis=0)

            # Add a switch to replace sample weights with None
            # since sample weights were already used in subsampling
            # and should not then be propagated to _find_binning_thresholds
            sample_weight = None

        if self.is_categorical is None:
            self.is_categorical_ = np.zeros(X.shape[1], dtype=np.uint8)
        else:
            self.is_categorical_ = np.asarray(self.is_categorical, dtype=np.uint8)

        n_features = X.shape[1]
        known_categories = self.known_categories
        if known_categories is None:
            known_categories = [None] * n_features

        # validate is_categorical and known_categories parameters
        for f_idx in range(n_features):
            is_categorical = self.is_categorical_[f_idx]
            known_cats = known_categories[f_idx]
            if is_categorical and known_cats is None:
                raise ValueError(
                    f"Known categories for feature {f_idx} must be provided."
                )
            if not is_categorical and known_cats is not None:
                raise ValueError(
                    f"Feature {f_idx} isn't marked as a categorical feature, "
                    "but categories were passed."
                )

        self.missing_values_bin_idx_ = self.n_bins - 1

        self.bin_thresholds_ = [None] * n_features
        n_bins_non_missing = [None] * n_features

        non_cat_thresholds = Parallel(n_jobs=self.n_threads, backend="threading")(
            delayed(_find_binning_thresholds)(
                X[:, f_idx], max_bins, sample_weight=sample_weight
            )
            for f_idx in range(n_features)
            if not self.is_categorical_[f_idx]
        )
        non_cat_idx = 0
        for f_idx in range(n_features):
            if self.is_categorical_[f_idx]:
                # Since categories are assumed to be encoded in
                # [0, n_cats] and since n_cats <= max_bins,
                # the thresholds *are* the unique categorical values. This will
                # lead to the correct mapping in transform()
                thresholds = known_categories[f_idx]
                n_bins_non_missing[f_idx] = thresholds.shape[0]
                self.bin_thresholds_[f_idx] = thresholds
            else:
                self.bin_thresholds_[f_idx] = non_cat_thresholds[non_cat_idx]
                n_bins_non_missing[f_idx] = self.bin_thresholds_[f_idx].shape[0] + 1
                non_cat_idx += 1

        self.n_bins_non_missing_ = np.array(n_bins_non_missing, dtype=np.uint32)
        return self

    def transform(self, X):
        """Bin data X.

        Missing values will be mapped to the last bin.

        For categorical features, the mapping will be incorrect for unknown
        categories. Since the BinMapper is given known_categories of the
        entire training data (i.e. before the call to train_test_split() in
        case of early-stopping), this never happens.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            The data to bin.

        Returns
        -------
        X_binned : array-like of shape (n_samples, n_features)
            The binned data (fortran-aligned).
        """
        X = check_array(X, dtype=[X_DTYPE], ensure_all_finite=False)
        check_is_fitted(self)
        if X.shape[1] != self.n_bins_non_missing_.shape[0]:
            raise ValueError(
                "This estimator was fitted with {} features but {} got passed "
                "to transform()".format(self.n_bins_non_missing_.shape[0], X.shape[1])
            )

        n_threads = _openmp_effective_n_threads(self.n_threads)
        binned = np.zeros_like(X, dtype=X_BINNED_DTYPE, order="F")
        _map_to_bins(
            X,
            self.bin_thresholds_,
            self.is_categorical_,
            self.missing_values_bin_idx_,
            n_threads,
            binned,
        )
        return binned

    def make_known_categories_bitsets(self):
        """Create bitsets of known categories.

        Returns
        -------
        - known_cat_bitsets : ndarray of shape (n_categorical_features, 8)
            Array of bitsets of known categories, for each categorical feature.
        - f_idx_map : ndarray of shape (n_features,)
            Map from original feature index to the corresponding index in the
            known_cat_bitsets array.
        """

        categorical_features_indices = np.flatnonzero(self.is_categorical_)

        n_features = self.is_categorical_.size
        n_categorical_features = categorical_features_indices.size

        f_idx_map = np.zeros(n_features, dtype=np.uint32)
        f_idx_map[categorical_features_indices] = np.arange(
            n_categorical_features, dtype=np.uint32
        )

        known_categories = self.bin_thresholds_

        known_cat_bitsets = np.zeros(
            (n_categorical_features, 8), dtype=X_BITSET_INNER_DTYPE
        )

        # TODO: complexity is O(n_categorical_features * 255). Maybe this is
        # worth cythonizing
        for mapped_f_idx, f_idx in enumerate(categorical_features_indices):
            for raw_cat_val in known_categories[f_idx]:
                set_bitset_memoryview(known_cat_bitsets[mapped_f_idx], raw_cat_val)

        return known_cat_bitsets, f_idx_map