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

Code before:
            metric_class = partial(FBeta, beta=beta)

        self.run_class_metric_test(
            ddp=ddp,
            preds=preds,
            target=target,
            metric_class=metric_class,
            sk_metric=partial(sk_metric, average=average, beta=beta),
            dist_sync_on_step=dist_sync_on_step,
            metric_args={
                "num_classes": num_classes,
                "average": average,
                "multilabel": multilabel,
                "threshold": THRESHOLD,
            },
            check_dist_sync_on_step=False,
            check_batch=False,
        )

    def test_fbeta_functional(
        self, preds, target, sk_metric, num_classes, multilabel, average, beta
    ):
        if beta == 1.0:
            metric_functional = f1
        else:
            metric_functional = partial(fbeta, beta=beta)

        self.run_functional_metric_test(preds=preds,
                                        target=target,
                                        metric_functional=metric_functional,
                                        sk_metric=partial(sk_metric, average=average, beta=beta),
                                        metric_args={
                                            "num_classes": num_classes,
                                            "average": average,
                                            "multilabel": multilabel,
                                            "threshold": THRESHOLD}
                                        )


@pytest.mark.parametrize(['pred', 'target', 'beta', 'exp_score'], [
    pytest.param([1., 0., 1., 0.], [0., 1., 1., 0.], 0.5, [0.5, 0.5]),
    pytest.param([1., 0., 1., 0.], [0., 1., 1., 0.], 1, [0.5, 0.5]),
    pytest.param([1., 0., 1., 0.], [0., 1., 1., 0.], 2, [0.5, 0.5]),
])
def test_fbeta_score(pred, target, beta, exp_score):
    score = fbeta(torch.tensor(pred), torch.tensor(target), num_classes=1, beta=beta, average='none')
    assert torch.allclose(score, torch.tensor(exp_score))


@pytest.mark.parametrize(['pred', 'target', 'exp_score'], [
    pytest.param([0., 0., 0., 0.], [1., 1., 1., 1.], [0.0, 0.0]),

Review comment (suggestion):
```suggestion
        metric_functional = f1 if beta == 1.0 else partial(fbeta, beta=beta)
```

Code after:
            ddp=ddp,
            preds=preds,
            target=target,
            metric_class=metric_class,
            sk_metric=partial(sk_metric, average=average, beta=beta),
            dist_sync_on_step=dist_sync_on_step,
            metric_args={
                "num_classes": num_classes,
                "average": average,
                "multilabel": multilabel,
                "threshold": THRESHOLD,
            },
            check_dist_sync_on_step=False,
            check_batch=False,
        )

    def test_fbeta_functional(
        self, preds, target, sk_metric, num_classes, multilabel, average, beta
    ):
        metric_functional = f1 if beta == 1.0 else partial(fbeta, beta=beta)

        self.run_functional_metric_test(preds=preds,
                                        target=target,
                                        metric_functional=metric_functional,
                                        sk_metric=partial(sk_metric, average=average, beta=beta),
                                        metric_args={
                                            "num_classes": num_classes,
                                            "average": average,
                                            "multilabel": multilabel,
                                            "threshold": THRESHOLD}
                                        )


@pytest.mark.parametrize(['pred', 'target', 'beta', 'exp_score'], [
    pytest.param([1., 0., 1., 0.], [0., 1., 1., 0.], 0.5, [0.5, 0.5]),
    pytest.param([1., 0., 1., 0.], [0., 1., 1., 0.], 1, [0.5, 0.5]),
    pytest.param([1., 0., 1., 0.], [0., 1., 1., 0.], 2, [0.5, 0.5]),
])
def test_fbeta_score(pred, target, beta, exp_score):
    score = fbeta(torch.tensor(pred), torch.tensor(target), num_classes=1, beta=beta, average='none')
    assert torch.allclose(score, torch.tensor(exp_score))


@pytest.mark.parametrize(['pred', 'target', 'exp_score'], [
    pytest.param([0., 0., 0., 0.], [1., 1., 1., 1.], [0.0, 0.0]),
    pytest.param([1., 0., 1., 0.], [0., 1., 1., 0.], [0.5, 0.5]),
    pytest.param([1., 0., 1., 0.], [1., 0., 1., 0.], [1.0, 1.0]),
])
def test_f1_score(pred, target, exp_score):
    score = f1(torch.tensor(pred), torch.tensor(target), num_classes=1, average='none')
    assert torch.allclose(score, torch.tensor(exp_score))