Repository: scikit-learn/scikit-learn
File: sklearn/tree/tests/test_split.py (Python)

Code before:
from dataclasses import dataclass
from itertools import product
from operator import itemgetter

import numpy as np
import pytest
from numpy.testing import assert_allclose
from scipy.sparse import csc_array
from scipy.special import xlogy

from sklearn.metrics import mean_poisson_deviance
from sklearn.tree import (
    DecisionTreeClassifier,
    DecisionTreeRegressor,
    ExtraTreeClassifier,
    ExtraTreeRegressor,
)
from sklearn.utils.stats import _weighted_percentile

CLF_CRITERIONS = ("gini", "log_loss")

REG_CRITERIONS = ("squared_error", "absolute_error", "poisson")

CLF_TREES = {
    "DecisionTreeClassifier": DecisionTreeClassifier,
    "ExtraTreeClassifier": ExtraTreeClassifier,
}

REG_TREES = {
    "DecisionTreeRegressor": DecisionTreeRegressor,
    "ExtraTreeRegressor": ExtraTreeRegressor,
}


@dataclass
class NaiveSplitter:
    criterion: str
    n_classes: int = 0

    def compute_node_impurity(self, y, w):
        sum_weights = np.sum(w)
        if sum_weights < 1e-7:
            return np.inf  # invalid split
        if self.criterion in ["gini", "entropy", "log_loss"]:
            p = np.bincount(y, weights=w, minlength=self.n_classes) / sum_weights
            if self.criterion == "gini":
                # 1 - sum(pk^2)
                loss = 1.0 - np.sum(p**2)
            else:
                # -sum(pk * log2(pk))
                loss = -np.sum(xlogy(p, p)) / np.log(2)
        elif self.criterion == "squared_error":
            mean_y = np.average(y, weights=w)
            loss = np.average((y - mean_y) ** 2, weights=w)
        elif self.criterion == "absolute_error":
            median = _weighted_percentile(y, w, percentile_rank=50, average=True)
            loss = np.average(np.abs(y - median), weights=w)
        elif self.criterion == "poisson":
            mean_y = np.average(y, weights=w)
            loss = mean_poisson_deviance(y, np.repeat(mean_y, y.size), sample_weight=w)
            loss *= 1 / 2
        else:
            raise ValueError(f"Unknown criterion: {self.criterion}")
        return loss * sum_weights

    def compute_split_impurity(
        self, X, y, w, feature, threshold=None, missing_left=False
    ):
        x = X[:, feature]
        go_left = x <= threshold
        if missing_left:
            go_left |= np.isnan(x)
        return (
            self.compute_node_impurity(y[go_left], w[go_left]),
            self.compute_node_impurity(y[~go_left], w[~go_left]),
        )

    def _generate_all_splits(self, X):
        for f in range(X.shape[1]):
            x = X[:, f]
            nan_mask = np.isnan(x)
            xu = np.unique(x[~nan_mask], sorted=True)
            thresholds = (xu[1:] + xu[:-1]) / 2
            for th in thresholds:
                yield {
                    "feature": f,
                    "threshold": th,
                    "missing_left": False,
                }
            if not nan_mask.any():
                continue
            for th in [*thresholds, -np.inf]:
                # include -inf to test all NaNs on the left
                yield {
                    "feature": f,
                    "threshold": th,
                    "missing_left": True,
                }

    def best_split_naive(self, X, y, w):
        splits = []
        splits = list(self._generate_all_splits(X))
        split_impurities = [
            sum(self.compute_split_impurity(X, y, w, **split)) for split in splits
        ]

        return min(zip(split_impurities, splits), key=itemgetter(0))


def sparsify(X, density):
    X -= 0.5
    th_low = np.quantile(X.ravel(), q=density / 2)
    th_up = np.quantile(X.ravel(), q=1 - density / 2)
    X[(th_low < X) & (X < th_up)] = 0
    return csc_array(X)


def make_simple_dataset(
    n,
    d,
    with_nans,
    is_sparse,
    is_clf,
    n_classes,
    rng: np.random.Generator,
):
    X_dense = rng.random((n, d))
    y = rng.random(n) + X_dense.sum(axis=1)
    w = rng.integers(0, 5, size=n) if rng.uniform() < 0.5 else rng.random(n)

    with_duplicates = rng.integers(2) == 0
    if with_duplicates:
        X_dense = X_dense.round(1 if n < 50 else 2)
    if with_nans:
        for i in range(d):
            step = rng.integers(2, 10)
            X_dense[i::step, i] = np.nan
    if is_sparse:
        density = rng.uniform(0.05, 0.99)
        X = sparsify(X_dense, density)
    else:
        X = X_dense

    if is_clf:
        q = np.linspace(0, 1, num=n_classes + 1)[1:-1]
        y = np.searchsorted(np.quantile(y, q), y)

    return X_dense.astype("float32"), X, y, w


