scikit-learn · EvgeniDubov · Oct 23, 2021 · Oct 23, 2021 · Oct 23, 2021 · Jan 10, 2022
diff --git a/doc/modules/tree.rst b/doc/modules/tree.rst
@@ -84,7 +84,8 @@ The disadvantages of decision trees include:
 
     - Decision tree learners create biased trees if some classes dominate.
       It is therefore recommended to balance the dataset prior to fitting
-      with the decision tree.
+      with the decision tree. When dealing with a binary classification problem,
+      using the Hellinger distance as a split criterion will alleviate this bias.
 
 
 .. _tree_classification:
@@ -394,6 +395,9 @@ Tips on practical use
     matrix input compared to a dense matrix when features have zero values in
     most of the samples.
 
+  * If the target is binary and imbalanced, it is recommended to use
+    Hellinger distance, `criterion="hellinger"`.
+
 
 .. _tree_algorithms:
 
@@ -531,6 +535,30 @@ Log Loss or Entropy:
 
       \mathrm{LL}(D, T) = \sum_{m \in T} \frac{n_m}{n} H(Q_m)
 
+**HellingerDistance**
+
+Hellinger distance is implemented for the binary case and a single label only.
+
+The score is calculated per split, not per population:
+
+.. math::
+
+    H(Q_m) = \sqrt{\left(\sqrt{\frac{N_{m,+}^{left}}{N_{m,+}}}-\sqrt{\frac{N_{m,-}^{left}}{N_{m,-}}}\right)^2+\left(\sqrt{\frac{N_{m,+}^{right}}{N_{m,+}}}-\sqrt{\frac{N_{m,-}^{right}}{N_{m,-}}}\right)^2 }
+
+where:
+     - :math:`N_{m,+}` is the number of positive samples for the :math:`m` node
+     - :math:`N_{m,-}` is the number of negative samples for the :math:`m` node
+     - :math:`N_{m,+}^{left}` and :math:`N_{m,+}^{right}` are the numbers of positive
+       samples respectively at the left and the right of the :math:`m` node split
+     - :math:`N_{m,-}^{left}` and :math:`N_{m,-}^{right}` are the numbers of negative
+       samples respectively at the left and the right of the :math:`m` node split
+
+
+
+.. topic:: References:
+
+    * :doi:`Cieslak, D. A., Hoens, T. R., Chawla, N. V., & Kegelmeyer, W. P., Hellinger distance decision trees are robust and skew-insensitive, 2012 <10.1007/s10618-011-0222-1>`
+
 Regression criteria
 -------------------
 

diff --git a/doc/whats_new/v1.3.rst b/doc/whats_new/v1.3.rst
@@ -485,6 +485,11 @@ Changelog
 
 :mod:`sklearn.tree`
 ...................
+- |Feature| :class:`tree.DecisionTreeClassifier`, :class:`tree.ExtraTreeClassifier`,
+  :class:`ensemble.RandomForestClassifier` and :class:`ensemble.ExtraTreesClassifier`
+  now support the Hellinger distance as a split criterion for
+  binary classification problem.
+  :pr:`16478` by :user:`Evgeni Dubov <EvgeniDubov>`.
 
 - |Enhancement| Adds a `class_names` parameter to
   :func:`tree.export_text`. This allows specifying the parameter `class_names`

