Source code for pytyche.analysis._calibrate

"""Attach an SBC-fitted calibration artifact to a posterior.

Implementation behind ``posterior.apply_calibration(calibration)`` on the
three posterior result types.  The fitted corrections are query-time
remappings — they never transform samples — so "applying" a calibration
means attaching the artifact to the posterior (``is_calibrated=True``),
not changing any array.  Sample arrays on the returned posterior are the
same objects as the original's.
"""

from __future__ import annotations

import dataclasses
from typing import TypeVar

from pytyche.analysis._contrasts import require_observed
from pytyche.bcf.config import (
    BinaryBCFResult,
    ContinuousBCFResult,
    HurdleBCFResult,
)
from pytyche.calibrate.artifact import Calibration

AnyBCFResultT = TypeVar(
    "AnyBCFResultT", HurdleBCFResult, ContinuousBCFResult, BinaryBCFResult
)




def apply_calibration(
    posterior: AnyBCFResultT, calibration: Calibration
) -> AnyBCFResultT:
    """Return a new posterior of the same type with *calibration* attached.

    Attach, don't transform: the artifact's corrections are query-time
    interval remappings, so the returned posterior shares every field of
    the original by reference (``dataclasses.replace`` with only
    ``is_calibrated=True`` and ``calibration`` changed) — no array is
    copied or re-derived, and the original posterior is untouched.

    The correction currently applies to intervals only: probabilities and
    expected losses are computed on the raw draws and are identical
    pre/post calibration; corrected CIs appear where the interval
    summaries are built (``analyze``).  K = 2 experiments only —
    per-contrast recalibration for K >= 3 is not yet implemented.

    Args:
        posterior: One of the three posterior result types, carrying
            observed data (raises otherwise).
        calibration: SBC-fitted artifact whose fitted regime (``metric``,
            ``n_treatments``) must match the posterior's observed data.

    Returns:
        A new instance of the same result type with ``is_calibrated=True``
        and ``calibration`` attached by identity.

    Raises:
        TypeError: When *calibration* is ``None`` — an artifact is
            required; an uncalibrated posterior is one you never applied
            a calibration to.  (The fit entry points' ``calibration=``
            kwarg only delegates here when non-None, so the default fit
            path never produces this error.)
        ValueError: When ``posterior.observed`` is ``None`` (no regime to
            check against), when ``calibration.applies_to(observed)``
            is ``False`` (the message names each mismatched dimension),
            or when a hurdle posterior's ``pooling`` differs from the
            artifact's (a correction fitted on one pooling mode's
            miscalibration profile must not silently apply to the other).
        NotImplementedError: When the (regime-matching) posterior has
            K >= 3 variants — per-contrast recalibration is not yet
            shipped.
    """
    if calibration is None:
        raise TypeError(
            "apply_calibration requires a Calibration artifact, got None. "
            "To leave a posterior uncalibrated, don't apply a calibration "
            "at all."
        )
    observed = require_observed(posterior)
    if not calibration.applies_to(observed):
        mismatches = []
        if observed.metric != calibration.metric:
            mismatches.append(
                f"metric (posterior {observed.metric!r}, "
                f"artifact {calibration.metric!r})"
            )
        if len(observed.variants) != calibration.n_treatments:
            mismatches.append(
                f"n_treatments (posterior K={len(observed.variants)}, "
                f"artifact K={calibration.n_treatments})"
            )
        raise ValueError(
            "calibration was fitted on a different regime than this "
            "posterior's observed data — mismatched " + "; ".join(mismatches)
        )
    if (
        isinstance(posterior, HurdleBCFResult)
        and posterior.pooling != calibration.pooling
    ):
        # Part of the regime match: the correction was fitted to one
        # pooling mode's miscalibration profile.  Binary/continuous
        # posteriors have no pooling concept and skip this check.
        raise ValueError(
            f"calibration was fitted on pooling="
            f"{calibration.pooling!r} sweeps but this posterior was fit "
            f"with pooling={posterior.pooling!r} — no silent cross-regime "
            "reuse."
        )
    if len(observed.variants) >= 3:
        raise NotImplementedError(
            "apply_calibration supports K = 2 only: per-contrast R(p) "
            "recalibration at K >= 3 requires per-contrast SBC machinery "
            "that is not yet implemented."
        )
    return dataclasses.replace(
        posterior, is_calibrated=True, calibration=calibration
    )