pipeline/evaluation/caption/protocol.py

"""Family-aware fixed-panel protocol for the balanced Caption diagnostic."""

from __future__ import annotations

from typing import Any


PROTOCOL_VERSION = "caption-family-panel-v2"
NATURAL_PROTOCOL_VERSION = "caption-family-panel-v3-natural"
PANEL_SIZE = 16
SPLIT_ORDER = ("univar", "bivar", "multivar")
FAMILY_ORDER = ("overall", "pattern", "stability", "risk")

# Selection depends only on the public split and prompt family. It must never
# depend on a record's ground-truth values or on whether a metric is defined.
PANELS: dict[str, dict[str, tuple[str, ...]]] = {
    "univar": {
        "overall": (
            "median",
            "std",
            "iqr",
            "min",
            "max",
            "robust_slope",
            "acf_lag1",
            "dominant_period",
            "seasonal_strength",
            "spectral_entropy",
            "mean_abs_change",
            "num_peaks",
            "num_troughs",
            "max_peak_prominence",
            "anomaly_count",
            "change_point_count",
        ),
        "pattern": (
            "std",
            "iqr",
            "min",
            "max",
            "robust_slope",
            "acf_lag1",
            "dominant_period",
            "seasonal_strength",
            "spectral_entropy",
            "mean_abs_change",
            "sign_changes",
            "num_peaks",
            "num_troughs",
            "max_peak_prominence",
            "main_peak_position",
            "main_trough_position",
        ),
        "stability": (
            "std",
            "iqr",
            "min",
            "max",
            "robust_slope",
            "acf_lag1",
            "dominant_period",
            "seasonal_strength",
            "spectral_entropy",
            "mean_abs_change",
            "sign_changes",
            "num_peaks",
            "num_troughs",
            "max_peak_prominence",
            "anomaly_count",
            "change_point_count",
        ),
        "risk": (
            "median",
            "std",
            "iqr",
            "min",
            "max",
            "robust_slope",
            "acf_lag1",
            "spectral_entropy",
            "mean_abs_change",
            "sign_changes",
            "max_peak_prominence",
            "main_peak_position",
            "main_trough_position",
            "anomaly_count",
            "strongest_anomaly_position",
            "change_point_count",
        ),
    },
    "bivar": {
        "overall": (
            "mean_x",
            "mean_y",
            "std_x",
            "std_y",
            "beta_x",
            "beta_y",
            "r2_x",
            "r2_y",
            "pearson_raw",
            "pearson_detrended",
            "period_x",
            "period_y",
            "power_ratio_x",
            "power_ratio_y",
            "vol_corr",
            "corr_structure_stability",
        ),
        "pattern": (
            "beta_x",
            "beta_y",
            "pearson_raw",
            "pearson_detrended",
            "period_x",
            "period_y",
            "power_ratio_x",
            "power_ratio_y",
            "spectral_entropy_x",
            "spectral_entropy_y",
            "acf_lag1_x",
            "acf_lag1_y",
            "ccf_peak_lag",
            "ccf_peak_value",
            "vol_corr",
            "corr_structure_stability",
        ),
        "stability": (
            "acf_lag1_x",
            "acf_lag1_y",
            "spectral_entropy_x",
            "spectral_entropy_y",
            "r2_x",
            "r2_y",
            "power_ratio_x",
            "power_ratio_y",
            "pearson_raw",
            "pearson_detrended",
            "change_point_count_x",
            "change_point_count_y",
            "corr_structure_stability",
            "vol_corr_std",
            "pc1_period_strength",
            "pc1_spectral_entropy",
        ),
        "risk": (
            "beta_x",
            "beta_y",
            "mean_abs_change_x",
            "mean_abs_change_y",
            "pearson_raw",
            "anomaly_ratio_x",
            "anomaly_ratio_y",
            "pc1_anomaly_ratio",
            "pc1_max_abs_z",
            "change_point_count_x",
            "change_point_count_y",
            "joint_extreme_ratio",
            "tail_concordance",
            "vol_corr",
            "vol_corr_std",
            "corr_structure_stability",
        ),
    },
    "multivar": {
        "overall": (
            "pca_k",
            "pca_expl_ratio_1",
            "effective_rank",
            "sys_sync_avg_r2_topk",
            "sys_sync_std_r2_topk",
            "sync_ratio",
            "leader_ratio",
            "lag_mean",
            "pc1_trend_r2",
            "pc1_regime_shift",
            "pc1_period_strength",
            "pc1_seasonality_strength",
            "pc1_mean_abs_change",
            "volatility_correlation_mean",
            "corr_structure_stability",
            "mahal_outlier_ratio",
        ),
        "pattern": (
            "pca_k",
            "pca_expl_ratio_1",
            "effective_rank",
            "loading_ipr_pc1",
            "r2_topk_skew",
            "sys_sync_avg_r2_topk",
            "sys_sync_std_r2_topk",
            "sync_ratio",
            "leader_ratio",
            "lag_mean",
            "lag_range",
            "pc1_trend_r2",
            "pc1_period_strength",
            "pc1_seasonality_strength",
            "seasonality_consistency",
            "pc1_mean_abs_change",
        ),
        "stability": (
            "pca_expl_ratio_1",
            "effective_rank",
            "sys_sync_avg_r2_topk",
            "sys_sync_std_r2_topk",
            "corr_structure_stability",
            "corr_dynamics_mean_std",
            "corr_frobenius_change_mean",
            "sync_stability",
            "volatility_correlation_mean",
            "volatility_correlation_std",
            "volatility_sync_index",
            "pc1_trend_r2",
            "pc1_regime_shift",
            "pc1_period_strength",
            "pc1_seasonality_strength",
            "seasonality_consistency",
        ),
        "risk": (
            "pca_expl_ratio_1",
            "effective_rank",
            "sys_sync_avg_r2_topk",
            "sys_sync_std_r2_topk",
            "pc1_trend_r2",
            "pc1_regime_shift",
            "pc1_anomaly_ratio",
            "pc1_max_abs_z",
            "mahal_outlier_ratio",
            "mahal_max_distance",
            "corr_structure_stability",
            "corr_dynamics_mean_std",
            "corr_frobenius_change_mean",
            "volatility_correlation_mean",
            "volatility_correlation_max",
            "high_vol_correlation_ratio",
        ),
    },
}

