analysis/thread_statistics/composition.py

"""Paired prompt-cluster inference for the rubric-first composition accuracy.

Only NumPy is required. Decisions and human tie masks are held fixed: this is
not a bootstrap of TEST standardization, DEV calibration, or model fitting.
"""

from __future__ import annotations

from numbers import Integral
from typing import Any

import numpy as np


def _ids(values: Any, size: int, name: str, *, unique: bool = True) -> np.ndarray:
    ids = np.asarray(values)
    if ids.ndim != 1 or len(ids) != size:
        raise ValueError(f"{name} must be one-dimensional with length {size}")
    # Missing IDs, floats (including NaN), and mixed-type coercions are unsafe
    # join keys. Object arrays from pandas are supported for strings/integers.
    items = list(values)
    kinds = set()
    for value in items:
        if isinstance(value, (str, np.str_)) and value:
            kinds.add("str")
        elif isinstance(value, Integral) and not isinstance(value, (bool, np.bool_)):
            kinds.add("int")
        else:
            raise ValueError(f"{name} must contain nonempty strings or integer IDs")
    if len(kinds) != 1:
        raise ValueError(f"{name} must not mix string and integer IDs")
    if unique and len(set(items)) != size:
        raise ValueError(f"{name} must contain unique IDs (no duplicate runs/axes)")
    return ids


def _check_axis_ids(a: Any, b: Any, size: int, axis: str) -> bool:
    if a is None and b is None:
        return False
    if a is None or b is None:
        raise ValueError(f"provide both {axis}_ids_a and {axis}_ids_b")
    left = _ids(a, size, f"{axis}_ids_a")
    right = _ids(b, size, f"{axis}_ids_b")
    if not np.array_equal(left, right):
        raise ValueError(f"{axis} axis alignment error: IDs must match in order")
    return True


def _validate_decisions(a_correct, b_correct, a_valid, b_valid):
    arrays = [np.asarray(x) for x in (a_correct, b_correct, a_valid, b_valid)]
    a, b, va, vb = arrays
    if a.ndim != 3 or any(x.shape != a.shape for x in arrays):
        raise ValueError("all correctness/valid arrays must have identical [seed, pair, rubric] shape")
    if any(size == 0 for size in a.shape):
        raise ValueError("seed, pair, and rubric axes must all be nonempty")
    for mask in (va, vb):
        if mask.dtype.kind not in "biuf" or not np.isin(mask, [0, 1]).all():
            raise ValueError("valid masks must contain only boolean or numeric 0/1 values")
    va = va.astype(bool, copy=False)
    vb = vb.astype(bool, copy=False)
    if not np.array_equal(va, vb) or not np.all(va == va[:1]):
        raise ValueError("human non-tie masks must be identical across methods and seeds; do not intersect missing predictions")
    cleaned = []
    for name, values in (("a_correct", a), ("b_correct", b)):
        if values.dtype.kind not in "biuf" or not np.isin(values[va], [0, 1]).all():
            raise ValueError(f"{name} must be finite binary correctness on every human-valid decision")
        # NaN or arbitrary numeric placeholders on ties never enter a sum.
        cleaned.append(np.where(va, values, 0).astype(np.float64))
    valid = va[0]
    denominators = valid.sum(axis=0)
    active = denominators > 0
    if not active.any():
        raise ValueError("composition is undefined: every rubric has only human ties")
    return *cleaned, valid, denominators, active


def _point(a, b, denominators, active):
    means_a = (a.sum(axis=1)[:, active] / denominators[active]).mean(axis=1)
    means_b = (b.sum(axis=1)[:, active] / denominators[active]).mean(axis=1)
    point_a, point_b = float(means_a.mean()), float(means_b.mean())
    return {
        "point_a": point_a,
        "point_b": point_b,
        "difference": point_b - point_a,
        "difference_direction": "b minus a",
        "per_seed_accuracy_a": means_a.tolist(),
        "per_seed_accuracy_b": means_b.tolist(),
        "per_seed_difference": (means_b - means_a).tolist(),
    }


