scripts/audit_chart_feature_sources.py

#!/usr/bin/env python
from __future__ import annotations

import argparse
import json
import math
import subprocess
import sys
from collections import Counter, defaultdict
from pathlib import Path
from typing import Any

PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

import numpy as np  # noqa: E402

from cil.chart_features import CHART_FEATURE_MODES, chart_feature_dim  # noqa: E402


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description="Audit deployment-visible chart feature sources in CIL chart shards."
    )
    parser.add_argument(
        "--indexes",
        nargs="+",
        type=Path,
        default=[
            Path("data/cil_charts/train/index.json"),
            Path("data/cil_charts/val/index.json"),
            Path("data/cil_charts/test/index.json"),
        ],
    )
    parser.add_argument("--out-dir", type=Path, default=Path("runs/chart_feature_audit"))
    args = parser.parse_args(argv)

    out_dir = args.out_dir
    out_dir.mkdir(parents=True, exist_ok=True)
    _write_provenance(out_dir, args)

    split_rows = []
    row_details = []
    for index_path in args.indexes:
        index = json.loads(index_path.read_text())
        split = str(index.get("split", index_path.parent.name))
        counts: Counter[str] = Counter()
        task_counts: Counter[str] = Counter()
        feature_dims: dict[str, int] = {}
        for shard in index.get("shards", []):
            shard_path = index_path.parent / shard["path"]
            with np.load(shard_path, allow_pickle=False) as data:
                metadata_values = data["metadata_json"]
                base_actions = data["base_action"]
                action_shapes = data["action_shape"]
                task_ids = data["task_id"]
                chart_ids = data["chart_id"]
                for row in range(metadata_values.shape[0]):
                    metadata = _json_loads(str(metadata_values[row]))
                    task_id = str(task_ids[row])
                    task_counts[task_id] += 1
                    counts["rows"] += 1
                    for key in (
                        "observation_embedding_path",
                        "object_embedding_path",
                        "observation_ref",
                        "scene_id",
                        "instruction",
                        "source_dataset",
                    ):
                        if metadata.get(key):
                            counts[f"{key}_present"] += 1
                    if counts["sample_details"] < 5:
                        shape = tuple(int(value) for value in action_shapes[row])
                        flat_count = int(math.prod(shape))
                        base = np.asarray(base_actions[row][:flat_count], dtype=np.float32).reshape(shape)
                        for mode in CHART_FEATURE_MODES:
                            feature_dims[mode] = chart_feature_dim(base, mode=mode)
                        row_details.append(
                            {
                                "split": split,
                                "chart_id": str(chart_ids[row]),
                                "task_id": task_id,
                                "has_observation_embedding_path": bool(
                                    metadata.get("observation_embedding_path")
                                ),
                                "has_observation_ref": bool(metadata.get("observation_ref")),
                                "has_scene_id": bool(metadata.get("scene_id")),
                                "has_instruction": bool(metadata.get("instruction")),
                            }
                        )
                        counts["sample_details"] += 1
        rows = int(counts["rows"])
        split_rows.append(
            {
                "split": split,
                "index": str(index_path),
                "rows": rows,
                "charts": int(index.get("num_groups_exported", 0)),
                "retrieval_index_allowed": bool(index.get("retrieval_index_allowed")),
                "include_outcomes": bool(index.get("include_outcomes")),
                "observation_embedding_path_present": int(
                    counts["observation_embedding_path_present"]
                ),
                "observation_embedding_path_rate": _rate(
                    counts["observation_embedding_path_present"], rows
                ),
                "observation_ref_present": int(counts["observation_ref_present"]),
                "observation_ref_rate": _rate(counts["observation_ref_present"], rows),
                "object_embedding_path_present": int(counts["object_embedding_path_present"]),
                "object_embedding_path_rate": _rate(
                    counts["object_embedding_path_present"], rows
                ),
                "scene_id_present": int(counts["scene_id_present"]),
                "scene_id_rate": _rate(counts["scene_id_present"], rows),
                "instruction_present": int(counts["instruction_present"]),
                "instruction_rate": _rate(counts["instruction_present"], rows),
                "source_dataset_present": int(counts["source_dataset_present"]),
                "source_dataset_rate": _rate(counts["source_dataset_present"], rows),
                "feature_dims": feature_dims,
                "task_counts": dict(sorted(task_counts.items())),
            }
        )

    metrics = {
        "report_type": "chart_feature_source_audit",
        "schema_version": 1,
        "indexes": [str(path) for path in args.indexes],
        "splits": split_rows,
        "data_hash": {row["split"]: _index_hash(row, "content_hash") for row in split_rows},
        "split_hash": {row["split"]: _index_hash(row, "split_hash") for row in split_rows},
        "sample_details": row_details,
        "conclusion": _conclusion(split_rows),
    }
    (out_dir / "metrics.json").write_text(json.dumps(metrics, indent=2, sort_keys=True) + "\n")
    (out_dir / "metrics_by_task.json").write_text(_metrics_by_task(split_rows) + "\n")
    (out_dir / "metrics_by_seed.json").write_text("{}\n")
    (out_dir / "table.tex").write_text(_table(split_rows) + "\n")
    (out_dir / "report.md").write_text(_report(metrics) + "\n")
    (out_dir / "train.log").write_text("not a training run; audited chart feature sources\n")
    (out_dir / "eval.log").write_text("audited chart feature sources in exported chart indexes\n")
    print(json.dumps({"out_dir": str(out_dir), "splits": len(split_rows)}, indent=2))
    return 0


