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

Code before:
    assert kappa == cohen_kappa_score(y2, y1)

    # Add spurious labels and ignore them.
    y1 = np.append(y1, [2] * 4)
    y2 = np.append(y2, [2] * 4)
    assert cohen_kappa_score(y1, y2, labels=[0, 1]) == kappa

    assert_almost_equal(cohen_kappa_score(y1, y1), 1.0)

    # Multiclass example: Artstein and Poesio, Table 4.
    y1 = np.array([0] * 46 + [1] * 44 + [2] * 10)
    y2 = np.array([0] * 52 + [1] * 32 + [2] * 16)
    assert_almost_equal(cohen_kappa_score(y1, y2), 0.8013, decimal=4)

    # Weighting example: none, linear, quadratic.
    y1 = np.array([0] * 46 + [1] * 44 + [2] * 10)
    y2 = np.array([0] * 50 + [1] * 40 + [2] * 10)
    assert_almost_equal(cohen_kappa_score(y1, y2), 0.9315, decimal=4)
    assert_almost_equal(cohen_kappa_score(y1, y2, weights="linear"), 0.9412, decimal=4)
    assert_almost_equal(
        cohen_kappa_score(y1, y2, weights="quadratic"), 0.9541, decimal=4
    )


@pytest.mark.parametrize("replace_undefined_by", [0.0, np.nan])
def test_cohen_kappa_zero_division(replace_undefined_by):
    """Test that cohen_kappa_score handles divisions by 0 correctly by returning the
    `replace_undefined_by` param. (The fist two tests cover the first possible location
    in the function for an occurrence of a division by zero, the second two tests 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

    # test case: empty inputs
    y1 = np.array([])
    y2 = np.array([])
    assert check_equal(
        cohen_kappa_score(y1, y2, replace_undefined_by=replace_undefined_by),
        replace_undefined_by,
    )

    # test case: annotator y2 does not assign any label specified in `labels` (note:
    # also applicable if labels is default and y2 does not contain any label that is in
    # y1)
    labels = [1, 2]
    y1 = np.array([1] * 5 + [2] * 5)
    y2 = np.array([3] * 10)
    assert check_equal(

Review comment (refactor):
The test function looks good to me overall—my comment is just a minor nitpick:
```suggestion
@pytest.mark.parametrize(
    "test_case",
    [
        ([], [], None, None, None),
        ([1] * 5 + [2] * 5, [3] * 10, [1, 2], None, None),
        ([3] * 10, [3] * 10, None, None, None),
        ([1] * 5 + [2] * 5, [3] * 10, [1, 2], "linear", None),
    ],
)
@pytest.mark.parametrize("replace_undefined_by", [0.0, np.nan])
def test_cohen_kappa_zero_division(replace_undefined_by):
```
This way, we could also simplify the unpacking slightly:
```python
    y1, y2, labels, weights, sample_weight = test_case
    y_1, y2 = np.array(y1), np.array(y2)
```
and only need one `assert`:
```python
    assert check_equal(
        cohen_kappa_score(
            y1,
            y2,
            labels=labels,
            weights=weights,
            replace_undefined_by=replace_undefined_by,
        ),
        replace_undefined_by,
    )
```

Totally optional though—happy for you to resolve this if you’d prefer to keep it as is.

Code after:
        ([1] * 5 + [2] * 5, [3] * 10, [1, 2], None),
        # both inputs (`y1` and `y2`) only have one label:
        ([3] * 10, [3] * 10, None, None),
        # both inputs only have one label in common that is also in `labels`:
        ([1] * 5 + [2] * 5, [1] * 5 + [3] * 5, [1, 2], None),
        # like the last test case, but with `weights="linear"` (note that
        # weights="linear" and weights="quadratic" are different branches, though the
        # latter is so similar to the former that the test case is skipped here):
        ([1] * 5 + [2] * 5, [1] * 5 + [3] * 5, [1, 2], "linear"),
    ],
)
@pytest.mark.parametrize("replace_undefined_by", [0.0, np.nan])
def test_cohen_kappa_undefined(test_case, replace_undefined_by):
    """Test that cohen_kappa_score handles divisions by 0 correctly by returning the
    `replace_undefined_by` param. (The first test case covers 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."""

    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)


def test_cohen_kappa_zero_division_warning():
    """Test that cohen_kappa_score raises UndefinedMetricWarning when a division by 0
    occurs."""

    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 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.",