infer.calibration¶

@runtime_checkable
class Calibrator(Protocol):
    """Minimal contract for a probability calibrator.

    Implementations must:

    * Preserve the shape of *core_probs*.
    * Ensure the output sums to 1.0 along the class axis.
    * Be deterministic (same input → same output).
    """

    def calibrate(self, core_probs: np.ndarray) -> np.ndarray:
        """Adjust a probability vector (or matrix) and return calibrated probabilities.

        Args:
            core_probs: Array of shape ``(n_classes,)`` or ``(n_rows, n_classes)``.

        Returns:
            Calibrated probabilities with the same shape.
        """
        ...  # pragma: no cover

@dataclass(frozen=True, eq=False)
class IdentityCalibrator:
    """No-op calibrator that returns probabilities unchanged."""

    def calibrate(self, core_probs: np.ndarray) -> np.ndarray:
        return core_probs

@dataclass(frozen=True, eq=False)
class TemperatureCalibrator:
    """Scale logits by a learned temperature before softmax.

    A temperature > 1 softens the distribution (less confident);
    a temperature < 1 sharpens it (more confident).

    Args:
        temperature: Positive scalar.  Defaults to 1.0 (identity).
    """

    temperature: float = 1.0

    def __post_init__(self) -> None:
        if self.temperature <= 0:
            raise ValueError(f"temperature must be positive, got {self.temperature}")

    def calibrate(self, core_probs: np.ndarray) -> np.ndarray:
        if self.temperature == 1.0:
            return core_probs

        eps = 1e-12
        logits = np.log(np.clip(core_probs, eps, None))
        scaled = logits / self.temperature

        # Numerically stable softmax
        shifted = scaled - scaled.max(axis=-1, keepdims=True)
        exp_vals = np.exp(shifted)
        return exp_vals / exp_vals.sum(axis=-1, keepdims=True)

@dataclass(frozen=True, eq=False)
class IsotonicCalibrator:
    """Per-class isotonic regression calibrator.

    Wraps one ``sklearn.isotonic.IsotonicRegression`` per class.  Each
    regressor maps the model's raw probability for that class to a
    calibrated value.  After per-class transformation the vector is
    renormalized to sum to 1.0.

    Args:
        regressors: List of fitted ``IsotonicRegression`` instances,
            one per class, ordered by label ID.
    """

    regressors: list

    def __post_init__(self) -> None:
        if not self.regressors:
            raise ValueError("regressors list must not be empty")

    @property
    def n_classes(self) -> int:
        return len(self.regressors)

    def calibrate(self, core_probs: np.ndarray) -> np.ndarray:
        single = core_probs.ndim == 1
        probs = core_probs[np.newaxis, :] if single else core_probs

        calibrated = np.empty_like(probs)
        for c, reg in enumerate(self.regressors):
            calibrated[:, c] = np.clip(reg.predict(probs[:, c]), 1e-12, None)

        row_sums = calibrated.sum(axis=1, keepdims=True)
        calibrated = calibrated / row_sums

        return calibrated[0] if single else calibrated

@dataclass(eq=False)
class CalibratorStore:
    """Per-user calibrator registry with global fallback.

    Holds a global calibrator (fitted on all users' validation data)
    and optional per-user calibrators for users that meet the
    personalization eligibility thresholds.

    Args:
        global_calibrator: Calibrator applied to users without a
            dedicated calibrator.
        user_calibrators: Mapping from ``user_id`` to a fitted
            calibrator.  ``None`` or empty dict means all users fall
            back to the global calibrator.
        method: Calibration method label (``"temperature"`` or
            ``"isotonic"``).
        model_bundle_id: Run directory name of the model bundle this
            store was fitted against.  ``None`` for legacy stores.
        model_schema_hash: Schema hash of the model bundle.  Used by
            :func:`~taskclf.core.inference_policy.validate_policy` to
            verify that calibrator and model are compatible.
        created_at: ISO-8601 timestamp of when the store was created.
    """

    global_calibrator: Calibrator
    user_calibrators: dict[str, Calibrator] = field(default_factory=dict)
    method: str = "temperature"
    model_bundle_id: str | None = None
    model_schema_hash: str | None = None
    created_at: str | None = None

    def get_calibrator(self, user_id: str) -> Calibrator:
        """Return the per-user calibrator if available, else the global one."""
        return self.user_calibrators.get(user_id, self.global_calibrator)

    def calibrate_batch(
        self,
        core_probs: np.ndarray,
        user_ids: Sequence[str],
    ) -> np.ndarray:
        """Apply per-user calibration row by row.

        Args:
            core_probs: Probability matrix ``(n_rows, n_classes)``.
            user_ids: Sequence of user_id strings aligned with rows.

        Returns:
            Calibrated probability matrix of the same shape.
        """
        result = np.empty_like(core_probs)
        for i, uid in enumerate(user_ids):
            cal = self.get_calibrator(uid)
            result[i] = cal.calibrate(core_probs[i : i + 1])[0]
        return result

    @property
    def user_ids(self) -> list[str]:
        """Return sorted list of user IDs with per-user calibrators."""
        return sorted(self.user_calibrators)