diff --git a/examples/tree/plot_imbalanced_data_classification.py b/examples/tree/plot_imbalanced_data_classification.py
@@ -0,0 +1,97 @@
+"""
+===================================================================
+Decision Tree imbalanced data classification
+===================================================================
+
+In this example we demonstrate the affect of different split criteria
+on decision tree classifier predictions on imbalanced data.
+
+You can see that Gini index is biased towards the majority class while
+Hellinger distance is sensitive to the ratio between the classes.
+
+"""
+
+# Import the necessary modules and libraries
+import numpy as np
+from sklearn import datasets
+from sklearn.tree import DecisionTreeClassifier
+import matplotlib.pyplot as plt
+
+# Create imbalanced dataset
+minority_class_ratio = 0.001
+n_classes = 2
+X, y = datasets.make_classification(
+    n_samples=1000,
+    n_features=2,
+    n_informative=2,
+    n_redundant=0,
+    n_repeated=0,
+    n_classes=n_classes,
+    n_clusters_per_class=1,
+    weights=[1 - minority_class_ratio],
+    shuffle=False,
+    random_state=0,
+)
+
+# Criteria to compare
+criterions = ["gini", "entropy", "hellinger"]
+
+# Plot parameters
+fig = plt.figure(figsize=(10, 6))
+max_subplots = len(criterions) * 2 - 1
+plot_colors = ["darkgrey", "yellow"]
+target_names = ["majority", "minority"]
+markers = [".", "o"]
+
+# Create mesh grid on feature space to draw classifier predictions
+x_min, x_max = X[:, 0].min() - 0.2, X[:, 0].max() + 0.2
+y_min, y_max = X[:, 1].min() - 2, X[:, 1].max() + 0.5
+xx, yy = np.meshgrid(np.arange(x_min, x_max, 0.1), np.arange(y_min, y_max, 0.01))
+
+# Create plot per criterion
+for sub_plot_idx, criterion in enumerate(criterions):
+    # Add subplot to figure
+    fig.add_subplot(1, max_subplots, sub_plot_idx * 2 + 1)
+
+    # Train classifier
+    clf = DecisionTreeClassifier(criterion=criterion)
+    clf.fit(X, y)
+
+    # Draw tree classifier prediction probability for minority class on feature space
+    Z = clf.predict_proba(np.c_[xx.ravel(), yy.ravel()])[:, 1].reshape(xx.shape)
+    plt.contourf(xx, yy, Z, cmap=plt.cm.rainbow)
+    plt.colorbar(label="minority class prediction probability")
+    # Set the range for probability on the colorbar.
+    plt.clim(0, 1)
+
+    # Draw axis labels
+    plt.xlabel("feature_1")
+    plt.ylabel("feature_2")
+
+    # Draw training points
+    for i, color, marker in zip(range(n_classes), plot_colors, markers):
+        idx = np.where(y == i)
+        plt.scatter(
+            X[:, 0][idx],
+            X[:, 1][idx],
+            c=color,
+            label=target_names[i],
+            cmap=plt.cm.RdYlBu,
+            marker=marker,
+            edgecolor="black",
+        )
+
+    # Draw subplot legend and title
+    plt.legend(
+        title="classes",
+        handletextpad=0,
+        loc="lower right",
+        borderpad=0,
+        scatterpoints=1,
+        labelspacing=1,
+    )
+    plt.title(criterion)
+
+    plt.plot()
+
+plt.show()
diff --git a/sklearn/ensemble/_forest.py b/sklearn/ensemble/_forest.py
@@ -1105,11 +1105,13 @@ class RandomForestClassifier(ForestClassifier):
            The default value of ``n_estimators`` changed from 10 to 100
            in 0.22.
 
-    criterion : {"gini", "entropy", "log_loss"}, default="gini"
+    criterion : {"gini", "entropy", "log_loss", "hellinger"}, default="gini"
         The function to measure the quality of a split. Supported criteria are
-        "gini" for the Gini impurity and "log_loss" and "entropy" both for the
-        Shannon information gain, see :ref:`tree_mathematical_formulation`.
-        Note: This parameter is tree-specific.
+         - "gini" for the Gini impurity
+         - "log_loss" and "entropy" both for the Shannon information gain
+         - "hellinger" for the Hellinger distance
+
+        Note: `criterion` is tree-specific, see :ref:`tree_mathematical_formulation`.
 
     max_depth : int, default=None
         The maximum depth of the tree. If None, then nodes are expanded until
@@ -1794,11 +1796,13 @@ class ExtraTreesClassifier(ForestClassifier):
            The default value of ``n_estimators`` changed from 10 to 100
            in 0.22.
 
-    criterion : {"gini", "entropy", "log_loss"}, default="gini"
+    criterion : {"gini", "entropy", "log_loss", "hellinger"}, default="gini"
         The function to measure the quality of a split. Supported criteria are
-        "gini" for the Gini impurity and "log_loss" and "entropy" both for the
-        Shannon information gain, see :ref:`tree_mathematical_formulation`.
-        Note: This parameter is tree-specific.
+         - "gini" for the Gini impurity
+         - "log_loss" and "entropy" both for the Shannon information gain
+         - "hellinger" for the Hellinger distance
+
+        Note: `criterion` is tree-specific, see :ref:`tree_mathematical_formulation`.
 
     max_depth : int, default=None
         The maximum depth of the tree. If None, then nodes are expanded until

diff --git a/sklearn/ensemble/tests/test_forest.py b/sklearn/ensemble/tests/test_forest.py
@@ -76,6 +76,19 @@
     random_state=0,
 )
 
+# Imbalanced classification sample used for testing imbalanced criteria
+IMBL_MINORITY_CLASS_RATIO = 0.05
+X_large_imbl, y_large_imbl = datasets.make_classification(
+    n_samples=500,
+    n_features=10,
+    n_informative=3,
+    n_redundant=0,
+    n_repeated=0,
+    weights=[IMBL_MINORITY_CLASS_RATIO],
+    shuffle=False,
+    random_state=42,
+)
+
 # also load the iris dataset
 # and randomly permute it
 iris = datasets.load_iris()
@@ -165,6 +178,27 @@ def test_iris(name, criterion):
     check_iris_criterion(name, criterion)
 
 
