Repository: wearepal/conduit

Issue:
TPR and TNR are missing from conduit.metrics
- [ ] TPR (aka Recall)
- [ ] TNR
- [ ] Precision


Discussion:

Resolving patch:
diff --git a/conduit/metrics.py b/conduit/metrics.py
--- a/conduit/metrics.py
+++ b/conduit/metrics.py
@@ -1,18 +1,7 @@
 from abc import abstractmethod
 from enum import Enum
 from functools import partial, wraps
-from typing import (
-    Callable,
-    List,
-    Literal,
-    Optional,
-    Protocol,
-    Tuple,
-    TypeVar,
-    Union,
-    cast,
-    overload,
-)
+from typing import Callable, List, Literal, Optional, Protocol, Tuple, Union, cast, overload
 
 from ranzen import some
 import torch
@@ -37,6 +26,8 @@
     "nanmax",
     "nanmin",
     "pdist_1d",
+    "pos_preds",
+    "pos_preds_per_subclass",
     "precision_at_k",
     "robust_accuracy",
     "robust_fscore",
@@ -47,7 +38,9 @@
     "robust_tpr",
     "subclass_balanced_accuracy",
     "subclasswise_metric",
+    "tnr",
     "tnr_per_subclass",
+    "tpr",
     "tpr_differences",
     "tpr_per_subclass",
     "weighted_nanmean",
@@ -80,12 +73,9 @@ def precision_at_k(
     return res
 
 
-R = TypeVar("R", Tensor, Tuple[Tensor, Union[Tensor, slice]], covariant=True)
-
-
-class Comparator(Protocol[R]):
+class Comparator(Protocol):
     @abstractmethod
-    def __call__(self, y_pred: Tensor, *, y_true: Tensor) -> R:
+    def __call__(self, y_pred: Tensor, *, y_true: Tensor) -> Tensor | tuple[Tensor, Tensor | slice]:
         """Compare.
 
         :param y_pred: Predicted labels or raw logits of a classifier.
@@ -98,10 +88,6 @@ def __call__(self, y_pred: Tensor, *, y_true: Tensor) -> R:
         ...
 
 
-C = TypeVar("C", bound=Comparator)
-C_co = TypeVar("C_co", bound=Comparator, covariant=True)
-
-
 @torch.no_grad()  # pyright: ignore
 def nanmax(x: Tensor) -> Tensor:
     return torch.amax(torch.nan_to_num(x, nan=-torch.inf))
@@ -297,25 +283,21 @@ def _apply_groupwise_metric(
     return comps.mean()
 
 
-A = TypeVar("A", Aggregator, None)
-A_co = TypeVar("A_co", Aggregator, None, covariant=True)
-
-
-class Metric(Protocol[C_co, A_co]):
+class Metric(Protocol):
     @staticmethod
     def __call__(y_pred: Tensor, *, y_true: Tensor) -> Tensor:
         ...
 
 
-class GroupwiseMetric(Protocol[C_co, A_co]):
+class GroupwiseMetric(Protocol):
     @staticmethod
     def __call__(y_pred: Tensor, *, y_true: Tensor, s: Tensor) -> Tensor:
         ...
 
 
 def groupwise_metric(
-    comparator: C, *, aggregator: A, cond_on_pred: bool = False
-) -> GroupwiseMetric[C, A]:
+    comparator: Comparator, *, aggregator: Aggregator | None, cond_on_pred: bool = False
+) -> GroupwiseMetric:
     """Converts a given ``comparator`` and ``aggregator`` into a group-wise metric.
 
     :param comparator: Function used to assess the correctness of ``y_pred`` with respect
@@ -342,10 +324,8 @@ def _decorated_comparator(y_pred: Tensor, *, y_true: Tensor, s: Tensor) -> Tenso
 
 
 def subclasswise_metric(
-    comparator: C,
-    *,
-    aggregator: A,
-) -> GroupwiseMetric[C, A]:
+    comparator: Comparator, *, aggregator: Aggregator | None
+) -> GroupwiseMetric:
     """Converts a given ``comparator`` and ``aggregator`` into a subclass-wise metric.
 
     :param comparator: Function used to assess the correctness of ``y_pred`` with respect
@@ -367,11 +347,8 @@ def _decorated_comparator(y_pred: Tensor, *, y_true: Tensor, s: Tensor) -> Tenso
 
 
 def classwise_metric(
-    comparator: C,
-    *,
-    aggregator: A,
-    cond_on_pred: bool = False,
-) -> Metric[C, A]:
+    comparator: Comparator, *, aggregator: Aggregator | None, cond_on_pred: bool = False
+) -> Metric:
     """Converts a given ``comparator`` and ``aggregator`` into a subclass-wise metric.
 
     :param comparator: Function used to assess the correctness of ``y_pred`` with respect
@@ -398,6 +375,7 @@ def _decorated_comparator(y_pred: Tensor, *, y_true: Tensor) -> Tensor:
 
 
 def equal(y_pred: Tensor, *, y_true: Tensor) -> Tensor:
+    """Comparator which tests whether ``y_pred`` and ``y_true`` are equal."""
     y_true = torch.atleast_1d(y_true.squeeze()).long()
     if len(y_pred) != len(y_true):
         raise ValueError("'y_pred' and 'y_true' must match in size at dimension 0.")
@@ -408,6 +386,12 @@ def equal(y_pred: Tensor, *, y_true: Tensor) -> Tensor:
     return (y_pred == y_true).float()
 
 
