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

Code before:
def test_deterministic_pickle():
    # Non-regression test for:
    # https://github.com/scikit-learn/scikit-learn/issues/27268
    # Uninitialised memory would lead to the two pickle strings being different.
    tree1 = DecisionTreeClassifier(random_state=0).fit(iris.data, iris.target)
    tree2 = DecisionTreeClassifier(random_state=0).fit(iris.data, iris.target)

    pickle1 = pickle.dumps(tree1)
    pickle2 = pickle.dumps(tree2)

    assert pickle1 == pickle2


@pytest.mark.parametrize("Tree", [DecisionTreeRegressor, ExtraTreeRegressor])
@pytest.mark.parametrize(
    "X",
    [
        # missing values will go left for greedy splits
        np.array([np.nan, 2, np.nan, 4, 5, 6]),
        np.array([np.nan, np.nan, 3, 4, 5, 6]),
        # missing values will go right for greedy splits
        np.array([1, 2, 3, 4, np.nan, np.nan]),
        np.array([1, 2, 3, np.nan, 6, np.nan]),
    ],
)
def test_regression_tree_missing_values_toy(Tree, X, global_random_seed):
    """Check that we properly handle missing values in regression trees using a toy
    dataset.

    The regression targeted by this test was that we were not reinitializing the
    criterion when it comes to the number of missing values. Therefore, the value
    of the critetion (i.e. MSE) was completely wrong.

    This test check that the MSE is null when there is a single sample in the leaf.

    Non-regression test for:
    https://github.com/scikit-learn/scikit-learn/issues/28254
    https://github.com/scikit-learn/scikit-learn/issues/28316
    """
    X = X.reshape(-1, 1)
    y = np.arange(6)

    tree = Tree(random_state=global_random_seed).fit(X, y)
    tree_ref = clone(tree).fit(y.reshape(-1, 1), y)

    impurity = tree.tree_.impurity
    assert all(impurity >= 0), impurity.min()  # MSE should always be positive

    # Note: the impurity matches after the first split only on greedy trees
    # see https://github.com/scikit-learn/scikit-learn/issues/32125
    if Tree is DecisionTreeRegressor:

Review comment (suggestion):
Actually, this test should rather stay unchanged but annotated with `@pytest.mark.filterwarnings("ignore:.*friedman_mse.*:FutureWarning")` and a `TODO(1.11)` comment.

Code after:

    # Create dataset with missing values
    X[rng.choice([False, True], size=X.shape, p=[0.9, 0.1])] = np.nan

    # Zero sample weight is the same as removing the sample
    sample_weight = np.ones(X.shape[0])
    sample_weight[::2] = 0.0

    tree_with_sw = Tree(random_state=0)
    tree_with_sw.fit(X, y, sample_weight=sample_weight)

    tree_samples_removed = Tree(random_state=0)
    tree_samples_removed.fit(X[1::2, :], y[1::2])

    assert_allclose(tree_samples_removed.predict(X), tree_with_sw.predict(X))


def test_deterministic_pickle():
    # Non-regression test for:
    # https://github.com/scikit-learn/scikit-learn/issues/27268
    # Uninitialised memory would lead to the two pickle strings being different.
    tree1 = DecisionTreeClassifier(random_state=0).fit(iris.data, iris.target)
    tree2 = DecisionTreeClassifier(random_state=0).fit(iris.data, iris.target)

    pickle1 = pickle.dumps(tree1)
    pickle2 = pickle.dumps(tree2)

    assert pickle1 == pickle2


# TODO(1.11): remove the deprecated friedman_mse criterion parametrization
@pytest.mark.filterwarnings("ignore:.*friedman_mse.*:FutureWarning")
@pytest.mark.parametrize("Tree", [DecisionTreeRegressor, ExtraTreeRegressor])
@pytest.mark.parametrize(
    "X",
    [
        # missing values will go left for greedy splits
        np.array([np.nan, 2, np.nan, 4, 5, 6]),
        np.array([np.nan, np.nan, 3, 4, 5, 6]),
        # missing values will go right for greedy splits
        np.array([1, 2, 3, 4, np.nan, np.nan]),
        np.array([1, 2, 3, np.nan, 6, np.nan]),
    ],
)
@pytest.mark.parametrize("criterion", ["squared_error", "friedman_mse"])
def test_regression_tree_missing_values_toy(Tree, X, criterion, global_random_seed):
    """Check that we properly handle missing values in regression trees using a toy
    dataset.

    The regression targeted by this test was that we were not reinitializing the
    criterion when it comes to the number of missing values. Therefore, the value