Repository: Project-MONAI/MONAI
File: tests/metrics/test_compute_average_precision.py (Python)

Code before:
        ]
    )
    def test_value(self, y_pred, y, softmax, to_onehot, average, expected_value):
        y_pred_trans = Compose([ToTensor(), Activations(softmax=softmax)])
        y_trans = Compose([ToTensor(), AsDiscrete(to_onehot=to_onehot)])
        y_pred = torch.stack([y_pred_trans(i) for i in decollate_batch(y_pred)], dim=0)
        y = torch.stack([y_trans(i) for i in decollate_batch(y)], dim=0)
        result = compute_average_precision(y_pred=y_pred, y=y, average=average)
        np.testing.assert_allclose(expected_value, result, rtol=1e-5)

    @parameterized.expand(
        [
            TEST_CASE_1,
            TEST_CASE_2,
            TEST_CASE_3,
            TEST_CASE_4,
            TEST_CASE_5,
            TEST_CASE_6,
            TEST_CASE_7,
            TEST_CASE_8,
            TEST_CASE_9,
            TEST_CASE_10,
            TEST_CASE_11,
            TEST_CASE_12,
        ]
    )
    def test_class_value(self, y_pred, y, softmax, to_onehot, average, expected_value):
        y_pred_trans = Compose([ToTensor(), Activations(softmax=softmax)])
        y_trans = Compose([ToTensor(), AsDiscrete(to_onehot=to_onehot)])
        y_pred = [y_pred_trans(i) for i in decollate_batch(y_pred)]
        y = [y_trans(i) for i in decollate_batch(y)]
        metric = AveragePrecisionMetric(average=average)
        metric(y_pred=y_pred, y=y)
        result = metric.aggregate()
        np.testing.assert_allclose(expected_value, result, rtol=1e-5)
        result = metric.aggregate(average=average)  # test optional argument
        metric.reset()
        np.testing.assert_allclose(expected_value, result, rtol=1e-5)


if __name__ == "__main__":
    unittest.main()


Review comment (performance):
```suggestion
]

ALL_TESTS = [
    TEST_CASE_1,
    TEST_CASE_2,
    TEST_CASE_3,
    TEST_CASE_4,
    TEST_CASE_5,
    TEST_CASE_6,
    TEST_CASE_7,
    TEST_CASE_8,
    TEST_CASE_9,
    TEST_CASE_10,
    TEST_CASE_11,
    TEST_CASE_12,
]

class TestComputeAveragePrecision(unittest.TestCase):

    @parameterized.expand(ALL_TESTS)
    def test_value(self, y_pred, y, softmax, to_onehot, average, expected_value):
        y_pred_trans = Compose([ToTensor(), Activations(softmax=softmax)])
        y_trans = Compose([ToTensor(), AsDiscrete(to_onehot=to_onehot)])
        y_pred = torch.stack([y_pred_trans(i) for i in decollate_batch(y_pred)], dim=0)
        y = torch.stack([y_trans(i) for i in decollate_batch(y)], dim=0)
        result = compute_average_precision(y_pred=y_pred, y=y, average=average)
        np.testing.assert_allclose(expected_value, result, rtol=1e-5)

    @parameterized.expand(ALL_TESTS)
```

Code after:

    @parameterized.expand(ALL_TESTS)
    def test_value(self, y_pred, y, softmax, to_onehot, average, expected_value):
        y_pred_trans = Compose([ToTensor(), Activations(softmax=softmax)])
        y_trans = Compose([ToTensor(), AsDiscrete(to_onehot=to_onehot)])
        y_pred = torch.stack([y_pred_trans(i) for i in decollate_batch(y_pred)], dim=0)
        y = torch.stack([y_trans(i) for i in decollate_batch(y)], dim=0)
        result = compute_average_precision(y_pred=y_pred, y=y, average=average)
        np.testing.assert_allclose(expected_value, result, rtol=1e-5)

    @parameterized.expand(ALL_TESTS)
    def test_class_value(self, y_pred, y, softmax, to_onehot, average, expected_value):
        y_pred_trans = Compose([ToTensor(), Activations(softmax=softmax)])
        y_trans = Compose([ToTensor(), AsDiscrete(to_onehot=to_onehot)])
        y_pred = [y_pred_trans(i) for i in decollate_batch(y_pred)]
        y = [y_trans(i) for i in decollate_batch(y)]
        metric = AveragePrecisionMetric(average=average)
        metric(y_pred=y_pred, y=y)
        result = metric.aggregate()
        np.testing.assert_allclose(expected_value, result, rtol=1e-5)
        result = metric.aggregate(average=average)  # test optional argument
        metric.reset()
        np.testing.assert_allclose(expected_value, result, rtol=1e-5)


if __name__ == "__main__":
    unittest.main()