@pytest.mark.filterwarnings("ignore:.*friedman_mse.*:FutureWarning")
@pytest.mark.parametrize(
    "Tree, criterion",
    [
        *product(REG_TREES.values(), REG_CRITERIONS),
        *product(CLF_TREES.values(), CLF_CRITERIONS),
    ],
)
@pytest.mark.parametrize("sparse", ["x", "sparse"])
@pytest.mark.parametrize("missing_values", ["x", "missing_values"])
def test_best_split_optimality(
    Tree, criterion, sparse, missing_values, global_random_seed
):
    is_clf = criterion in CLF_CRITERIONS
    with_nans = missing_values != "x"
    is_sparse = sparse != "x"
    if is_sparse and with_nans:
        pytest.skip("Sparse + missing values not supported yet")

    # TODO: (remove in PR #32119)
    if with_nans and criterion == "absolute_error":
        pytest.skip("AE + missing values not supported yet")
    if with_nans and criterion == "poisson":
        pytest.xfail("Poisson criterion is buggy for now")

    rng = np.random.default_rng(global_random_seed)

    ns = [5] * 5 + [10] * 5 + [20, 30, 50, 100]

    for it, n in enumerate(ns):
        d = rng.integers(1, 4)
        n_classes = rng.integers(2, 5)  # only used for classification
        X_dense, X, y, w = make_simple_dataset(
            n, d, with_nans, is_sparse, is_clf, n_classes, rng
        )

        naive_splitter = NaiveSplitter(criterion, n_classes)

        tree = Tree(
            criterion=criterion,
            max_depth=1,
            random_state=global_random_seed,
        )
        tree.fit(X, y, sample_weight=w)
        actual_impurity = tree.tree_.impurity * tree.tree_.weighted_n_node_samples

        # Check root's impurity:
        # The root is 0, left child is 1 and right child is 2.
        root_impurity = naive_splitter.compute_node_impurity(y, w)
        assert_allclose(root_impurity, actual_impurity[0], atol=1e-12)

        if tree.tree_.node_count == 1:
            # if no splits was made assert that either:
            assert (
                "Extra" in Tree.__name__
                or root_impurity < 1e-12  # root impurity is 0
                # or no valid split can be made:
                or naive_splitter.best_split_naive(X_dense, y, w)[0] == np.inf
            )
            continue

        # Check children impurity:
        actual_split = {
            "feature": tree.tree_.feature[0],
            "threshold": tree.tree_.threshold[0],
            "missing_left": bool(tree.tree_.missing_go_to_left[0]),
        }
        left_right_impurity = naive_splitter.compute_split_impurity(
            X_dense, y, w, **actual_split
        )
        assert_allclose(left_right_impurity, actual_impurity[1:], atol=1e-12)

        if "Extra" in Tree.__name__:
            continue

        # check that the selected split was really the best possible split:
        best_impurity, excepted_split = naive_splitter.best_split_naive(X_dense, y, w)
        actual_split_impurity = actual_impurity[1:].sum()
        assert np.isclose(best_impurity, actual_split_impurity), (
            excepted_split,
            actual_split,
        )


Review comment (suggestion):
There is not necessarily a unique minimizer. I think we should instead check that the split found by `Tree.fit` belongs to the set of minimizers found by the brute force exhaustive search of `NaiveSplitter`.

Code after:
        }
        nodes = naive_splitter.compute_split_nodes(X_dense, y, w, **actual_split)
        (left_val, left_impurity), (right_val, right_impurity) = nodes
        assert_allclose(left_impurity, actual_impurity[1], atol=1e-12)
        assert_allclose(right_impurity, actual_impurity[2], atol=1e-12)
        assert_allclose(left_val, actual_value[1], atol=1e-12)
        assert_allclose(right_val, actual_value[2], atol=1e-12)

        if "Extra" in Tree.__name__:
            # The remainder of the test checks for optimality of the found split.
            # However, randomized trees are not guaranteed to find an optimal split
            # but only a "better-than-nothing" split.
            # Therefore, end the test here for these models.
            continue

        # Check that the selected split has the same impurity as the best split
        # found by the naive splitter. Note that there could exist multiple splits
        # with the same optimal impurity, so the assertion is made on the impurity
        # value: the split value is only displayed to help debugging in case
        # of assertion failure.
        best_impurity, best_split = naive_splitter.best_split_naive(X_dense, y, w)
        actual_split_impurity = actual_impurity[1:].sum()
        assert np.isclose(best_impurity, actual_split_impurity), (
            best_split,
            actual_split,
        )