"""Post-hoc analyses for the downstream plan-validity story.""" from __future__ import annotations import argparse import csv import math from collections import defaultdict from pathlib import Path import numpy as np from sklearn.linear_model import LogisticRegression from sklearn.metrics import ( accuracy_score, average_precision_score, balanced_accuracy_score, confusion_matrix, f1_score, precision_score, recall_score, roc_auc_score, ) from sklearn.preprocessing import StandardScaler from code.downstream.features import load_feature_matrix DEFAULT_FAMILIES = [ "ad_lstm_wl_delta", "ad_xgb_wl_delta", "dd_lstm_shortest_path_delta", "dd_xgb_wl_delta", ] DEFAULT_SEEDS = [13, 23, 37] DEFAULT_SPLITS = ["train", "validation", "test-interpolation", "test-extrapolation"] def compute_metrics(y_true, probs, threshold=0.5) -> dict: labels = np.asarray(y_true, dtype=np.int64) scores = np.asarray(probs, dtype=np.float64) preds = (scores >= threshold).astype(np.int64) unique = np.unique(labels) auroc = float("nan") auprc = float("nan") bal_acc = float("nan") if len(unique) == 2: auroc = float(roc_auc_score(labels, scores)) auprc = float(average_precision_score(labels, scores)) bal_acc = float(balanced_accuracy_score(labels, preds)) tn, fp, fn, tp = confusion_matrix(labels, preds, labels=[0, 1]).ravel() return { "num_examples": int(len(labels)), "positive_rate": float(labels.mean()) if len(labels) else float("nan"), "threshold": float(threshold), "accuracy": float(accuracy_score(labels, preds)) if len(labels) else float("nan"), "balanced_accuracy": bal_acc, "auroc": auroc, "auprc": auprc, "f1": float(f1_score(labels, preds, zero_division=0)), "precision": float(precision_score(labels, preds, zero_division=0)), "recall": float(recall_score(labels, preds, zero_division=0)), "tn": int(tn), "fp": int(fp), "fn": int(fn), "tp": int(tp), } def tune_threshold(y_true, probs, objective: str) -> float: """Choose a threshold on validation predictions.""" scores = np.asarray(probs, dtype=np.float64) candidates = sorted(set([0.0, 0.5, 1.0] + scores.tolist())) best_threshold = 0.5 best_score = -float("inf") for threshold in candidates: metrics = compute_metrics(y_true, probs, threshold) score = float(metrics[objective]) if math.isnan(score): continue if score > best_score: best_score = score best_threshold = float(threshold) return best_threshold def read_predictions(path: Path) -> dict[str, dict[str, list]]: by_split: dict[str, dict[str, list]] = defaultdict(lambda: {"y": [], "probs": []}) with path.open(newline="", encoding="utf-8") as f: for row in csv.DictReader(f): split = row["split"] by_split[split]["y"].append(int(row["label_valid"])) by_split[split]["probs"].append(float(row["prob_valid"])) return by_split def analyze_mlp_thresholds(args) -> list[dict]: rows: list[dict] = [] mlp_root = Path(args.mlp_results_dir) if args.mlp_results_dir else Path(args.dataset_dir) / "mlp_results" for family in args.families: for seed in args.seeds: pred_path = ( mlp_root / family / f"source_seed_{seed}" / f"head_seed_{seed}" / "predictions.csv" ) if not pred_path.exists(): continue preds = read_predictions(pred_path) if "validation" not in preds: continue threshold = tune_threshold( preds["validation"]["y"], preds["validation"]["probs"], args.threshold_objective, ) for split, values in preds.items(): metrics = compute_metrics(values["y"], values["probs"], threshold) metrics.update( { "method": f"mlp_threshold_tuned_{args.threshold_objective}", "family": family, "seed": seed, "split": split, } ) rows.append(metrics) return rows def evaluate_baselines(args) -> list[dict]: rows: list[dict] = [] for seed in args.seeds: split_data = { split: load_feature_matrix( Path(args.dataset_dir) / "features" / args.baseline_reference_family / f"seed_{seed}" / f"{split}.npz" ) for split in DEFAULT_SPLITS } rows.extend(evaluate_majority(seed, split_data)) rows.extend(evaluate_plan_length_logreg(seed, split_data)) rows.extend(evaluate_corruption_diagnostic(seed, split_data)) return rows def evaluate_majority(seed: int, split_data: dict) -> list[dict]: train_y = np.asarray(split_data["train"]["y"], dtype=np.int64) positive_rate = float(train_y.mean()) majority_prob = 1.0 if positive_rate >= 0.5 else 0.0 rows = [] for split, data in split_data.items(): y = np.asarray(data["y"], dtype=np.int64) probs = np.full(len(y), majority_prob, dtype=np.float64) metrics = compute_metrics(y, probs, threshold=0.5) metrics.update( { "method": "majority", "family": "baseline", "seed": seed, "split": split, } ) rows.append(metrics) return rows def evaluate_plan_length_logreg(seed: int, split_data: dict) -> list[dict]: feature_names = [str(name) for name in split_data["train"]["feature_names"]] cols = [ idx for idx, name in enumerate(feature_names) if name in {"plan_len", "log_plan_len", "plan_to_budget_ratio"} ] if not cols: raise RuntimeError("Plan-length feature columns were not found.") train_X = np.asarray(split_data["train"]["X"], dtype=np.float32)[:, cols] train_y = np.asarray(split_data["train"]["y"], dtype=np.int64) scaler = StandardScaler() train_X = scaler.fit_transform(train_X) model = LogisticRegression(class_weight="balanced", max_iter=1000, random_state=seed) model.fit(train_X, train_y) rows = [] for split, data in split_data.items(): X = scaler.transform(np.asarray(data["X"], dtype=np.float32)[:, cols]) y = np.asarray(data["y"], dtype=np.int64) probs = model.predict_proba(X)[:, 1] metrics = compute_metrics(y, probs, threshold=0.5) metrics.update( { "method": "plan_length_logreg", "family": "baseline", "seed": seed, "split": split, } ) rows.append(metrics) return rows def evaluate_corruption_diagnostic(seed: int, split_data: dict) -> list[dict]: train_types = sorted(set(str(item) for item in split_data["train"]["corruption_types"])) def one_hot(data): values = [str(item) for item in data["corruption_types"]] arr = np.zeros((len(values), len(train_types)), dtype=np.float32) index = {name: idx for idx, name in enumerate(train_types)} for row_idx, value in enumerate(values): if value in index: arr[row_idx, index[value]] = 1.0 return arr train_X = one_hot(split_data["train"]) train_y = np.asarray(split_data["train"]["y"], dtype=np.int64) model = LogisticRegression(class_weight="balanced", max_iter=1000, random_state=seed) model.fit(train_X, train_y) rows = [] for split, data in split_data.items(): y = np.asarray(data["y"], dtype=np.int64) probs = model.predict_proba(one_hot(data))[:, 1] metrics = compute_metrics(y, probs, threshold=0.5) metrics.update( { "method": "corruption_type_diagnostic", "family": "diagnostic", "seed": seed, "split": split, } ) rows.append(metrics) return rows def write_rows(path: Path, rows: list[dict]) -> None: path.parent.mkdir(parents=True, exist_ok=True) fields = [ "method", "family", "seed", "split", "num_examples", "positive_rate", "threshold", "accuracy", "balanced_accuracy", "auroc", "auprc", "f1", "precision", "recall", "tn", "fp", "fn", "tp", ] with path.open("w", newline="", encoding="utf-8") as f: writer = csv.DictWriter(f, fieldnames=fields) writer.writeheader() for row in rows: writer.writerow({field: row.get(field, "") for field in fields}) def mean_std(values): vals = [float(value) for value in values if not math.isnan(float(value))] if not vals: return float("nan"), float("nan") mean = sum(vals) / len(vals) var = sum((value - mean) ** 2 for value in vals) / len(vals) return mean, math.sqrt(var) def write_summary(path: Path, rows: list[dict]) -> None: grouped: dict[tuple[str, str, str], list[dict]] = defaultdict(list) for row in rows: grouped[(row["method"], row["family"], row["split"])].append(row) fields = [ "method", "family", "split", "num_runs", "accuracy_mean", "accuracy_std", "balanced_accuracy_mean", "balanced_accuracy_std", "auroc_mean", "auroc_std", "auprc_mean", "auprc_std", "f1_mean", "f1_std", "precision_mean", "precision_std", "recall_mean", "recall_std", ] path.parent.mkdir(parents=True, exist_ok=True) with path.open("w", newline="", encoding="utf-8") as f: writer = csv.DictWriter(f, fieldnames=fields) writer.writeheader() for key in sorted(grouped): method, family, split = key group = grouped[key] out = { "method": method, "family": family, "split": split, "num_runs": len(group), } for metric in ["accuracy", "balanced_accuracy", "auroc", "auprc", "f1", "precision", "recall"]: mean, std = mean_std([row[metric] for row in group]) out[f"{metric}_mean"] = mean out[f"{metric}_std"] = std writer.writerow(out) def main() -> None: parser = argparse.ArgumentParser(description="Analyze downstream validity experiments.") parser.add_argument("--dataset_dir", default="outputs/downstream_validity/frozen_transition_validity") parser.add_argument( "--mlp_results_dir", default=None, help="Optional alternate MLP results directory, e.g. an ablation folder.", ) parser.add_argument("--output_dir", default=None) parser.add_argument("--families", nargs="+", default=DEFAULT_FAMILIES) parser.add_argument("--seeds", nargs="+", type=int, default=DEFAULT_SEEDS) parser.add_argument("--threshold_objective", choices=["balanced_accuracy", "f1"], default="balanced_accuracy") parser.add_argument("--baseline_reference_family", default="ad_lstm_wl_delta") parser.add_argument("--skip_mlp_thresholds", action="store_true") parser.add_argument("--skip_baselines", action="store_true") args = parser.parse_args() args.dataset_dir = str(Path(args.dataset_dir).resolve()) output_dir = Path(args.output_dir).resolve() if args.output_dir else Path(args.dataset_dir) / "analysis" threshold_rows = [] if args.skip_mlp_thresholds else analyze_mlp_thresholds(args) baseline_rows = [] if args.skip_baselines else evaluate_baselines(args) all_rows = threshold_rows + baseline_rows write_rows(output_dir / "tuned_threshold_metrics.csv", threshold_rows) write_summary(output_dir / "tuned_threshold_summary.csv", threshold_rows) write_rows(output_dir / "baseline_metrics.csv", baseline_rows) write_summary(output_dir / "baseline_summary.csv", baseline_rows) write_rows(output_dir / "story_metrics.csv", all_rows) write_summary(output_dir / "story_summary.csv", all_rows) print(f"Wrote downstream story analysis to {output_dir}") if __name__ == "__main__": main()