From 4526191fe1542578fe9d5af05e48c4fbb577587f Mon Sep 17 00:00:00 2001 From: Jiarui Li Date: Sat, 25 Jul 2026 11:15:06 +0800 Subject: [PATCH] Report AUCs in Delphi2M format --- delphi2m_auc_report.py | 183 +++++++++++++++++++++++++++++++++++++++++ evaluate_auc.py | 52 +++++++----- evaluate_auc_v2.py | 131 ++++++----------------------- 3 files changed, 239 insertions(+), 127 deletions(-) create mode 100644 delphi2m_auc_report.py diff --git a/delphi2m_auc_report.py b/delphi2m_auc_report.py new file mode 100644 index 0000000..a8b49c9 --- /dev/null +++ b/delphi2m_auc_report.py @@ -0,0 +1,183 @@ +"""Build Delphi2M-style sex-specific AUC reports. + +The Delphi2M evaluation code uses 0.1 years for the no-gap evaluation. The +published report displays that point as 0 months, while retaining the actual +0.1-year evaluation period in this project's report output. +""" +from __future__ import annotations + +from pathlib import Path +from typing import Dict, Optional + +import numpy as np +import pandas as pd + + +DEFAULT_DELPHI2M_PERIODS_YEARS = (0.1, 1.0, 5.0, 10.0) + +_CHAPTER_SHORT_NAMES = { + "I": "I. Infectious Diseases", + "II": "II. Neoplasms", + "III": "III. Blood & Immune Disorders", + "IV": "IV. Metabolic Diseases", + "V": "V. Mental Disorders", + "VI": "VI. Nervous System Diseases", + "VII": "VII. Eye Diseases", + "VIII": "VIII. Ear Diseases", + "IX": "IX. Circulatory Diseases", + "X": "X. Respiratory Diseases", + "XI": "XI. Digestive Diseases", + "XII": "XII. Skin Diseases", + "XIII": "XIII. Musculoskeletal Diseases", + "XIV": "XIV. Genitourinary Diseases", + "XV": "XV. Pregnancy & Childbirth", + "XVI": "XVI. Perinatal Conditions", + "XVII": "XVII. Congenital Abnormalities", + "XVIII": "XVIII. Symptoms & Signs", + "XIX": "XIX. Injury & Poisoning", + "XX": "XX. External Causes", + "XXI": "XXI. Health Services", + "XXII": "XXII. Special Purposes", + "Death": "Death", + "Unmapped": "Unmapped", +} + + +def _is_no_gap(period_years: float) -> bool: + return bool(np.isclose(float(period_years), 0.1, rtol=0.0, atol=1e-8)) + + +def _canonical_period_years(period_years: float) -> float: + value = float(period_years) + for canonical in DEFAULT_DELPHI2M_PERIODS_YEARS: + if np.isclose(value, canonical, rtol=0.0, atol=1e-6): + return float(canonical) + return value + + +def _gap_months(period_years: float) -> int: + if _is_no_gap(period_years): + return 0 + return int(round(float(period_years) * 12.0)) + + +def _gap_label(period_years: float) -> str: + if _is_no_gap(period_years): + return "No gap" + value = float(period_years) + value_text = f"{value:g}" + unit = "year" if np.isclose(value, 1.0) else "years" + return f"{value_text} {unit}" + + +def _load_chapter_by_code( + chapter_mapping_path: Optional[str | Path] = None, +) -> Dict[str, str]: + if chapter_mapping_path is None: + chapter_mapping_path = Path(__file__).with_name( + "icd10_chapter_organ_mapping.csv" + ) + path = Path(chapter_mapping_path) + if not path.exists(): + return {} + + mapping = pd.read_csv( + path, + usecols=["code", "icd10_chapter"], + dtype={"code": str, "icd10_chapter": str}, + ) + mapping["code"] = mapping["code"].str.strip() + mapping["chapter"] = ( + mapping["icd10_chapter"] + .str.strip() + .map(_CHAPTER_SHORT_NAMES) + .fillna("Unmapped") + ) + return dict(zip(mapping["code"], mapping["chapter"])) + + +def build_delphi2m_auc_report( + df_unpooled: pd.DataFrame, + *, + period_col: str, + chapter_mapping_path: Optional[str | Path] = None, +) -> pd.DataFrame: + """Aggregate age strata by sex and return a Delphi2M-style AUC report. + + Required input columns are ``token``, ``label_code``, ``sex``, + ``auc_delong``, and the supplied ``period_col`` (``offset`` or + ``horizon``). The output begins with the five columns used by Delphi2M + Fig. 2e and then records the actual evaluation period and ICD-10 code. + """ + required = {"token", "label_code", "sex", "auc_delong", period_col} + missing = sorted(required - set(df_unpooled.columns)) + if missing: + raise ValueError( + "Cannot build Delphi2M AUC report; missing columns: " + + ", ".join(missing) + ) + + source = df_unpooled.loc[ + :, + ["token", "label_code", "sex", "auc_delong", period_col], + ].copy() + source["sex"] = source["sex"].astype(str).str.strip().str.lower() + source = source[source["sex"].isin(["female", "male"])] + source["auc_delong"] = pd.to_numeric( + source["auc_delong"], errors="coerce" + ) + source[period_col] = pd.to_numeric(source[period_col], errors="coerce") + source = source.dropna(subset=[period_col, "auc_delong"]) + source[period_col] = source[period_col].map(_canonical_period_years) + + if source.empty: + raise ValueError("Cannot build Delphi2M AUC report from empty AUC data.") + + grouped = ( + source.groupby( + ["token", "label_code", period_col, "sex"], + dropna=False, + as_index=False, + ) + .agg(auc=("auc_delong", "mean")) + ) + report = ( + grouped.pivot( + index=["token", "label_code", period_col], + columns="sex", + values="auc", + ) + .reset_index() + .rename_axis(columns=None) + .rename(columns={"female": "Female", "male": "Male"}) + ) + for col in ["Female", "Male"]: + if col not in report.columns: + report[col] = np.nan + + chapter_by_code = _load_chapter_by_code(chapter_mapping_path) + report["chapter"] = ( + report["label_code"].astype(str).map(chapter_by_code).fillna("Unmapped") + ) + report["Gap, months"] = report[period_col].map(_gap_months).astype("Int64") + report["Gap label"] = report[period_col].map(_gap_label) + report["icd10"] = pd.to_numeric(report["token"], errors="coerce").astype( + "Int64" + ) + + report = report.sort_values( + ["icd10", period_col], kind="stable", ignore_index=True + ) + return report.loc[ + :, + [ + "Gap, months", + "chapter", + "icd10", + "Female", + "Male", + period_col, + "Gap label", + "label_code", + ], + ] diff --git a/evaluate_auc.py b/evaluate_auc.py index 96d05b7..ce5fd4f 100644 --- a/evaluate_auc.py +++ b/evaluate_auc.py @@ -7,7 +7,8 @@ This script follows the logic of the Delphi evaluation script supplied by the us at least `offset` years before the target time; 3. run model inference by disease chunks to avoid materializing all logits; 4. compute AUC separately by sex and age bracket; - 5. aggregate age brackets with DeLong variance. + 5. average age-bracket AUCs within each sex and write a Delphi2M-style + Female/Male report. Efficiency notes: - transformer/readout inference is executed once and cached; @@ -39,6 +40,10 @@ from torch.utils.data import DataLoader, Subset from tqdm.auto import tqdm from dataset import HealthDataset +from delphi2m_auc_report import ( + DEFAULT_DELPHI2M_PERIODS_YEARS, + build_delphi2m_auc_report, +) from eval_data import load_sequence_eval_dataset, sequence_eval_collate_fn from models import DeepHealth from readouts import build_readout @@ -1158,30 +1163,23 @@ def evaluate_auc_pipeline( df_auc_unpooled["label_code"] = df_auc_unpooled["token"].map( dataset.label_id_to_code) - print("Using DeLong method to calculate AUC confidence intervals.") - grouped = df_auc_unpooled.groupby( - ["token", "label_code", "offset"], dropna=False, as_index=False) - df_auc = grouped.agg( - auc=("auc_delong", "mean"), - n_strata=("auc_delong", "size"), - n_diseased=("n_diseased", "sum"), - n_healthy=("n_healthy", "sum"), - auc_variance_sum=("auc_variance_delong", "sum"), + print( + "Building Delphi2M-style report: mean AUC across age strata, " + "reported separately for Female and Male." ) - df_auc["auc_variance_delong"] = ( - df_auc["auc_variance_sum"] - / (df_auc["n_strata"].clip(lower=1).astype(np.float64) ** 2) + df_report = build_delphi2m_auc_report( + df_auc_unpooled, + period_col="offset", ) - df_auc = df_auc.drop(columns=["auc_variance_sum"]) if output_path is not None: out_dir = Path(output_path) out_dir.mkdir(parents=True, exist_ok=True) - df_auc.to_csv(out_dir / "df_both.csv", index=False) - df_auc_unpooled.to_csv( - out_dir / "df_auc_unpooled.csv", index=False) + report_path = out_dir / "df_auc_delphi2m_report.csv" + df_report.to_csv(report_path, index=False) + print(f"Saved Delphi2M-style AUC report: {report_path}") - return df_auc_unpooled, df_auc + return df_auc_unpooled, df_report # --------------------------------------------------------------------------- @@ -1231,8 +1229,18 @@ def make_auc_offsets(args: argparse.Namespace, cfg: Dict[str, Any]) -> List[floa if explicit_offsets is not None: base_offsets = explicit_offsets else: - next_token_offset = float(cfg_get(args, cfg, "offset", 0.1)) - base_offsets = [next_token_offset, 1.0, 5.0, 10.0] + next_token_offset = float( + cfg_get( + args, + cfg, + "offset", + DEFAULT_DELPHI2M_PERIODS_YEARS[0], + ) + ) + base_offsets = [ + next_token_offset, + *DEFAULT_DELPHI2M_PERIODS_YEARS[1:], + ] offsets: List[float] = [] seen = set() @@ -1280,9 +1288,9 @@ def main() -> None: parser.add_argument("--filter_min_total", type=int, default=None, help="Minimum metadata count for disease selection; default 0.") parser.add_argument("--offset", type=float, default=None, - help="Next-token prediction offset in years; preserved and evaluated alongside 1, 5, and 10 years by default.") + help="Next-token prediction offset in years; 0.1 is Delphi2M no gap and is evaluated alongside 1, 5, and 10 years by default.") parser.add_argument("--offsets", type=str, default=None, - help="Comma-separated prediction offsets in years. Overrides the default set of offset,1,5,10.") + help="Comma-separated prediction offsets in years. Overrides the default set of 0.1,1,5,10.") parser.add_argument("--age_start", type=float, default=None) parser.add_argument("--age_stop", type=float, default=None) parser.add_argument("--age_step", type=float, default=None) diff --git a/evaluate_auc_v2.py b/evaluate_auc_v2.py index 952d5f1..111518e 100644 --- a/evaluate_auc_v2.py +++ b/evaluate_auc_v2.py @@ -3,6 +3,9 @@ This script supports DeepHealth fixed-horizon risk scores for exponential, Weibull, and mixed all-future distributions. +The default horizons are 0.1, 1, 5, and 10 years. As in Delphi2M, 0.1 years +is reported as the no-gap evaluation. + Landmark querying depends on the model target mode saved in train_config.json: - next_token: insert a token at landmark age and read it out; - all_future: pass landmark age directly as t_query. @@ -28,6 +31,10 @@ from torch.utils.data import DataLoader, Dataset from tqdm.auto import tqdm from dataset import HealthDataset +from delphi2m_auc_report import ( + DEFAULT_DELPHI2M_PERIODS_YEARS, + build_delphi2m_auc_report, +) from eval_data import load_sequence_eval_dataset from models import DeepHealth from readouts import build_readout @@ -324,44 +331,6 @@ def _first_existing_column(df: pd.DataFrame, candidates: Sequence[str]) -> Optio return None -def build_metadata_for_merge(dataset: HealthDataset, labels_meta: Optional[pd.DataFrame]) -> pd.DataFrame: - base_rows = [] - for token, code in dataset.label_id_to_code.items(): - token = int(token) - code_text = str(code) - if token in SPECIAL_TOKENS or code_text.startswith("<"): - continue - base_rows.append({"token": token, "label_code": code_text}) - base = pd.DataFrame(base_rows) - if labels_meta is None or labels_meta.empty: - return base - - meta = labels_meta.copy() - code_col = _first_existing_column( - meta, ["Name", "code", "ICD10", "icd10", "label", "token", "disease_code"]) - if code_col is not None: - meta["_label_code"] = meta[code_col].astype( - str).map(lambda s: s.split()[0].strip()) - merged = base.merge(meta, left_on="label_code", - right_on="_label_code", how="left") - return merged.drop(columns=["_label_code"], errors="ignore") - - if "index" in meta.columns: - idx = pd.to_numeric(meta["index"], errors="coerce") - has_no_event = ( - NO_EVENT_IDX in dataset.label_id_to_code - and dataset.label_id_to_code.get(NO_EVENT_IDX) == "" - ) - if has_no_event: - idx = idx.where(idx < NO_EVENT_IDX, idx + 1) - meta["_index_int"] = idx.astype("Int64") - merged = base.merge(meta, left_on="token", - right_on="_index_int", how="left") - return merged.drop(columns=["_index_int"], errors="ignore") - - return base - - def _metadata_count_map(dataset: HealthDataset, labels_meta: Optional[pd.DataFrame]) -> Dict[int, float]: if labels_meta is None or labels_meta.empty or "count" not in labels_meta.columns: return {} @@ -1101,7 +1070,6 @@ def evaluate_landmark_auc( loader: DataLoader, landmark_dataset: LandmarkDataset, output_path: Path, - labels_meta: Optional[pd.DataFrame], disease_ids: Sequence[int], disease_chunk_size: int, score_mode: str, @@ -1118,7 +1086,6 @@ def evaluate_landmark_auc( use_amp: bool, hidden_cache_dtype: str, logit_batch_size: int, - meta_info: Dict[str, Any], ) -> Tuple[pd.DataFrame, pd.DataFrame]: model.eval().to(device) @@ -1235,54 +1202,21 @@ def evaluate_landmark_auc( df_unpooled["label_code"] = df_unpooled["token"].map( landmark_dataset.dataset.label_id_to_code) - for k, v in meta_info.items(): - df_unpooled[k] = v - - meta_table = build_metadata_for_merge(landmark_dataset.dataset, labels_meta) - df_unpooled = df_unpooled.merge( - meta_table, on=["token", "label_code"], how="left") - - grouped = df_unpooled.groupby( - ["token", "label_code", "horizon"], dropna=False, as_index=False) - df_merged = grouped.agg( - auc=("auc_delong", "mean"), - n_strata=("auc_delong", "size"), - n_diseased=("n_diseased", "sum"), - n_healthy=("n_healthy", "sum"), - auc_variance_sum=("auc_variance_delong", "sum"), + print( + "Building Delphi2M-style report: mean AUC across landmark-age " + "strata, reported separately for Female and Male." ) - df_merged["auc_variance_delong"] = ( - df_merged["auc_variance_sum"] - / (df_merged["n_strata"].clip(lower=1).astype(np.float64) ** 2) + df_report = build_delphi2m_auc_report( + df_unpooled, + period_col="horizon", ) - df_merged = df_merged.drop(columns=["auc_variance_sum"]) - - keep_meta = [ - c for c in [ - "model_ckpt_path", - "config_path", - "target_mode", - "model_target_mode", - "dist_mode", - "time_mode", - "attn_mask_mode", - "readout_name", - "landmark_query_mode", - "landmark_token_mode", - "score_mode", - "eval_split", - ] - if c in df_unpooled.columns - ] - for col in keep_meta: - df_merged[col] = meta_info[col] output_path.mkdir(parents=True, exist_ok=True) - df_unpooled.to_csv( - output_path / "df_auc_landmark_unpooled.csv", index=False) - df_merged.to_csv(output_path / "df_auc_landmark.csv", index=False) + report_path = output_path / "df_auc_landmark_delphi2m_report.csv" + df_report.to_csv(report_path, index=False) + print(f"Saved Delphi2M-style landmark AUC report: {report_path}") - return df_unpooled, df_merged + return df_unpooled, df_report def main() -> None: @@ -1308,7 +1242,12 @@ def main() -> None: parser.add_argument("--landmark_start", type=float, default=None) parser.add_argument("--landmark_stop", type=float, default=None) parser.add_argument("--landmark_step", type=float, default=None) - parser.add_argument("--horizons", type=str, default=None) + parser.add_argument( + "--horizons", + type=str, + default=None, + help="Comma-separated horizons in years; defaults to 0.1,1,5,10, where 0.1 is Delphi2M no gap.", + ) parser.add_argument("--min_cases", type=int, default=None) parser.add_argument("--min_history_events", type=int, default=None) @@ -1428,8 +1367,9 @@ def main() -> None: "Landmark ages are empty. Check landmark_start/landmark_stop/landmark_step.") horizons = np.asarray( - parse_float_list(cfg_get(args, cfg, "horizons", "1,5,10")) or [ - 1.0, 5.0, 10.0], + parse_float_list( + cfg_get(args, cfg, "horizons", "0.1,1,5,10") + ) or list(DEFAULT_DELPHI2M_PERIODS_YEARS), dtype=np.float32, ) if horizons.size == 0: @@ -1520,8 +1460,6 @@ def main() -> None: if model_target_mode == "next_token" else "direct_t_query" ) - score_mode_out = f"{landmark_query_mode}_{score_mode}" - num_workers_auc = int( cfg_get(args, cfg, "num_workers_auc", max(1, (os.cpu_count() or 2) - 1))) auc_task_chunk_size = int(cfg_get(args, cfg, "auc_task_chunk_size", 0)) @@ -1553,27 +1491,11 @@ def main() -> None: print(f"AUC workers: {num_workers_auc}") print(f"Output path: {output_path}") - meta_info = { - "score_mode": score_mode_out, - "eval_split": eval_split, - "model_ckpt_path": str(model_ckpt_path), - "config_path": str(config_path), - "target_mode": str(target_mode), - "model_target_mode": str(model_target_mode), - "dist_mode": str(dist_mode), - "time_mode": str(time_mode), - "attn_mask_mode": str(attn_mask_mode), - "readout_name": str(readout_name), - "landmark_query_mode": landmark_query_mode, - "landmark_token_mode": "no_event" if model_target_mode == "next_token" else "none", - } - evaluate_landmark_auc( model=model, loader=loader, landmark_dataset=landmark_dataset, output_path=output_path, - labels_meta=labels_meta, disease_ids=disease_ids, disease_chunk_size=disease_chunk_size, score_mode=score_mode, @@ -1590,7 +1512,6 @@ def main() -> None: use_amp=use_amp, hidden_cache_dtype=hidden_cache_dtype, logit_batch_size=logit_batch_size, - meta_info=meta_info, )