eval/audit_split_manifest.py

#!/usr/bin/env python3
"""Audit a MaizeGuard CSV manifest for group leakage and source concentration.

The manifest is expected to contain ``split``, ``label``, ``source_id``,
``record_id`` and ``group_id`` columns. The command prints a deterministic JSON
report and exits non-zero when a biological/collection group or record appears
in more than one split.
"""

from __future__ import annotations

import argparse
import csv
import json
from collections import Counter, defaultdict
from pathlib import Path

REQUIRED_COLUMNS = {"split", "label", "source_id", "record_id", "group_id"}


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("manifest", type=Path, help="CSV split manifest to audit")
    parser.add_argument(
        "--max-source-share",
        type=float,
        default=0.80,
        help="Flag classes where one source supplies more than this share (default: 0.80)",
    )
    return parser.parse_args()


def audit(path: Path, max_source_share: float) -> dict[str, object]:
    groups: dict[tuple[str, str], set[str]] = defaultdict(set)
    records: dict[str, set[str]] = defaultdict(set)
    split_counts: Counter[str] = Counter()
    label_counts: Counter[str] = Counter()
    label_sources: dict[str, Counter[str]] = defaultdict(Counter)

    with path.open(newline="", encoding="utf-8") as handle:
        reader = csv.DictReader(handle)
        missing = REQUIRED_COLUMNS - set(reader.fieldnames or ())
        if missing:
            raise ValueError(f"manifest is missing required columns: {sorted(missing)}")
        for row in reader:
            split = row["split"].strip()
            label = row["label"].strip()
            source = row["source_id"].strip()
            record = row["record_id"].strip()
            group = row["group_id"].strip()
            if not all((split, label, source, record, group)):
                raise ValueError("manifest contains an empty required field")
            groups[(source, group)].add(split)
            records[record].add(split)
            split_counts[split] += 1
            label_counts[label] += 1
            label_sources[label][source] += 1

    cross_split_groups = sorted(
        (source, group, sorted(splits))
        for (source, group), splits in groups.items()
        if len(splits) > 1
    )
    cross_split_records = sorted(
        (record, sorted(splits)) for record, splits in records.items() if len(splits) > 1
    )
    concentration: dict[str, object] = {}
    concentrated_labels: list[str] = []
    for label in sorted(label_sources):
        total = label_counts[label]
        source, count = label_sources[label].most_common(1)[0]
        share = count / total
        concentration[label] = {
            "record_count": total,
            "largest_source": source,
            "largest_source_count": count,
            "largest_source_share": round(share, 6),
        }
        if share > max_source_share:
            concentrated_labels.append(label)

    return {
        "schema_version": "1",
        "manifest": path.name,
        "record_count": sum(split_counts.values()),
        "group_count": len(groups),
        "split_counts": dict(sorted(split_counts.items())),
        "cross_split_group_count": len(cross_split_groups),
        "cross_split_group_examples": [
            {"source_id": source, "group_id": group, "splits": splits}
            for source, group, splits in cross_split_groups[:20]
        ],
        "cross_split_record_count": len(cross_split_records),
        "cross_split_record_examples": [
            {"record_id": record, "splits": splits}
            for record, splits in cross_split_records[:20]
        ],
        "class_source_concentration": concentration,
        "classes_above_source_share_threshold": concentrated_labels,
        "max_source_share_threshold": max_source_share,
        "leakage_gate_passed": not cross_split_groups and not cross_split_records,
    }


def main() -> int:
    args = parse_args()
    if not 0 < args.max_source_share <= 1:
        raise SystemExit("--max-source-share must be greater than 0 and at most 1")
    report = audit(args.manifest, args.max_source_share)
    print(json.dumps(report, indent=2, sort_keys=True))
    return 0 if report["leakage_gate_passed"] else 2


if __name__ == "__main__":
    raise SystemExit(main())