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

Code before:
        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, missing_values",
    [(False, False), (True, False), (False, True)],
    ids=["dense-without_missing", "sparse-without_missing", "dense-with_missing"],
)
def test_split_impurity(Tree, criterion, sparse, missing_values, global_random_seed):
    is_clf = criterion in CLF_CRITERIONS

    # TODO: (remove in PR #32119)
    if missing_values and criterion == "absolute_error":
        pytest.skip("AE + missing values not supported yet")
    if missing_values 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, missing_values, 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
        actual_value = tree.tree_.value[:, 0]

        # 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)

Review comment (bug):
```suggestion
        pytest.xfail("Poisson criterion is faulty for now")
```

Code after:

@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, missing_values",
    [(False, False), (True, False), (False, True)],
    ids=["dense-without_missing", "sparse-without_missing", "dense-with_missing"],
)
def test_split_impurity(Tree, criterion, sparse, missing_values, global_random_seed):
    is_clf = criterion in CLF_CRITERIONS

    # TODO: (remove in PR #32119)
    if missing_values and criterion == "absolute_error":
        pytest.skip("AE + missing values not supported yet")
    if missing_values and criterion == "poisson":
        pytest.xfail("Poisson criterion is faulty 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, missing_values, 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
        actual_value = tree.tree_.value[:, 0]

        # 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: