Add AUROC metric option to ClassificationTask.

rslearn/train/tasks/classification.py

@@ -10,6 +10,8 @@
 from torchmetrics import Metric, MetricCollection
 from torchmetrics.classification import (
     MulticlassAccuracy,
+    MulticlassAUROC,
+    MulticlassAveragePrecision,
     MulticlassF1Score,
     MulticlassPrecision,
     MulticlassRecall,
@@ -43,6 +45,10 @@ def __init__(
         metric_kwargs: dict[str, Any] = {},
         enable_f1_metric: bool = False,
         f1_metric_kwargs: dict[str, Any] = {},
+        enable_auroc: bool = False,
+        auroc_metric_kwargs: dict[str, Any] = {},
+        enable_prauc: bool = False,
+        prauc_metric_kwargs: dict[str, Any] = {},
         positive_class: str | None = None,
         positive_class_threshold: float = 0.5,
         enable_confusion_matrix: bool = False,
@@ -68,6 +74,13 @@ def __init__(
                 torchmetrics.classification.MulticlassAccuracy.
             enable_f1_metric: whether to compute F1 (default false)
             f1_metric_kwargs: extra arguments to pass to F1 metric.
+            enable_auroc: whether to compute AUROC (default false).
+            auroc_metric_kwargs: extra arguments to pass to the AUROC metric, see
+                torchmetrics.classification.MulticlassAUROC.
+            enable_prauc: whether to compute PRAUC, i.e. average precision / area under
+                the precision-recall curve (default false).
+            prauc_metric_kwargs: extra arguments to pass to the PRAUC metric, see
+                torchmetrics.classification.MulticlassAveragePrecision.
             positive_class: positive class name.
             positive_class_threshold: threshold for classifying the positive class in
                 binary classification (default 0.5).
@@ -86,6 +99,10 @@ def __init__(
         self.metric_kwargs = metric_kwargs
         self.enable_f1_metric = enable_f1_metric
         self.f1_metric_kwargs = f1_metric_kwargs
+        self.enable_auroc = enable_auroc
+        self.auroc_metric_kwargs = auroc_metric_kwargs
+        self.enable_prauc = enable_prauc
+        self.prauc_metric_kwargs = prauc_metric_kwargs
         self.positive_class = positive_class
         self.positive_class_threshold = positive_class_threshold
         self.enable_confusion_matrix = enable_confusion_matrix
@@ -283,6 +300,18 @@ def get_metrics(self) -> MetricCollection:
                 )
                 metrics["f1"] = ClassificationMetric(MulticlassF1Score(**kwargs))
 
+        if self.enable_auroc:
+            auroc_kwargs = {"num_classes": len(self.classes)}
+            auroc_kwargs.update(self.auroc_metric_kwargs)
+            metrics["auroc"] = ClassificationMetric(MulticlassAUROC(**auroc_kwargs))
+
+        if self.enable_prauc:
+            prauc_kwargs = {"num_classes": len(self.classes)}
+            prauc_kwargs.update(self.prauc_metric_kwargs)
+            metrics["prauc"] = ClassificationMetric(
+                MulticlassAveragePrecision(**prauc_kwargs)
+            )
+
         if self.enable_confusion_matrix:
             metrics["confusion_matrix"] = ClassificationMetric(
                 ConfusionMatrixMetric(

tests/unit/train/tasks/test_classification.py

@@ -109,3 +109,172 @@ def test_per_class_f1() -> None:
     assert results["1_recall"] == pytest.approx(1)
     assert results["2_precision"] == pytest.approx(0)
     assert results["2_recall"] == pytest.approx(0)
+
+
+def test_auroc_partial() -> None:
+    # Binary case with hand-computable macro AUROC.
+    # Labels are [0, 0, 1, 1] with class-1 probabilities [0.1, 0.4, 0.35, 0.8].
+    # For 2 classes, macro one-vs-rest AUROC equals the class-1 AUROC, i.e. the
+    # fraction of (positive, negative) score pairs that are correctly ranked:
+    #   positives (label 1) scores: 0.35, 0.8
+    #   negatives (label 0) scores: 0.1, 0.4
+    #   correctly ranked: (0.35>0.1), (0.8>0.1), (0.8>0.4) -> 3 of 4 -> 0.75
+    targets = [
+        {
+            "class": torch.tensor(0, dtype=torch.int32),
+            "valid": torch.tensor(1, dtype=torch.int32),
+        },
+        {
+            "class": torch.tensor(0, dtype=torch.int32),
+            "valid": torch.tensor(1, dtype=torch.int32),
+        },
+        {
+            "class": torch.tensor(1, dtype=torch.int32),
+            "valid": torch.tensor(1, dtype=torch.int32),
+        },
+        {
+            "class": torch.tensor(1, dtype=torch.int32),
+            "valid": torch.tensor(1, dtype=torch.int32),
+        },
+    ]
+    preds = torch.tensor(
+        [
+            [0.9, 0.1],
+            [0.6, 0.4],
+            [0.65, 0.35],
+            [0.2, 0.8],
+        ],
+        dtype=torch.float32,
+    )
+
+    task = ClassificationTask(
+        property_name="ignored", classes=["0", "1"], enable_auroc=True
+    )
+    metrics = task.get_metrics()
+    metrics.update(preds, targets)
+    results = metrics.compute()
+    assert results["auroc"] == pytest.approx(0.75)
+
+
+def test_auroc_perfect() -> None:
+    # Perfectly separable predictions should yield AUROC of 1.0.
+    targets = [
+        {
+            "class": torch.tensor(0, dtype=torch.int32),
+            "valid": torch.tensor(1, dtype=torch.int32),
+        },
+        {
+            "class": torch.tensor(0, dtype=torch.int32),
+            "valid": torch.tensor(1, dtype=torch.int32),
+        },
+        {
+            "class": torch.tensor(1, dtype=torch.int32),
+            "valid": torch.tensor(1, dtype=torch.int32),
+        },
+        {
+            "class": torch.tensor(1, dtype=torch.int32),
+            "valid": torch.tensor(1, dtype=torch.int32),
+        },
+    ]
+    preds = torch.tensor(
+        [
+            [0.9, 0.1],
+            [0.8, 0.2],
+            [0.2, 0.8],
+            [0.1, 0.9],
+        ],
+        dtype=torch.float32,
+    )
+
+    task = ClassificationTask(
+        property_name="ignored", classes=["0", "1"], enable_auroc=True
+    )
+    metrics = task.get_metrics()
+    metrics.update(preds, targets)
+    results = metrics.compute()
+    assert results["auroc"] == pytest.approx(1.0)
+
+
+def test_prauc_partial() -> None:
+    # Binary case with hand-computable macro PRAUC (average precision).
+    # Labels are [0, 0, 1, 1] with class-1 probabilities [0.1, 0.4, 0.35, 0.8].
+    # Average precision uses AP = sum_n (R_n - R_{n-1}) * P_n over score thresholds.
+    #   Class 1: ranking by class-1 prob gives 0.8(pos), 0.4(neg), 0.35(pos), 0.1(neg)
+    #     -> (1/2)*1 + (1/2)*(2/3) = 5/6.
+    #   Class 0: ranking by class-0 prob gives 0.9(pos), 0.65(neg), 0.6(pos), 0.2(neg)
+    #     -> (1/2)*1 + (1/2)*(2/3) = 5/6.
+    #   Macro average = 5/6 ~= 0.8333.
+    targets = [
+        {
+            "class": torch.tensor(0, dtype=torch.int32),
+            "valid": torch.tensor(1, dtype=torch.int32),
+        },
+        {
+            "class": torch.tensor(0, dtype=torch.int32),
+            "valid": torch.tensor(1, dtype=torch.int32),
+        },
+        {
+            "class": torch.tensor(1, dtype=torch.int32),
+            "valid": torch.tensor(1, dtype=torch.int32),
+        },
+        {
+            "class": torch.tensor(1, dtype=torch.int32),
+            "valid": torch.tensor(1, dtype=torch.int32),
+        },
+    ]
+    preds = torch.tensor(
+        [
+            [0.9, 0.1],
+            [0.6, 0.4],
+            [0.65, 0.35],
+            [0.2, 0.8],
+        ],
+        dtype=torch.float32,
+    )
+
+    task = ClassificationTask(
+        property_name="ignored", classes=["0", "1"], enable_prauc=True
+    )
+    metrics = task.get_metrics()
+    metrics.update(preds, targets)
+    results = metrics.compute()
+    assert results["prauc"] == pytest.approx(5 / 6)
+
+
+def test_prauc_perfect() -> None:
+    # Perfectly separable predictions should yield PRAUC of 1.0.
+    targets = [
+        {
+            "class": torch.tensor(0, dtype=torch.int32),
+            "valid": torch.tensor(1, dtype=torch.int32),
+        },
+        {
+            "class": torch.tensor(0, dtype=torch.int32),
+            "valid": torch.tensor(1, dtype=torch.int32),
+        },
+        {
+            "class": torch.tensor(1, dtype=torch.int32),
+            "valid": torch.tensor(1, dtype=torch.int32),
+        },
+        {
+            "class": torch.tensor(1, dtype=torch.int32),
+            "valid": torch.tensor(1, dtype=torch.int32),
+        },
+    ]
+    preds = torch.tensor(
+        [
+            [0.9, 0.1],
+            [0.8, 0.2],
+            [0.2, 0.8],
+            [0.1, 0.9],
+        ],
+        dtype=torch.float32,
+    )
+
+    task = ClassificationTask(
+        property_name="ignored", classes=["0", "1"], enable_prauc=True
+    )
+    metrics = task.get_metrics()
+    metrics.update(preds, targets)
+    results = metrics.compute()
+    assert results["prauc"] == pytest.approx(1.0)