Fixes sample weight with sklearn 1.9

- Stores sample_weight before sorting
- Updates _get_n_samples_bootstrap and _generate_sample_indices. - Adds unit tests.

Author: reidjohnson

quantile_forest/_quantile_forest.py

@@ -174,12 +174,16 @@ def fit(self, X, y, sample_weight=None, sparse_pickle=False):
         for i in range(y.shape[1]):
             y_sorted[:, i] = y[sorter[:, i], i]
 
+        self.sample_weight_ = sample_weight
+
         if sample_weight is not None:
-            sample_weight = np.asarray(sample_weight)[sorter]
+            sample_weight_sorted = np.asarray(sample_weight)[sorter]
+        else:
+            sample_weight_sorted = None
 
         # Get map of tree leaf nodes to training indices.
         y_train_leaves = self._get_y_train_leaves(
-            X, y_sorted, sorter=sorter, sample_weight=sample_weight
+            X, y_sorted, sorter=sorter, sample_weight=sample_weight_sorted
         )
 
         # Get map of tree leaf nodes to target value bounds (for monotonicity constraints).
@@ -352,17 +356,28 @@ def _get_y_train_leaves(self, X, y, sorter=None, sample_weight=None):
             X_leaves = self.apply(X)
 
         if self.bootstrap:
-            n_samples_bootstrap = _get_n_samples_bootstrap(n_samples, self.max_samples)
+            params = {
+                "n_samples": n_samples,
+                "max_samples": self.max_samples,
+            }
+            if sklearn_version >= parse_version("1.9.0"):
+                params["sample_weight"] = self.sample_weight_
+            n_samples_bootstrap = _get_n_samples_bootstrap(**params)
 
         shape = (n_samples if not self.bootstrap else n_samples_bootstrap, self.n_estimators)
         bootstrap_indices = np.empty(shape, dtype=np.int64)
         X_leaves_bootstrap = np.empty(shape, dtype=np.int64)
         for i, estimator in enumerate(self.estimators_):
             # Get bootstrap indices.
             if self.bootstrap:
-                bootstrap_indices[:, i] = _generate_sample_indices(
-                    estimator.random_state, n_samples, n_samples_bootstrap
-                )
+                params = {
+                    "random_state": estimator.random_state,
+                    "n_samples": n_samples,
+                    "n_samples_bootstrap": n_samples_bootstrap,
+                }
+                if sklearn_version >= parse_version("1.9.0"):
+                    params["sample_weight"] = self.sample_weight_
+                bootstrap_indices[:, i] = _generate_sample_indices(**params)
             else:
                 bootstrap_indices[:, i] = np.arange(n_samples)
 
@@ -534,7 +549,9 @@ def _oob_samples(self, X, indices=None, duplicates=None):
         for i, estimator in enumerate(self.estimators_):
             # Get the indices excluded from the bootstrapping process.
             if self.unsampled_indices_ is None:
-                unsampled_indices = self._get_unsampled_indices(estimator, duplicates=duplicates)
+                unsampled_indices = self._get_unsampled_indices(
+                    estimator, duplicates=duplicates, sample_weight=self.sample_weight_
+                )
             else:
                 # Avoid generating unsampled indices if they are precomputed.
                 unsampled_indices = np.asarray(self.unsampled_indices_[i])
@@ -564,7 +581,7 @@ def _oob_samples(self, X, indices=None, duplicates=None):
 
         return X_leaves, X_indices
 
