PRF: Optimize binning in HGBT (scikit-learn#34248)

Author: cakedev0

doc/whats_new/upcoming_changes/sklearn.ensemble/34248.efficiency.rst

@@ -0,0 +1,6 @@
+- Improved the performance of binning in
+  :class:`ensemble.HistGradientBoostingClassifier` and
+  :class:`ensemble.HistGradientBoostingRegressor`. Fitting with sample weights
+  could be slow for datasets with fewer than 200,000 samples; this is now fixed.
+  Binning without sample weights was also optimized.
+  By :user:`Arthur Lacote <cakedev0>`.

sklearn/ensemble/_hist_gradient_boosting/binning.py

@@ -24,7 +24,7 @@
 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
+from sklearn.utils.stats import _weighted_percentile_1d_sorted
 from sklearn.utils.validation import check_is_fitted
 
 
@@ -50,35 +50,42 @@ def _find_binning_thresholds(col_data, max_bins, sample_weight=None):
         A given value x will be mapped into bin value i iff
         bining_thresholds[i - 1] < x <= binning_thresholds[i]
     """
-    # ignore missing values when computing bin thresholds
-    missing_mask = np.isnan(col_data)
-    any_missing = missing_mask.any()
-    if any_missing:
-        col_data = col_data[~missing_mask]
-
-    # If sample_weight is not None and 0-weighted values exist, we need to
-    # remove those before calculating the distinct points.
-    if sample_weight is not None:
-        if any_missing:
+    # 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]
 
-    # The data will be sorted anyway in np.unique and again in percentile, so we do it
-    # here. Sorting also returns a contiguous array.
-    sort_idx = np.argsort(col_data)
-    col_data = col_data[sort_idx]
-    if sample_weight is not None:
+        sort_idx = np.argsort(col_data)
+        col_data = col_data[sort_idx]
         sample_weight = sample_weight[sort_idx]
 
-    distinct_values = np.unique(col_data).astype(X_DTYPE)
+    # 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 len(distinct_values) == 1:
+    if n_distincts == 1:
         return np.asarray([])
-
-    if len(distinct_values) <= max_bins:
+    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
@@ -93,17 +100,12 @@ def _find_binning_thresholds(col_data, max_bins, sample_weight=None):
     else:
         percentiles = np.linspace(0, 100, num=max_bins + 1)
         percentiles = percentiles[1:-1]
-        bin_thresholds = np.array(
-            [
-                _weighted_percentile(col_data, sample_weight, percentile, average=True)
-                for percentile in percentiles
-            ]
+        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.
-    unique_bin_values = np.unique(bin_thresholds)
-    if unique_bin_values.shape[0] != bin_thresholds.shape[0]:
-        bin_thresholds = unique_bin_values
+    bin_thresholds = np.unique(bin_thresholds)
 
     # We avoid having +inf thresholds: +inf thresholds are only allowed in
     # a "split on nan" situation.

sklearn/ensemble/_hist_gradient_boosting/tests/test_binning.py

@@ -220,17 +220,27 @@ def test_bin_mapper_repeated_values_invariance(n_distinct):
 
 
 @pytest.mark.parametrize("n_bins", [50, 250])
-def test_binmapper_weighted_vs_repeated_equivalence(global_random_seed, n_bins):
+@pytest.mark.parametrize("with_missing_values", [False, True])
+def test_binmapper_weighted_vs_repeated_equivalence(
+    global_random_seed, n_bins, with_missing_values
+):
     rng = np.random.RandomState(global_random_seed)
 
     n_samples = 200
     X = rng.randn(n_samples, 3)
+    if with_missing_values:
+        missing_mask = rng.rand(*X.shape) < 0.1
+        X[missing_mask] = np.nan
+
     sw = rng.randint(0, 5, size=n_samples)
     X_repeated = np.repeat(X, sw, axis=0)
 
     est_weighted = _BinMapper(n_bins=n_bins).fit(X, sample_weight=sw)
     est_repeated = _BinMapper(n_bins=n_bins).fit(X_repeated, sample_weight=None)
-    assert_allclose(est_weighted.bin_thresholds_, est_repeated.bin_thresholds_)
+    for thresholds_weighted, thresholds_repeated in zip(
+        est_weighted.bin_thresholds_, est_repeated.bin_thresholds_
+    ):
+        assert_allclose(thresholds_weighted, thresholds_repeated)
 
     X_trans_weighted = est_weighted.transform(X)
     X_trans_repeated = est_repeated.transform(X)

sklearn/utils/stats.py

@@ -1,6 +1,8 @@
 # Authors: The scikit-learn developers
 # SPDX-License-Identifier: BSD-3-Clause
 
+import numpy as np
+
 from sklearn.utils._array_api import (
     _find_matching_floating_dtype,
     get_namespace_and_device,
@@ -214,3 +216,30 @@ def _weighted_percentile(
         result = result[..., 0]
 
     return result[0, ...] if n_dim == 1 else result
+
+
+def _weighted_percentile_1d_sorted(array, sample_weight, percentile_rank):
+    """Compute weighted percentiles for sorted 1D data and percentile ranks.
+
+    This implements the "averaged_inverted_cdf" method as computed by
+    `_weighted_percentile(..., average=True)`, but it expects:
+    - `array` is sorted
+    - `array` has no NaN values
+    - `sample_weight` has strictly positive values
+    - `percentile_rank` is a 1D array
+    """
+    weight_cdf = np.cumsum(sample_weight, dtype=array.dtype)
+    adjusted_percentile_rank = percentile_rank / 100 * weight_cdf[-1]
+
+    percentile_indices = np.searchsorted(weight_cdf, adjusted_percentile_rank)
+    percentile_indices = np.clip(percentile_indices, 0, array.shape[0] - 1)
+
+    fraction_above = weight_cdf[percentile_indices] - adjusted_percentile_rank
+    is_fraction_above = fraction_above > np.finfo(array.dtype).eps
+    percentile_plus_one_indices = np.clip(percentile_indices + 1, 0, array.shape[0] - 1)
+
+    return np.where(
+        is_fraction_above,
+        array[percentile_indices],
+        (array[percentile_indices] + array[percentile_plus_one_indices]) / 2,
+    )

sklearn/utils/tests/test_stats.py

@@ -12,7 +12,7 @@
 )
 from sklearn.utils.estimator_checks import _array_api_for_tests
 from sklearn.utils.fixes import np_version, parse_version
-from sklearn.utils.stats import _weighted_percentile
+from sklearn.utils.stats import _weighted_percentile, _weighted_percentile_1d_sorted
 
 
 @pytest.mark.parametrize("average", [True, False])
@@ -183,6 +183,27 @@ def test_weighted_percentile_constant_multiplier(
     assert percentile == approx(percentile_multiplier)
 
 
+@pytest.mark.parametrize("percentile_rank", [np.array([0]), np.array([20, 50, 100])])
+def test_weighted_percentile_1d_sorted_matches_weighted_percentile(
+    global_random_seed, percentile_rank
+):
+    """Check sorted 1D helper against `_weighted_percentile`."""
+    rng = np.random.RandomState(global_random_seed)
+    for x in [
+        np.sort(rng.uniform(size=100)),
+        np.array([0, 1, 1, 2], dtype=float),
+        np.repeat([0, 1, 2], rng.choice(10, size=3)).astype(np.float32),
+    ]:
+        sample_weight = rng.uniform(low=0.1, high=10, size=x.shape[0])
+
+        percentile = _weighted_percentile_1d_sorted(x, sample_weight, percentile_rank)
+        expected_percentile = _weighted_percentile(
+            x, sample_weight, percentile_rank, average=True
+        )
+
+        assert_allclose(percentile, expected_percentile)
+
+
 @pytest.mark.parametrize("percentile_rank", [50, [20, 35, 50]])
 @pytest.mark.parametrize("average", [True, False])
 def test_weighted_percentile_2d(global_random_seed, percentile_rank, average):