utility.py

"""Paired optimizer outcomes used by GATE-A and later utility supervision."""

from __future__ import annotations

from dataclasses import asdict, dataclass
from enum import Enum


class UtilityLabel(str, Enum):
    HELPFUL = "helpful"
    NEUTRAL = "neutral"
    HARMFUL = "harmful"


@dataclass(frozen=True)
class OptimizerOutcome:
    """Metrics from one frozen optimizer run, in the Gate-A evaluation units."""

    gt_edge_distance_mm: float
    contact_f1_30mm: float
    mpjpe_mm: float
    pa_mpjpe_mm: float
    pve_translation_aligned_mm: float
    relative_translation_error_mm: float
    penetration_max_mm: float
    penetration_mean_mm: float
    reprojection_error_px: float
    optimizer_success: bool = True


@dataclass(frozen=True)
class FrozenDecisionThresholds:
    primary_contact_improvement_mm: float
    contact_f1_change: float
    mpjpe_noninferiority_mm: float
    pa_mpjpe_noninferiority_mm: float
    pve_translation_aligned_noninferiority_mm: float
    relative_translation_error_noninferiority_mm: float
    penetration_max_noninferiority_mm: float
    penetration_mean_noninferiority_mm: float
    reprojection_noninferiority_px: float


SAFETY_DELTA_TO_MARGIN = {
    "mpjpe_mm": "mpjpe_noninferiority_mm",
    "pa_mpjpe_mm": "pa_mpjpe_noninferiority_mm",
    "pve_translation_aligned_mm": "pve_translation_aligned_noninferiority_mm",
    "relative_translation_error_mm": "relative_translation_error_noninferiority_mm",
    "penetration_max_mm": "penetration_max_noninferiority_mm",
    "penetration_mean_mm": "penetration_mean_noninferiority_mm",
    "reprojection_error_px": "reprojection_noninferiority_px",
}


def paired_deltas(base: OptimizerOutcome, intervention: OptimizerOutcome) -> dict[str, float]:
    """Return intervention-minus-base deltas; lower is better except contact F1."""
    return {
        "gt_edge_distance_mm": intervention.gt_edge_distance_mm - base.gt_edge_distance_mm,
        "contact_f1_30mm": intervention.contact_f1_30mm - base.contact_f1_30mm,
        "mpjpe_mm": intervention.mpjpe_mm - base.mpjpe_mm,
        "pa_mpjpe_mm": intervention.pa_mpjpe_mm - base.pa_mpjpe_mm,
        "pve_translation_aligned_mm": (
            intervention.pve_translation_aligned_mm - base.pve_translation_aligned_mm
        ),
        "relative_translation_error_mm": (
            intervention.relative_translation_error_mm - base.relative_translation_error_mm
        ),
        "penetration_max_mm": intervention.penetration_max_mm - base.penetration_max_mm,
        "penetration_mean_mm": intervention.penetration_mean_mm - base.penetration_mean_mm,
        "reprojection_error_px": intervention.reprojection_error_px - base.reprojection_error_px,
    }


def safety_violations(
    deltas: dict[str, float], thresholds: FrozenDecisionThresholds
) -> list[str]:
    """Return safety metrics whose intervention delta exceeds its frozen margin."""
    return [
        metric
        for metric, margin_name in SAFETY_DELTA_TO_MARGIN.items()
        if deltas[metric] > getattr(thresholds, margin_name)
    ]


def classify_utility(
    base: OptimizerOutcome,
    intervention: OptimizerOutcome,
    thresholds: FrozenDecisionThresholds,
) -> tuple[UtilityLabel, dict[str, float]]:
    """Classify exactly as preregistered for the 50-event Gate-A experiment.

    Contact F1 is retained as a diagnostic and paper-claim check.  It does not
    override the primary GT-edge-distance definition used for the utility label.
    """
    delta = paired_deltas(base, intervention)
    if not intervention.optimizer_success:
        return UtilityLabel.HARMFUL, delta

    primary_delta = delta["gt_edge_distance_mm"]
    primary_help = primary_delta < -thresholds.primary_contact_improvement_mm
    primary_harm = primary_delta > thresholds.primary_contact_improvement_mm

    if safety_violations(delta, thresholds) or primary_harm:
        return UtilityLabel.HARMFUL, delta
    if primary_help:
        return UtilityLabel.HELPFUL, delta
    return UtilityLabel.NEUTRAL, delta


def serialize_pair(
    event_id: str,
    condition: str,
    base: OptimizerOutcome,
    intervention: OptimizerOutcome,
    thresholds: FrozenDecisionThresholds,
    *,
    scientific_result: bool = False,
) -> dict:
    label, deltas = classify_utility(base, intervention, thresholds)
    return {
        "event_id": event_id,
        "condition": condition,
        "scientific_result": scientific_result,
        "label": label.value,
        "base": asdict(base),
        "intervention": asdict(intervention),
        "deltas": deltas,
    }