scripts/audit_labels.py

#!/usr/bin/env python3
"""Label variance audit for PIMT pyramid targets.

Computes per-tier (Top/Mid/Base) and per-descriptor variance and positive-label
frequency across the full dataset.  Flags tiers whose mean label variance is
below a threshold and descriptors with <5 positive examples.

Additionally computes and caches training-set mean/variance of intensity scalars
(log₁₀ΣOAV per tier snapshot, Σx_liquid) to artifacts/scalar_stats_v1.json.

Usage:
    python scripts/audit_labels.py [--data data/empirical_dataset_v4.jsonl]
"""
from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np


VARIANCE_THRESHOLD = 0.01
MIN_POSITIVE_EXAMPLES = 5

# Step→tier mapping for 49-step trajectories (~3600s total):
#   top: 0-12 (~0-15min), mid: 13-33 (~15min-1hr), base: 34-48 (1hr+)
STEP_RANGES = {"top": (0, 13), "mid": (13, 34), "base": (34, 49)}
TIER_NAMES = ["top", "mid", "base"]


def load_dataset(path: str) -> list[dict]:
    with open(path) as f:
        return [json.loads(line) for line in f]


def load_vocabulary(path: str = "data/pyrfume_vocabulary.json") -> list[str]:
    with open(path) as f:
        data = json.load(f)
    return data["vocabulary"] if isinstance(data, dict) else data


def audit_pyramid_labels(
    records: list[dict], vocab: list[str], variance_threshold: float = VARIANCE_THRESHOLD
) -> dict:
    """Compute per-tier variance, positive-label frequency, and descriptor stats."""
    mixtures = [r for r in records if not r.get("is_control")]
    pyramids = []
    for r in mixtures:
        if "pyramid_targets" in r:
            pt = np.array(r["pyramid_targets"], dtype=np.float32)
            if pt.shape == (3, 138):
                pyramids.append(pt)

    if not pyramids:
        raise ValueError("No records with pyramid_targets found in dataset")

    all_pyramids = np.stack(pyramids)  # (N, 3, 138)
    N = len(all_pyramids)
    report = {"N": N, "tiers": {}}

    for i, name in enumerate(TIER_NAMES):
        tier = all_pyramids[:, i, :]  # (N, 138)
        positive_freq = tier.mean(axis=0)
        per_desc_var = tier.var(axis=0)
        mean_var = float(per_desc_var.mean())
        records_positive = int((tier.sum(axis=1) > 0).sum())
        descs_active = int(np.count_nonzero(positive_freq))
        descs_below_min = int(np.sum(positive_freq * N < MIN_POSITIVE_EXAMPLES))

        tier_report = {
            "records_with_any_positive": records_positive,
            "records_with_any_positive_pct": round(100 * records_positive / N, 2),
            "descriptors_active": descs_active,
            "descriptors_active_pct": round(100 * descs_active / 138, 2),
            "mean_variance": round(mean_var, 6),
            "variance_threshold": variance_threshold,
            "passes_variance_threshold": mean_var >= variance_threshold,
            "descriptors_below_min_examples": descs_below_min,
            "mean_label_frequency": round(float(tier.mean()), 6),
            "top_descriptors": [],
        }

        top_idx = np.argsort(positive_freq)[::-1][:20]
        for idx in top_idx:
            n_pos = int(positive_freq[idx] * N)
            desc_name = vocab[idx] if idx < len(vocab) else f"desc_{idx}"
            tier_report["top_descriptors"].append({
                "idx": int(idx),
                "name": desc_name,
                "positive_count": n_pos,
                "frequency": round(float(positive_freq[idx]), 6),
                "variance": round(float(per_desc_var[idx]), 6),
            })

        report["tiers"][name] = tier_report

    return report


def compute_scalar_stats(records: list[dict]) -> dict:
    """Compute training-set mean/variance of intensity scalars per tier."""
    mixtures = [r for r in records if not r.get("is_control")]
    log_oav = {t: [] for t in TIER_NAMES}
    x_liq = {t: [] for t in TIER_NAMES}

    for r in mixtures:
        traj = r.get("trajectory", [])
        for tier_name, (start, end) in STEP_RANGES.items():
            for step in traj[start:end]:
                oavs = step.get("OAV", {})
                xliqs = step.get("x_liquid", {})
                log_oav[tier_name].append(
                    sum(np.log10(max(v, 1e-10)) for v in oavs.values())
                )
                x_liq[tier_name].append(sum(xliqs.values()))

    stats = {}
    for tier_name in TIER_NAMES:
        lo = np.array(log_oav[tier_name])
        xl = np.array(x_liq[tier_name])
        stats[f"log10_oav_sum_{tier_name}"] = {
            "mean": round(float(lo.mean()), 6),
            "std": round(float(lo.std()), 6),
            "min": round(float(lo.min()), 6),
            "max": round(float(lo.max()), 6),
            "count": len(lo),
        }
        stats[f"x_liquid_sum_{tier_name}"] = {
            "mean": round(float(xl.mean()), 6),
            "std": round(float(xl.std()), 6),
            "min": round(float(xl.min()), 6),
            "max": round(float(xl.max()), 6),
            "count": len(xl),
        }

    all_lo = np.concatenate([np.array(log_oav[t]) for t in TIER_NAMES])
    all_xl = np.concatenate([np.array(x_liq[t]) for t in TIER_NAMES])
    stats["log10_oav_sum_overall"] = {
        "mean": round(float(all_lo.mean()), 6),
        "std": round(float(all_lo.std()), 6),
        "min": round(float(all_lo.min()), 6),
        "max": round(float(all_lo.max()), 6),
    }
    stats["x_liquid_sum_overall"] = {
        "mean": round(float(all_xl.mean()), 6),
        "std": round(float(all_xl.std()), 6),
        "min": round(float(all_xl.min()), 6),
        "max": round(float(all_xl.max()), 6),
    }
    return stats


