Add simple conformal prediction classification representation and cov…

.pre-commit-config.yaml

@@ -14,3 +14,4 @@ repos:
     hooks:
       - id: ruff
         args: [--fix, --exit-non-zero-on-fix, --no-cache]
+      - id: ruff-format

src/probly/metrics.py

@@ -29,3 +29,47 @@ def expected_calibration_error(probs: np.ndarray, labels: np.ndarray, num_bins:
         ece += weight * np.abs(acc_bin - conf_bin)
 
     return ece
+
+
+def coverage(preds: np.ndarray, targets: np.ndarray) -> float:
+    """
+    Compute the coverage of set-valued predictions.
+    Args:
+        preds: numpy.ndarray of shape (n_instances, n_classes) or
+        (n_instances, n_samples, n_classes)
+        targets: numpy.ndarray of shape (n_instances,) or (n_instances, n_classes)
+    Returns:
+        cov: float, coverage of the set-valued predictions
+    """
+    if preds.ndim == 2:
+        cov = np.mean(preds[np.arange(preds.shape[0]), targets])
+    elif preds.ndim == 3:
+        probs_lower = np.round(np.nanmin(preds, axis=1), decimals=3)
+        probs_upper = np.round(np.nanmax(preds, axis=1), decimals=3)
+        covered = np.all((probs_lower <= targets) & (targets <= probs_upper), axis=1)
+        cov = np.mean(covered)
+    else:
+        raise ValueError(f"Expected 2D or 3D array, got {preds.ndim}D")
+    return cov
+
+
+def efficiency(preds: np.ndarray) -> float:
+    """
+    Compute the efficiency of set-valued predictions. In the case of a set over classes this
+    is the mean of the number of classes in the set. In the case of a credal set, this is computed
+    by the mean difference between the upper and lower probabilities.
+    Args:
+        preds: numpy.ndarray of shape (n_instances, n_classes) or
+        (n_instances, n_samples, n_classes)
+    Returns:
+        eff: float, efficiency of the set-valued predictions
+    """
+    if preds.ndim == 2:
+        eff = np.mean(preds)
+    elif preds.ndim == 3:
+        probs_lower = np.round(np.nanmin(preds, axis=1), decimals=3)
+        probs_upper = np.round(np.nanmax(preds, axis=1), decimals=3)
+        eff = np.mean(probs_upper - probs_lower)
+    else:
+        raise ValueError(f"Expected 2D or 3D array, got {preds.ndim}D")
+    return eff

src/probly/representation/conformalprediction.py

@@ -0,0 +1,64 @@
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+
+class ConformalPrediction:
+    """
+    This class implements conformal prediction for a given model.
+    Args:
+        base: torch.nn.Module, The base model to be used for conformal prediction.
+        alpha: float, The error rate for conformal prediction.
+
+    Attributes:
+        model: torch.nn.Module, The base model.
+        alpha: float, The error rate for conformal prediction.
+        q: float, The quantile value for conformal prediction.
+    """
+
+    def __init__(self, base: torch.nn.Module, alpha: float = 0.05) -> None:
+        self.model = base
+        self.alpha = alpha
+
+    def __call__(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass of the model without conformal prediction.
+        Args:
+            x: torch.Tensor, input data
+        Returns:
+            torch.Tensor, model output
+        """
+        return self.model(x)
+
+    def represent_uncertainty(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Represent the uncertainty of the model by a conformal prediction set.
+        Args:
+            x: torch.Tensor, input data
+        Returns:
+            torch.Tensor of shape (n_instances, n_classes), the conformal prediction set,
+            where each element is a boolean indicating whether the class is included in the set.
+        """
+        with torch.no_grad():
+            outputs = self.model(x)
+            scores = F.softmax(outputs, dim=1)
+        sets = scores >= (1 - self.q)
+        return sets
+
+    def calibrate(self, loader: torch.utils.data.DataLoader) -> None:
+        """
+        Perform the calibration step for conformal prediction.
+        Args:
+            loader: DataLoader, The data loader for the calibration set.
+        """
+        self.model.eval()
+        scores = []
+        with torch.no_grad():
+            for inputs, targets in loader:
+                outputs = self.model(inputs)
+                score = 1 - F.softmax(outputs, dim=1).numpy()
+                score = score[torch.arange(score.shape[0]), targets]
+                scores.append(score)
+        scores = np.concatenate(scores)
+        n = scores.shape[0]
+        self.q = np.quantile(scores, np.ceil((n + 1) * (1 - self.alpha)) / n, method="inverted_cdf")