import torch
import numpy as np
import pandas as pd
import warnings
from tqdm import tqdm
from typing import Tuple, List, Dict, Union, Optional
from .utils import _check_multihot_labels, _is_tensor
from . import constants
PredictionItem = Union[Tuple[torch.Tensor, float], torch.Tensor]
SingleAlphaOutput = List[PredictionItem]
MultiAlphaOutput = Dict[float, SingleAlphaOutput]
PredictionOutput = Union[SingleAlphaOutput, MultiAlphaOutput]
[docs]
class PredictionRegions:
"""
A container class for conformal prediction results.
This class wraps the p-values generated by the Inductive Conformal Predictor
and provides methods to extract prediction sets at specific significance levels
and evaluate performance metrics.
Parameters
----------
p_values : torch.Tensor
The p-values associated with every possible label combination for each test sample.
Shape: (t_samples, 2^(c_classes)).
combinations : torch.Tensor
The matrix of all possible label combinations corresponding to the columns of `p_values`.
Shape: (2^(c_classes), c_classes).
non_empty_prediction_regions : bool, optional
If True (default), the combination with the highest p-value is returned to ensure non-empty predictions.
Example
--------
>>> import torch
>>> from mulaconf.prediction_regions import PredictionRegions
>>>
>>> # 1. Generate dummy data (100 samples, 5 classes -> 32 combinations)
>>> combinations = torch.cartesian_prod(*[torch.tensor([0, 1])] * 5)
>>> p_values = torch.rand(100, 2**5)
>>>
>>> # 2. Initialize container
>>> prediction_regions_obj = PredictionRegions(p_values, combinations)
"""
def __init__(self, p_values: torch.Tensor, combinations: torch.Tensor, non_empty_prediction_regions:bool=True):
self.device = p_values.device
self.p_values = p_values
self.combinations = combinations.to(self.device)
self.non_empty_prediction_regions = non_empty_prediction_regions
self.valid_tuples_size = None
def get_valid_tuples(self, alpha: float) -> List[torch.Tensor]:
"""
Extracts valid label combinations for a specific significance level, using a batched approach
to prevent System RAM overflow on high-label datasets.
A combination is considered valid if its p-value is greater than the significance level `alpha`.
Parameters
----------
alpha : float
The significance level (error rate). Confidence level = 1 - alpha.
Returns
-------
List[torch.Tensor]
A list where each element corresponds to a test sample. Each element is a Tensor
containing the valid label combinations for that sample.
"""
all_valid_tuples = []
n_samples, n_combinations = self.p_values.shape
batch_size = max(500, constants._REGION_BATCH_SIZE// n_combinations)
for i in range(0, n_samples, batch_size):
batch_p_values = self.p_values[i: i + batch_size]
mask = batch_p_values > alpha
row_counts = mask.sum(dim=1)
if self.non_empty_prediction_regions:
empty_mask = (row_counts == 0)
if empty_mask.any():
max_indices = batch_p_values.argmax(dim=1)
mask[empty_mask, max_indices[empty_mask]] = True
row_counts[empty_mask] = 1
del empty_mask
valid_indices = mask.nonzero()
comb_indices = valid_indices[:, 1]
flat_predictions = self.combinations[comb_indices]
chunk_tuples = list(torch.split(flat_predictions.int(), row_counts.cpu().tolist()))
all_valid_tuples.extend(chunk_tuples)
self.valid_tuples_size = len(all_valid_tuples)
del batch_p_values, mask, row_counts, valid_indices, comb_indices, flat_predictions
return all_valid_tuples
def _get_true_labelsets_p_values_and_indices(self, true_labelsets: torch.Tensor) -> torch.Tensor:
"""
Retrieves the exact column indices of the actual ground truth label combinations.
Parameters
----------
true_labelsets : torch.Tensor
The ground truth labels for the test samples.
Shape: (t_samples, c_classes).
Returns
-------
torch.Tensor
A 1D tensor of p-values and a 1D tensor of indices corresponding to the true labels.
Shape: (t_samples).
"""
device = self.p_values.device
n_classes = self.combinations.shape[1]
n_samples = true_labelsets.shape[0]
powers_of_two = 2 ** torch.arange(n_classes, device=device)
comb_hashes = (self.combinations * powers_of_two).sum(dim=1)
true_hashes = (true_labelsets * powers_of_two).sum(dim=1)
sorted_comb_hashes, sorted_indices = torch.sort(comb_hashes)
positions = torch.searchsorted(sorted_comb_hashes, true_hashes)
positions = torch.clamp(positions, max=len(sorted_comb_hashes) - 1)
matched_indices = sorted_indices[positions]
return self.p_values[torch.arange(n_samples), matched_indices], matched_indices
def _parse_significance_level(
self,
significance_level: Union[float, List[float], np.ndarray, pd.Series, torch.Tensor]
) -> Tuple[List[float], bool]:
"""
Parse significance level input into a list of alphas.
Returns
-------
Tuple[List[float], bool]
A tuple of (alphas list, is_scalar flag).
"""
is_scalar = False
if torch.is_tensor(significance_level):
if significance_level.ndim == 0:
is_scalar = True
alphas = [significance_level.item()]
else:
alphas = significance_level.tolist()
elif isinstance(significance_level, np.ndarray):
if significance_level.ndim == 0:
is_scalar = True
alphas = [significance_level.item()]
else:
alphas = significance_level.tolist()
elif isinstance(significance_level, pd.Series):
alphas = significance_level.tolist()
elif isinstance(significance_level, (list, tuple)):
alphas = list(significance_level)
else:
is_scalar = True
alphas = [significance_level]
alphas = [round(float(a), 3) for a in alphas]
for a in alphas:
if not (0 <= a <= 1):
raise ValueError(f"Significance level must be strictly between 0 and 1, got {a}")
return alphas, is_scalar
@torch.no_grad()
def __call__(self,
significance_level: Optional[Union[float, List[float], np.ndarray, pd.Series, torch.Tensor]] = None
) -> PredictionOutput:
"""
Extracts prediction sets for one or multiple significance levels.
Parameters
----------
significance_level : float or List[float]
The significance level(s) (alpha) for which to return prediction regions.
Must be strictly between 0 and 1.
Returns
-------
Union[List[torch.Tensor], Dict[float, List[torch.Tensor]]]
- If ``significance_level`` is a scalar: Returns a list of valid tuples.
- If ``significance_level`` is a list: Returns a dictionary mapping alpha values to lists of valid tuples.
Raises
------
ValueError
If ``significance_level`` is None or out of range [0, 1].
Example
--------
>>> # Scalar alpha
>>> prediction_sets_lst = prediction_regions_obj(0.1)
>>>
>>> # Multiple alphas
>>> prediction_sets_dict = prediction_regions_obj([0.05, 0.1])
"""
if significance_level is None:
raise ValueError("significance_level must be specified")
else:
print(f'Returning prediction regions for significance level {significance_level}.')
alphas, is_scalar = self._parse_significance_level(significance_level)
iterator = tqdm(alphas, desc="Extracting Prediction Sets") if len(alphas) > 1 else alphas
results = {}
for alpha in iterator:
results[alpha] = self.get_valid_tuples(alpha)
if is_scalar:
return results[alphas[0]]
else:
return results
@torch.no_grad()
def evaluate(self,
return_true_label_p_value: bool = True,
return_coverage: bool = True,
return_n_criterion: bool = True,
return_s_criterion: bool = True,
return_observed_fuzziness: bool = True,
return_observed_excess: bool = True,
*,
true_labelsets: Optional[Union[np.ndarray, torch.Tensor, pd.DataFrame, pd.Series, List]]=None,
significance_level: Optional[Union[float, List[float]]] = None,
) -> Optional[Dict]:
"""
Evaluates the conformal predictor using standard metrics.
Calculates validity (coverage) and efficiency (N-criterion, S-criterion, Observed fuzziness, Observed excess) metrics.
Parameters
----------
return_true_label_p_value : bool, default=True
Whether to return the p-values of the true class (requires `true labelsets`).
return_coverage : bool, default=True
Whether to calculate empirical coverage (requires `true labelsets` and `significance_level`).
return_n_criterion : bool, default=True
Whether to calculate average set size (N-criterion) (requires `true labelsets` and `significance_level`).
return_s_criterion : bool, default=True
Whether to calculate the S-criterion (sum of p-values), a threshold-independent efficiency metric.
return_observed_fuzziness : bool, default=True
Whether to calculate the observed fuzziness (requires `true labelsets`).
return_observed_excess : bool, default=True
Whether to calculate the observed excess (requires `true labelsets` and `significance_level`).
true_labelsets : torch.Tensor or array-like, optional
The ground truth labels. Required for coverage and p-value metrics.
significance_level : float or List[float], optional
The alpha level(s) to evaluate coverage and N-criterion.
.. note::
The ``true_labelsets`` and the ``significance_level`` must be passed as a keyword argument.
Returns
-------
Dict
A dictionary containing the requested evaluation metrics.
Depending on the boolean flags provided, the returned dictionary will contain
the following specific string keys:
- ``'true_labels_p_values'`` : A 1D tensor of the exact p-values for the true labels.
- ``'coverage'`` : The empirical coverage (validity).
- ``'n_criterion'`` : The average prediction set size.
- ``'s_criterion'`` : The average sum of all p-values.
- ``'observed_fuzziness'`` : The average sum of the p-values of the false labels.
- ``'observed_excess'`` : The average sum of false labels included in the prediction region.
If ``significance_level`` is a list of multiple alphas, it returns a nested dictionary
where the top-level keys are the alpha values (floats), and the values are metric
dictionaries containing the keys listed above.
Raises
------
ValueError
If `true_labelsets`` and the ``significance_level`` is not provided.
ValueError
If ``true_labelsets`` shape does not match the number of classes.
Example
--------
>>> # Generate dummy test labels
>>> y_test = torch.rand(10, 5)
>>>
>>> # Evaluate specific alpha
>>> metrics = prediction_regions_obj.evaluate(true_labelsets=y_test, significance_level=0.1)
>>>
>>> # Evaluate multiple alphas
>>> metrics_multi = prediction_regions_obj.evaluate(true_labelsets=y_test, significance_level=[0.05, 0.1])
"""
if true_labelsets is None and significance_level is None:
raise ValueError("You must provide either true_labelsets or significance_level, or both.")
if true_labelsets is None:
out = {}
if return_s_criterion: out['s_criterion'] = self.p_values.sum(dim=1).mean().item()
if return_true_label_p_value or return_coverage or return_n_criterion or return_observed_fuzziness or return_observed_excess:
warnings.warn("Missing true_labelsets. Only s_criterion calculated.", RuntimeWarning)
return out
true_labelsets = _check_multihot_labels(true_labelsets)
true_labelsets = _is_tensor(true_labelsets).to(self.device)
if true_labelsets.shape[1] != self.combinations.shape[1]:
raise ValueError("True labels must have the same number of columns as the number of classes.")
true_labelsets_p_values, true_labelsets_indices = self._get_true_labelsets_p_values_and_indices(true_labelsets)
if significance_level is not None:
alphas, is_scalar = self._parse_significance_level(significance_level)
def _evaluate_metrics(alpha):
out = {}
if return_true_label_p_value:
out['true_labels_p_values'] = true_labelsets_p_values
if return_coverage:
out['coverage'] = (true_labelsets_p_values > alpha).float().mean().item()
if return_n_criterion or return_observed_excess:
n_samples, n_combinations = self.p_values.shape
batch_size = max(1, constants._REGION_BATCH_SIZE // n_combinations)
total_sizes = 0
total_excess = 0
for i in range(0, n_samples, batch_size):
batch_p_values = self.p_values[i: i + batch_size]
batch_true_indices = true_labelsets_indices[i: i + batch_size]
mask = batch_p_values > alpha
counts = mask.sum(dim=1)
if self.non_empty_prediction_regions:
empty_mask = (counts == 0)
if empty_mask.any():
max_indices = batch_p_values.argmax(dim=1)
mask[empty_mask, max_indices[empty_mask]] = True
counts[empty_mask] = 1
total_sizes += counts.sum().item()
if return_observed_excess:
true_label_in_set = mask[torch.arange(len(batch_p_values)), batch_true_indices].int()
excess = counts - true_label_in_set
total_excess += excess.sum().item()
if return_n_criterion:
out['n_criterion'] = total_sizes / n_samples
if return_observed_excess:
out['observed_excess'] = total_excess / n_samples
if return_s_criterion:
out['s_criterion'] = self.p_values.sum(dim=1).mean().item()
if return_observed_fuzziness:
sum_p = self.p_values.sum(dim=1)
of_vals = sum_p - true_labelsets_p_values
out['observed_fuzziness'] = of_vals.mean().item()
return out
if is_scalar:
return _evaluate_metrics(alphas[0])
else:
return {a: _evaluate_metrics(a) for a in tqdm(alphas, desc="Evaluating Metrics")}
else:
out = {}
if return_true_label_p_value: out['true_labelsets_p_values'] = true_labelsets_p_values
if return_s_criterion: out['s_criterion'] = self.p_values.sum(dim=1).mean().item()
if return_coverage or return_n_criterion:
warnings.warn("Coverage/N-criterion require a significance_level.", RuntimeWarning)
return out