Add CV+ Adaptive & Simple Prediction Conformal Classifier

docs/source/usage_modes/flax_models.rst

@@ -195,7 +195,8 @@ but a new one could be used.
     conformal_sets = AdaptivePredictionConformalClassifier().conformal_set(
         val_probs=calib_means,
         test_probs=test_means,
-        val_targets=calib_targets
+        val_targets=calib_targets,
+        error=0.05
     )
 
 .. _flax_models_regression:

docs/source/usage_modes/model_outputs.rst

@@ -116,7 +116,8 @@ and :code:`val_targets` to be the corresponding validation target variables.
     conformal_sets = AdaptivePredictionConformalClassifier().conformal_set(
         val_probs=val_means,
         test_probs=test_means,
-        val_targets=val_targets
+        val_targets=val_targets,
+        error=0.05
     )
 
 .. _model_outputs_regression:

docs/source/usage_modes/uncertainty_estimates.rst

@@ -36,7 +36,8 @@ Please check :class:`~fortuna.conformal.classification.adaptive_prediction.Adapt
     conformal_sets = AdaptivePredictionConformalClassifier().conformal_set(
         val_probs=val_probs,
         test_probs=test_probs,
-        val_targets=val_targets
+        val_targets=val_targets,
+        error=0.05
     )
 
 You should usually expect your test predictions to be included in the conformal sets, as they contain the most probable

examples/mnist_classification.pct.py

@@ -153,6 +153,7 @@ def download(split_range, shuffle=False):
     val_probs=val_means,
     test_probs=test_means,
     val_targets=val_data_loader.to_array_targets(),
+    error=0.05
 )
 
 # %% [markdown]

examples/mnist_classification_sghmc.pct.py

@@ -143,6 +143,7 @@ def download(split_range, shuffle=False):
     val_probs=val_means,
     test_probs=test_means,
     val_targets=val_data_loader.to_array_targets(),
+    error=0.05
 )
 
 # %% [markdown]

fortuna/conformal/__init__.py

@@ -3,11 +3,13 @@
 )
 from fortuna.conformal.classification.adaptive_prediction import (
     AdaptivePredictionConformalClassifier,
+    CVPlusAdaptivePredictionConformalClassifier,
 )
 from fortuna.conformal.classification.maxcovfixprec_binary_classfication import (
     MaxCoverageFixedPrecisionBinaryClassificationCalibrator,
 )
 from fortuna.conformal.classification.simple_prediction import (
+    CVPlusSimplePredictionConformalClassifier,
     SimplePredictionConformalClassifier,
 )
 from fortuna.conformal.multivalid.iterative.classification.binary_multicalibrator import (

fortuna/conformal/classification/adaptive_prediction.py

@@ -1,134 +1,40 @@
-from typing import (
-    List,
-    Optional,
-)
-
 from jax import vmap
 import jax.numpy as jnp
-import numpy as np
 
-from fortuna.conformal.classification.base import ConformalClassifier
+from fortuna.conformal.classification.base import (
+    CVPlusConformalClassifier,
+    SplitConformalClassifier,
+)
 from fortuna.typing import Array
 
 
-class AdaptivePredictionConformalClassifier(ConformalClassifier):
-    def score(
-        self,
-        val_probs: Array,
-        val_targets: Array,
-    ) -> jnp.ndarray:
-        """
-        Compute score function.
-
-        Parameters
-        ----------
-        val_probs: Array
-            A two-dimensional array of class probabilities for each validation data point.
-        val_targets: Array
-            A one-dimensional array of validation target variables.
-
-        Returns
-        -------
-        jnp.ndarray
-            The conformal scores.
-        """
-        if val_probs.ndim != 2:
-            raise ValueError(
-                """`val_probs` must be a two-dimensional array. The first dimension is over the validation
-            inputs. The second is over the classes."""
-            )
+@vmap
+def _score_fn(probs: Array, perm: Array, inv_perm: Array, targets: Array):
+    return jnp.cumsum(probs[perm])[inv_perm[targets]]
 
-        perms = jnp.argsort(val_probs, axis=1)[:, ::-1]
-        inv_perms = jnp.argsort(perms, axis=1)
 
-        @vmap
-        def score_fn(prob, perm, inv_perm, target):
-            sorted_prob = prob[perm]
-            return jnp.cumsum(sorted_prob)[inv_perm[target]]
+def score_fn(
+    probs: Array,
+    targets: Array,
+):
+    perms = jnp.argsort(probs, axis=1)[:, ::-1]
+    inv_perms = jnp.argsort(perms, axis=1)
+    return _score_fn(probs, perms, inv_perms, targets)
 
-        return score_fn(val_probs, perms, inv_perms, val_targets)
 
-    def quantile(
+class AdaptivePredictionConformalClassifier(SplitConformalClassifier):
+    def score_fn(
         self,
-        val_probs: Array,
-        val_targets: Array,
-        error: float = 0.05,
-        scores: Optional[Array] = None,
-    ) -> Array:
-        """
-        Compute a quantile of the scores.
+        probs: Array,
+        targets: Array,
+    ):
+        return score_fn(probs=probs, targets=targets)
 
-        Parameters
-        ----------
-        val_probs: Array
-            A two-dimensional array of class probabilities for each validation data point.
-        val_targets: Array
-            A one-dimensional array of validation target variables.
-        error: float
-            Coverage error. This must be a scalar between 0 and 1, extremes included.
-        scores: Optional[Array]
-            The conformal scores. This should be the output of
-            :meth:`~fortuna.conformal.classification.adaptive_prediction.AdaptivePredictionConformalClassifier.score`.
 
-        Returns
-        -------
-        float
-            The conformal quantiles.
-        """
-        if error < 0 or error > 1:
-            raise ValueError("""`error` must be a scalar between 0 and 1.""")
-
-        if scores is None:
-            scores = self.score(val_probs, val_targets)
-        n = scores.shape[0]
-        return jnp.quantile(scores, jnp.ceil((n + 1) * (1 - error)) / n)
-
-    def conformal_set(
+class CVPlusAdaptivePredictionConformalClassifier(CVPlusConformalClassifier):
+    def score_fn(
         self,
-        val_probs: Array,
-        test_probs: Array,
-        val_targets: Array,
-        error: float = 0.05,
-        quantile: Optional[float] = None,
-    ) -> List[List[int]]:
-        """
-        Coverage set of each of the test inputs, at the desired coverage error.
-
-        Parameters
-        ----------
-        val_probs: Array
-            A two-dimensional array of class probabilities for each validation data point.
-        test_probs: Array
-            A two-dimensional array of class probabilities for each test data point.
-        val_targets: Array
-            A one-dimensional array of validation target variables.
-        error: float
-            The coverage error. This must be a scalar between 0 and 1, extremes included.
-        quantile: Optional[float]
-            Conformal quantiles. This should be the output of
-            :meth:`~fortuna.conformal.classification.adaptive_prediction.AdaptivePredictionConformalClassifier.quantile`.
-
-        Returns
-        -------
-        List[List[int, ...]]
-            The coverage sets.
-        """
-        if test_probs.ndim != 2:
-            raise ValueError(
-                """`test_probs` must be a two-dimensional array. The first dimension is over the validation
-            inputs. The second is over the classes."""
-            )
-
-        if quantile is None:
-            quantile = self.quantile(val_probs, val_targets, error)
-        test_perms = jnp.argsort(test_probs, axis=1)[:, ::-1]
-        test_sorted_probs = vmap(lambda prob, perm: prob[perm])(test_probs, test_perms)
-        sizes = (
-            (test_sorted_probs.cumsum(axis=1) > quantile).astype("int32").argmax(axis=1)
-        )
-
-        sets = np.zeros(len(sizes), dtype=object)
-        for s in jnp.unique(sizes):
-            idx = jnp.where(sizes == s)[0]
-            sets[idx] = test_perms[idx, : s + 1].tolist()
-        return sets.tolist()
+        probs: Array,
+        targets: Array,
+    ):
+        return score_fn(probs=probs, targets=targets)

fortuna/conformal/classification/base.py

@@ -1,6 +1,12 @@
-from typing import List
+import abc
+from typing import (
+    List,
+    Tuple,
+)
 
+from jax import vmap
 import jax.numpy as jnp
+import numpy as np
 
 from fortuna.typing import Array
 
@@ -27,3 +33,100 @@ def is_in(self, values: Array, conformal_sets: List) -> Array:
             An array of ones or zero, indicating whether the values lie within their respective conformal sets.
         """
         return jnp.array([v in s for v, s in zip(values.tolist(), conformal_sets)])
+
+    @abc.abstractmethod
+    def score_fn(
+        self,
+        probs: Array,
+        targets: Array,
+    ):
+        pass
+
+    @staticmethod
+    def _get_conformal_sets_from_scores(
+        val_scores: Array,
+        test_scores: Array,
+        error: float,
+    ) -> List[List[int]]:
+        conds = jnp.sum(val_scores[:, None, None] > test_scores[None], axis=0) < (
+            1 - error
+        ) * (len(val_scores) + 1)
+        sizes = conds.sum(1)
+
+        sets = np.zeros(len(test_scores), dtype=object)
+        for us in jnp.unique(sizes):
+            idx = jnp.where(sizes == us)[0]
+            if us == 0:
+                sets[idx] = [len(idx) * []]
+            else:
+                sets[idx] = np.where(conds[idx])[1].reshape(-1, us).tolist()
+
+        return sets.tolist()
+
+    @abc.abstractmethod
+    def get_scores(self, *args, **kwargs) -> Tuple[Array, Array]:
+        pass
+
+
+class SplitConformalClassifier(ConformalClassifier, abc.ABC):
+    def get_scores(
+        self, val_probs: Array, val_targets: Array, test_probs: Array
+    ) -> Tuple[Array, Array]:
+        val_scores = self.score_fn(val_probs, val_targets)
+        test_scores = vmap(
+            lambda i: self.score_fn(
+                test_probs, i * jnp.ones(len(test_probs), dtype="int32")
+            ),
+            out_axes=1,
+        )(jnp.arange(val_probs.shape[1]))
+        return val_scores, test_scores
+
+    def conformal_set(
+        self, val_probs: Array, val_targets: Array, test_probs: Array, error: float
+    ) -> List[List[int]]:
+        val_scores, test_scores = self.get_scores(
+            val_probs=val_probs, val_targets=val_targets, test_probs=test_probs
+        )
+        return super()._get_conformal_sets_from_scores(
+            val_scores=val_scores, test_scores=test_scores, error=error
+        )
+
+
+class CVPlusConformalClassifier(ConformalClassifier):
+    def conformal_set(
+        self,
+        cross_val_probs: List[Array],
+        cross_val_targets: List[Array],
+        cross_test_probs: List[Array],
+        error: float,
+    ) -> List[List[int]]:
+        val_scores, test_scores = self.get_scores(
+            cross_val_probs=cross_val_probs,
+            cross_val_targets=cross_val_targets,
+            cross_test_probs=cross_test_probs,
+        )
+        return super()._get_conformal_sets_from_scores(
+            val_scores=val_scores, test_scores=test_scores, error=error
+        )
+
+    def get_scores(
+        self,
+        cross_val_probs: List[Array],
+        cross_val_targets: List[Array],
+        cross_test_probs: List[Array],
+    ) -> Tuple[Array, Array]:
+        val_scores, test_scores = [], []
+        for val_probs, val_targets, test_probs in zip(
+            cross_val_probs, cross_val_targets, cross_test_probs
+        ):
+            val_scores.append(self.score_fn(val_probs, val_targets))
+            test_scores.append(
+                vmap(
+                    lambda i: self.score_fn(
+                        test_probs, i * jnp.ones(len(test_probs), dtype="int32")
+                    ),
+                    out_axes=1,
+                )(jnp.arange(cross_val_probs[0].shape[1]))
+            )
+
+        return jnp.concatenate(val_scores), jnp.concatenate(test_scores, axis=0)