+def check_imbalanced_criterion(name, criterion):
+    ForestClassifier = FOREST_CLASSIFIERS[name]
+
+    clf = ForestClassifier(n_estimators=10, criterion=criterion, random_state=1)
+    clf.fit(X_large_imbl, y_large_imbl)
+
+    # score is a mean of minority class predict_proba
+    score = clf.predict_proba(X_large_imbl)[:, 1].mean()
+
+    assert score > IMBL_MINORITY_CLASS_RATIO, (
+        f"Failed with imbalanced criterion {criterion}, score = {score}, minority class"
+        f" ratio = {IMBL_MINORITY_CLASS_RATIO}"
+    )
+
+
+@pytest.mark.parametrize("name", FOREST_CLASSIFIERS)
+@pytest.mark.parametrize("criterion", ["hellinger"])
+def test_imbalanced_criteria(name, criterion):
+    check_imbalanced_criterion(name, criterion)
+
+
 def check_regression_criterion(name, criterion):
     # Check consistency on regression dataset.
     ForestRegressor = FOREST_REGRESSORS[name]
@@ -308,11 +342,7 @@ def test_probability(name):
     check_probability(name)
 
 
-def check_importances(name, criterion, dtype, tolerance):
-    # cast as dype
-    X = X_large.astype(dtype, copy=False)
-    y = y_large.astype(dtype, copy=False)
-
+def check_importances(name, criterion, tolerance, X, y):
     ForestEstimator = FOREST_ESTIMATORS[name]
 
     est = ForestEstimator(n_estimators=10, criterion=criterion, random_state=0)
@@ -350,15 +380,25 @@ def check_importances(name, criterion, dtype, tolerance):
 @pytest.mark.parametrize(
     "name, criterion",
     itertools.chain(
-        product(FOREST_CLASSIFIERS, ["gini", "log_loss"]),
+        product(FOREST_CLASSIFIERS, ["gini", "log_loss", "hellinger"]),
         product(FOREST_REGRESSORS, ["squared_error", "friedman_mse", "absolute_error"]),
     ),
 )
 def test_importances(dtype, name, criterion):
     tolerance = 0.01
     if name in FOREST_REGRESSORS and criterion == "absolute_error":
         tolerance = 0.05
-    check_importances(name, criterion, dtype, tolerance)
+
+    # cast as dtype
+    # use imbalanced data for testing imbalanced criterion
+    if criterion == "hellinger":
+        X = X_large_imbl.astype(dtype, copy=False)
+        y = y_large_imbl.astype(dtype, copy=False)
+    else:
+        X = X_large.astype(dtype, copy=False)
+        y = y_large.astype(dtype, copy=False)
+
+    check_importances(name, criterion, tolerance, X, y)
 
 
 def test_importances_asymptotic():

diff --git a/sklearn/tree/_classes.py b/sklearn/tree/_classes.py
@@ -68,6 +68,7 @@
     "gini": _criterion.Gini,
     "log_loss": _criterion.Entropy,
     "entropy": _criterion.Entropy,
+    "hellinger": _criterion.HellingerDistance,
 }
 CRITERIA_REG = {
     "squared_error": _criterion.MSE,
@@ -604,10 +605,13 @@ class DecisionTreeClassifier(ClassifierMixin, BaseDecisionTree):
 
     Parameters
     ----------
-    criterion : {"gini", "entropy", "log_loss"}, default="gini"
+    criterion : {"gini", "entropy", "log_loss", "hellinger"}, default="gini"
         The function to measure the quality of a split. Supported criteria are
-        "gini" for the Gini impurity and "log_loss" and "entropy" both for the
-        Shannon information gain, see :ref:`tree_mathematical_formulation`.
+         - "gini" for the Gini impurity
+         - "log_loss" and "entropy" both for the Shannon information gain
+         - "hellinger" for the Hellinger distance
+
+        Note: `criterion` is tree-specific, see :ref:`tree_mathematical_formulation`.
 
     splitter : {"best", "random"}, default="best"
         The strategy used to choose the split at each node. Supported
@@ -821,7 +825,10 @@ class DecisionTreeClassifier(ClassifierMixin, BaseDecisionTree):
 
     _parameter_constraints: dict = {
         **BaseDecisionTree._parameter_constraints,
-        "criterion": [StrOptions({"gini", "entropy", "log_loss"}), Hidden(Criterion)],
+        "criterion": [
+            StrOptions({"gini", "entropy", "log_loss", "hellinger"}),
+            Hidden(Criterion),
+        ],
         "class_weight": [dict, list, StrOptions({"balanced"}), None],
     }
 
@@ -1291,10 +1298,13 @@ class ExtraTreeClassifier(DecisionTreeClassifier):
 
     Parameters
     ----------
-    criterion : {"gini", "entropy", "log_loss"}, default="gini"
+    criterion : {"gini", "entropy", "log_loss", "hellinger"}, default="gini"
         The function to measure the quality of a split. Supported criteria are
-        "gini" for the Gini impurity and "log_loss" and "entropy" both for the
-        Shannon information gain, see :ref:`tree_mathematical_formulation`.
+         - "gini" for the Gini impurity
+         - "log_loss" and "entropy" both for the Shannon information gain
+         - "hellinger" for the Hellinger distance
+
+        Note: `criterion` is tree-specific, see :ref:`tree_mathematical_formulation`.
 
     splitter : {"random", "best"}, default="random"
         The strategy used to choose the split at each node. Supported