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

Code before:
    dtc = DecisionTreeRegressor(random_state=42, max_depth=1)
    dtc.fit(X, y)

    # Goes to right node because it has the most data points
    y_pred = dtc.predict([[np.nan]])
    assert_allclose(y_pred, [np.mean(y[-5:])])

    # equal number of elements in both nodes
    X_equal = X[:-1]
    y_equal = y[:-1]

    dtc = DecisionTreeRegressor(random_state=42, max_depth=1)
    dtc.fit(X_equal, y_equal)

    # Goes to right node because the implementation sets:
    # missing_go_to_left = n_left > n_right, which is False
    y_pred = dtc.predict([[np.nan]])
    assert_allclose(y_pred, [np.mean(y_equal[-4:])])


# TODO(1.11): remove the deprecated friedman_mse criterion parametrization and
# use the default instead. 
@pytest.mark.filterwarnings("ignore:.*friedman_mse.*:FutureWarning")
@pytest.mark.parametrize("seed", range(3))
@pytest.mark.parametrize("criterion", ["squared_error", "friedman_mse"])
def test_missing_values_random_splitter_on_equal_nodes_no_missing(seed, criterion):
    """Check missing values go to the correct node during predictions for ExtraTree.

    Since ETC use random splits, we use different seeds to verify that the
    left/right node is chosen correctly when the splits occur.
    """
    X = np.array([[0, 1, 2, 3, 8, 9, 11, 12, 15]]).T
    y = np.array([0.1, 0.2, 0.3, 0.2, 1.4, 1.4, 1.5, 1.6, 2.6])

    etr = ExtraTreeRegressor(random_state=seed, max_depth=1, criterion=criterion)
    etr.fit(X, y)

    # Get the left and right children of the root node
    left_child = etr.tree_.children_left[0]
    right_child = etr.tree_.children_right[0]

    # Get the number of samples for the left and right children
    left_samples = etr.tree_.weighted_n_node_samples[left_child]
    right_samples = etr.tree_.weighted_n_node_samples[right_child]
    went_left = left_samples > right_samples

    # predictions
    y_pred_left = etr.tree_.value[left_child][0]
    y_pred_right = etr.tree_.value[right_child][0]

    # Goes to node with the most data points

Review comment (suggestion):
```suggestion
def test_missing_values_random_splitter_on_equal_nodes_no_missing(criterion, seed):
```

Code after:
    msg = (
        r"'min_samples_split' .* must be an int in the range \[2, inf\) "
        r"or a float in the range \(0.0, 1.0\]"
    )
    with pytest.raises(ValueError, match=msg):
        tree.fit(X, y)


# TODO(1.11): remove the deprecated friedman_mse criterion parametrization
@pytest.mark.filterwarnings("ignore:.*friedman_mse.*:FutureWarning")
@pytest.mark.parametrize("criterion", ["squared_error", "friedman_mse"])
def test_missing_values_best_splitter_on_equal_nodes_no_missing(criterion):
    """Check missing values goes to correct node during predictions."""
    X = np.array([[0, 1, 2, 3, 8, 9, 11, 12, 15]]).T
    y = np.array([0.1, 0.2, 0.3, 0.2, 1.4, 1.4, 1.5, 1.6, 2.6])

    dtc = DecisionTreeRegressor(random_state=42, max_depth=1, criterion=criterion)
    dtc.fit(X, y)

    # Goes to right node because it has the most data points
    y_pred = dtc.predict([[np.nan]])
    assert_allclose(y_pred, [np.mean(y[-5:])])

    # equal number of elements in both nodes
    X_equal = X[:-1]
    y_equal = y[:-1]

    dtc = DecisionTreeRegressor(random_state=42, max_depth=1, criterion=criterion)
    dtc.fit(X_equal, y_equal)

    # Goes to right node because the implementation sets:
    # missing_go_to_left = n_left > n_right, which is False
    y_pred = dtc.predict([[np.nan]])
    assert_allclose(y_pred, [np.mean(y_equal[-4:])])


# TODO(1.11): remove the deprecated friedman_mse criterion parametrization
@pytest.mark.filterwarnings("ignore:.*friedman_mse.*:FutureWarning")
@pytest.mark.parametrize("seed", range(3))
@pytest.mark.parametrize("criterion", ["squared_error", "friedman_mse"])
def test_missing_values_random_splitter_on_equal_nodes_no_missing(criterion, seed):
    """Check missing values go to the correct node during predictions for ExtraTree.

    Since ETC use random splits, we use different seeds to verify that the
    left/right node is chosen correctly when the splits occur.
    """
    X = np.array([[0, 1, 2, 3, 8, 9, 11, 12, 15]]).T
    y = np.array([0.1, 0.2, 0.3, 0.2, 1.4, 1.4, 1.5, 1.6, 2.6])

    etr = ExtraTreeRegressor(random_state=seed, max_depth=1, criterion=criterion)
    etr.fit(X, y)