Selective prediction

neureptrace.decoding.selective_prediction converts probability traces into predictions plus an abstention mask.

Fixed confidence, entropy, and margin thresholds are ordinary post-hoc prediction rules. A requested target coverage sets the confidence threshold from the unlabeled probability batch and should be reported as an unlabeled target-adaptive thresholding step.

The public API intentionally has no label-vector argument.

Typical usage:

from neureptrace.decoding.selective_prediction import selective_predict

result = selective_predict(
    probabilities,
    classes=classes,
    confidence_threshold=0.75,
)

predictions = result.predictions
selected = result.selected_mask

neureptrace.decoding.selective_prediction

Selective prediction utilities for probability traces.

The helpers in this module turn classifier probability rows into predictions plus an abstention mask. Fixed thresholds are ordinary post-hoc prediction rules. A coverage-calibrated threshold uses the unlabeled target probability distribution itself and is therefore marked as an unlabeled target-adaptive Category-2 step.

No label vector is accepted by the public prediction API.

SelectivePredictionResult dataclass

Predictions, retained-row mask, and uncertainty diagnostics.

Source code in src/neureptrace/decoding/selective_prediction.py

@dataclass(frozen=True, slots=True)
class SelectivePredictionResult:
    """Predictions, retained-row mask, and uncertainty diagnostics."""

    predictions: np.ndarray
    selected_mask: np.ndarray
    probabilities: np.ndarray
    confidence: np.ndarray
    entropy: np.ndarray
    margin: np.ndarray
    threshold: float | None
    metadata: dict[str, Any] = field(default_factory=dict)

    @property
    def coverage(self) -> float:
        """Fraction of rows retained for prediction."""

        return float(np.mean(self.selected_mask)) if self.selected_mask.size else 0.0

coverage property

Fraction of rows retained for prediction.

selective_predict(probabilities, *, classes=None, confidence_threshold=None, max_entropy=None, min_margin=None, target_coverage=None, epsilon=DEFAULT_EPSILON)

Predict with optional abstention from probability rows.

Parameters:

Name	Type	Description	Default
probabilities	Sequence[Sequence[float]] | ndarray	Row-wise class probabilities or non-negative scores. Rows are normalized defensively before metrics are computed.	required
classes	Sequence[Any] | ndarray | None	Optional class labels for prediction output. If omitted, dense integer class ids are used.	None
confidence_threshold	float | str | None	Fixed minimum top-class probability required for selection.	None
max_entropy	float | str | None	Fixed maximum entropy allowed for selection.	None
min_margin	float | str | None	Fixed minimum gap between the largest and second-largest probability.	None
target_coverage	float | str | None	Optional desired retained fraction in (0, 1]. When supplied, the confidence threshold is set from the unlabeled probability batch so that approximately this coverage is retained. This is a Category-2 threshold selection rule.	None
epsilon	float | str	Numerical floor for row normalization and entropy.	DEFAULT_EPSILON

Returns:

Type	Description
SelectivePredictionResult	Predicted labels, selection mask, and uncertainty metrics.

Source code in src/neureptrace/decoding/selective_prediction.py