+def nonzero_preds(y_pred: Tensor, *, y_true: Tensor) -> Tensor:
+    if y_pred.is_floating_point():
+        y_pred = hard_prediction(y_pred)
+    return (y_pred != 0).float()
+
+
 def conditional_equal(
     y_pred: Tensor,
     *,
@@ -436,10 +420,29 @@ def accuracy(y_pred: Tensor, *, y_true: Tensor) -> Tensor:
     return Aggregator.MEAN(equal(y_pred=y_pred.squeeze(), y_true=y_true.squeeze()))
 
 
+@torch.no_grad()  # pyright: ignore
+def pos_preds(y_pred: Tensor, *, y_true: Tensor) -> Tensor:
+    """Computes the fraction of positive predictions in ``y_pred``."""
+    return nonzero_preds(y_pred=y_pred.squeeze(), y_true=y_true.squeeze()).mean()
+
+
+@torch.no_grad()  # pyright: ignore
+def tpr(y_pred: Tensor, *, y_true: Tensor) -> Tensor:
+    """Computes the true positive rate (TPR) given predicted and ground-truth labels."""
+    return conditional_equal(y_pred.squeeze(), y_true=y_true.squeeze(), y_true_cond=1)[0].mean()
+
+
+@torch.no_grad()  # pyright: ignore
+def tnr(y_pred: Tensor, *, y_true: Tensor) -> Tensor:
+    """Computes the true negative rate (TNR) given predicted and ground-truth labels."""
+    return conditional_equal(y_pred.squeeze(), y_true=y_true.squeeze(), y_true_cond=0)[0].mean()
+
+
 robust_accuracy = subclasswise_metric(comparator=equal, aggregator=Aggregator.MIN)
 accuracy_per_subclass = subclasswise_metric(comparator=equal, aggregator=None)
 subclass_balanced_accuracy = subclasswise_metric(comparator=equal, aggregator=Aggregator.MEAN)
 robust_gap = subclasswise_metric(comparator=equal, aggregator=Aggregator.MAX_DIFF)
+pos_preds_per_subclass = subclasswise_metric(comparator=nonzero_preds, aggregator=None)
 
 group_balanced_accuracy = groupwise_metric(comparator=equal, aggregator=Aggregator.MEAN)
 accuracy_per_group = groupwise_metric(comparator=equal, aggregator=None)
@@ -532,18 +535,10 @@ def fscore(
         target_size = card_y * card_s
 
     precs = _apply_groupwise_metric(
-        prec_ids,
-        comparator=equal,
-        y_pred=y_pred,
-        y_true=y_true,
-        aggregator=None,
+        prec_ids, comparator=equal, y_pred=y_pred, y_true=y_true, aggregator=None
     )
     recs = _apply_groupwise_metric(
-        rec_ids,
-        comparator=equal,
-        y_pred=y_pred,
-        y_true=y_true,
-        aggregator=None,
+        rec_ids, comparator=equal, y_pred=y_pred, y_true=y_true, aggregator=None
     )
 
     if len(precs) < target_size:


Test patch:
diff --git a/tests/metrics_test.py b/tests/metrics_test.py
--- a/tests/metrics_test.py
+++ b/tests/metrics_test.py
@@ -82,6 +82,41 @@ def test_groupwise_metrics(
         tps_ref = cdtm.nans(card_s)
     assert_close(tps_ref, tps)
 
+    pp = cdtm.pos_preds(y_pred=y_pred, y_true=y_true)
+    match card_y:
+        case 1:
+            if disjoint:  # all predictions are 1
+                assert_close(pp, 1.0)
+            else:  # all predictions are 0
+                assert_close(pp, 0.0)
+        case 2:
+            if disjoint:  # all predictions are 2 and 3
+                assert_close(pp, 1.0)
+            else:
+                assert_close(pp, y_pred.eq(1).float().mean())
+        case _:
+            assert_close(pp, 1 - y_pred.eq(0).float().mean())
+
+    ppps = cdtm.pos_preds_per_subclass(y_pred=y_pred, y_true=y_true, s=s)
+    match ppps.tolist():
+        case [pp]:
+            assert card_s == 1
+            if disjoint:
+                assert_close(pp, 1.0)
+        case [pp1, pp2]:
+            assert card_s == 2
+            if disjoint:
+                assert_close(pp1, 1.0)
+                assert_close(pp2, 1.0)
+        case [pp1, pp2, pp3]:
+            assert card_s == 3
+            if disjoint:
+                assert_close(pp1, 1.0)
+                assert_close(pp2, 1.0)
+                assert_close(pp3, 1.0)
+        case _:
+            raise ValueError("Unexpected number of subclasses")
+
     apc = cdtm.accuracy_per_class(y_pred=y_pred, y_true=y_true)
     _, y_counts = y_true.unique(return_counts=True)
     cw_hits = cdtm.nans(NUM_SAMPLES, card_y).scatter_(-1, index=y_true, src=hits)
@@ -102,6 +137,16 @@ def test_groupwise_metrics(
     td_ref = (tps - tps.unsqueeze(-1)).abs().squeeze()
     assert_close(td_ref, td)
 
+    tpr = cdtm.tpr(y_pred=y_pred, y_true=y_true)
+    match card_y:
+        case 1:
+            assert tpr.isnan()  # TPR is undefined because `y_true` is 0 everywhere.
+        case _:
+            if disjoint:
+                assert_close(tpr, 0.0)  # All predictions are wrong in the disjoint case.
+            else:
+                assert 0.0 < tpr.item() < 1.0
+
     apg = cdtm.accuracy_per_group(y_pred=y_pred, y_true=y_true, s=s)
     assert len(apg) == card_sy