Repository: Lightning-AI/pytorch-lightning
File: tests/metrics/functional/test_classification.py (Python)

Code before:
    pytest.param([1, 1, 0, 0], [0, 0, 1, 1], None, 0.),
    pytest.param([1, 1, 1, 1], [1, 1, 0, 0], 0.8, 0.5),
    pytest.param([0.5, 0.5, 0.5, 0.5], [1, 1, 0, 0], 0.2, 0.5),
    pytest.param([1, 1, 0, 0], [1, 1, 0, 0], 0.5, 1.),
])
def test_auroc(pred, target, max_fpr, expected):
    if max_fpr is not None and LooseVersion(torch.__version__) < LooseVersion('1.6.0'):
        pytest.skip('requires torch v1.6 or higher to test max_fpr argument')

    score = auroc(torch.tensor(pred), torch.tensor(target), max_fpr=max_fpr).item()
    assert score == expected


@pytest.mark.skipif(LooseVersion(torch.__version__) < LooseVersion('1.6.0'),
                    reason='requires torch v1.6 or higher to test max_fpr argument')
@pytest.mark.parametrize(['max_fpr'], [
    pytest.param(None),
    pytest.param(1),
    pytest.param(0.99),
    pytest.param(0.9),
    pytest.param(0.75),
    pytest.param(0.5),
    pytest.param(0.25),
    pytest.param(0.1),
    pytest.param(0.01),
    pytest.param(0.001),
])
def test_auroc_with_max_fpr_against_sklearn(max_fpr):
    device = 'cuda' if torch.cuda.is_available() else 'cpu'

    pred = torch.rand((300,), device=device)
    # Supports only binary classification
    target = torch.randint(2, (300,), dtype=torch.float64, device=device)
    sk_score = sk_roc_auc_score(target.cpu().detach().numpy(),
                                pred.cpu().detach().numpy(),
                                max_fpr=max_fpr)
    pl_score = auroc(pred, target, max_fpr=max_fpr)

    sk_score = torch.tensor(sk_score, dtype=torch.float, device=device)
    assert torch.allclose(sk_score, pl_score)


def test_multiclass_auroc():
    with pytest.raises(ValueError,
                       match=r".*probabilities, i.e. they should sum up to 1.0 over classes"):
        _ = multiclass_auroc(pred=torch.tensor([[0.9, 0.9],
                                                [1.0, 0]]),
                             target=torch.tensor([0, 1]))

    with pytest.raises(ValueError,
                       match=r".*not defined when all of the classes do not occur in the target.*"):

Review comment (suggestion):
```suggestion
    None, 1, 0.99, 0.9, 0.75, 0.5, 0.25, 0.1, 0.01, 0.001,
```

Code after:
    pytest.param([1, 1, 0, 0], [0, 0, 1, 1], None, 0.),
    pytest.param([1, 1, 1, 1], [1, 1, 0, 0], 0.8, 0.5),
    pytest.param([0.5, 0.5, 0.5, 0.5], [1, 1, 0, 0], 0.2, 0.5),
    pytest.param([1, 1, 0, 0], [1, 1, 0, 0], 0.5, 1.),
])
def test_auroc(pred, target, max_fpr, expected):
    if max_fpr is not None and LooseVersion(torch.__version__) < LooseVersion('1.6.0'):
        pytest.skip('requires torch v1.6 or higher to test max_fpr argument')

    score = auroc(torch.tensor(pred), torch.tensor(target), max_fpr=max_fpr).item()
    assert score == expected


@pytest.mark.skipif(LooseVersion(torch.__version__) < LooseVersion('1.6.0'),
                    reason='requires torch v1.6 or higher to test max_fpr argument')
@pytest.mark.parametrize('max_fpr', [
    None, 1, 0.99, 0.9, 0.75, 0.5, 0.25, 0.1, 0.01, 0.001,
])
def test_auroc_with_max_fpr_against_sklearn(max_fpr):
    device = 'cuda' if torch.cuda.is_available() else 'cpu'

    pred = torch.rand((300,), device=device)
    # Supports only binary classification
    target = torch.randint(2, (300,), dtype=torch.float64, device=device)
    sk_score = sk_roc_auc_score(target.cpu().detach().numpy(),
                                pred.cpu().detach().numpy(),
                                max_fpr=max_fpr)
    pl_score = auroc(pred, target, max_fpr=max_fpr)

    sk_score = torch.tensor(sk_score, dtype=torch.float, device=device)
    assert torch.allclose(sk_score, pl_score)


def test_multiclass_auroc():
    with pytest.raises(ValueError,
                       match=r".*probabilities, i.e. they should sum up to 1.0 over classes"):
        _ = multiclass_auroc(pred=torch.tensor([[0.9, 0.9],
                                                [1.0, 0]]),
                             target=torch.tensor([0, 1]))

    with pytest.raises(ValueError,
                       match=r".*not defined when all of the classes do not occur in the target.*"):
        _ = multiclass_auroc(pred=torch.rand((4, 3)).softmax(dim=1),
                             target=torch.tensor([1, 0, 1, 0]))

    with pytest.raises(ValueError,
                       match=r".*does not equal the number of classes passed in 'num_classes'.*"):
        _ = multiclass_auroc(pred=torch.rand((5, 4)).softmax(dim=1),
                             target=torch.tensor([0, 1, 2, 2, 3]),
                             num_classes=6)