def selective_predict(
    probabilities: Sequence[Sequence[float]] | np.ndarray,
    *,
    classes: Sequence[Any] | np.ndarray | None = None,
    confidence_threshold: float | str | None = None,
    max_entropy: float | str | None = None,
    min_margin: float | str | None = None,
    target_coverage: float | str | None = None,
    epsilon: float | str = DEFAULT_EPSILON,
) -> SelectivePredictionResult:
    """Predict with optional abstention from probability rows.

    Parameters
    ----------
    probabilities:
        Row-wise class probabilities or non-negative scores.  Rows are normalized
        defensively before metrics are computed.
    classes:
        Optional class labels for prediction output.  If omitted, dense integer
        class ids are used.
    confidence_threshold:
        Fixed minimum top-class probability required for selection.
    max_entropy:
        Fixed maximum entropy allowed for selection.
    min_margin:
        Fixed minimum gap between the largest and second-largest probability.
    target_coverage:
        Optional desired retained fraction in ``(0, 1]``.  When supplied, the
        confidence threshold is set from the unlabeled probability batch so that
        approximately this coverage is retained.  This is a Category-2 threshold
        selection rule.
    epsilon:
        Numerical floor for row normalization and entropy.

    Returns
    -------
    SelectivePredictionResult
        Predicted labels, selection mask, and uncertainty metrics.
    """

    eps = _positive_float(epsilon, name="epsilon")
    probs = normalize_probability_rows(probabilities, epsilon=eps)
    class_values = _classes(classes, n_classes=probs.shape[1])
    confidence = np.max(probs, axis=1)
    prediction_indices = np.argmax(probs, axis=1)
    predictions = class_values[prediction_indices]
    entropy = probability_entropy(probs, epsilon=eps)
    margin = probability_margin(probs)

    threshold = None if confidence_threshold is None else _unit_interval_float(confidence_threshold, name="confidence_threshold")
    coverage_requested = None if target_coverage is None else _open_unit_interval_float(target_coverage, name="target_coverage", include_one=True)
    if coverage_requested is not None:
        threshold = confidence_threshold_for_coverage(confidence, target_coverage=coverage_requested)
    max_entropy_threshold = None if max_entropy is None else _nonnegative_float(max_entropy, name="max_entropy")
    min_margin_threshold = None if min_margin is None else _unit_interval_float(min_margin, name="min_margin")

    selected = np.ones(probs.shape[0], dtype=bool)
    if threshold is not None:
        selected &= confidence >= threshold
    if max_entropy_threshold is not None:
        selected &= entropy <= max_entropy_threshold
    if min_margin_threshold is not None:
        selected &= margin >= min_margin_threshold

    metadata = _metadata(
        n_rows=probs.shape[0],
        n_classes=probs.shape[1],
        selected_count=int(np.count_nonzero(selected)),
        confidence_threshold=threshold,
        max_entropy=max_entropy_threshold,
        min_margin=min_margin_threshold,
        target_coverage=coverage_requested,
    )
    return SelectivePredictionResult(
        predictions=predictions,
        selected_mask=selected,
        probabilities=probs.astype(np.float32, copy=False),
        confidence=confidence.astype(float, copy=False),
        entropy=entropy.astype(float, copy=False),
        margin=margin.astype(float, copy=False),
        threshold=threshold,
        metadata=metadata,
    )

normalize_probability_rows(probabilities, *, epsilon=DEFAULT_EPSILON)

Return finite row-normalized probabilities.

Source code in src/neureptrace/decoding/selective_prediction.py

def normalize_probability_rows(probabilities: Sequence[Sequence[float]] | np.ndarray, *, epsilon: float = DEFAULT_EPSILON) -> np.ndarray:
    """Return finite row-normalized probabilities."""

    matrix = np.asarray(probabilities, dtype=float)
    if matrix.ndim != 2:
        raise ValueError("probabilities must be a two-dimensional matrix.")
    if matrix.shape[0] < 1 or matrix.shape[1] < 2:
        raise ValueError("probabilities must contain at least one row and two classes.")
    if not np.all(np.isfinite(matrix)) or np.any(matrix < 0.0):
        raise ValueError("probabilities must contain finite non-negative values.")
    eps = _positive_float(epsilon, name="epsilon")
    matrix = np.maximum(matrix, eps)
    row_sums = np.sum(matrix, axis=1, keepdims=True)
    if np.any(row_sums <= 0.0):
        raise ValueError("probability rows must have positive mass.")
    return matrix / row_sums

probability_entropy(probabilities, *, epsilon=DEFAULT_EPSILON)

Return row-wise Shannon entropy.

Source code in src/neureptrace/decoding/selective_prediction.py

def probability_entropy(probabilities: Sequence[Sequence[float]] | np.ndarray, *, epsilon: float = DEFAULT_EPSILON) -> np.ndarray:
    """Return row-wise Shannon entropy."""

    eps = _positive_float(epsilon, name="epsilon")
    probs = normalize_probability_rows(probabilities, epsilon=eps)
    return -np.sum(probs * np.log(np.maximum(probs, eps)), axis=1)

probability_margin(probabilities)

Return top-one minus top-two probability margin per row.

Source code in src/neureptrace/decoding/selective_prediction.py

def probability_margin(probabilities: Sequence[Sequence[float]] | np.ndarray) -> np.ndarray:
    """Return top-one minus top-two probability margin per row."""

    probs = normalize_probability_rows(probabilities)
    sorted_probs = np.sort(probs, axis=1)
    return sorted_probs[:, -1] - sorted_probs[:, -2]

confidence_threshold_for_coverage(confidence, *, target_coverage)

Return the smallest confidence among the top-coverage rows.

Source code in src/neureptrace/decoding/selective_prediction.py

def confidence_threshold_for_coverage(confidence: Sequence[float] | np.ndarray, *, target_coverage: float | str) -> float:
    """Return the smallest confidence among the top-coverage rows."""

    values = np.asarray(confidence, dtype=float).reshape(-1)
    if values.size < 1:
        raise ValueError("confidence must contain at least one row.")
    if not np.all(np.isfinite(values)):
        raise ValueError("confidence must be finite.")
    coverage = _open_unit_interval_float(target_coverage, name="target_coverage", include_one=True)
    keep = max(1, int(np.ceil(values.size * coverage)))
    ordered = np.sort(values)[::-1]
    return float(ordered[keep - 1])