"""Select a deterministic, diverse panel of errors for the final TFM analysis.""" from __future__ import annotations import argparse import json import re from collections import Counter from pathlib import Path import pandas as pd def _model_name(path: Path) -> str: value = str(path).lower() if "b2_1_random_forest" in value: return "B2.1 Random Forest" if "b2_2_xgboost" in value: return "B2.2 XGBoost" if "baseline3" in value: return "B3 EdgeMLP" if "graphsage" in value: return "GraphSAGE" raise ValueError(f"Cannot infer model from prediction path: {path}") def _seed(path: Path) -> int: match = re.search(r"seed[-_](\d+)", str(path)) if match: return int(match.group(1)) raise ValueError(f"Prediction path does not identify its seed: {path}") def _range_pattern(value: str) -> str: clauses = [clause.strip() for clause in value.split(",") if clause.strip()] has_lower = any(clause.startswith((">", ">=")) for clause in clauses) has_upper = any(clause.startswith(("<", "<=")) for clause in clauses) if any(clause.startswith(("==", "===")) for clause in clauses): return "exact_or_wildcard" if has_lower and has_upper: return "bounded" if has_lower: return "lower_only" if has_upper: return "upper_only" if any(clause.startswith("~=") for clause in clauses): return "compatible_release" return "other" def _persistence_level(error_count: int) -> str: """Classify recurrence across the 40 comparable model/seed runs.""" if error_count >= 30: return "high" if error_count >= 15: return "medium" return "low" def _prediction_paths(roots: list[Path]) -> list[Path]: paths: list[Path] = [] for root in roots: paths.extend(root.glob("**/predictions/*.parquet")) return sorted(set(paths)) def _load_runs(paths: list[Path]) -> pd.DataFrame: frames: list[pd.DataFrame] = [] required = {"pair_id", "label", "risk_incompatible", "prediction"} for path in paths: frame = pd.read_parquet(path) missing = required.difference(frame.columns) if missing: raise ValueError(f"{path} is missing columns: {sorted(missing)}") if len(frame) != 1_200 or frame["pair_id"].duplicated().any(): raise ValueError(f"Unexpected test prediction contract in {path}") frame = frame[list(required)].copy() frame["model"] = _model_name(path) frame["seed"] = _seed(path) frames.append(frame) if not frames: raise ValueError("No prediction files were found") result = pd.concat(frames, ignore_index=True) result["error_kind"] = "" result.loc[(result["label"] == 0) & (result["prediction"] == 1), "error_kind"] = "FP" result.loc[(result["label"] == 1) & (result["prediction"] == 0), "error_kind"] = "FN" return result[result["error_kind"] != ""].copy() def _aggregate_errors(errors: pd.DataFrame, dataset: pd.DataFrame) -> pd.DataFrame: grouped = errors.groupby(["pair_id", "error_kind", "label"], as_index=False).agg( error_count=("prediction", "size"), mean_risk=("risk_incompatible", "mean"), min_risk=("risk_incompatible", "min"), max_risk=("risk_incompatible", "max"), models=("model", lambda values: ", ".join(sorted(set(values)))), seeds=("seed", lambda values: ", ".join(str(value) for value in sorted(set(values)))), ) test = dataset.loc[dataset["split"] == "test"].copy() if len(test) != 1_200 or test["pair_id"].duplicated().any(): raise ValueError("Dataset test split does not contain 1,200 unique pairs") merged = grouped.merge(test, on="pair_id", how="left", validate="one_to_one", suffixes=("", "_data")) if merged["pa"].isna().any(): raise ValueError("Some prediction pair IDs do not exist in the dataset") if not (merged["label"] == merged["label_data"]).all(): raise ValueError("Prediction labels disagree with the dataset") merged = merged.drop(columns=["label_data"]) merged["package_pair"] = merged["pa"] + " -> " + merged["pb"] merged["range_pattern"] = merged["declared_range"].astype(str).map(_range_pattern) merged["decision_impact"] = merged["error_kind"].map( { "FN": "high: incompatible pair accepted", "FP": "moderate: compatible pair rejected", } ) merged["persistence"] = merged["error_count"].map(_persistence_level) merged["temporal_gap_days"] = ( pd.to_datetime(merged["va_upload_time"], utc=True) - pd.to_datetime(merged["vb_upload_time"], utc=True) ).abs().dt.total_seconds() / 86_400 return merged def _select_diverse(frame: pd.DataFrame, kind: str, size: int, max_per_parent: int) -> pd.DataFrame: candidates = frame.loc[frame["error_kind"] == kind].copy() ascending_risk = kind == "FN" candidates = candidates.sort_values( ["error_count", "mean_risk", "pair_id"], ascending=[False, ascending_risk, True], kind="mergesort", ) selected: list[int] = [] parent_counts: Counter[str] = Counter() for index, row in candidates.iterrows(): if parent_counts[row["pa"]] >= max_per_parent: continue selected.append(index) parent_counts[row["pa"]] += 1 if len(selected) == size: break if len(selected) < size: for index in candidates.index: if index not in selected: selected.append(index) if len(selected) == size: break if len(selected) < size: raise ValueError(f"Only {len(selected)} distinct {kind} cases are available; requested {size}") result = candidates.loc[selected].copy() result["selection_rank"] = range(1, len(result) + 1) return result def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--dataset", type=Path, required=True) parser.add_argument("--prediction-root", type=Path, action="append", required=True) parser.add_argument("--output-dir", type=Path, required=True) parser.add_argument("--per-kind", type=int, default=30) parser.add_argument("--max-per-parent", type=int, default=5) args = parser.parse_args() dataset = pd.read_parquet(args.dataset) paths = _prediction_paths(args.prediction_root) errors = _load_runs(paths) aggregate = _aggregate_errors(errors, dataset) panel = pd.concat( [ _select_diverse(aggregate, "FP", args.per_kind, args.max_per_parent), _select_diverse(aggregate, "FN", args.per_kind, args.max_per_parent), ], ignore_index=True, ) args.output_dir.mkdir(parents=True, exist_ok=True) aggregate.to_parquet(args.output_dir / "all_errors.parquet", index=False) panel.to_parquet(args.output_dir / "selected_60.parquet", index=False) panel.to_csv(args.output_dir / "selected_60.csv", index=False) summary = { "dataset": str(args.dataset), "prediction_files": [str(path) for path in paths], "runs": int(len(paths)), "models": sorted(errors["model"].unique()), "seeds": sorted(int(seed) for seed in errors["seed"].unique()), "error_events": int(len(errors)), "distinct_error_pairs": int(aggregate["pair_id"].nunique()), "selected": panel["error_kind"].value_counts().sort_index().to_dict(), "decision_impact": panel["decision_impact"].value_counts().sort_index().to_dict(), "persistence": panel["persistence"].value_counts().sort_index().to_dict(), "selection_rule": { "primary": "descending error frequency across model/seed runs", "secondary_fp": "descending mean predicted risk", "secondary_fn": "ascending mean predicted risk", "parent_cap": args.max_per_parent, "tie_break": "pair_id ascending", }, } (args.output_dir / "summary.json").write_text( json.dumps(summary, indent=2, ensure_ascii=False, sort_keys=True), encoding="utf-8" ) print(json.dumps(summary, indent=2, ensure_ascii=False, sort_keys=True)) if __name__ == "__main__": main()