def composition_point_difference(
    a_correct, b_correct, a_valid, b_valid, *, seed_ids_a, seed_ids_b,
    pair_ids_a=None, pair_ids_b=None, rubric_ids_a=None, rubric_ids_b=None,
) -> dict:
    """Return the observed rubric-first, then seed-averaged difference, B - A.

    All four arrays must be [seed, pair, rubric]. Human-valid masks must match
    across methods and seeds. Exact zeros in correctness count as incorrect;
    only the supplied human tie mask excludes observations. Rubrics with no
    valid pairs anywhere are omitted, matching eval.metrics.composition's
    nanmean; their indices are disclosed. This function makes no CI claim.

    Seed IDs are required and must match in order, without duplicates. Pair
    and rubric IDs are optional paired checks; absent IDs mean the caller
    explicitly assumes positional alignment. Labels cannot detect falsely
    labeled/permuted contents. No axis is silently transposed or reordered.
    """
    a, b, valid, den, active = _validate_decisions(a_correct, b_correct, a_valid, b_valid)
    s, p, r = a.shape
    alignment = {"seed": _check_axis_ids(seed_ids_a, seed_ids_b, s, "seed")}
    if not alignment["seed"]:
        raise ValueError("seed_ids_a and seed_ids_b are required")
    alignment["pair"] = _check_axis_ids(pair_ids_a, pair_ids_b, p, "pair")
    alignment["rubric"] = _check_axis_ids(rubric_ids_a, rubric_ids_b, r, "rubric")
    return {
        **_point(a, b, den, active),
        "estimand": "mean over seeds of equal-rubric means of valid-pair accuracy",
        "alignment_checked": alignment,
        "counts": {
            "seeds": s, "pairs": p, "rubrics": r,
            "active_rubrics": int(active.sum()),
            "excluded_all_tie_rubric_indices": np.flatnonzero(~active).tolist(),
            "valid_pairs_per_rubric": den.tolist(),
            "human_valid_pair_rubric_decisions": int(valid.sum()),
            "human_tie_pair_rubric_decisions": int(valid.size - valid.sum()),
            "valid_model_decisions_per_method": int(s * valid.sum()),
        },
    }


