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

Code before:

@pytest.mark.parametrize(
    "TreeEstimator, criterion",
    [
        *product(REG_TREES.values(), REG_CRITERIONS),
        *product(CLF_TREES.values(), CLF_CRITERIONS),
    ],
)
def test_min_impurity_decrease(TreeEstimator, criterion, global_random_seed):
    # test if min_impurity_decrease ensure that a split is made only if
    # if the impurity decrease is at least that value
    X, y = datasets.make_classification(n_samples=100, random_state=global_random_seed)

    # test both DepthFirstTreeBuilder and BestFirstTreeBuilder
    # by setting max_leaf_nodes
    for max_leaf_nodes in [None, 1000]:
        for expected_decrease in [0.05, 0.0001, 0.1]:
            est = TreeEstimator(
                criterion=criterion,
                max_leaf_nodes=max_leaf_nodes,
                min_impurity_decrease=expected_decrease,
                random_state=global_random_seed,
            )
            est.fit(X, y)
            tree = est.tree_
            weighted_impurity = tree.impurity * tree.n_node_samples / X.shape[0]

            for node in range(tree.node_count):
                # If current node is a not leaf node, check if the split was
                # justified w.r.t the min_impurity_decrease
                if tree.children_left[node] != TREE_LEAF:
                    left = tree.children_left[node]
                    right = tree.children_right[node]

                    actual_decrease = weighted_impurity[node] - (
                        weighted_impurity[left] + weighted_impurity[right]
                    )

                    assert actual_decrease >= expected_decrease


def test_pickle():
    """Test pickling preserves Tree properties and performance."""
    for name, TreeEstimator in ALL_TREES.items():
        if "Classifier" in name:
            X, y = iris.data, iris.target
        else:
            X, y = diabetes.data, diabetes.target

        est = TreeEstimator(random_state=0)
        est.fit(X, y)

Review comment (suggestion):
```suggestion
            weighted_impurity = tree.impurity * tree.weighted_n_node_samples / X.shape[0]
```

Code after:
    name, csc_container
):
    check_min_weight_fraction_leaf_with_min_samples_leaf(
        name, "multilabel", sparse_container=csc_container
    )


# TODO(1.11): remove the deprecated friedman_mse criterion parametrization
@pytest.mark.filterwarnings("ignore:.*friedman_mse.*:FutureWarning")
@pytest.mark.parametrize(
    "TreeEstimator, criterion",
    [
        *product(REG_TREES.values(), REG_CRITERIONS),
        *product(CLF_TREES.values(), CLF_CRITERIONS),
    ],
)
def test_min_impurity_decrease(TreeEstimator, criterion, global_random_seed):
    # test if min_impurity_decrease ensure that a split is made only if
    # if the impurity decrease is at least that value
    X, y = datasets.make_classification(n_samples=100, random_state=global_random_seed)

    # test both DepthFirstTreeBuilder and BestFirstTreeBuilder
    # by setting max_leaf_nodes
    for max_leaf_nodes in [None, 1000]:
        for expected_decrease in [0.05, 0.0001, 0.1]:
            est = TreeEstimator(
                criterion=criterion,
                max_leaf_nodes=max_leaf_nodes,
                min_impurity_decrease=expected_decrease,
                random_state=global_random_seed,
            )
            est.fit(X, y)
            tree = est.tree_
            weighted_impurity = (
                tree.impurity * tree.weighted_n_node_samples / X.shape[0]
            )

            for node in range(tree.node_count):
                # If current node is a not leaf node, check if the split was
                # justified w.r.t the min_impurity_decrease
                if tree.children_left[node] != TREE_LEAF:
                    left = tree.children_left[node]
                    right = tree.children_right[node]

                    actual_decrease = weighted_impurity[node] - (
                        weighted_impurity[left] + weighted_impurity[right]
                    )

                    # Allow a tiny slack to account for floating-point rounding errors:
                    assert actual_decrease > expected_decrease - 1e-10