paired_test.py

#!/usr/bin/env python3
import argparse
import glob
import json
import os
from collections import defaultdict

import networkx as nx
from scipy.stats import ttest_rel, wilcoxon

from experiment_paths import REPO_ROOT


SOURCE_DIR = "challenge-source"
INPUT_FILE = os.path.join(SOURCE_DIR, "TaskA-Flagship", "train_task_a.json")
TAXONOMY_RELATIONS = {"is-a"}


def parse_tuple(line: str) -> tuple[str, str, str]:
    line = line.strip()
    if line.startswith("(") and line.endswith(")"):
        line = line[1:-1]
    parts = [part.strip() for part in line.split(",")]
    if len(parts) != 3:
        raise ValueError(f"Invalid tuple format: {line}")
    return (parts[0], parts[1], parts[2])


def load_output_file(path: str) -> list[tuple[str, str, str]]:
    with open(path, "r", encoding="utf-8") as f:
        return [parse_tuple(line.strip()) for line in f if line.strip()]


def load_gold() -> dict[str, list[tuple[str, str, str]]]:
    source_path = os.path.join(REPO_ROOT, INPUT_FILE)
    with open(source_path, "r", encoding="utf-8") as f:
        source = json.load(f)

    gold: dict[str, list[tuple[str, str, str]]] = {}
    for sample in source:
        gold[sample["id"]] = [tuple(triple) for triple in sample["primitive-ontology-triples"]]
    return gold


def normalize_sample_id(sample_id: str) -> str:
    normalized = sample_id
    while normalized.endswith(".txt"):
        normalized = normalized[:-4]
    return normalized


def normalize(text: str) -> str:
    return " ".join(str(text).lower().strip().split())


def normalize_triples(triples: list[tuple[str, str, str]]) -> list[tuple[str, str, str]]:
    return [tuple(normalize(x) for x in triple) for triple in triples]


def edge_f1(gold_edges: list[tuple[str, str, str]], pred_edges: list[tuple[str, str, str]]) -> float:
    gold_set = set(gold_edges)
    pred_set = set(pred_edges)

    intersection = gold_set & pred_set
    if len(gold_set) == 0 and len(pred_set) == 0:
        return 1.0
    if len(pred_set) == 0 or len(gold_set) == 0:
        return 0.0

    precision = len(intersection) / len(pred_set)
    recall = len(intersection) / len(gold_set)
    if precision + recall == 0:
        return 0.0
    return 2 * precision * recall / (precision + recall)


def get_neighborhood(triples: list[tuple[str, str, str]]) -> dict[str, set[tuple[str, str, str]]]:
    neigh: dict[str, set[tuple[str, str, str]]] = defaultdict(set)
    for subject, predicate, obj in triples:
        neigh[subject].add(("OUT", predicate, obj))
        neigh[obj].add(("IN", predicate, subject))
    return neigh


def neighborhood_similarity(gold_triples: list[tuple[str, str, str]], pred_triples: list[tuple[str, str, str]]) -> float:
    gold_neigh = get_neighborhood(gold_triples)
    pred_neigh = get_neighborhood(pred_triples)

    all_nodes = set(gold_neigh.keys()) | set(pred_neigh.keys())
    scores: list[float] = []

    for node in all_nodes:
        gold_node = gold_neigh.get(node, set())
        pred_node = pred_neigh.get(node, set())
        union = gold_node | pred_node
        if len(union) == 0:
            continue
        jaccard = len(gold_node & pred_node) / len(union)
        scores.append(jaccard)

    return sum(scores) / len(scores) if scores else 0.0


def build_taxonomy_graph(triples: list[tuple[str, str, str]]) -> nx.DiGraph:
    graph = nx.DiGraph()
    for subject, predicate, obj in triples:
        if predicate in TAXONOMY_RELATIONS:
            graph.add_edge(subject, obj)
    return graph


def taxonomy_similarity(gold_triples: list[tuple[str, str, str]], pred_triples: list[tuple[str, str, str]]) -> float:
    gold_graph = build_taxonomy_graph(gold_triples)
    pred_graph = build_taxonomy_graph(pred_triples)

    all_nodes = set(gold_graph.nodes()) | set(pred_graph.nodes())
    scores: list[float] = []

    for node in all_nodes:
        gold_anc = nx.ancestors(gold_graph, node) if node in gold_graph else set()
        pred_anc = nx.ancestors(pred_graph, node) if node in pred_graph else set()
        gold_desc = nx.descendants(gold_graph, node) if node in gold_graph else set()
        pred_desc = nx.descendants(pred_graph, node) if node in pred_graph else set()

        gold_rel = gold_anc | gold_desc
        pred_rel = pred_anc | pred_desc
        union = gold_rel | pred_rel
        if len(union) == 0:
            continue
        jaccard = len(gold_rel & pred_rel) / len(union)
        scores.append(jaccard)

    return sum(scores) / len(scores) if scores else 0.0


