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

Code before:

    assert_allclose(tree_leaf_weighted_median, weighted_median)


def test_splitting_with_missing_values():
    # Non regression test for https://github.com/scikit-learn/scikit-learn/issues/32178
    X = (
        np.vstack([[0, 0, 0, 0, 1, 2, 3, 4], [1, 2, 1, 2, 1, 2, 1, 2]])
        .swapaxes(0, 1)
        .astype(float)
    )
    y = [0, 0, 0, 0, 1, 1, 1, 1]
    X[X == 0] = np.nan

    # The important thing here is that we try several trees, where each one tries
    # one of the two features first. The resulting tree should be the same in all
    # cases. The way to control which feature is tried first is `random_state`.
    # Twenty trees is a good guess for how many we need to try to make sure we get
    # both orders of features at least once.
    for i in range(20):
        tree = DecisionTreeRegressor(max_depth=1, random_state=i).fit(X, y)
        assert_array_equal(tree.tree_.impurity, np.array([0.25, 0.0, 0.0]))


def test_missing_values_and_constant_toy():
    # Non regression test for https://github.com/scikit-learn/scikit-learn/issues/32272
    X = [0, 0, 0, np.nan, np.nan]
    y = [0, 0, 0, 1, 1]
    X = np.array(X).reshape(-1, 1)
    tree = DecisionTreeClassifier().fit(X, y)
    assert_array_equal(tree.predict(X), y)
    assert tree.tree_.node_count == 3


Review comment (suggestion):
Is it possible to explain in a few words why we expect to predict perfectly in this case and why you expect 3 nodes, which I guess means one split was made rather than 0?

It's probably obvious to you since you have this fresh in your head but try to put yourself in the shoes of a busy maintainer that comes back to this in 6 months time and would like to understand without having to think too much.

Code after:

    assert_allclose(tree_leaf_weighted_median, weighted_median)


def test_splitting_with_missing_values():
    # Non regression test for https://github.com/scikit-learn/scikit-learn/issues/32178
    X = (
        np.vstack([[0, 0, 0, 0, 1, 2, 3, 4], [1, 2, 1, 2, 1, 2, 1, 2]])
        .swapaxes(0, 1)
        .astype(float)
    )
    y = [0, 0, 0, 0, 1, 1, 1, 1]
    X[X == 0] = np.nan

    # The important thing here is that we try several trees, where each one tries
    # one of the two features first. The resulting tree should be the same in all
    # cases. The way to control which feature is tried first is `random_state`.
    # Twenty trees is a good guess for how many we need to try to make sure we get
    # both orders of features at least once.
    for i in range(20):
        tree = DecisionTreeRegressor(max_depth=1, random_state=i).fit(X, y)
        assert_array_equal(tree.tree_.impurity, np.array([0.25, 0.0, 0.0]))


def test_missing_values_and_constant_toy():
    # Non regression test for https://github.com/scikit-learn/scikit-learn/issues/32272
    # This test ensures that a feature with constant non-missing values plus some
    # missing values is correctly identified as splittable (not constant).
    X = [0, 0, 0, np.nan, np.nan]  # constant non-missing values (all 0s)
    y = [0, 0, 0, 1, 1]  # perfectly separable by missingness
    X = np.array(X).reshape(-1, 1)
    tree = DecisionTreeClassifier().fit(X, y)
    # We expect perfect predictions because the missing value pattern perfectly
    # separates the two classes (non-missing -> class 0, missing -> class 1)
    assert_array_equal(tree.predict(X), y)
    # with just one split (-> three nodes: the root + 2 leaves)
    assert tree.tree_.node_count == 3