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

Code before:

        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 (bug):
```suggestion
        # 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.
```

Code after:

        # Check root's impurity:
        # The root is 0, left child is 1 and right child is 2.
        root_val, root_impurity = naive_splitter.compute_node_value_and_impurity(y, w)
        assert_allclose(root_impurity, actual_impurity[0], atol=1e-12)
        assert_allclose(root_val, actual_value[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": int(tree.tree_.feature[0]),
            "threshold": tree.tree_.threshold[0],
            "missing_left": bool(tree.tree_.missing_go_to_left[0]),
        }
        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,
        )