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

Code before:
        X_dense -= 0.5
        mask = rng.random(X_dense.shape) > density
        X_dense[mask] = 0
        X = csc_array(X_dense)
    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, missing_values", [("x", "x"), ("sparse", "x"), ("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"

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

Review comment (suggestion):
Best split optimality is only guaranteed for non-randomized trees but this test is parametrized to run for all tree estimators.

```suggestion
def test_split_impurity(
```

Code after:
        mask = rng.random(X_dense.shape) < nan_density
        X_dense[mask] = np.nan
    if is_sparse:
        density = rng.uniform(0.05, 0.99)
        X_dense -= 0.5
        mask = rng.random(X_dense.shape) > density
        X_dense[mask] = 0
        X = csc_array(X_dense)
    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)

    # Trees cast X to float32 internally; match that dtype here to avoid
    # routing/impurity mismatches from rounding with `<=`.
    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 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
        )