staplebridge/hydrocarbon/plan_validation.py

"""Exact-SB validation diagnostics for the hydrocarbon plan-control head.

The lead-local Gibbs target is computed with exactly the training definition,
by delegating to :func:`configured_plan_level_objective` with
``exact_sb_objective=True``. Nothing here participates in candidate
generation, plan selection, or reranking: ``q*`` is a diagnostic and a
model-selection signal only.
"""

from __future__ import annotations

from collections import Counter
from typing import Any, Sequence

import torch

from staplebridge.hydrocarbon.plan_control import (
    HydrocarbonPlanControlConfig,
    configured_plan_level_objective,
)

# The validation metrics this module contributes. They are additive: no
# existing validation metric is replaced or removed.
Q_STAR_VALIDATION_METRICS = (
    "q_star_vs_q_theta_kl",
    "q_star_top1_agreement",
    "q_star_spearman",
    "q_star_entropy",
    "q_theta_entropy",
    "exact_sb_forward_objective",
)

S5_MODE = "S5-S5/i,i+4"
R8_MODE = "R8-S5/i,i+7"


def exact_sb_validation_enabled(root_config: dict[str, Any] | None) -> bool:
    """Only ``exact_sb_objective`` **and** ``validation_all_plans`` opt in."""
    cfg = HydrocarbonPlanControlConfig.from_config(root_config)
    return bool(cfg.enabled and cfg.exact_sb_objective and cfg.validation_all_plans)


def _average_ranks(values: torch.Tensor) -> torch.Tensor:
    """Ascending ranks with ties averaged, as Spearman's rho requires."""
    order = torch.argsort(values)
    sorted_values = values[order]
    ranks = torch.empty_like(sorted_values, dtype=torch.float64)
    start = 0
    for index in range(1, sorted_values.numel() + 1):
        if index == sorted_values.numel() or sorted_values[index] != sorted_values[start]:
            ranks[start:index] = 0.5 * float(start + index - 1) + 1.0
            start = index
    out = torch.empty_like(ranks)
    out[order] = ranks
    return out


def spearman_correlation(left: torch.Tensor, right: torch.Tensor) -> float | None:
    """Rank correlation, or ``None`` when it is undefined for this lead.

    Undefined means fewer than two plans, or one side constant (every plan
    tied), in which case there is no ordering to agree or disagree with.
    """
    if left.shape != right.shape:
        raise ValueError("spearman inputs must cover the same plans")
    if left.numel() < 2:
        return None
    left_ranks = _average_ranks(left.detach().to(dtype=torch.float64).flatten())
    right_ranks = _average_ranks(right.detach().to(dtype=torch.float64).flatten())
    left_centered = left_ranks - left_ranks.mean()
    right_centered = right_ranks - right_ranks.mean()
    denominator = float(
        torch.sqrt((left_centered**2).sum() * (right_centered**2).sum()).item()
    )
    if denominator <= 0.0:
        return None
    return float((left_centered * right_centered).sum().item() / denominator)


