scripts/analyze_drill_graph_subgroups.py

#!/usr/bin/env python3
"""Post-hoc subgroup audit for the frozen DRiLL graph OOF predictions.

This script does not fit or select a model.  It asks whether the observed graph
gain is confined to the easy document-clustering proxy, and whether queries
whose candidate pools expose VBPL-derived causal-labelled links behave
differently from queries for which those links are unavailable.
"""

from __future__ import annotations

import argparse
import json
import platform
import sys
from datetime import datetime, timezone
from pathlib import Path
from typing import Mapping, Sequence

import numpy as np

REPO_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(REPO_ROOT / "src"))

from scripts.run_drill_graph_evaluation import (  # noqa: E402
    _ranking_observations,
    _stable_candidate_rankings,
)
from telen.drill_graph import (  # noqa: E402
    DrillLegalGraph,
    build_drill_legal_graph,
    load_drill_graph_dataset,
)
from telen.static_benchmark import (  # noqa: E402
    paired_bootstrap_difference,
    sha256_file,
)

CAUSAL_RELATIONS = ("AMENDS", "REPEALS", "REPLACES", "TERMINATES_VERSION")


def causal_candidate_masks(
    graph: DrillLegalGraph,
    candidates: np.ndarray,
    relevant_indices: Sequence[frozenset[int]],
) -> tuple[np.ndarray, np.ndarray]:
    """Return query masks for exposed causal links and gold-incident links."""

    if candidates.ndim != 2 or len(candidates) != len(relevant_indices):
        raise ValueError("candidate matrix and qrels must align")
    candidate_exposed = np.zeros(len(candidates), dtype=bool)
    gold_incident = np.zeros(len(candidates), dtype=bool)
    reverse: dict[str, dict[str, set[str]]] = {
        relation: {} for relation in CAUSAL_RELATIONS
    }
    for relation in CAUSAL_RELATIONS:
        for source, targets in graph.document_neighbors[relation].items():
            for target in targets:
                reverse[relation].setdefault(target, set()).add(source)

    for query_index, (row, relevant) in enumerate(
        zip(candidates, relevant_indices, strict=True)
    ):
        candidate_documents = {
            graph.article_document_ids[int(aid)] for aid in row
        }
        gold_documents = {
            graph.article_document_ids[int(aid)] for aid in relevant
        }
        for relation in CAUSAL_RELATIONS:
            outgoing = graph.document_neighbors[relation]
            incoming = reverse[relation]
            for document in candidate_documents:
                linked_candidates = (
                    set(outgoing.get(document, ()))
                    | incoming.get(document, set())
                ) & candidate_documents
                if linked_candidates:
                    candidate_exposed[query_index] = True
                    if document in gold_documents or linked_candidates & gold_documents:
                        gold_incident[query_index] = True
            if candidate_exposed[query_index] and gold_incident[query_index]:
                break
    return candidate_exposed, gold_incident


def build_subgroup_masks(
    article_document_ids: Sequence[str],
    relevant_indices: Sequence[frozenset[int]],
    causal_exposed: np.ndarray,
    causal_gold_incident: np.ndarray,
) -> Mapping[str, np.ndarray]:
    """Create mutually interpretable, label-audited query strata."""

    if not (
        len(relevant_indices) == len(causal_exposed) == len(causal_gold_incident)
    ):
        raise ValueError("subgroup inputs must align")
    qrel_counts = np.asarray(list(map(len, relevant_indices)), dtype=np.int64)
    gold_document_counts = np.asarray(
        [
            len({article_document_ids[int(aid)] for aid in relevant})
            for relevant in relevant_indices
        ],
        dtype=np.int64,
    )
    all_queries = np.ones(len(relevant_indices), dtype=bool)
    return {
        "all_queries": all_queries,
        "single_qrel": qrel_counts == 1,
        "multiple_qrels": qrel_counts > 1,
        "multiple_qrels_single_gold_document":
            (qrel_counts > 1) & (gold_document_counts == 1),
        "multiple_gold_documents": gold_document_counts > 1,
        "causal_link_exposed_in_candidate_pool": causal_exposed,
        "no_causal_link_exposed_in_candidate_pool": ~causal_exposed,
        "gold_document_incident_to_exposed_causal_link": causal_gold_incident,
        "gold_document_not_incident_to_exposed_causal_link": ~causal_gold_incident,
    }


def summarize_subgroups(
    observations: Mapping[str, np.ndarray],
    masks: Mapping[str, np.ndarray],
    *,
    bootstrap_samples: int,
    bootstrap_seed: int,
) -> dict[str, object]:
    """Summarize fixed graph/text/sham nDCG observations by subgroup."""

    required = {"graph", "text", "sham", "causal_ablation"}
    if set(observations) != required:
        raise ValueError(f"observations must contain {sorted(required)}")
    size = len(observations["graph"])
    if any(len(values) != size for values in observations.values()):
        raise ValueError("observation vectors must align")
    output: dict[str, object] = {}
    for name, mask in masks.items():
        indices = np.flatnonzero(mask)
        if not len(indices):
            output[name] = {"query_count": 0}
            continue
        graph = observations["graph"][indices]
        text = observations["text"][indices]
        sham = observations["sham"][indices]
        causal_ablation = observations["causal_ablation"][indices]
        output[name] = {
            "query_count": int(len(indices)),
            "mean_ndcg_at_10": {
                "learned_text_residual": float(text.mean()),
                "learned_graph_residual": float(graph.mean()),
                "learned_sham_graph_residual": float(sham.mean()),
                "graph_vbpl_causal_labeled_edges_removed": float(
                    causal_ablation.mean()
                ),
            },
            "graph_minus_text": paired_bootstrap_difference(
                text,
                graph,
                samples=bootstrap_samples,
                seed=bootstrap_seed,
            ),
            "graph_minus_sham": paired_bootstrap_difference(
                sham,
                graph,
                samples=bootstrap_samples,
                seed=bootstrap_seed,
            ),
            "graph_minus_causal_edge_ablation": paired_bootstrap_difference(
                causal_ablation,
                graph,
                samples=bootstrap_samples,
                seed=bootstrap_seed,
            ),
        }
    return output


