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

Code before:
    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:])])


@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
    y_pred = etr.predict([[np.nan]])
    if went_left:
        assert_allclose(y_pred_left, y_pred)
    else:
        assert_allclose(y_pred_right, y_pred)


@pytest.mark.parametrize("criterion", ["entropy", "gini"])
def test_missing_values_best_splitter_three_classes(criterion):
    """Test when missing values are uniquely present in a class among 3 classes."""
    missing_values_class = 0
    X = np.array([[np.nan] * 4 + [0, 1, 2, 3, 8, 9, 11, 12]]).T
    y = np.array([missing_values_class] * 4 + [1] * 4 + [2] * 4)
    dtc = DecisionTreeClassifier(random_state=42, max_depth=2, criterion=criterion)

Review comment (suggestion):
Let's make it easier to `git grep` TODOs before the 1.10 release.

```suggestion
# TODO(1.11): remove the deprecated friedman_mse criterion parametrization and
# use the default instead. 
@pytest.mark.filterwarnings("ignore:.*friedman_mse.*:FutureWarning")
```

Code after:
    """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)

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