-    def _get_unsampled_indices(self, estimator, duplicates=None):
+    def _get_unsampled_indices(self, estimator, duplicates=None, sample_weight=None):
         """Get the unsampled indices for a base estimator.
 
         Parameters
@@ -575,6 +592,13 @@ def _get_unsampled_indices(self, estimator, duplicates=None):
         duplicates : list of lists, default=None
             List of sets of functionally identical indices.
 
+        sample_weight : array-like of shape (n_samples, n_outputs), \
+                default=None
+            Sample weights. If None, then samples are equally weighted. Splits
+            that would create child nodes with net zero or negative weight are
+            ignored while searching for a split in each node. For each output,
+            the ordering of the weights correspond to the sorted samples.
+
         Returns
         -------
         unsampled_indices : array of shape (n_unsampled,)
@@ -584,9 +608,12 @@ def _get_unsampled_indices(self, estimator, duplicates=None):
             warn("Unsampled indices only exist if bootstrap=True.")
             return np.array([])
         n_train_samples = self.n_train_samples_
-        n_samples_bootstrap = _get_n_samples_bootstrap(n_train_samples, self.max_samples)
+        sample_weight = self.sample_weight_
+        params = {} if sklearn_version < parse_version("1.9") else {"sample_weight": sample_weight}
+        n_samples_bootstrap = _get_n_samples_bootstrap(n_train_samples, self.max_samples, **params)
+        params = {} if sklearn_version < parse_version("1.9") else {"sample_weight": sample_weight}
         sample_indices = _generate_sample_indices(
-            estimator.random_state, n_train_samples, n_samples_bootstrap
+            estimator.random_state, n_train_samples, n_samples_bootstrap, **params
         )
         unsampled_indices = generate_unsampled_indices(
             sample_indices,

quantile_forest/tests/test_quantile_forest.py

@@ -7,6 +7,7 @@
 
 import numpy as np
 import pytest
+import sklearn
 from scipy.stats import percentileofscore
 from sklearn import datasets
 from sklearn.ensemble import ExtraTreesRegressor, RandomForestRegressor
@@ -19,6 +20,7 @@
     assert_array_almost_equal,
     assert_array_equal,
 )
+from sklearn.utils.fixes import parse_version
 from sklearn.utils.validation import check_is_fitted
 
 from quantile_forest import ExtraTreesQuantileRegressor, RandomForestQuantileRegressor
@@ -30,6 +32,8 @@
     calc_weighted_quantile,
 )
 
+sklearn_version = parse_version(sklearn.__version__)
+
 np.random.seed(0)
 
 FOREST_REGRESSORS: dict[str, Any] = {
@@ -765,8 +769,10 @@ def check_proximity_counts(name):
     n_samples = len(X)
     proximities = est.proximity_counts(X)
     X_leaves = est.apply(X)
+    params = (n_samples, n_samples)
+    params = params if sklearn_version < parse_version("1.9") else params + (None,)
     bootstrap_indices = np.array(
-        [_generate_sample_indices(e.random_state, n_samples, n_samples) for e in est.estimators_]
+        [_generate_sample_indices(e.random_state, *params) for e in est.estimators_]
     )
     for train_idx, train_idx_prox in enumerate(proximities):
         for proximity_idx, proximity_count in train_idx_prox:
@@ -1218,8 +1224,10 @@ def check_proximity_counts_oob(name):
     # Check that OOB proximity counts match OOB bootstrap counts.
     X_leaves = est.apply(X)
     n_samples = len(X)
+    params = (n_samples, n_samples)
+    params = params if sklearn_version < parse_version("1.9") else params + (None,)
     bootstrap_indices = np.array(
-        [_generate_sample_indices(e.random_state, n_samples, n_samples) for e in est.estimators_]
+        [_generate_sample_indices(e.random_state, *params) for e in est.estimators_]
     )
     for train_idx, train_idx_prox in enumerate(proximities):
         for proximity_idx, proximity_count in train_idx_prox:
@@ -1384,6 +1392,115 @@ def test_serialization(name, sparse_pickle, monotonic_cst, multi_target):
     check_serialization(name, sparse_pickle, monotonic_cst, multi_target)
 
 
+@pytest.mark.skipif(
+    sklearn_version < parse_version("1.9"),
+    reason="sample_weight in bootstrap functions requires scikit-learn >= 1.9",
+)
+@pytest.mark.parametrize("name", FOREST_REGRESSORS)
+def test_sample_weight_stored_on_fit(name):
+    """Check that sample_weight is stored during fit."""
+    X, y = datasets.make_regression(n_samples=100, n_features=4, random_state=0)
+    sample_weight = np.random.default_rng(0).uniform(0.1, 2.0, size=len(y))
+
+    ForestRegressor = FOREST_REGRESSORS[name]
+
+    est = ForestRegressor(n_estimators=5, random_state=0)
+    est.fit(X, y, sample_weight=sample_weight)
+    assert est.sample_weight_ is not None
+    assert est.sample_weight_.shape[0] == len(y)
+
+    est_no_weight = ForestRegressor(n_estimators=5, random_state=0)
+    est_no_weight.fit(X, y)
+    assert est_no_weight.sample_weight_ is None
+
+
+@pytest.mark.skipif(
+    sklearn_version < parse_version("1.9"),
+    reason="sample_weight in bootstrap functions requires scikit-learn >= 1.9",
+)
+@pytest.mark.parametrize("name", FOREST_REGRESSORS)
+def test_sample_weight_predict(name):
+    """Check that predictions work with sample_weight on sklearn 1.9+."""
+    X, y = datasets.make_regression(n_samples=200, n_features=4, random_state=0)
+    sample_weight = np.random.default_rng(0).uniform(0.1, 2.0, size=len(y))
+
+    ForestRegressor = FOREST_REGRESSORS[name]
+
+    est = ForestRegressor(n_estimators=10, bootstrap=True, random_state=0)
+    est.fit(X, y, sample_weight=sample_weight)
+
+    y_pred = est.predict(X, quantiles=[0.25, 0.5, 0.75])
+    assert y_pred.shape == (len(X), 3)
+    assert not np.isnan(y_pred).any()
+
+    # Predictions with sample_weight should differ from uniform weights.
+    est_uniform = ForestRegressor(n_estimators=10, bootstrap=True, random_state=0)
+    est_uniform.fit(X, y)
+    y_pred_uniform = est_uniform.predict(X, quantiles=[0.25, 0.5, 0.75])
+    assert not np.allclose(y_pred, y_pred_uniform)
+
+
+@pytest.mark.skipif(
+    sklearn_version < parse_version("1.9"),
+    reason="sample_weight in bootstrap functions requires scikit-learn >= 1.9",
+)
+@pytest.mark.parametrize("name", FOREST_REGRESSORS)
+def test_sample_weight_oob(name):
+    """Check that OOB predictions work with sample_weight on sklearn 1.9+."""
+    X, y = datasets.make_regression(n_samples=500, n_features=8, noise=0.1, random_state=0)
+    sample_weight = np.random.default_rng(0).uniform(0.1, 2.0, size=len(y))
+
+    ForestRegressor = FOREST_REGRESSORS[name]
+
+    est = ForestRegressor(n_estimators=20, bootstrap=True, oob_score=True, random_state=0)
+    est.fit(X, y, sample_weight=sample_weight)
+
+    y_pred = est.predict(X, quantiles=[0.2, 0.5, 0.8], oob_score=True)
+    assert y_pred.shape == (len(X), 3)
+
+
+@pytest.mark.skipif(
+    sklearn_version < parse_version("1.9"),
+    reason="sample_weight in bootstrap functions requires scikit-learn >= 1.9",
+)
+@pytest.mark.parametrize("name", FOREST_REGRESSORS)
+def test_sample_weight_unsampled_indices(name):
+    """Check that _get_unsampled_indices uses sample_weight on sklearn 1.9+."""
+    X, y = datasets.make_regression(n_samples=200, n_features=4, random_state=0)
+    sample_weight = np.random.default_rng(0).uniform(0.1, 2.0, size=len(y))
+
+    ForestRegressor = FOREST_REGRESSORS[name]
+
+    est = ForestRegressor(n_estimators=5, bootstrap=True, random_state=0)
+    est.fit(X, y, sample_weight=sample_weight)
+
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", UserWarning)
+        unsampled = est._get_unsampled_indices(est.estimators_[0])
+
+    assert len(unsampled) > 0
+    assert all(idx < len(y) for idx in unsampled)
+
+
+@pytest.mark.skipif(
+    sklearn_version < parse_version("1.9"),
+    reason="sample_weight in bootstrap functions requires scikit-learn >= 1.9",
+)
+@pytest.mark.parametrize("name", FOREST_REGRESSORS)
+def test_sample_weight_proximity_counts(name):
+    """Check that proximity_counts works with sample_weight on sklearn 1.9+."""
+    X, y = datasets.make_regression(n_samples=100, n_features=4, random_state=0)
+    sample_weight = np.random.default_rng(0).uniform(0.1, 2.0, size=len(y))
+
+    ForestRegressor = FOREST_REGRESSORS[name]
+
+    est = ForestRegressor(n_estimators=5, bootstrap=True, random_state=0)
+    est.fit(X, y, sample_weight=sample_weight)
+
+    proximities = est.proximity_counts(X)
+    assert len(proximities) == len(X)
+
+
 def test_calc_quantile():
     """Check quantile calculations."""
     quantiles = [0.0, 0.25, 0.5, 0.75, 1.0]