def generate_histograms(report: dict, output_dir: Path) -> None:
    """Generate histogram plots of per-descriptor positive-label frequency."""
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    output_dir.mkdir(parents=True, exist_ok=True)

    fig, axes = plt.subplots(1, 3, figsize=(18, 5))
    fig.suptitle("Per-Descriptor Positive-Label Frequency by Tier", fontsize=14)

    for ax, tier_name in zip(axes, TIER_NAMES):
        tier_data = report["tiers"][tier_name]
        freqs = [d["frequency"] for d in tier_data["top_descriptors"]]
        names = [d["name"] for d in tier_data["top_descriptors"]]

        status = "PASS" if tier_data["passes_variance_threshold"] else "FAIL"
        ax.barh(range(len(freqs)), freqs, color="steelblue")
        ax.set_yticks(range(len(names)))
        ax.set_yticklabels(names, fontsize=8)
        ax.set_xlabel("Frequency")
        ax.set_title(f"{tier_name.upper()} ({status}, var={tier_data['mean_variance']:.4f})")
        ax.invert_yaxis()

    plt.tight_layout()
    path = output_dir / "label_frequency_histograms.png"
    fig.savefig(path, dpi=150)
    plt.close(fig)
    print(f"📊 Histogram saved to {path}")


def main():
    parser = argparse.ArgumentParser(description="Label variance audit for PIMT pyramid targets")
    parser.add_argument("--data", default="data/empirical_dataset_v4.jsonl")
    parser.add_argument("--variance-threshold", type=float, default=VARIANCE_THRESHOLD)
    parser.add_argument("--output", default="reports/label_variance_audit_v4.json")
    args = parser.parse_args()

    print(f"Loading dataset from {args.data}...")
    records = load_dataset(args.data)
    vocab = load_vocabulary()

    print(f"Auditing pyramid labels ({len(records)} records)...")
    label_report = audit_pyramid_labels(records, vocab, args.variance_threshold)

    print("Computing scalar statistics...")
    scalar_stats = compute_scalar_stats(records)

    # Build full report
    report = {
        "dataset": args.data,
        "total_records": len(records),
        "controls": sum(1 for r in records if r.get("is_control")),
        "mixtures": sum(1 for r in records if not r.get("is_control")),
        "variance_threshold": args.variance_threshold,
        "min_positive_examples": MIN_POSITIVE_EXAMPLES,
        **label_report,
        "scalar_stats": scalar_stats,
    }

    # Save JSON report
    report_path = Path(args.output)
    report_path.parent.mkdir(parents=True, exist_ok=True)
    report_path.write_text(json.dumps(report, indent=2, ensure_ascii=False))
    print(f"📄 Report saved to {report_path}")

    # Save scalar stats artifact
    scalar_artifact = {
        "version": "v1",
        "dataset": args.data,
        "description": "Training-set statistics for standardization in ConcentrationAwarePyramidHead",
        "log10_oav_sum": scalar_stats["log10_oav_sum_overall"],
        "x_liquid_sum": scalar_stats["x_liquid_sum_overall"],
        "per_tier": {
            t: {
                "log10_oav_sum": scalar_stats[f"log10_oav_sum_{t}"],
                "x_liquid_sum": scalar_stats[f"x_liquid_sum_{t}"],
            }
            for t in TIER_NAMES
        },
    }
    scalar_path = Path("artifacts/scalar_stats_v1.json")
    scalar_path.parent.mkdir(parents=True, exist_ok=True)
    scalar_path.write_text(json.dumps(scalar_artifact, indent=2, ensure_ascii=False))
    print(f"📄 Scalar stats artifact saved to {scalar_path}")

    # Generate histogram plots
    generate_histograms(report, Path("reports/figures"))

    # Print verdict
    print(f"\n{'='*60}")
    print("VARIANCE AUDIT VERDICT")
    print(f"{'='*60}")
    all_pass = True
    for name in TIER_NAMES:
        t = report["tiers"][name]
        status = "✅ PASS" if t["passes_variance_threshold"] else "❌ FAIL"
        print(f"  {name}: {status} (variance={t['mean_variance']:.6f}, threshold={args.variance_threshold})")
        print(f"    Records with any positive: {t['records_with_any_positive']}/{report['N']} ({t['records_with_any_positive_pct']}%)")
        print(f"    Active descriptors: {t['descriptors_active']}/138 ({t['descriptors_active_pct']}%)")
        if not t["passes_variance_threshold"]:
            all_pass = False

    if not all_pass:
        print(f"\n⚠️  STOP CONDITION TRIGGERED")
        print(f"One or more tiers fail the variance threshold ({args.variance_threshold}).")
        print(f"The top/mid tier labels are too sparse for reliable training.")
        print(f"Recommended: enrich tier labels before proceeding to Phase 1.")
    else:
        print(f"\n✅ All tiers pass variance threshold.")


if __name__ == "__main__":
    main()