from __future__ import annotations import argparse import json from dataclasses import asdict, dataclass from pathlib import Path from typing import Any import pandas as pd DEFAULT_METHODS = ("standard", "missingness_aware", "weighted_missingness_aware") @dataclass class AppendixPhase2Config: alpha_sweep_dirs: tuple[Path, ...] grouping_sweep_dir: Path selection_sweep_dir: Path output_dir: Path methods: tuple[str, ...] = DEFAULT_METHODS model_type: str = "xgboost" def _load_csv(path: Path) -> pd.DataFrame: return pd.read_csv(path) if path.exists() else pd.DataFrame() def build_alpha_sweep_summary( *, repeated_summary_frames: list[pd.DataFrame], methods: tuple[str, ...], model_type: str, ) -> pd.DataFrame: if not repeated_summary_frames: return pd.DataFrame() combined = pd.concat(repeated_summary_frames, ignore_index=True) filtered = combined[ combined["method"].isin(methods) & combined["model_type"].eq(model_type) ].copy() if filtered.empty: return filtered preferred_columns = [ "alpha", "method", "run_count", "empirical_coverage_mean", "empirical_coverage_std", "max_group_coverage_gap_mean", "max_group_coverage_gap_std", "average_set_size_mean", "average_set_size_std", "worst_hospital_coverage_mean", "worst_hospital_coverage_std", "wcr", "source_dir", ] filtered = filtered.sort_values(["alpha", "method"]).reset_index(drop=True) return filtered[[column for column in preferred_columns if column in filtered.columns]] def build_grouping_baselines_summary( *, repeated_summary: pd.DataFrame, methods: tuple[str, ...], model_type: str, ) -> pd.DataFrame: filtered = repeated_summary[ repeated_summary["method"].isin(methods) & repeated_summary["model_type"].eq(model_type) ].copy() if filtered.empty: return filtered base_columns = [ "missingness_grouping_strategy", "model_type", "alpha", "selection_fraction", "mask_strategy", "mask_rate", "weighted_shrinkage_lambda", ] standard = ( filtered[filtered["method"] == "standard"][base_columns + ["max_group_coverage_gap_mean"]] .rename(columns={"max_group_coverage_gap_mean": "standard_gap_mean"}) ) summary = filtered.merge(standard, on=base_columns, how="left", validate="many_to_one") summary["gap_reduction_vs_standard_mean"] = ( summary["standard_gap_mean"] - summary["max_group_coverage_gap_mean"] ) preferred_columns = [ "missingness_grouping_strategy", "method", "alpha", "run_count", "empirical_coverage_mean", "max_group_coverage_gap_mean", "gap_reduction_vs_standard_mean", "average_set_size_mean", "worst_hospital_coverage_mean", "wcr", ] return summary.sort_values( ["missingness_grouping_strategy", "method"] ).reset_index(drop=True)[preferred_columns] def build_selection_stability_summary( *, selected_variable_stability: pd.DataFrame, selected_variable_jaccard: pd.DataFrame, repeated_summary: pd.DataFrame, model_type: str, ) -> pd.DataFrame: if selected_variable_stability.empty: return pd.DataFrame() filtered = selected_variable_stability[selected_variable_stability["model_type"].eq(model_type)].copy() if filtered.empty: return filtered base_columns = [ "experiment", "method", "min_hospital_admissions", "alpha", "selection_fraction", "model_type", "missingness_grouping_strategy", "mask_strategy", "mask_rate", "selective_feature_group", "weighted_shrinkage_lambda", ] grouped = filtered.groupby(base_columns, dropna=False) rows: list[dict[str, Any]] = [] for key, frame in grouped: top = frame.sort_values(["frequency", "count", "selected_variable"], ascending=[False, False, True]).iloc[0] top3_mass = float(frame.sort_values(["frequency", "count"], ascending=[False, False]).head(3)["frequency"].sum()) row = dict(zip(base_columns, key if isinstance(key, tuple) else (key,), strict=False)) row.update( { "top_selected_variable": str(top["selected_variable"]), "top_selected_variable_frequency": float(top["frequency"]), "top3_frequency_mass": top3_mass, "unique_selected_variables": int(frame["selected_variable"].nunique()), } ) rows.append(row) summary = pd.DataFrame(rows) if not selected_variable_jaccard.empty: summary = summary.merge( selected_variable_jaccard[ base_columns + ["run_count", "mean_pairwise_jaccard", "min_pairwise_jaccard"] ], on=base_columns, how="left", validate="one_to_one", ) if not repeated_summary.empty: repeated_filtered = repeated_summary[repeated_summary["model_type"].eq(model_type)].copy() keep = base_columns + [ "empirical_coverage_mean", "max_group_coverage_gap_mean", "average_set_size_mean", "wcr", ] summary = summary.merge( repeated_filtered[keep], on=base_columns, how="left", validate="one_to_one", ) return summary.sort_values(["method", "alpha"]).reset_index(drop=True) def build_selection_performance_by_variable( *, overall_summary: pd.DataFrame, model_type: str, ) -> pd.DataFrame: if overall_summary.empty or "selected_variable" not in overall_summary.columns: return pd.DataFrame() filtered = overall_summary[ overall_summary["model_type"].eq(model_type) & overall_summary["selected_variable"].notna() ].copy() if filtered.empty: return filtered group_columns = ["method", "selected_variable"] aggregated = filtered.groupby(group_columns, dropna=False, as_index=False).agg( run_count=("run_id", "nunique"), empirical_coverage_mean=("empirical_coverage", "mean"), empirical_coverage_std=("empirical_coverage", "std"), max_group_coverage_gap_mean=("max_group_coverage_gap", "mean"), max_group_coverage_gap_std=("max_group_coverage_gap", "std"), average_set_size_mean=("average_set_size", "mean"), average_set_size_std=("average_set_size", "std"), worst_hospital_coverage_mean=("worst_hospital_coverage", "mean"), worst_hospital_coverage_std=("worst_hospital_coverage", "std"), ) aggregated["empirical_coverage_std"] = aggregated["empirical_coverage_std"].fillna(0.0) aggregated["max_group_coverage_gap_std"] = aggregated["max_group_coverage_gap_std"].fillna(0.0) aggregated["average_set_size_std"] = aggregated["average_set_size_std"].fillna(0.0) aggregated["worst_hospital_coverage_std"] = aggregated["worst_hospital_coverage_std"].fillna(0.0) return aggregated.sort_values(["method", "max_group_coverage_gap_mean", "selected_variable"]).reset_index(drop=True) def run_appendix_phase2_analysis(config: AppendixPhase2Config) -> dict[str, Path]: config.output_dir.mkdir(parents=True, exist_ok=True) alpha_repeated_frames: list[pd.DataFrame] = [] for sweep_dir in config.alpha_sweep_dirs: frame = _load_csv(sweep_dir / "repeated_summary.csv") if frame.empty: continue frame = frame.copy() frame["source_dir"] = str(sweep_dir) alpha_repeated_frames.append(frame) grouping_repeated_summary = _load_csv(config.grouping_sweep_dir / "repeated_summary.csv") selection_stability = _load_csv(config.selection_sweep_dir / "selected_variable_stability.csv") selection_jaccard = _load_csv(config.selection_sweep_dir / "selected_variable_jaccard_summary.csv") selection_repeated = _load_csv(config.selection_sweep_dir / "repeated_summary.csv") selection_overall = _load_csv(config.selection_sweep_dir / "overall_summary.csv") alpha_summary = build_alpha_sweep_summary( repeated_summary_frames=alpha_repeated_frames, methods=config.methods, model_type=config.model_type, ) grouping_summary = build_grouping_baselines_summary( repeated_summary=grouping_repeated_summary, methods=config.methods, model_type=config.model_type, ) stability_summary = build_selection_stability_summary( selected_variable_stability=selection_stability, selected_variable_jaccard=selection_jaccard, repeated_summary=selection_repeated, model_type=config.model_type, ) performance_by_variable = build_selection_performance_by_variable( overall_summary=selection_overall, model_type=config.model_type, ) alpha_path = config.output_dir / "alpha_sweep_summary.csv" grouping_path = config.output_dir / "grouping_baselines_summary.csv" stability_path = config.output_dir / "selection_stability_summary.csv" performance_path = config.output_dir / "selection_performance_by_variable.csv" config_path = config.output_dir / "config.json" manifest_path = config.output_dir / "manifest.json" alpha_summary.to_csv(alpha_path, index=False) grouping_summary.to_csv(grouping_path, index=False) stability_summary.to_csv(stability_path, index=False) performance_by_variable.to_csv(performance_path, index=False) config_path.write_text(json.dumps(asdict(config), indent=2, default=str), encoding="utf-8") manifest_path.write_text( json.dumps( { "alpha_sweep_summary": str(alpha_path), "grouping_baselines_summary": str(grouping_path), "selection_stability_summary": str(stability_path), "selection_performance_by_variable": str(performance_path), "config": str(config_path), }, indent=2, sort_keys=True, ), encoding="utf-8", ) return { "alpha_sweep_summary": alpha_path, "grouping_baselines_summary": grouping_path, "selection_stability_summary": stability_path, "selection_performance_by_variable": performance_path, "config": config_path, "manifest": manifest_path, } def _parse_path_list(value: str) -> tuple[Path, ...]: return tuple(Path(item.strip()) for item in value.split(",") if item.strip()) def _parse_methods(value: str) -> tuple[str, ...]: return tuple(item.strip() for item in value.split(",") if item.strip()) def build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser(prog="sepsis-mcp-appendix-phase2-analysis") parser.add_argument("--alpha-sweep-dirs", type=_parse_path_list, required=True) parser.add_argument("--grouping-sweep-dir", type=Path, required=True) parser.add_argument("--selection-sweep-dir", type=Path, required=True) parser.add_argument("--methods", type=_parse_methods, default=",".join(DEFAULT_METHODS)) parser.add_argument("--model-type", default="xgboost") parser.add_argument("--output-dir", type=Path, required=True) return parser def main(argv: list[str] | None = None) -> int: parser = build_parser() args = parser.parse_args(argv) run_appendix_phase2_analysis( AppendixPhase2Config( alpha_sweep_dirs=args.alpha_sweep_dirs, grouping_sweep_dir=args.grouping_sweep_dir, selection_sweep_dir=args.selection_sweep_dir, output_dir=args.output_dir, methods=args.methods, model_type=args.model_type, ) ) return 0 if __name__ == "__main__": raise SystemExit(main())