FEA Support missing-values in ExtraTrees* (#28268)

adam2392 · web-flow · commit 775587b6b6bb · 2024-07-10T11:53:49.000Z
diff --git a/doc/whats_new/v1.6.rst b/doc/whats_new/v1.6.rst
@@ -148,6 +148,11 @@ Changelog
   :pr:`28622` by :user:`Adam Li <adam2392>` and
   :user:`Sérgio Pereira <sergiormpereira>`.
 
+- |Feature| :class:`ensemble.ExtraTreesClassifier` and :class:`ensemble.ExtraTreesRegressor` now support
+  missing-values in the data matrix `X`. Missing-values are handled by randomly moving all of
+  the samples to the left, or right child node as the tree is traversed. 
+  :pr:`28268` by :user:`Adam Li <adam2392>`.
+
 :mod:`sklearn.impute`
 .....................
 
diff --git a/sklearn/ensemble/tests/test_forest.py b/sklearn/ensemble/tests/test_forest.py
@@ -1767,6 +1767,8 @@ def test_estimators_samples(ForestClass, bootstrap, seed):
     [
         (datasets.make_regression, RandomForestRegressor),
         (datasets.make_classification, RandomForestClassifier),
+        (datasets.make_regression, ExtraTreesRegressor),
+        (datasets.make_classification, ExtraTreesClassifier),
     ],
 )
 def test_missing_values_is_resilient(make_data, Forest):
@@ -1800,12 +1802,21 @@ def test_missing_values_is_resilient(make_data, Forest):
     assert score_with_missing >= 0.80 * score_without_missing
 
 
-@pytest.mark.parametrize("Forest", [RandomForestClassifier, RandomForestRegressor])
+@pytest.mark.parametrize(
+    "Forest",
+    [
+        RandomForestClassifier,
+        RandomForestRegressor,
+        ExtraTreesRegressor,
+        ExtraTreesClassifier,
+    ],
+)
 def test_missing_value_is_predictive(Forest):
     """Check that the forest learns when missing values are only present for
     a predictive feature."""
     rng = np.random.RandomState(0)
     n_samples = 300
+    expected_score = 0.75
 
     X_non_predictive = rng.standard_normal(size=(n_samples, 10))
     y = rng.randint(0, high=2, size=n_samples)
@@ -1835,19 +1846,20 @@ def test_missing_value_is_predictive(Forest):
 
     predictive_test_score = forest_predictive.score(X_predictive_test, y_test)
 
-    assert predictive_test_score >= 0.75
+    assert predictive_test_score >= expected_score
     assert predictive_test_score >= forest_non_predictive.score(
         X_non_predictive_test, y_test
     )
 
 
-def test_non_supported_criterion_raises_error_with_missing_values():
+@pytest.mark.parametrize("Forest", FOREST_REGRESSORS.values())
+def test_non_supported_criterion_raises_error_with_missing_values(Forest):
     """Raise error for unsupported criterion when there are missing values."""
     X = np.array([[0, 1, 2], [np.nan, 0, 2.0]])
     y = [0.5, 1.0]
 
-    forest = RandomForestRegressor(criterion="absolute_error")
+    forest = Forest(criterion="absolute_error")
 
-    msg = "RandomForestRegressor does not accept missing values"
+    msg = ".*does not accept missing values"
     with pytest.raises(ValueError, match=msg):
         forest.fit(X, y)
diff --git a/sklearn/tree/_classes.py b/sklearn/tree/_classes.py
@@ -1686,6 +1686,16 @@ def __init__(
             monotonic_cst=monotonic_cst,
         )
 
+    def _more_tags(self):
+        # XXX: nan is only supported for dense arrays, but we set this for the
+        # common test to pass, specifically: check_estimators_nan_inf
+        allow_nan = self.splitter == "random" and self.criterion in {
+            "gini",
+            "log_loss",
+            "entropy",
+        }
+        return {"multilabel": True, "allow_nan": allow_nan}
+
 
 class ExtraTreeRegressor(DecisionTreeRegressor):
     """An extremely randomized tree regressor.
@@ -1929,3 +1939,13 @@ def __init__(
             ccp_alpha=ccp_alpha,
             monotonic_cst=monotonic_cst,
         )
+
+    def _more_tags(self):
+        # XXX: nan is only supported for dense arrays, but we set this for the
+        # common test to pass, specifically: check_estimators_nan_inf
+        allow_nan = self.splitter == "random" and self.criterion in {
+            "squared_error",
+            "friedman_mse",
+            "poisson",
+        }
+        return {"allow_nan": allow_nan}
