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

Code before:
    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:
        nan_density = rng.uniform(0.05, 0.8)
        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)

    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]


Review comment (suggestion):
One more question, why do we need to convert to `float32` here? If it is really required then maybe a short comment above might help make sense of it.

Code after:
    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:
        nan_density = rng.uniform(0.05, 0.8)
        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)