def exact_match_metrics(
    gold_triples: list[tuple[str, str, str]],
    pred_triples: list[tuple[str, str, str]],
) -> dict[str, float]:
    gold_triples = normalize_triples(gold_triples)
    pred_triples = normalize_triples(pred_triples)

    metric_edge_f1 = edge_f1(gold_triples, pred_triples)
    metric_neighborhood = neighborhood_similarity(gold_triples, pred_triples)
    metric_taxonomy = taxonomy_similarity(gold_triples, pred_triples)
    metric_graph = (metric_edge_f1 + metric_neighborhood + metric_taxonomy) / 3

    return {
        "exact_match.edge_f1": metric_edge_f1,
        "exact_match.neighborhood_similarity": metric_neighborhood,
        "exact_match.taxonomy_similarity": metric_taxonomy,
        "exact_match.graph_similarity": metric_graph,
    }


def list_output_files(experiment: str) -> dict[str, str]:
    pattern = os.path.join(REPO_ROOT, "experiments", experiment, "output", "*.txt")
    files = sorted(glob.glob(pattern))
    return {
        normalize_sample_id(os.path.splitext(os.path.basename(path))[0]): path
        for path in files
    }


def paired_values(
    experiment_a: str,
    experiment_b: str,
    metric_name: str,
) -> tuple[list[float], list[float], int]:
    files_a = list_output_files(experiment_a)
    files_b = list_output_files(experiment_b)
    ids_a = set(files_a.keys())
    ids_b = set(files_b.keys())

    if ids_a != ids_b:
        missing_in_b = sorted(ids_a - ids_b)
        missing_in_a = sorted(ids_b - ids_a)
        detail_lines = []
        if missing_in_b:
            detail_lines.append(f"Missing in {experiment_b}: {missing_in_b[:5]}")
        if missing_in_a:
            detail_lines.append(f"Missing in {experiment_a}: {missing_in_a[:5]}")
        detail = " | ".join(detail_lines)
        raise ValueError(f"Experiment outputs do not contain the same ids. {detail}")

    gold = load_gold()
    experiment_ids = sorted(ids_a)

    values_a: list[float] = []
    values_b: list[float] = []
    skipped = 0

    for sample_id in experiment_ids:
        gold_triples = gold.get(sample_id)
        if not gold_triples:
            skipped += 1
            continue

        pred_a = load_output_file(files_a[sample_id])
        pred_b = load_output_file(files_b[sample_id])

        metrics_a = exact_match_metrics(gold_triples, pred_a)
        metrics_b = exact_match_metrics(gold_triples, pred_b)

        values_a.append(metrics_a[metric_name])
        values_b.append(metrics_b[metric_name])

    return values_a, values_b, skipped


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Run a paired statistical test between two experiments and print p-value."
    )
    parser.add_argument("experiment_a", help="First experiment name (folder under experiments/)")
    parser.add_argument("experiment_b", help="Second experiment name (folder under experiments/)")
    parser.add_argument(
        "--metric",
        default="exact_match.graph_similarity",
        choices=[
            "exact_match.graph_similarity",
            "exact_match.edge_f1",
            "exact_match.neighborhood_similarity",
            "exact_match.taxonomy_similarity",
        ],
        help="Per-record metric used for paired comparison.",
    )
    parser.add_argument(
        "--test",
        default="ttest",
        choices=["ttest", "wilcoxon"],
        help="Paired test to run.",
    )
    args = parser.parse_args()

    values_a, values_b, skipped = paired_values(args.experiment_a, args.experiment_b, args.metric)

    if len(values_a) == 0:
        raise ValueError("No comparable records found.")

    if args.test == "ttest":
        _, p_value = ttest_rel(values_a, values_b)
    else:
        _, p_value = wilcoxon(values_a, values_b)

    if skipped:
        print(f"warning: skipped {skipped} records with missing gold triples")
    print(p_value)


if __name__ == "__main__":
    main()