def _json_loads(value: str) -> dict[str, Any]:
    try:
        payload = json.loads(value)
    except json.JSONDecodeError:
        return {}
    return payload if isinstance(payload, dict) else {}


def _rate(count: int, total: int) -> float:
    return float(count) / float(total) if total else 0.0


def _index_hash(row: dict[str, Any], key: str) -> Any:
    path = Path(str(row["index"]))
    if not path.exists():
        return None
    return json.loads(path.read_text()).get(key)


def _conclusion(rows: list[dict[str, Any]]) -> str:
    if all(float(row["observation_embedding_path_rate"]) > 0.0 for row in rows):
        if all(float(row["object_embedding_path_rate"]) > 0.0 for row in rows):
            return (
                "Current chart exports expose observation and object-layout embeddings "
                "in every split; visual/object-layout chart features can be evaluated "
                "with leakage-audited indexes."
            )
        return (
            "Current chart exports expose observation embeddings in every split; "
            "visual chart features can be evaluated with leakage-audited indexes."
        )
    if all(float(row["observation_ref_rate"]) > 0.0 for row in rows):
        return (
            "Current chart exports expose raw observation refs in every split but "
            "not observation embeddings; run scripts/export_chart_observation_embeddings.py "
            "before claiming embedded visual chart tokens."
        )
    if all(float(row["observation_embedding_path_rate"]) == 0.0 for row in rows) and all(
        float(row["observation_ref_rate"]) == 0.0 for row in rows
    ):
        return (
            "Current chart exports do not contain observation embeddings or raw observation refs; "
            "visual/object-centric chart tokens require a new export or embedding pass."
        )
    return (
        "Observation feature availability is partial across splits; inspect per-split "
        "embedding/ref rates before running visual chart-token experiments."
    )


def _metrics_by_task(rows: list[dict[str, Any]]) -> str:
    payload: dict[str, dict[str, int]] = defaultdict(dict)
    for row in rows:
        for task, count in row["task_counts"].items():
            payload[task][str(row["split"])] = int(count)
    return json.dumps(payload, indent=2, sort_keys=True)


def _table(rows: list[dict[str, Any]]) -> str:
    lines = [
        "% Auto-generated by scripts/audit_chart_feature_sources.py",
        "\\begin{tabular}{lrrrrr}",
        "\\toprule",
        "Split & Rows & ObsEmbed & ObjEmbed & ObsRef & Instruction \\\\",
        "\\midrule",
    ]
    for row in rows:
        lines.append(
            f"{row['split']} & {row['rows']} & "
            f"{row['observation_embedding_path_present']} & "
            f"{row['object_embedding_path_present']} & "
            f"{row['observation_ref_present']} & "
            f"{row['instruction_present']} \\\\"
        )
    lines.extend(["\\bottomrule", "\\end{tabular}"])
    return "\n".join(lines)


def _report(metrics: dict[str, Any]) -> str:
    lines = [
        "# Chart Feature Source Audit",
        "",
        metrics["conclusion"],
        "",
        "| Split | Rows | Obs embedding path | Object embedding path | Obs ref | Scene id | Instruction | Feature dims |",
        "| --- | ---: | ---: | ---: | ---: | ---: | ---: | --- |",
    ]
    for row in metrics["splits"]:
        lines.append(
            f"| {row['split']} | {row['rows']} | "
            f"{row['observation_embedding_path_present']} ({row['observation_embedding_path_rate']:.2%}) | "
            f"{row['object_embedding_path_present']} ({row['object_embedding_path_rate']:.2%}) | "
            f"{row['observation_ref_present']} ({row['observation_ref_rate']:.2%}) | "
            f"{row['scene_id_present']} ({row['scene_id_rate']:.2%}) | "
            f"{row['instruction_present']} ({row['instruction_rate']:.2%}) | "
            f"{json.dumps(row['feature_dims'], sort_keys=True)} |"
        )
    return "\n".join(lines)


def _write_provenance(out_dir: Path, args: argparse.Namespace) -> None:
    (out_dir / "config.yaml").write_text(
        "\n".join(f"{key}: {value}" for key, value in sorted(vars(args).items())) + "\n"
    )
    (out_dir / "command.txt").write_text(
        "python scripts/audit_chart_feature_sources.py " + " ".join(sys.argv[1:]) + "\n"
    )
    (out_dir / "git_hash.txt").write_text(_run(["git", "rev-parse", "HEAD"]) + "\n")
    hashes = {}
    for index_path in args.indexes:
        index = json.loads(index_path.read_text())
        hashes[str(index_path)] = {
            "content_hash": index.get("content_hash"),
            "split_hash": index.get("split_hash"),
        }
    (out_dir / "data_hash.txt").write_text(json.dumps(hashes, indent=2, sort_keys=True) + "\n")
    (out_dir / "split_hash.txt").write_text(json.dumps(hashes, indent=2, sort_keys=True) + "\n")


def _run(command: list[str]) -> str:
    try:
        return subprocess.check_output(command, cwd=PROJECT_ROOT, text=True).strip()
    except (subprocess.CalledProcessError, FileNotFoundError):
        return ""


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