processed/covtoken_code/arch/conformal_head.py

"""Phase 6 — conformal coverage certificate (highest novelty, lowest compute).

Turns the per-image coverage certificate (gate/certificate.py) into a CALIBRATED,
distribution-free guarantee via split conformal prediction. No pretraining; pure post-hoc
calibration on a held-out split over the frozen-backbone inference-time pruner.

Setup. For each image we prune to a retained set S and obtain delta_C = C*(x) - C(S;x).
We want a guarantee about a downstream lesion-coverage quantity Y(x) in [0,1] (e.g. the
fraction of lesion mass retained, or a probe's lesion-detection score on S). Using a
nonconformity score s(x) = 1 - Y(x) (higher = worse lesion preservation), split conformal
gives a threshold q_hat (the ceil((n+1)(1-alpha))/n empirical quantile of calibration
scores) such that, for an exchangeable test point,
        P( Y(x_test) >= 1 - q_hat ) >= 1 - alpha .
So the certificate emits a guaranteed_coverage_lowerbound = 1 - q_hat with nominal coverage
1-alpha. Gate 6 checks the empirical coverage lands in [1-alpha-tol, 1].

This is label-free at INFERENCE: q_hat is fixed once on a calibration split; masks are used
ONLY to compute Y on the calibration set (eval-only), never in subspace construction.
"""
from __future__ import annotations

from dataclasses import dataclass

import numpy as np


def conformal_quantile(cal_scores: np.ndarray, alpha: float = 0.1) -> float:
    """Split-conformal threshold q_hat: the ceil((n+1)(1-alpha))/n empirical quantile."""
    cal_scores = np.asarray(cal_scores, float)
    n = len(cal_scores)
    if n == 0:
        return 1.0
    level = min(1.0, np.ceil((n + 1) * (1 - alpha)) / n)
    return float(np.quantile(cal_scores, level, method="higher"))


@dataclass
class ConformalCertificate:
    alpha: float                      # miscoverage level (nominal coverage = 1-alpha)
    q_hat: float                      # calibrated nonconformity threshold
    guaranteed_coverage: float        # 1 - q_hat : guaranteed lesion-coverage lower bound
    n_cal: int

    def certify(self, y: float) -> dict:
        """Per-image: does the observed lesion-coverage y meet the guaranteed lower bound?"""
        return {"y": float(y), "guaranteed_coverage": self.guaranteed_coverage,
                "alpha": self.alpha, "covered": bool(y >= self.guaranteed_coverage)}


def calibrate(cal_y: np.ndarray, alpha: float = 0.1) -> ConformalCertificate:
    """Fit the conformal certificate on calibration lesion-coverage values y in [0,1]."""
    cal_y = np.asarray(cal_y, float)
    scores = 1.0 - cal_y                      # nonconformity = lesion mass LOST
    q_hat = conformal_quantile(scores, alpha)
    return ConformalCertificate(alpha=alpha, q_hat=q_hat,
                                guaranteed_coverage=float(1.0 - q_hat), n_cal=len(cal_y))


def empirical_coverage(test_y: np.ndarray, cert: ConformalCertificate) -> float:
    """Fraction of test images whose lesion-coverage meets the guaranteed lower bound."""
    test_y = np.asarray(test_y, float)
    return float(np.mean(test_y >= cert.guaranteed_coverage))