def lead_q_star_diagnostics(
    controlled_log_probabilities: torch.Tensor,
    reference_log_probabilities: torch.Tensor,
    terminal_energies: torch.Tensor,
    beta: float,
    plan_modes: Sequence[str] | None = None,
    target_support_mask: torch.Tensor | None = None,
) -> dict[str, Any]:
    """Per-lead ``q*`` diagnostics over the hard-supported legal plans.

    ``q*`` and the forward objective come from the same
    :func:`configured_plan_level_objective` call the training loop uses, so a
    validation KL can never drift from the trained objective's definition.
    """
    reverse_kl, forward_objective, log_q_star = configured_plan_level_objective(
        controlled_log_probabilities,
        [],
        [],
        # Unused by the exact-SB branch; passed only to satisfy the shared
        # signature. Any positive value leaves the exact-SB result untouched.
        1.0,
        exact_sb_objective=True,
        reference_log_probabilities=reference_log_probabilities,
        terminal_energies=terminal_energies,
        exact_sb_beta=beta,
        target_support_mask=target_support_mask,
    )
    if forward_objective is None or log_q_star is None:
        raise RuntimeError("exact-SB validation diagnostics require the exact-SB branch")
    q_star = log_q_star.exp()
    q_theta = controlled_log_probabilities.exp()
    star_top = int(torch.argmax(log_q_star).item())
    theta_top = int(torch.argmax(controlled_log_probabilities).item())
    if target_support_mask is None:
        # Historical false/fallback branch is deliberately bit-for-bit.
        q_star_entropy = -(q_star * log_q_star).sum()
    else:
        q_star_entropy = -torch.where(
            q_star > 0.0, q_star * log_q_star, torch.zeros_like(q_star)
        ).sum()
    result: dict[str, Any] = {
        "n_plans": int(controlled_log_probabilities.numel()),
        "q_star_vs_q_theta_kl": float(reverse_kl.detach().cpu().item()),
        "exact_sb_forward_objective": float(forward_objective.detach().cpu().item()),
        "q_star_entropy": float(q_star_entropy.detach().cpu().item()),
        "q_theta_entropy": float(
            (-(q_theta * controlled_log_probabilities).sum()).detach().cpu().item()
        ),
        "q_star_top1_index": star_top,
        "q_theta_top1_index": theta_top,
        "q_star_top1_agreement": bool(star_top == theta_top),
        "q_star_spearman": spearman_correlation(
            log_q_star.detach(), controlled_log_probabilities.detach()
        ),
        "exact_sb_beta": float(beta),
        "q_star_support_size": int(
            target_support_mask.to(dtype=torch.bool).sum().item()
            if target_support_mask is not None
            else controlled_log_probabilities.numel()
        ),
    }
    if plan_modes is not None:
        if len(plan_modes) != controlled_log_probabilities.numel():
            raise ValueError("plan_modes must cover the same plans as q* and q_theta")
        for distribution_name, probabilities in (
            ("q_star", q_star),
            ("q_theta", q_theta),
        ):
            for metric_name, mode in (
                ("s5_s5_i4", S5_MODE),
                ("r8_s5_i7", R8_MODE),
            ):
                indices = [index for index, value in enumerate(plan_modes) if value == mode]
                mass = (
                    float(probabilities[indices].sum().detach().cpu().item())
                    if indices
                    else 0.0
                )
                result[f"{distribution_name}_{metric_name}_probability_mass"] = mass
        result["q_star_top1_mode"] = str(plan_modes[star_top])
        result["q_theta_top1_mode"] = str(plan_modes[theta_top])
    return result


def _mean(values: Sequence[float]) -> float | None:
    return float(sum(values) / len(values)) if values else None


def aggregate_q_star_diagnostics(rows: Sequence[dict[str, Any]]) -> dict[str, Any]:
    """Average per-lead diagnostics into the validation summary metrics."""
    spearmans = [
        float(row["q_star_spearman"])
        for row in rows
        if row.get("q_star_spearman") is not None
    ]
    result: dict[str, Any] = {
        "q_star_vs_q_theta_kl": _mean([float(row["q_star_vs_q_theta_kl"]) for row in rows]),
        "q_star_top1_agreement": _mean(
            [float(bool(row["q_star_top1_agreement"])) for row in rows]
        ),
        "q_star_spearman": _mean(spearmans),
        "q_star_entropy": _mean([float(row["q_star_entropy"]) for row in rows]),
        "q_theta_entropy": _mean(
            [float(row["q_theta_entropy"]) for row in rows if row.get("q_theta_entropy") is not None]
        ),
        "exact_sb_forward_objective": _mean(
            [float(row["exact_sb_forward_objective"]) for row in rows]
        ),
        "q_star_leads_scored": len(rows),
        "q_star_plans_scored": int(sum(int(row["n_plans"]) for row in rows)),
        "q_star_spearman_defined_leads": len(spearmans),
        # Stated explicitly because it is the invariant this module must hold:
        # diagnostics never touch the decoded candidate set or its ranking.
        "q_star_used_for_candidate_reranking": False,
    }
    for distribution_name in ("q_star", "q_theta"):
        for metric_name in ("s5_s5_i4", "r8_s5_i7"):
            key = f"{distribution_name}_{metric_name}_probability_mass"
            values = [float(row[key]) for row in rows if row.get(key) is not None]
            if values:
                result[key] = _mean(values)
        top1_key = f"{distribution_name}_top1_mode"
        modes = Counter(str(row[top1_key]) for row in rows if row.get(top1_key))
        if modes:
            result[f"{distribution_name}_top1_counts"] = dict(modes)
            result[f"{distribution_name}_top1_mix"] = {
                mode: count / len(rows) for mode, count in modes.items()
            }
            result[f"{distribution_name}_s5_s5_i4_top1_rate"] = (
                modes.get(S5_MODE, 0) / len(rows)
            )
            result[f"{distribution_name}_r8_s5_i7_top1_rate"] = (
                modes.get(R8_MODE, 0) / len(rows)
            )
    return result