def paired_composition_bootstrap(
    a_correct, b_correct, a_valid, b_valid, pair_prompt_ids, *,
    seed_ids_a, seed_ids_b,
    pair_ids_a=None, pair_ids_b=None, rubric_ids_a=None, rubric_ids_b=None,
    pair_prompt_ids_b=None,
    n_bootstrap: int = 10_000, random_seed: int = 190,
    resample_seeds: bool = True, batch_size: int = 128,
    return_bootstrap_distribution: bool = False,
) -> dict:
    """Paired percentile 95% CI matching the observed rubric-first estimand.

    Prompts are independent sampling clusters. Each draw samples G prompts
    uniformly with replacement, keeping *all* their pair/rubric/seed cells.
    A prompt selected twice contributes its pair counts twice. Within each
    rubric, sum correct decisions and divide by its resampled valid-pair
    count; then average rubrics equally, then seeds equally. Thus prompts
    have equal sampling probability, not equal accuracy weight.

    By default S paired seed IDs are also sampled with replacement, once per
    replicate and shared across all prompts and both methods. This crossed
    (pigeonhole) bootstrap is an approximate random-seed inference; with
    resample_seeds=False the interval is conditional on the supplied models.
    Paired seeds must represent the same planned runs in both methods.

    Rubrics, fitted models, DEV assignments, TEST standardization, and human
    tie masks are fixed. Arrays alone cannot bootstrap upstream operations.
    Sparse support: if any draw loses an originally active rubric, ci95 is
    None with explicit counts/status; no zero filling, redraw conditioning,
    or silently changing the target rubric set. The same applies to G < 2.
    An observed all-tie rubric is excluded just as in the main metric.

    Results are JSON serializable. Optional bootstrap_distribution contains
    one float (or None for an undefined draw) per requested replicate.
    """
    for name, value, minimum in (("n_bootstrap", n_bootstrap, 2), ("batch_size", batch_size, 1)):
        if isinstance(value, bool) or not isinstance(value, Integral) or value < minimum:
            raise ValueError(f"{name} must be an integer >= {minimum}")
    if isinstance(random_seed, bool) or not isinstance(random_seed, Integral) or random_seed < 0:
        raise ValueError("random_seed must be a nonnegative integer")
    if not isinstance(resample_seeds, (bool, np.bool_)):
        raise ValueError("resample_seeds must be boolean")
    result = composition_point_difference(
        a_correct, b_correct, a_valid, b_valid,
        seed_ids_a=seed_ids_a, seed_ids_b=seed_ids_b,
        pair_ids_a=pair_ids_a, pair_ids_b=pair_ids_b,
        rubric_ids_a=rubric_ids_a, rubric_ids_b=rubric_ids_b,
    )
    a, b, valid, den, active = _validate_decisions(a_correct, b_correct, a_valid, b_valid)
    s, p, r = a.shape
    pids = _ids(pair_prompt_ids, p, "pair_prompt_ids", unique=False)
    if pair_prompt_ids_b is not None:
        right = _ids(pair_prompt_ids_b, p, "pair_prompt_ids_b", unique=False)
        if not np.array_equal(pids, right):
            raise ValueError("prompt axis alignment error: pair prompt IDs must match in order")
    result["alignment_checked"]["prompt_b"] = pair_prompt_ids_b is not None
    unique, inverse = np.unique(pids, return_inverse=True)
    g = len(unique)

    # Sufficient statistics preserve every original pair while avoiding
    # repeated materialization of [seed, resampled pairs, rubric] arrays.
    cluster_valid = np.zeros((g, r), dtype=np.float64)
    np.add.at(cluster_valid, inverse, valid)
    cluster_diff = np.zeros((g, s, int(active.sum())), dtype=np.float64)
    for si in range(s):
        np.add.at(cluster_diff[:, si], inverse, (b[si] - a[si])[:, active])
    flat_diff = cluster_diff.reshape(g, -1)
    active_valid = cluster_valid[:, active]
    seed_stream, prompt_stream = np.random.SeedSequence(int(random_seed)).spawn(2)
    seed_rng, prompt_rng = np.random.default_rng(seed_stream), np.random.default_rng(prompt_stream)
    boot = np.full(n_bootstrap, np.nan)
    missing_by_rubric = np.zeros(int(active.sum()), dtype=np.int64)
    for start in range(0, n_bootstrap, batch_size):
        size = min(batch_size, n_bootstrap - start)
        pw = prompt_rng.multinomial(g, np.full(g, 1.0 / g), size=size)
        sw = seed_rng.multinomial(s, np.full(s, 1.0 / s), size=size) if resample_seeds else np.ones((size, s))
        denom = pw @ active_valid
        missing = denom == 0
        missing_by_rubric += missing.sum(axis=0)
        good = ~missing.any(axis=1)
        numer = (pw @ flat_diff).reshape(size, s, -1)
        per_seed = (numer[good] / denom[good, None, :]).mean(axis=2)
        boot[start + np.flatnonzero(good)] = (per_seed * sw[good]).sum(axis=1) / s

    undefined = int((~np.isfinite(boot)).sum())
    status = "ok"
    if g < 2:
        status = "insufficient_prompt_clusters"
    elif undefined:
        status = "undefined_rubric_in_bootstrap"
    notes = [
        "Independent exchangeable prompt clusters; arbitrary dependence within each prompt is retained.",
        "Prompt universe is the prompts present in the supplied pair table, including all-tie prompts; absent prompts cannot be inferred.",
        "All pairs must be within their stated prompt; no response may be shared across purportedly independent prompt clusters.",
        "Fixed rubric set and fixed human tie masks shared across methods/seeds; model ties count as incorrect.",
        "Conditional on saved decisions, TEST standardization, DEV calibration, training data, and fitted models apart from optional empirical seed resampling.",
        "Percentile interval is approximate, not a finite-sample coverage guarantee; increasing replicates does not replace more prompts or seeds.",
    ]
    if resample_seeds:
        notes.append("Training seed pairs are exchangeable, independent of sampled prompts; crossed seed/prompt resampling can be conservative and is weak with few seeds.")
    else:
        notes.append("Seed resampling disabled: interval conditional on these trained models, with a shared prompt draw for every seed.")
    if s == 1:
        notes.append("Only one paired seed: training-seed uncertainty cannot be estimated, even when seed resampling is requested.")
    if not active.all():
        notes.append("Observed all-tie rubrics are omitted to match the main nanmean; this is not an identified full-rubric population target.")
    if undefined:
        notes.append("CI withheld: at least one replicate lacks a valid pair for an originally active rubric; no failed draws were discarded or redrawn.")
    if g < 2:
        notes.append("CI withheld: fewer than two prompt clusters cannot estimate prompt sampling uncertainty.")
    result["counts"].update({
        "prompts": g,
        "pairs_per_prompt": np.bincount(inverse, minlength=g).tolist(),
        "prompts_with_valid_pairs": int((cluster_valid.sum(axis=1) > 0).sum()),
        "valid_prompts_per_rubric": (cluster_valid > 0).sum(axis=0).tolist(),
        "bootstrap_undefined_by_active_rubric": missing_by_rubric.tolist(),
    })
    result.update({
        "ci95": np.quantile(boot, [0.025, 0.975]).tolist() if status == "ok" else None,
        "interval_status": status,
        "confidence_level": 0.95,
        "interval_method": "percentile; numpy linear quantiles",
        "bootstrap_samples": int(n_bootstrap),
        "bootstrap_defined_samples": int(n_bootstrap - undefined),
        "bootstrap_undefined_samples": undefined,
        "bootstrap_mean": float(boot.mean()) if not undefined else None,
        "random_seed": int(random_seed),
        "resample_seeds": bool(resample_seeds),
        "units": "paired prompt clusters and paired training seeds" if resample_seeds else "paired prompt clusters conditional on supplied seeds",
        "prompt_weighting": "uniform cluster sampling; valid-pair weighting within each rubric",
        "assumptions": notes,
    })
    if return_bootstrap_distribution:
        result["bootstrap_distribution"] = [float(v) if np.isfinite(v) else None for v in boot]
    return result