PERCENT_TO_FRACTION_KEYS = frozenset(
    {
        "acf_lag1",
        "seasonal_strength",
        "spectral_entropy",
        "r2_x",
        "r2_y",
        "pearson_raw",
        "pearson_detrended",
        "power_ratio_x",
        "power_ratio_y",
        "vol_corr",
        "corr_structure_stability",
        "spectral_entropy_x",
        "spectral_entropy_y",
        "acf_lag1_x",
        "acf_lag1_y",
        "ccf_peak_value",
        "vol_corr_std",
        "pc1_period_strength",
        "pc1_spectral_entropy",
        "anomaly_ratio_x",
        "anomaly_ratio_y",
        "pc1_anomaly_ratio",
        "joint_extreme_ratio",
        "tail_concordance",
        "pca_expl_ratio_1",
        "sys_sync_avg_r2_topk",
        "sys_sync_std_r2_topk",
        "sync_ratio",
        "leader_ratio",
        "pc1_trend_r2",
        "pc1_regime_shift",
        "pc1_seasonality_strength",
        "volatility_correlation_mean",
        "mahal_outlier_ratio",
        "seasonality_consistency",
        "sync_stability",
        "volatility_correlation_std",
        "volatility_sync_index",
        "volatility_correlation_max",
        "high_vol_correlation_ratio",
    }
)

# These clarify collisions in the short shared metric-definitions file without
# changing the underlying scorer key or ground truth.
DEFINITION_OVERRIDES = {
    "corr_structure_stability": "滚动相关结构的稳定性",
    "sync_stability": "系统同步关系随时间的稳定性",
    "vol_corr_std": "滚动波动率相关系数的标准差",
}

PROMPT_APPENDIX = """

为回答上述问题，请先计算下面全部 {panel_size} 个相关统计属性（key — 含义）：
{metric_lines}{scope_note}

输出要求：
1. 先输出“统计属性”，每个 key 单独一行并严格使用：
   [STAT]<metric_key> = <number>[/STAT]
2. 必须依次报告上面全部 {panel_size} 个 key；metric_key 必须逐字复制。
3. 若属性对该序列不适用或无法定义，使用：
   [STAT]<metric_key> = NA[/STAT]
4. number 只能使用十进制或科学计数法，不写单位；比例统一写成 0–1 小数。
5. 随后输出“分析描述”，围绕开头问题综合解释关键现象；不要逐项复述统计量。
6. 不要复述原始时间序列。
"""

NATURAL_PROMPT_APPENDIX = """

请围绕开头的问题，从{analysis_aspects}等角度分析上述时间序列。{scope_note}请直接输出连贯的分析描述，并将有助于判断的关键数值自然融入正文，解释其反映的主要现象和变量关系；不要逐项列出统计量，不要使用表格或清单，不要展示计算过程，也不要复述原始时间序列。
"""


def panel_for(split: str, family: str) -> tuple[str, ...]:
    try:
        panel = PANELS[split][family]
    except KeyError as error:
        raise ValueError(f"Unknown split/family: {split!r}/{family!r}") from error
    if len(panel) != PANEL_SIZE or len(set(panel)) != PANEL_SIZE:
        raise AssertionError(f"Invalid panel size or duplicate key: {split}/{family}")
    return panel


def metric_description(key: str, definitions: dict[str, str]) -> str:
    description = DEFINITION_OVERRIDES.get(key, definitions.get(key))
    if not isinstance(description, str) or not description.strip():
        raise ValueError(f"Missing metric definition for {key!r}")
    return description.strip()


def candidate_specs(
    split: str,
    family: str,
    definitions: dict[str, str],
) -> list[dict[str, Any]]:
    return [
        {
            "key": key,
            "description": metric_description(key, definitions),
            "percent_to_fraction": key in PERCENT_TO_FRACTION_KEYS,
            "allow_na": True,
        }
        for key in panel_for(split, family)
    ]


def build_prompt(
    source_prompt: str,
    split: str,
    family: str,
    definitions: dict[str, str],
) -> str:
    if "<ts> </ts>" not in source_prompt and "<ts></ts>" not in source_prompt:
        raise ValueError("Source prompt does not contain a time-series placeholder")
    specs = candidate_specs(split, family, definitions)
    metric_lines = "\n".join(
        f"- `{spec['key']}` — {spec['description']}" for spec in specs
    )
    scope_note = ""
    if split == "bivar":
        scope_note = "\n其中 X=Series1，Y=Series2。"
    elif split == "multivar":
        scope_note = "\nSeries1、Series2、……对应输入中的各个通道。"
    appendix = PROMPT_APPENDIX.format(
        panel_size=PANEL_SIZE,
        metric_lines=metric_lines,
        scope_note=scope_note,
    )
    return source_prompt.rstrip() + appendix


def build_natural_prompt(
    source_prompt: str,
    split: str,
    family: str,
    definitions: dict[str, str],
) -> str:
    if "<ts> </ts>" not in source_prompt and "<ts></ts>" not in source_prompt:
        raise ValueError("Source prompt does not contain a time-series placeholder")
    specs = candidate_specs(split, family, definitions)
    analysis_aspects = "、".join(spec["description"] for spec in specs)
    scope_note = ""
    if split == "bivar":
        scope_note = "其中 X=Series1，Y=Series2。"
    elif split == "multivar":
        scope_note = "Series1、Series2、……对应输入中的各个通道。"
    appendix = NATURAL_PROMPT_APPENDIX.format(
        analysis_aspects=analysis_aspects,
        scope_note=scope_note,
    )
    return source_prompt.rstrip() + appendix


def validate_protocol(definitions: dict[str, str]) -> None:
    if tuple(PANELS) != SPLIT_ORDER:
        raise AssertionError("Split order is not frozen")
    for split in SPLIT_ORDER:
        if tuple(PANELS[split]) != FAMILY_ORDER:
            raise AssertionError(f"Family order is not frozen for {split}")
        for family in FAMILY_ORDER:
            for key in panel_for(split, family):
                metric_description(key, definitions)