def run(args: argparse.Namespace) -> dict[str, object]:
    dataset = load_drill_graph_dataset(args.corpus, args.questions)
    graph = build_drill_legal_graph(
        dataset,
        metadata_path=args.metadata,
        canonical_relationships_path=args.relationships,
    )
    predictions = np.load(args.predictions)
    candidates = np.asarray(predictions["candidates"], dtype=np.int32)
    causal_exposed, causal_gold_incident = causal_candidate_masks(
        graph, candidates, dataset.retrieval.relevant_indices
    )
    masks = build_subgroup_masks(
        dataset.article_document_ids,
        dataset.retrieval.relevant_indices,
        causal_exposed,
        causal_gold_incident,
    )
    system_keys = {
        "graph": "scores_learned_graph_residual",
        "text": "scores_learned_text_residual",
        "sham": "scores_learned_sham_graph_residual",
        "causal_ablation": "scores_graph_vbpl_causal_labeled_edges_removed",
    }
    ndcg_observations: dict[str, np.ndarray] = {}
    for role, key in system_keys.items():
        rankings = _stable_candidate_rankings(
            candidates,
            np.asarray(predictions[key], dtype=np.float32),
            dataset.retrieval.corpus_ids,
        )
        _summary, per_query = _ranking_observations(
            rankings, dataset.retrieval.relevant_indices
        )
        ndcg_observations[role] = per_query["ndcg_at_10"]

    subgroups = summarize_subgroups(
        ndcg_observations,
        masks,
        bootstrap_samples=args.bootstrap_samples,
        bootstrap_seed=args.bootstrap_seed,
    )
    receipt: dict[str, object] = {
        "schema_version": 1,
        "artifact_type": "drill_graph_fixed_prediction_posthoc_subgroup_audit",
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "claim_status": "exploratory_subgroup_analysis_not_model_selection",
        "protocol": {
            "models_refit": False,
            "hyperparameters_selected": False,
            "predictions_frozen_before_subgroup_definition": True,
            "multiplicity_adjustment": None,
            "interpretation": (
                "descriptive heterogeneity audit; subgroup intervals do not replace "
                "the predeclared all-query graph credit gate"
            ),
        },
        "dataset_audit": dataset.retrieval.audit,
        "graph_audit": graph.audit,
        "subgroups": subgroups,
        "artifacts": {
            "corpus": {"path": str(args.corpus), "sha256": sha256_file(args.corpus)},
            "questions": {
                "path": str(args.questions),
                "sha256": sha256_file(args.questions),
            },
            "metadata": {
                "path": str(args.metadata),
                "sha256": sha256_file(args.metadata),
            },
            "relationships": {
                "path": str(args.relationships),
                "sha256": sha256_file(args.relationships),
            },
            "predictions": {
                "path": str(args.predictions),
                "sha256": sha256_file(args.predictions),
            },
        },
        "limitations": [
            "Subgroups were defined after observing the primary aggregate result and are exploratory.",
            "VBPL-derived causal labels and canonical directions come from a pinned community mirror and are not expert-verified event truth.",
            "Candidate exposure depends on the shared RRF top-100 pool; absence of an exposed edge is not absence from the full legal graph.",
        ],
        "environment": {"python": platform.python_version(), "numpy": np.__version__},
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(
        json.dumps(receipt, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
    )
    return receipt


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--corpus",
        type=Path,
        default=REPO_ROOT
        / "tmp/telen/drill-public/legal-corpus/data/train-00000-of-00001.parquet",
    )
    parser.add_argument(
        "--questions",
        type=Path,
        default=REPO_ROOT
        / "tmp/telen/drill-public/train/data/train-00000-of-00001.parquet",
    )
    parser.add_argument(
        "--metadata",
        type=Path,
        default=REPO_ROOT
        / "tmp/telen/vietnamese-legal-documents-source/data/metadata.parquet",
    )
    parser.add_argument(
        "--relationships",
        type=Path,
        default=REPO_ROOT
        / "tmp/telen/bootstrap-audit/canonical_relationships.parquet",
    )
    parser.add_argument(
        "--predictions",
        type=Path,
        default=REPO_ROOT
        / "tmp/telen/static-benchmarks/cache/drill/drill-rrf-graph-semantic-rank64-oof-predictions.npz",
    )
    parser.add_argument("--bootstrap-samples", type=int, default=10_000)
    parser.add_argument("--bootstrap-seed", type=int, default=20260829)
    parser.add_argument(
        "--output",
        type=Path,
        default=REPO_ROOT
        / "research/results/telen_drill_graph_subgroup_audit_2026-08-29.json",
    )
    return parser.parse_args()


if __name__ == "__main__":
    result = run(parse_args())
    print(json.dumps(result, ensure_ascii=False, indent=2))