gate/certificate.py

"""Per-image coverage certificate (formalization §5) — the clinical differentiator.

Each inference emits a label-free certificate, not just a throughput number:
    Cert(x) = < delta_C = C*(x) - C(S;x),  k = |S|,  mu,  retained lesion-subspace dirs >
delta_C <= epsilon is the audited guarantee that pruning did not collapse lesion-relevant
directions for THIS image. The same importance map that gated compute is the audited record,
so compute-optimality and audit-faithfulness are tied by construction.
"""
from __future__ import annotations

from dataclasses import dataclass, field

import torch

from .lagrangian import PrunerResult


@dataclass
class Certificate:
    delta_C: float                       # coverage drop under the applied mask
    k: int                               # retained budget |S|
    mu: float                            # dual value (marginal token cost of coverage)
    epsilon: float                       # coverage floor
    satisfied: bool                      # delta_C <= epsilon (the audited guarantee)
    n_tokens: int                        # original token count
    retained_dirs: torch.Tensor | None = None  # lesion-subspace dirs preserved by S
    # Phase 6 conformal head will add: guaranteed_coverage_prob, alpha
    extra: dict = field(default_factory=dict)

    def as_dict(self) -> dict:
        return {"delta_C": self.delta_C, "k": self.k, "mu": self.mu,
                "epsilon": self.epsilon, "satisfied": self.satisfied,
                "n_tokens": self.n_tokens, "retention_ratio": self.k / max(1, self.n_tokens),
                **self.extra}


def certificate_from_result(res: PrunerResult, epsilon: float, n_tokens: int,
                            Z: torch.Tensor | None = None,
                            P_L: torch.Tensor | None = None,
                            top_dirs: int = 8) -> Certificate:
    """Build a Certificate from a pruner result; optionally record the top retained
    lesion-subspace directions (principal axes of P_L applied to the retained tokens)."""
    retained = None
    if Z is not None and P_L is not None and res.k > 0:
        Z_S = (Z.float() * res.mask[:, None]) @ P_L.to(Z.device).float().T
        try:
            _, _, Vt = torch.linalg.svd(Z_S, full_matrices=False)
            retained = Vt[:top_dirs].detach().cpu()
        except Exception:
            retained = None
    return Certificate(
        delta_C=res.delta_C, k=res.k, mu=res.mu, epsilon=epsilon,
        satisfied=res.satisfied, n_tokens=n_tokens, retained_dirs=retained,
    )