Repository: scikit-learn/scikit-learn
File: sklearn/metrics/tests/test_classification.py (Python)

Code before:
@pytest.mark.parametrize("replace_undefined_by", [0.0, np.nan])
def test_cohen_kappa_zero_division(test_case, replace_undefined_by):
    """Test that cohen_kappa_score handles divisions by 0 correctly by returning the
    `replace_undefined_by` param. (The fist two test cases cover the first possible
    location in the function for an occurrence of a division by zero, the last three
    test cases cover a zero division in the the second possible location in the
    function."""

    def _check_equal(res, exp):
        if np.isnan(res) and np.isnan(exp):
            return True
        return res == exp

    y1, y2, labels, weights = test_case
    y1, y2 = np.array(y1), np.array(y2)

    assert _check_equal(
        cohen_kappa_score(
            y1,
            y2,
            labels=labels,
            weights=weights,
            replace_undefined_by=replace_undefined_by,
        ),
        replace_undefined_by,
    )


# TODO(1.9): remove the @ignore_warnings of the FutureWarning
@ignore_warnings(category=FutureWarning)
def test_cohen_kappa_zero_division_warning():
    """Test that cohen_kappa_score raises UndefinedMetricWarning when a division by 0
    occurs."""

    # test first place to raise warning
    labels = [1, 2]
    y1 = np.array([1] * 5 + [2] * 5)
    y2 = np.array([3] * 10)
    with pytest.warns(
        UndefinedMetricWarning,
        match="`y2` contains no labels that are presented in both `y1` and `labels`.",
    ):
        cohen_kappa_score(y1, y2, labels=labels)

    # test second place to raise warning
    labels = [1, 2]
    y1 = np.array([1] * 5 + [2] * 5)
    y2 = np.array([1] * 5 + [3] * 5)
    with pytest.warns(
        UndefinedMetricWarning,
        match="`y1`, `y2` and `labels` have only one label in common.",

Review comment (suggestion):
```suggestion
    y1, y2, labels, weights = test_case
    y1, y2 = np.array(y1), np.array(y2)

    score = cohen_kappa_score(
            y1,
            y2,
            labels=labels,
            weights=weights,
            replace_undefined_by=replace_undefined_by,
    )
    assert_allclose(score, replace_undefined_by, equal_nan=True)
```

Code after:
    y1 = np.array([1] * 5 + [2] * 5)
    y2 = np.array([3] * 10)
    with pytest.warns(
        UndefinedMetricWarning,
        match="`y2` contains no labels that are present in both `y1` and `labels`.",
    ):
        cohen_kappa_score(y1, y2, labels=labels)

    labels = [1, 2]
    y1 = np.array([1] * 5 + [2] * 5)
    y2 = np.array([1] * 5 + [3] * 5)
    with pytest.warns(
        UndefinedMetricWarning,
        match="`y1`, `y2` and `labels` have only one label in common.",
    ):
        cohen_kappa_score(y1, y2, labels=labels)


def test_cohen_kappa_score_error_wrong_label():
    """Test that correct error is raised when users pass labels that are not in y1."""
    labels = [1, 2]
    y1 = np.array(["a"] * 5 + ["b"] * 5)
    y2 = np.array(["b"] * 10)
    with pytest.raises(
        ValueError, match="At least one label in `labels` must be present in `y1`"
    ):
        cohen_kappa_score(y1, y2, labels=labels)


@pytest.mark.parametrize("zero_division", [0, 1, np.nan])
@pytest.mark.parametrize("y_true, y_pred", [([0], [0])])
@pytest.mark.parametrize(
    "metric",
    [
        f1_score,
        partial(fbeta_score, beta=1),
        precision_score,
        recall_score,
    ],
)
def test_zero_division_nan_no_warning(metric, y_true, y_pred, zero_division):
    """Check the behaviour of `zero_division` when setting to 0, 1 or np.nan.
    No warnings should be raised.
    """
    with warnings.catch_warnings():
        warnings.simplefilter("error")
        result = metric(y_true, y_pred, zero_division=zero_division)

    if np.isnan(zero_division):
        assert np.isnan(result)
    else: