scripts/run_power_analysis.py

#!/usr/bin/env python3
"""Event-cluster power scenarios for the TELEN temporal evaluation.

This is a design aid, not a claim that five proxy events identify the variance
of the future expert benchmark. It reports both the plug-in standard deviation
and a conservative one-sided 95% upper confidence bound, then requires the
calculation to be repeated after the first frozen gold pilot.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import math
from collections import defaultdict
from datetime import datetime, timezone
from pathlib import Path

import numpy as np
from scipy.stats import chi2, nct, t

REPO_ROOT = Path(__file__).resolve().parents[1]


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for block in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()


def achieved_power(
    cluster_count: int, *, effect: float, standard_deviation: float, alpha: float
) -> float:
    if cluster_count < 2:
        return 0.0
    if standard_deviation <= 0:
        return 1.0
    degrees = cluster_count - 1
    critical = t.ppf(1.0 - alpha / 2.0, degrees)
    noncentrality = effect * math.sqrt(cluster_count) / standard_deviation
    return float(
        nct.cdf(-critical, degrees, noncentrality)
        + 1.0
        - nct.cdf(critical, degrees, noncentrality)
    )


def required_clusters(
    *,
    effect: float,
    standard_deviation: float,
    alpha: float,
    target_power: float,
    maximum: int = 10_000,
) -> int | None:
    for cluster_count in range(2, maximum + 1):
        if achieved_power(
            cluster_count,
            effect=effect,
            standard_deviation=standard_deviation,
            alpha=alpha,
        ) >= target_power:
            return cluster_count
    return None


def parse_floats(value: str) -> tuple[float, ...]:
    values = tuple(float(item.strip()) for item in value.split(",") if item.strip())
    if not values or any(item <= 0 for item in values):
        raise argparse.ArgumentTypeError("expected positive comma-separated values")
    return values


def run(args: argparse.Namespace) -> dict:
    oracle = json.loads(args.oracle_result.read_text(encoding="utf-8"))
    detail_path = args.detail or Path(oracle["per_query_path"])
    rows = [
        json.loads(line)
        for line in detail_path.read_text(encoding="utf-8").splitlines()
        if line.strip()
    ]
    baseline = oracle["oracle_gate"]["strongest_non_oracle"]
    treatment = "correct_event_edge_oracle"
    by_event: dict[str, list[float]] = defaultdict(list)
    for row in rows:
        if row["phase"] != "post":
            continue
        by_event[row["candidate_id"]].append(
            row["metrics"][treatment]["ndcg_at_10"]
            - row["metrics"][baseline]["ndcg_at_10"]
        )
    event_differences = np.asarray(
        [np.mean(values) for _, values in sorted(by_event.items())],
        dtype=np.float64,
    )
    if len(event_differences) < 3:
        raise ValueError("power analysis requires at least three event clusters")
    observed_mean = float(event_differences.mean())
    observed_sd = float(event_differences.std(ddof=1))
    degrees = len(event_differences) - 1
    sd_upper = float(
        observed_sd
        * math.sqrt(degrees / chi2.ppf(args.sd_upper_tail, degrees))
    )

    scenarios = []
    for effect in args.effects:
        for alpha in args.alphas:
            for target_power in args.powers:
                scenarios.append(
                    {
                        "minimum_detectable_difference": effect,
                        "two_sided_alpha": alpha,
                        "target_power": target_power,
                        "required_event_clusters_plugin_sd": required_clusters(
                            effect=effect,
                            standard_deviation=observed_sd,
                            alpha=alpha,
                            target_power=target_power,
                        ),
                        "required_event_clusters_conservative_sd": required_clusters(
                            effect=effect,
                            standard_deviation=sd_upper,
                            alpha=alpha,
                            target_power=target_power,
                        ),
                    }
                )
    result = {
        "schema_version": 1,
        "artifact_type": "event_cluster_power_scenarios",
        "created_at": datetime.now(timezone.utc).isoformat(),
        "inputs": {
            "oracle_result": str(args.oracle_result),
            "oracle_result_sha256": sha256_file(args.oracle_result),
            "per_query_detail": str(detail_path),
            "per_query_detail_sha256": sha256_file(detail_path),
            "baseline": baseline,
            "treatment": treatment,
        },
        "proxy_estimate": {
            "event_cluster_count": len(event_differences),
            "event_macro_mean_difference": observed_mean,
            "event_difference_sample_sd": observed_sd,
            "sd_one_sided_confidence_level": 1.0 - args.sd_upper_tail,
            "event_difference_sd_upper_bound": sd_upper,
            "event_differences": {
                event_id: float(np.mean(values))
                for event_id, values in sorted(by_event.items())
            },
        },
        "scenarios": scenarios,
        "decision_rule": (
            "Use event/law clusters as the sampling unit. Re-estimate variance "
            "once a frozen expert gold pilot exists; do not size the final study "
            "from query-row count or treat this replay proxy as definitive."
        ),
        "limitations": [
            "Only five mechanistic replay clusters estimate variance.",
            "Oracle-minus-baseline variation may differ from learned-model variation.",
            "The normal/noncentral-t approximation does not replace simulation on the gold design.",
            "Relation-stratified coverage may require more events than statistical power alone.",
        ],
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(
        json.dumps(result, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
    )
    return result


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--oracle-result",
        type=Path,
        default=(
            REPO_ROOT / "research/results/telen_oracle_utility_pilot_2026-08-29.json"
        ),
    )
    parser.add_argument("--detail", type=Path)
    parser.add_argument(
        "--output",
        type=Path,
        default=REPO_ROOT / "research/results/telen_power_analysis_2026-08-29.json",
    )
    parser.add_argument("--effects", type=parse_floats, default=(0.05, 0.075, 0.10))
    parser.add_argument("--alphas", type=parse_floats, default=(0.05, 0.01))
    parser.add_argument("--powers", type=parse_floats, default=(0.80, 0.90))
    parser.add_argument("--sd-upper-tail", type=float, default=0.05)
    args = parser.parse_args()
    if not 0 < args.sd_upper_tail < 1:
        raise ValueError("sd upper tail must lie strictly between zero and one")
    result = run(args)
    print(json.dumps(result, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()