| from __future__ import annotations |
|
|
| import json |
| from pathlib import Path |
|
|
| import pandas as pd |
|
|
| from sepsis_mcp.paper_artifacts import ( |
| build_physionet_appendix_table, |
| build_physionet_gap_plot_frame, |
| build_physionet_model_auroc_frame, |
| build_gossis_validation_subgroup_frame, |
| build_gossis_validation_table, |
| build_physionet_main_table, |
| main, |
| ) |
|
|
|
|
| def _write_physionet_dropout_run(base_dir: Path, experiment_suffix: str, coverage_shift: float) -> None: |
| base_dir.mkdir(parents=True, exist_ok=True) |
| metrics = { |
| "config": {"alpha": 0.1}, |
| "experiments": { |
| f"A_to_A{experiment_suffix}": { |
| "standard": { |
| "empirical_coverage": 0.91 - coverage_shift, |
| "average_set_size": 0.95 - coverage_shift, |
| "target_coverage": 0.9, |
| }, |
| "missingness_aware": { |
| "empirical_coverage": 0.89 + coverage_shift, |
| "average_set_size": 0.96, |
| "target_coverage": 0.9, |
| }, |
| "weighted_missingness_aware": { |
| "empirical_coverage": 0.9, |
| "average_set_size": 0.94, |
| "target_coverage": 0.9, |
| }, |
| "shrunk_weighted_missingness_aware": { |
| "empirical_coverage": 0.905, |
| "average_set_size": 0.93, |
| "target_coverage": 0.9, |
| }, |
| }, |
| f"A_to_B{experiment_suffix}": { |
| "standard": { |
| "empirical_coverage": 0.9 - coverage_shift, |
| "average_set_size": 0.91 - coverage_shift, |
| "target_coverage": 0.9, |
| }, |
| "missingness_aware": { |
| "empirical_coverage": 0.88 + coverage_shift, |
| "average_set_size": 0.915 + coverage_shift, |
| "target_coverage": 0.9, |
| }, |
| "weighted_missingness_aware": { |
| "empirical_coverage": 0.885 + coverage_shift / 2, |
| "average_set_size": 0.9, |
| "target_coverage": 0.9, |
| }, |
| "shrunk_weighted_missingness_aware": { |
| "empirical_coverage": 0.89 + coverage_shift / 3, |
| "average_set_size": 0.905, |
| "target_coverage": 0.9, |
| }, |
| }, |
| }, |
| } |
| (base_dir / "metrics.json").write_text(json.dumps(metrics), encoding="utf-8") |
|
|
|
|
| def _write_dropout_summary(path: Path) -> Path: |
| pd.DataFrame( |
| [ |
| {"condition": "none", "mask_rate": 0.0, "method": "standard", "empirical_coverage": 0.8976, "target_coverage": 0.9, "absolute_coverage_gap": 0.0024, "average_set_size": 0.9126}, |
| {"condition": "none", "mask_rate": 0.0, "method": "missingness_aware", "empirical_coverage": 0.8772, "target_coverage": 0.9, "absolute_coverage_gap": 0.0228, "average_set_size": 0.8914}, |
| {"condition": "drop=0.3", "mask_rate": 0.3, "method": "standard", "empirical_coverage": 0.8847, "target_coverage": 0.9, "absolute_coverage_gap": 0.0153, "average_set_size": 0.8971}, |
| {"condition": "drop=0.3", "mask_rate": 0.3, "method": "missingness_aware", "empirical_coverage": 0.8962, "target_coverage": 0.9, "absolute_coverage_gap": 0.0038, "average_set_size": 0.9104}, |
| {"condition": "drop=0.5", "mask_rate": 0.5, "method": "standard", "empirical_coverage": 0.8661, "target_coverage": 0.9, "absolute_coverage_gap": 0.0339, "average_set_size": 0.8794}, |
| {"condition": "drop=0.5", "mask_rate": 0.5, "method": "missingness_aware", "empirical_coverage": 0.8993, "target_coverage": 0.9, "absolute_coverage_gap": 0.0007, "average_set_size": 0.9149}, |
| ] |
| ).to_csv(path, index=False) |
| return path |
|
|
|
|
| def _write_model_dropout_summary(path: Path) -> Path: |
| pd.DataFrame( |
| [ |
| {"condition": "none", "mask_rate": 0.0, "model_type": "xgboost", "auroc": 0.81, "auprc": 0.10}, |
| {"condition": "drop=0.3", "mask_rate": 0.3, "model_type": "xgboost", "auroc": 0.78, "auprc": 0.09}, |
| {"condition": "none", "mask_rate": 0.0, "model_type": "sklearn_gbdt", "auroc": 0.79, "auprc": 0.08}, |
| {"condition": "drop=0.3", "mask_rate": 0.3, "model_type": "sklearn_gbdt", "auroc": 0.75, "auprc": 0.07}, |
| {"condition": "none", "mask_rate": 0.0, "model_type": "logistic_regression", "auroc": 0.72, "auprc": 0.05}, |
| {"condition": "drop=0.3", "mask_rate": 0.3, "model_type": "logistic_regression", "auroc": 0.69, "auprc": 0.04}, |
| ] |
| ).to_csv(path, index=False) |
| return path |
|
|
|
|
| def _write_gossis_summary_files(base_dir: Path) -> tuple[Path, Path]: |
| base_dir.mkdir(parents=True, exist_ok=True) |
| overall_path = base_dir / "overall_summary.csv" |
| subgroup_path = base_dir / "subgroup_summary.csv" |
| pd.DataFrame( |
| [ |
| { |
| "method": "standard", |
| "min_hospital_admissions": 500, |
| "empirical_coverage": 0.895, |
| "average_set_size": 0.95, |
| "max_group_coverage_gap": 0.08, |
| "worst_hospital_coverage": 0.84, |
| }, |
| { |
| "method": "standard", |
| "min_hospital_admissions": 500, |
| "empirical_coverage": 0.885, |
| "average_set_size": 0.94, |
| "max_group_coverage_gap": 0.10, |
| "worst_hospital_coverage": 0.83, |
| }, |
| { |
| "method": "missingness_aware", |
| "min_hospital_admissions": 500, |
| "empirical_coverage": 0.898, |
| "average_set_size": 0.96, |
| "max_group_coverage_gap": 0.02, |
| "worst_hospital_coverage": 0.835, |
| }, |
| { |
| "method": "weighted_missingness_aware", |
| "min_hospital_admissions": 500, |
| "empirical_coverage": 0.9, |
| "average_set_size": 0.955, |
| "max_group_coverage_gap": 0.05, |
| "worst_hospital_coverage": 0.845, |
| }, |
| { |
| "method": "shrunk_weighted_missingness_aware", |
| "min_hospital_admissions": 500, |
| "empirical_coverage": 0.897, |
| "average_set_size": 0.952, |
| "max_group_coverage_gap": 0.07, |
| "worst_hospital_coverage": 0.844, |
| }, |
| ] |
| ).to_csv(overall_path, index=False) |
| pd.DataFrame( |
| [ |
| {"method": "standard", "min_hospital_admissions": 500, "group_label": "low", "coverage": 0.82}, |
| {"method": "standard", "min_hospital_admissions": 500, "group_label": "medium", "coverage": 0.92}, |
| {"method": "standard", "min_hospital_admissions": 500, "group_label": "high", "coverage": 0.94}, |
| {"method": "missingness_aware", "min_hospital_admissions": 500, "group_label": "low", "coverage": 0.89}, |
| {"method": "missingness_aware", "min_hospital_admissions": 500, "group_label": "medium", "coverage": 0.90}, |
| {"method": "missingness_aware", "min_hospital_admissions": 500, "group_label": "high", "coverage": 0.91}, |
| {"method": "weighted_missingness_aware", "min_hospital_admissions": 500, "group_label": "low", "coverage": 0.86}, |
| {"method": "weighted_missingness_aware", "min_hospital_admissions": 500, "group_label": "medium", "coverage": 0.91}, |
| {"method": "weighted_missingness_aware", "min_hospital_admissions": 500, "group_label": "high", "coverage": 0.93}, |
| ] |
| ).to_csv(subgroup_path, index=False) |
| return overall_path, subgroup_path |
|
|
|
|
| def _write_gossis_summary_files_with_opt_in_methods(base_dir: Path) -> tuple[Path, Path]: |
| overall_path, subgroup_path = _write_gossis_summary_files(base_dir) |
| overall = pd.read_csv(overall_path) |
| subgroup = pd.read_csv(subgroup_path) |
| pd.concat( |
| [ |
| overall, |
| pd.DataFrame( |
| [ |
| { |
| "method": "learned_partition", |
| "min_hospital_admissions": 500, |
| "empirical_coverage": 0.901, |
| "average_set_size": 0.958, |
| "max_group_coverage_gap": 0.03, |
| "worst_hospital_coverage": 0.84, |
| }, |
| { |
| "method": "gibbs_general", |
| "min_hospital_admissions": 500, |
| "empirical_coverage": 0.902, |
| "average_set_size": 0.959, |
| "max_group_coverage_gap": 0.04, |
| "worst_hospital_coverage": 0.845, |
| }, |
| ] |
| ), |
| ], |
| ignore_index=True, |
| ).to_csv(overall_path, index=False) |
| pd.concat( |
| [ |
| subgroup, |
| pd.DataFrame( |
| [ |
| {"method": "learned_partition", "min_hospital_admissions": 500, "group_label": "low", "coverage": 0.88}, |
| {"method": "gibbs_general", "min_hospital_admissions": 500, "group_label": "low", "coverage": 0.87}, |
| ] |
| ), |
| ], |
| ignore_index=True, |
| ).to_csv(subgroup_path, index=False) |
| return overall_path, subgroup_path |
|
|
|
|
| def test_build_physionet_tables_and_plot_frame(tmp_path: Path) -> None: |
| _write_physionet_dropout_run(tmp_path / "none", "", 0.0) |
| _write_physionet_dropout_run(tmp_path / "rd03", "__random_drop", 0.02) |
| _write_physionet_dropout_run(tmp_path / "rd05", "__random_drop", 0.04) |
| summary_path = _write_dropout_summary(tmp_path / "dropout_grid_summary.csv") |
| model_summary_path = _write_model_dropout_summary(tmp_path / "model_dropout_grid_summary.csv") |
|
|
| table = build_physionet_main_table(tmp_path / "none", tmp_path / "rd03", tmp_path / "rd05") |
| appendix = build_physionet_appendix_table(tmp_path / "none", tmp_path / "rd03", tmp_path / "rd05") |
| stress = build_physionet_gap_plot_frame(summary_path) |
| model_auroc = build_physionet_model_auroc_frame(model_summary_path) |
|
|
| assert len(table) == 18 |
| assert list(table.columns) == [ |
| "setting", |
| "experiment", |
| "condition", |
| "method", |
| "empirical_coverage", |
| "average_set_size", |
| "absolute_coverage_gap", |
| ] |
| assert set(table["method"]) == {"standard", "missingness_aware", "weighted_missingness_aware"} |
| assert len(appendix) == 24 |
| assert "shrunk_weighted_missingness_aware" in set(appendix["method"]) |
| assert set(stress["method"]) == {"standard", "missingness_aware"} |
| assert set(stress["condition"]) == {"none", "drop=0.3", "drop=0.5"} |
| assert list(stress.columns) == ["condition", "mask_rate", "method", "empirical_coverage", "target_coverage", "absolute_coverage_gap", "average_set_size"] |
| assert (stress["absolute_coverage_gap"] >= 0).all() |
| assert set(model_auroc["model_type"]) == {"xgboost", "sklearn_gbdt", "logistic_regression"} |
| assert list(model_auroc.columns) == ["condition", "mask_rate", "model_type", "auroc", "auprc"] |
|
|
|
|
| def test_build_gossis_validation_frames(tmp_path: Path) -> None: |
| overall_path, subgroup_path = _write_gossis_summary_files(tmp_path) |
|
|
| table = build_gossis_validation_table(overall_path, min_hospital_admissions=500) |
| subgroup = build_gossis_validation_subgroup_frame(subgroup_path, min_hospital_admissions=500) |
|
|
| assert list(table["method"]) == [ |
| "standard", |
| "missingness_aware", |
| "weighted_missingness_aware", |
| "shrunk_weighted_missingness_aware", |
| ] |
| assert table.loc[table["method"] == "standard", "empirical_coverage"].item() == 0.89 |
| assert set(subgroup["method"]) == {"standard", "missingness_aware", "weighted_missingness_aware"} |
| assert set(subgroup["group_label"]) == {"low", "medium", "high"} |
|
|
|
|
| def test_build_gossis_validation_frames_ignore_opt_in_methods(tmp_path: Path) -> None: |
| overall_path, subgroup_path = _write_gossis_summary_files_with_opt_in_methods(tmp_path) |
|
|
| table = build_gossis_validation_table(overall_path, min_hospital_admissions=500) |
| subgroup = build_gossis_validation_subgroup_frame(subgroup_path, min_hospital_admissions=500) |
|
|
| assert "learned_partition" not in set(table["method"]) |
| assert "gibbs_general" not in set(table["method"]) |
| assert "learned_partition" not in set(subgroup["method"]) |
| assert "gibbs_general" not in set(subgroup["method"]) |
|
|
|
|
| def test_external_baseline_dependency_is_documented_as_optional() -> None: |
| pyproject_text = Path("pyproject.toml").read_text(encoding="utf-8") |
| readme_text = Path("README.md").read_text(encoding="utf-8") |
|
|
| assert "[project.optional-dependencies]" in pyproject_text |
| assert "conditionalconformal==0.0.5" in pyproject_text |
| assert "external-baselines" in pyproject_text |
| assert "conditionalconformal==0.0.5" in readme_text |
|
|
|
|
| def test_main_writes_requested_paper_artifacts(tmp_path: Path) -> None: |
| _write_physionet_dropout_run(tmp_path / "none", "", 0.0) |
| _write_physionet_dropout_run(tmp_path / "rd03", "__random_drop", 0.02) |
| _write_physionet_dropout_run(tmp_path / "rd05", "__random_drop", 0.04) |
| summary_path = _write_dropout_summary(tmp_path / "dropout_grid_summary.csv") |
| model_summary_path = _write_model_dropout_summary(tmp_path / "model_dropout_grid_summary.csv") |
| overall_path, subgroup_path = _write_gossis_summary_files(tmp_path / "gossis") |
| output_dir = tmp_path / "paper" |
|
|
| main( |
| [ |
| "--physionet-none-dir", |
| str(tmp_path / "none"), |
| "--physionet-rd03-dir", |
| str(tmp_path / "rd03"), |
| "--physionet-rd05-dir", |
| str(tmp_path / "rd05"), |
| "--physionet-dropout-summary-path", |
| str(summary_path), |
| "--physionet-model-dropout-summary-path", |
| str(model_summary_path), |
| "--gossis-overall-path", |
| str(overall_path), |
| "--gossis-subgroup-path", |
| str(subgroup_path), |
| "--output-dir", |
| str(output_dir), |
| ] |
| ) |
|
|
| expected_files = { |
| "figure1_method_schematic.png", |
| "figure1_method_schematic.pdf", |
| "figure2_physionet_stress_test.png", |
| "figure2_physionet_stress_test.pdf", |
| "figureA2_physionet_model_auroc_dropout.png", |
| "figureA2_physionet_model_auroc_dropout.pdf", |
| "figure3_gossis_subgroup_coverage.png", |
| "figure3_gossis_subgroup_coverage.pdf", |
| "table1_physionet_main_results.csv", |
| "table1_physionet_main_results.png", |
| "table1_physionet_main_results.tex", |
| "tableA1_physionet_full_results.csv", |
| "tableA1_physionet_full_results.png", |
| "tableA1_physionet_full_results.tex", |
| "table2_gossis_validation_summary.csv", |
| "table2_gossis_validation_summary.png", |
| "table2_gossis_validation_summary.tex", |
| "artifact_manifest.json", |
| } |
| assert expected_files <= {path.name for path in output_dir.iterdir()} |
|
|
