Binary risk control - multi risk (#762)

* test_check_risks_targets_same_len working

* __init__ and _set_best_predict_param_choice handle multi risk

* update and rename _get_risks_and_effective_sample_sizes_per_param for multi risk

* update _set_best_predict_param for multi risk

* fix alpha

* self._risk is now always a list

* lists of len 1 = mono risk

* test updated to new behavior for binary: r_hat and n_obs now always have 2 dimensions (even in mono risk)

* simplify ltt for binary

* keep returning lists of lists for ltt

* update ltt_procedure and its calls and tests for multi risk

* ltt now fails when bad shape of inputs for multi risk

* add unit tests ltt multi risk

* ensure compatibility with python<3.10

Author: allglc

mapie/control_risk/ltt.py

@@ -1,5 +1,5 @@
 import warnings
-from typing import Any, List, Tuple, Union
+from typing import Any, List, Tuple
 
 import numpy as np
 
@@ -12,7 +12,7 @@ def ltt_procedure(
     r_hat: NDArray,
     alpha_np: NDArray,
     delta: float,
-    n_obs: Union[int, NDArray],
+    n_obs: NDArray,
     binary: bool = False,
 ) -> List[List[Any]]:
     """
@@ -24,28 +24,36 @@ def ltt_procedure(
         - Apply a family wise error rate algorithm, here Bonferonni correction
         - Return the index lambdas that give you the control at alpha level
 
+    Note that in the case of multi-risk, the arrays r_hat, alpha_np, and n_obs
+    should have the same length for the first dimension which corresponds
+    to the number of risks. In the case of a single risk, the length should be 1.
+
     Parameters
     ----------
-    r_hat: NDArray of shape (n_lambdas, ).
+    r_hat: NDArray of shape (n_risks, n_lambdas).
         Empirical risk with respect to the lambdas.
         Here lambdas are thresholds that impact decision-making,
         therefore empirical risk.
 
-    alpha_np: NDArray of shape (n_alpha, ).
+    alpha_np: NDArray of shape (n_risks, n_alpha).
         Contains the different alphas control level.
         The empirical risk should be less than alpha with
         probability 1-delta.
+        Note: MAPIE 1.2 does not support multiple risks and multiple alphas
+        simultaneously.
+        For PrecisionRecallController, the shape should be (1, n_alpha).
+        For BinaryClassificationController, the shape should be (n_risks, 1).
 
     delta: float.
         Probability of not controlling empirical risk.
         Correspond to proportion of failure we don't
         want to exceed.
 
-    n_obs: Union[int, NDArray]
+    n_obs: NDArray of shape (n_risks, n_lambdas).
         Correspond to the number of observations used to compute the risk.
         In the case of a conditional loss, n_obs must be the
         number of effective observations used to compute the empirical risk
-        for each lambda, hence of shape (n_lambdas, ).
+        for each lambda.
 
     binary: bool, default=False
         Must be True if the loss associated to the risk is binary.
@@ -62,11 +70,19 @@ def ltt_procedure(
     M. I., & Lei, L. (2021). Learn then test:
     "Calibrating predictive algorithms to achieve risk control".
     """
-    p_values = compute_hoeffding_bentkus_p_value(r_hat, n_obs, alpha_np, binary)
+    if not (r_hat.shape[0] == n_obs.shape[0] == alpha_np.shape[0]):
+        raise ValueError(
+            "r_hat, n_obs, and alpha_np must have the same length."
+        )
+    p_values = np.array([
+        compute_hoeffding_bentkus_p_value(r_hat_i, n_obs_i, alpha_np_i, binary)
+        for r_hat_i, n_obs_i, alpha_np_i in zip(r_hat, n_obs, alpha_np)
+    ])
+    p_values = p_values.max(axis=0)  # take max over risks (no effect if mono risk)
     N = len(p_values)
     valid_index = []
-    for i in range(len(alpha_np)):
-        l_index = np.where(p_values[:, i] <= delta/N)[0].tolist()
+    for i in range(alpha_np.shape[1]):
+        l_index = np.nonzero(p_values[:, i] <= delta/N)[0].tolist()
         valid_index.append(l_index)
     return valid_index
 

mapie/risk_control.py

@@ -471,7 +471,7 @@ def _transform_pred_proba(
             y_pred_proba_array = y_pred_proba
         else:
             y_pred_proba_stacked = np.stack(
-                y_pred_proba,  # type: ignore
+                y_pred_proba,
                 axis=0
             )[:, :, 1]
             y_pred_proba_array = np.moveaxis(y_pred_proba_stacked, 0, -1)
@@ -669,7 +669,10 @@ def predict(
             self.n_obs = len(self.risks)
             self.r_hat = self.risks.mean(axis=0)
             self.valid_index = ltt_procedure(
-                self.r_hat, alpha_np, cast(float, delta), self.n_obs
+                np.expand_dims(self.r_hat, axis=0),
+                np.expand_dims(alpha_np, axis=0),
+                cast(float, delta),
+                np.expand_dims(np.array([self.n_obs]), axis=0)
             )
             self._check_valid_index(alpha_np)
             self.lambdas_star, self.r_star = find_lambda_control_star(
@@ -865,16 +868,20 @@ class BinaryClassificationController:
         predict_proba method of a fitted binary classifier.
         Its output signature must be of shape (len(X), 2)
 
-    risk : BinaryClassificationRisk
+    risk : Union[BinaryClassificationRisk, List[BinaryClassificationRisk]]
         The risk or performance metric to control.
         Valid options:
 
         - An existing risk defined in `mapie.risk_control` (e.g. precision, recall,
           accuracy, false_positive_rate)
         - A custom instance of BinaryClassificationRisk object
 
-    target_level : float
+        Can be a list of risks in the case of multi risk control.
+
+    target_level : Union[float, List[float]]
         The maximum risk level (or minimum performance level). Must be between 0 and 1.
+        Can be a list of target levels in the case of multi risk control (length should
+        match the length of the risks list).
 
     confidence_level : float, default=0.9
         The confidence level with which the risk (or performance) is controlled.
@@ -950,18 +957,19 @@ class BinaryClassificationController:
     def __init__(
         self,
         predict_function: Callable[[ArrayLike], NDArray],
-        risk: BinaryClassificationRisk,
-        target_level: float,
+        risk: Union[BinaryClassificationRisk, List[BinaryClassificationRisk]],
+        target_level: Union[float, List[float]],
         confidence_level: float = 0.9,
         best_predict_param_choice: Union[
             Literal["auto"], BinaryClassificationRisk] = "auto",
     ):
+        self.is_multi_risk = self._check_if_multi_risk_control(risk, target_level)
         self._predict_function = predict_function
-        self._risk = risk
-        if self._risk.higher_is_better:
-            self._alpha = 1 - target_level
-        else:
-            self._alpha = target_level
+        self._risk = risk if isinstance(risk, list) else [risk]
+        target_level_list = (
+            target_level if isinstance(target_level, list) else [target_level]
+        )
+        self._alpha = self._convert_target_level_to_alpha(target_level_list)
         self._delta = 1 - confidence_level
 
         self._best_predict_param_choice = self._set_best_predict_param_choice(
@@ -1006,20 +1014,16 @@ def calibrate(  # pragma: no cover
             self._predict_params
         )
 
-        risks_and_eff_sizes = self._get_risks_and_effective_sample_sizes_per_param(
+        risk_values, eff_sample_sizes = self._get_risk_values_and_eff_sample_sizes(
             y_calibrate_,
             predictions_per_param,
             self._risk
         )
-
-        risks_per_param = risks_and_eff_sizes[:, 0]
-        eff_sample_sizes_per_param = risks_and_eff_sizes[:, 1]
-
         valid_params_index = ltt_procedure(
-            risks_per_param,
-            np.array([self._alpha]),
+            risk_values,
+            np.expand_dims(self._alpha, axis=1),
             self._delta,
-            eff_sample_sizes_per_param,
+            eff_sample_sizes,
             True,
         )[0]
 
@@ -1072,16 +1076,20 @@ def _set_best_predict_param_choice(
             Literal["auto"], BinaryClassificationRisk] = "auto",
     ) -> BinaryClassificationRisk:
         if best_predict_param_choice == "auto":
-            try:
-                return self._best_predict_param_choice_map[
-                    self._risk
-                ]
-            except KeyError:
-                raise ValueError(
-                    "When best_predict_param_choice is 'auto', "
-                    "risk must be one of the risks defined in mapie.risk_control"
-                    "(e.g. precision, accuracy, false_positive_rate)."
-                )
+            if self.is_multi_risk:
+                # when multi risk, we minimize the first risk in the list
+                return self._risk[0]
+            else:
+                try:
+                    return self._best_predict_param_choice_map[
+                        self._risk[0]
+                    ]
+                except KeyError:
+                    raise ValueError(
+                        "When best_predict_param_choice is 'auto', "
+                        "risk must be one of the risks defined in mapie.risk_control"
+                        "(e.g. precision, accuracy, false_positive_rate)."
+                    )
         else:
             return best_predict_param_choice
 
@@ -1099,29 +1107,37 @@ def _set_best_predict_param(
         predictions_per_param: NDArray,
         valid_params_index: List[Any],
     ):
-        secondary_risks_per_param = \
-            self._get_risks_and_effective_sample_sizes_per_param(
+        secondary_risks_per_param, _ = self._get_risk_values_and_eff_sample_sizes(
                 y_calibrate_,
                 predictions_per_param[valid_params_index],
-                self._best_predict_param_choice
-            )[:, 0]
+                [self._best_predict_param_choice]
+            )
 
         self.best_predict_param = self.valid_predict_params[
             np.argmin(secondary_risks_per_param)
         ]
 
     @staticmethod
-    def _get_risks_and_effective_sample_sizes_per_param(
+    def _get_risk_values_and_eff_sample_sizes(
         y_true: NDArray,
         predictions_per_param: NDArray,
-        risk: BinaryClassificationRisk,
-    ) -> NDArray:
-        return np.array(
-            [risk.get_value_and_effective_sample_size(
-                y_true,
-                predictions
-            ) for predictions in predictions_per_param]
-        )
+        risks: List[BinaryClassificationRisk],
+    ) -> Tuple[NDArray, NDArray]:
+        """
+        Compute the values of risks and effective sample sizes for multiple risks
+        and for multiple parameter values.
+        Returns arrays with shape (n_risks, n_params).
+        """
+        risks_values_and_eff_sizes = np.array([
+            [risk.get_value_and_effective_sample_size(y_true, predictions)
+             for predictions in predictions_per_param]
+            for risk in risks
+        ])
+
+        risk_values = risks_values_and_eff_sizes[:, :, 0]
+        effective_sample_sizes = risks_values_and_eff_sizes[:, :, 1]
+
+        return risk_values, effective_sample_sizes
 
     def _get_predictions_per_param(self, X: ArrayLike, params: NDArray) -> NDArray:
         try:
@@ -1148,3 +1164,42 @@ def _get_predictions_per_param(self, X: ArrayLike, params: NDArray) -> NDArray:
             else:
                 raise
         return (predictions_proba[:, np.newaxis] >= params).T.astype(int)
+
+    def _convert_target_level_to_alpha(self, target_level: List[float]) -> NDArray:
+        alpha = []
+        for risk, target in zip(self._risk, target_level):
+            if risk.higher_is_better:
+                alpha.append(1 - target)
+            else:
+                alpha.append(target)
+        return np.array(alpha)
+
+    @staticmethod
+    def _check_if_multi_risk_control(
+        risk: Union[BinaryClassificationRisk, List[BinaryClassificationRisk]],
+        target_level: Union[float, List[float]],
+    ) -> bool:
+        """
+        Check if we are in a multi risk setting and if inputs types are correct.
+        """
+        if (
+            isinstance(risk, list) and isinstance(target_level, list)
+            and len(risk) == len(target_level)
+            and len(risk) > 0
+        ):
+            if len(risk) == 1:
+                return False
+            else:
+                return True
+        elif (
+            isinstance(risk, BinaryClassificationRisk)
+            and isinstance(target_level, float)
+        ):
+            return False
+        else:
+            raise ValueError(
+                "If you provide a list of risks, you must provide "
+                "a list of target levels of the same length and vice versa. "
+                "If you provide a single BinaryClassificationRisk risk, "
+                "you must provide a single float target level."
+            )