diff --git a/export_weibull_death_parameter_stats.py b/export_weibull_death_parameter_stats.py new file mode 100644 index 0000000..79914cf --- /dev/null +++ b/export_weibull_death_parameter_stats.py @@ -0,0 +1,502 @@ +"""Export Weibull shape-parameter statistics on the test split. + +The script is intended for all_future checkpoints with dist_mode="weibull" or +dist_mode="mixed". For full Weibull models it reads rho_head[Death]; for mixed +models it reads rho_death_head. For full Weibull models it also exports disease +token rho summaries, which are the main evidence for whether risk/hazard changes +with horizon instead of following an exponential shape. +""" +from __future__ import annotations + +import argparse +import contextlib +import json +from pathlib import Path +from typing import Any, Dict, Iterable, List, Optional + +import numpy as np +import pandas as pd +import torch +import torch.nn.functional as F +from torch.utils.data import DataLoader +from tqdm.auto import tqdm + +from eval_data import load_sequence_eval_dataset +from evaluate_auc_v2 import ( + LandmarkDataset, + _build_first_occurrence_maps, + _get_death_token_ids, + build_model_from_dataset, + cfg_get, + collate_landmark_fn, + load_checkpoint_state_dict, + load_json_config, + load_model_state, + make_eval_indices, + parse_float_list, + parse_int_list, + resolve_dist_mode_for_checkpoint, + resolve_eval_device, + validate_dataset_metadata, +) + + +def quantile_summary(df: pd.DataFrame, group_cols: List[str], value_cols: List[str]) -> pd.DataFrame: + probs = [0.01, 0.05, 0.25, 0.50, 0.75, 0.95, 0.99] + rows: List[Dict[str, Any]] = [] + grouped = [((), df)] if not group_cols else df.groupby(group_cols, dropna=False) + + for key, g in grouped: + if not isinstance(key, tuple): + key = (key,) + base = {col: val for col, val in zip(group_cols, key)} + base["n"] = int(len(g)) + for col in value_cols: + x = pd.to_numeric(g[col], errors="coerce").to_numpy(dtype=np.float64) + x = x[np.isfinite(x)] + if x.size == 0: + continue + row = dict(base) + row["variable"] = col + row["mean"] = float(np.mean(x)) + row["std"] = float(np.std(x, ddof=1)) if x.size > 1 else 0.0 + row["min"] = float(np.min(x)) + row["max"] = float(np.max(x)) + for p in probs: + row[f"p{int(p * 100):02d}"] = float(np.quantile(x, p)) + rows.append(row) + return pd.DataFrame(rows) + + +def load_labels_meta(path: Optional[str]) -> Optional[pd.DataFrame]: + if path is None: + return None + fp = Path(path) + if not fp.exists(): + return None + return pd.read_csv(fp) + + +@torch.inference_mode() +def infer_landmark_hidden_local( + model, + loader: DataLoader, + device: torch.device, + use_amp: bool, + hidden_cache_dtype: str, +) -> tuple[np.ndarray, Dict[str, np.ndarray]]: + """Minimal all_future landmark hidden inference for parameter export.""" + out_dtype = np.float32 if str(hidden_cache_dtype).lower() == "float32" else np.float16 + hidden_parts: List[np.ndarray] = [] + arrays: Dict[str, List[np.ndarray]] = { + "patient_id": [], + "sex": [], + "landmark_age": [], + "followup_end_time": [], + "death_time": [], + } + amp_enabled = bool(use_amp and device.type == "cuda") + + for batch in tqdm(loader, desc="Landmark hidden", dynamic_ncols=True): + batch_dev = { + k: (v.to(device, non_blocking=True) if isinstance(v, torch.Tensor) else v) + for k, v in batch.items() + } + amp_ctx = ( + torch.autocast(device_type=device.type, dtype=torch.float16) + if amp_enabled + else contextlib.nullcontext() + ) + with amp_ctx: + hidden = model( + event_seq=batch_dev["event_seq"], + time_seq=batch_dev["time_seq"], + sex=batch_dev["sex"], + padding_mask=batch_dev["padding_mask"], + t_query=batch_dev["t_query"], + other_type=batch_dev["other_type"], + other_value=batch_dev["other_value"], + other_value_kind=batch_dev["other_value_kind"], + other_time=batch_dev["other_time"], + target_mode="all_future", + ) + hidden_parts.append(hidden.detach().cpu().numpy().astype(out_dtype, copy=False)) + for key in arrays: + arrays[key].append(batch[key].cpu().numpy()) + + hidden_all = np.concatenate(hidden_parts, axis=0) + row_arrays = {key: np.concatenate(parts, axis=0) for key, parts in arrays.items()} + return hidden_all, row_arrays + + +@torch.inference_mode() +def project_death_params( + model, + hidden_all: np.ndarray, + dist_mode: str, + device: torch.device, + batch_size: int, + use_amp: bool, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + death_idx = int(getattr(model, "death_idx", getattr(model, "vocab_size", hidden_all.shape[0]) - 1)) + if not hasattr(model, "vocab_size"): + death_idx = int(model.risk_head.out_features - 1) + + compute_dtype = torch.float16 if (device.type == "cuda" and use_amp) else torch.float32 + risk_w = model.risk_head.weight[death_idx: death_idx + 1].detach().to(device=device, dtype=compute_dtype) + risk_b = None + if model.risk_head.bias is not None: + risk_b = model.risk_head.bias[death_idx: death_idx + 1].detach().to(device=device, dtype=compute_dtype) + + if dist_mode == "weibull": + rho_w = model.rho_head.weight[death_idx: death_idx + 1].detach().to(device=device, dtype=compute_dtype) + rho_b = model.rho_head.bias[death_idx: death_idx + 1].detach().to(device=device, dtype=compute_dtype) + elif dist_mode == "mixed": + rho_w = model.rho_death_head.weight.detach().to(device=device, dtype=compute_dtype) + rho_b = model.rho_death_head.bias.detach().to(device=device, dtype=compute_dtype) + else: + raise ValueError("Death Weibull parameter export requires dist_mode='weibull' or 'mixed'.") + + logits_out: List[np.ndarray] = [] + rate_out: List[np.ndarray] = [] + rho_out: List[np.ndarray] = [] + + for start in tqdm(range(0, hidden_all.shape[0], batch_size), desc="Death eta/rho", dynamic_ncols=True): + end = min(start + batch_size, hidden_all.shape[0]) + h = torch.from_numpy(hidden_all[start:end]).to(device=device, dtype=compute_dtype, non_blocking=True) + logits = F.linear(h, risk_w, risk_b).squeeze(-1) + rate = F.softplus(logits) + 1e-8 + rho = F.softplus(F.linear(h, rho_w, rho_b).squeeze(-1)) + 1e-6 + logits_out.append(logits.float().cpu().numpy()) + rate_out.append(rate.float().cpu().numpy()) + rho_out.append(rho.float().cpu().numpy()) + del h, logits, rate, rho + + return ( + np.concatenate(logits_out).astype(np.float32, copy=False), + np.concatenate(rate_out).astype(np.float32, copy=False), + np.concatenate(rho_out).astype(np.float32, copy=False), + ) + + +@torch.inference_mode() +def export_all_token_rho_summary( + model, + hidden_all: np.ndarray, + dataset, + device: torch.device, + output_dir: Path, + token_chunk_size: int, + row_batch_size: int, + use_amp: bool, + horizons: np.ndarray, +) -> None: + if not hasattr(model, "rho_head"): + print("[INFO] Skipping all-token rho summary because this is not a full Weibull model.") + return + + special = {0, 1, 2} + token_ids = [ + int(t) + for t, code in dataset.label_id_to_code.items() + if int(t) not in special and not str(code).startswith("<") + ] + token_ids = sorted(set(token_ids)) + death_idx = int(getattr(model, "death_idx", getattr(model, "vocab_size", len(token_ids)) - 1)) + if not hasattr(model, "vocab_size"): + death_idx = int(model.risk_head.out_features - 1) + compute_dtype = torch.float16 if (device.type == "cuda" and use_amp) else torch.float32 + rows: List[Dict[str, Any]] = [] + + for chunk_start in tqdm(range(0, len(token_ids), token_chunk_size), desc="All-token rho chunks", dynamic_ncols=True): + chunk = token_ids[chunk_start: chunk_start + token_chunk_size] + w = model.rho_head.weight[chunk].detach().to(device=device, dtype=compute_dtype) + b = model.rho_head.bias[chunk].detach().to(device=device, dtype=compute_dtype) + vals_parts: List[np.ndarray] = [] + for row_start in range(0, hidden_all.shape[0], row_batch_size): + row_end = min(row_start + row_batch_size, hidden_all.shape[0]) + h = torch.from_numpy(hidden_all[row_start:row_end]).to(device=device, dtype=compute_dtype, non_blocking=True) + rho = F.softplus(F.linear(h, w, b)) + 1e-6 + vals_parts.append(rho.float().cpu().numpy()) + del h, rho + vals = np.concatenate(vals_parts, axis=0) + for j, token in enumerate(chunk): + x = vals[:, j].astype(np.float64, copy=False) + row = { + "token": int(token), + "label_code": dataset.label_id_to_code.get(int(token), ""), + "endpoint_type": "death" if int(token) == int(death_idx) else "disease", + "n_landmark_rows": int(x.size), + "rho_mean": float(np.mean(x)), + "rho_std": float(np.std(x, ddof=1)) if x.size > 1 else 0.0, + "rho_minus_one_mean": float(np.mean(x - 1.0)), + "frac_rho_gt_1": float(np.mean(x > 1.0)), + "frac_rho_lt_1": float(np.mean(x < 1.0)), + "frac_rho_gt_1_1": float(np.mean(x > 1.1)), + "frac_rho_lt_0_9": float(np.mean(x < 0.9)), + "rho_p01": float(np.quantile(x, 0.01)), + "rho_p05": float(np.quantile(x, 0.05)), + "rho_p25": float(np.quantile(x, 0.25)), + "rho_p50": float(np.quantile(x, 0.50)), + "rho_p75": float(np.quantile(x, 0.75)), + "rho_p95": float(np.quantile(x, 0.95)), + "rho_p99": float(np.quantile(x, 0.99)), + } + for horizon in horizons.tolist(): + h = float(horizon) + if h <= 0: + continue + # Shape-only time scaling. For rho=1 this equals 1, i.e. an + # exponential model with constant instantaneous hazard. + inst_scale = np.power(h, x - 1.0) + cumhaz_scale = np.power(h, x) + row[f"instant_hazard_scale_h{h:g}y_vs_1y_mean"] = float(np.mean(inst_scale)) + row[f"instant_hazard_scale_h{h:g}y_vs_1y_p50"] = float(np.quantile(inst_scale, 0.50)) + row[f"cumhaz_scale_h{h:g}y_mean"] = float(np.mean(cumhaz_scale)) + row[f"cumhaz_scale_h{h:g}y_p50"] = float(np.quantile(cumhaz_scale, 0.50)) + rows.append(row) + del vals, vals_parts + + out = pd.DataFrame(rows) + out.to_csv(output_dir / "all_token_weibull_shape_summary.csv", index=False) + out[out["endpoint_type"] == "disease"].to_csv( + output_dir / "disease_token_weibull_shape_summary.csv", index=False + ) + out[out["endpoint_type"] == "death"].to_csv( + output_dir / "death_token_weibull_shape_summary.csv", index=False + ) + + disease = out[out["endpoint_type"] == "disease"].copy() + if not disease.empty: + pd.DataFrame([ + { + "n_tokens": int(len(disease)), + "rho_mean_across_tokens": float(disease["rho_mean"].mean()), + "rho_median_across_tokens": float(disease["rho_p50"].median()), + "tokens_with_mean_rho_gt_1": int((disease["rho_mean"] > 1.0).sum()), + "tokens_with_mean_rho_lt_1": int((disease["rho_mean"] < 1.0).sum()), + "frac_tokens_with_mean_rho_gt_1": float((disease["rho_mean"] > 1.0).mean()), + "frac_tokens_with_mean_rho_lt_1": float((disease["rho_mean"] < 1.0).mean()), + "tokens_with_mean_rho_gt_1_1": int((disease["rho_mean"] > 1.1).sum()), + "tokens_with_mean_rho_lt_0_9": int((disease["rho_mean"] < 0.9).sum()), + } + ]).to_csv(output_dir / "disease_weibull_shape_overall_summary.csv", index=False) + + +def main() -> None: + parser = argparse.ArgumentParser(description="Export test-split Weibull shape parameter statistics.") + parser.add_argument("--run_path", type=str, required=True) + parser.add_argument("--output_path", type=str, default=None) + parser.add_argument("--eval_split", type=str, default="test", choices=["test", "val", "valid", "validation", "train", "all"]) + 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("--batch_size", type=int, default=None) + parser.add_argument("--num_workers", type=int, default=None) + parser.add_argument("--device", type=str, default=None) + parser.add_argument("--use_amp", action=argparse.BooleanOptionalAction, default=None) + parser.add_argument("--hidden_cache_dtype", type=str, default="float32", choices=["float16", "float32"]) + parser.add_argument( + "--include_all_token_rho_summary", + action=argparse.BooleanOptionalAction, + default=True, + help=( + "For full Weibull models, export disease/death token rho summaries. " + "Use --no-include_all_token_rho_summary to skip the heavier token projection." + ), + ) + parser.add_argument("--token_chunk_size", type=int, default=32) + parser.add_argument("--row_batch_size", type=int, default=512) + args = parser.parse_args() + + run_path = Path(args.run_path) + config_path = run_path / "train_config.json" + ckpt_path = run_path / "best_model.pt" + if not config_path.exists(): + raise FileNotFoundError(config_path) + if not ckpt_path.exists(): + raise FileNotFoundError(ckpt_path) + + cfg = load_json_config(config_path) + model_target_mode = str(cfg.get("model_target_mode", "next_token")).lower() + if model_target_mode != "all_future": + raise ValueError("This export is intended for all_future checkpoints.") + + data_prefix = cfg.get("data_prefix", "ukb") + labels_file = cfg.get("labels_file", "labels.csv") + no_event_interval_years = cfg.get("no_event_interval_years", 5.0) + include_no_event_in_uts_target = cfg.get("include_no_event_in_uts_target", False) + + dataset = load_sequence_eval_dataset( + model_target_mode=model_target_mode, + data_prefix=data_prefix, + labels_file=labels_file, + no_event_interval_years=float(no_event_interval_years), + include_no_event_in_uts_target=bool(include_no_event_in_uts_target), + min_history_events=int(cfg.get("all_future_min_history_events", 1)), + min_future_events=int(cfg.get("all_future_min_future_events", 1)), + extra_info_types=parse_int_list(cfg.get("extra_info_types", None)), + ) + validate_dataset_metadata(dataset, cfg) + + subset_indices = make_eval_indices(dataset, args, cfg) + first_occurrence_by_token, _, _, _ = _build_first_occurrence_maps(dataset, subset_indices) + + landmark_start = float(cfg_get(args, cfg, "landmark_start", 40.0)) + landmark_stop = float(cfg_get(args, cfg, "landmark_stop", 80.0)) + landmark_step = float(cfg_get(args, cfg, "landmark_step", 5.0)) + landmark_ages = np.arange(landmark_start, landmark_stop, landmark_step, dtype=np.float32) + if landmark_ages.size == 0: + raise ValueError("No landmark ages produced.") + + horizons = np.asarray( + parse_float_list(cfg_get(args, cfg, "horizons", "1,5,10")) or [1.0, 5.0, 10.0], + dtype=np.float32, + ) + if horizons.size == 0: + raise ValueError("No horizons provided.") + + state_dict = load_checkpoint_state_dict(ckpt_path, map_location="cpu") + dist_mode = resolve_dist_mode_for_checkpoint(str(cfg.get("dist_mode", "exponential")), state_dict) + if dist_mode not in {"weibull", "mixed"}: + raise ValueError( + f"Resolved dist_mode={dist_mode!r}; expected 'weibull' or 'mixed' for Weibull shape export." + ) + + cfg_model = dict(cfg) + cfg_model["dist_mode"] = dist_mode + device = resolve_eval_device(args.device) + model = build_model_from_dataset(args, cfg_model, dataset).to(device) + load_model_state(model, state_dict) + model.eval() + + death_token_ids = _get_death_token_ids(dataset, None) + death_idx = int(death_token_ids[0]) + attn_mask_mode = str(cfg.get("attn_mask_mode", "target_aware")) + landmark_dataset = LandmarkDataset( + dataset=dataset, + subset_indices=subset_indices, + landmark_ages=landmark_ages, + attn_mask_mode=attn_mask_mode, + model_target_mode=model_target_mode, + min_history_events=int(cfg_get(args, cfg, "min_history_events", 1)), + first_occurrence_by_token=first_occurrence_by_token, + death_token_ids=death_token_ids, + ) + + batch_size = int(cfg_get(args, cfg, "batch_size", 128)) + num_workers = int(cfg_get(args, cfg, "num_workers", 4)) + loader = DataLoader( + landmark_dataset, + batch_size=batch_size, + shuffle=False, + collate_fn=collate_landmark_fn, + num_workers=num_workers, + pin_memory=device.type == "cuda", + persistent_workers=num_workers > 0, + prefetch_factor=2 if num_workers > 0 else None, + ) + + use_amp = bool(cfg_get(args, cfg, "use_amp", False)) + hidden_all, row_arrays = infer_landmark_hidden_local( + model=model, + loader=loader, + device=device, + use_amp=use_amp, + hidden_cache_dtype=str(args.hidden_cache_dtype), + ) + eta, rate, rho = project_death_params( + model=model, + hidden_all=hidden_all, + dist_mode=dist_mode, + device=device, + batch_size=int(args.row_batch_size), + use_amp=use_amp, + ) + + rows = pd.DataFrame({ + "patient_id": row_arrays["patient_id"].astype(np.int64), + "sex": row_arrays["sex"].astype(np.int64), + "sex_label": np.where(row_arrays["sex"].astype(np.int64) == 0, "female", "male"), + "landmark_age": row_arrays["landmark_age"].astype(np.float32), + "followup_end_time": row_arrays["followup_end_time"].astype(np.float32), + "death_time": row_arrays["death_time"].astype(np.float32), + "death_eta": eta, + "death_rate": rate, + "death_rho": rho, + }) + for horizon in horizons.tolist(): + h = float(horizon) + cumulative_hazard = rows["death_rate"].to_numpy(dtype=np.float64) * np.power(h, rows["death_rho"].to_numpy(dtype=np.float64)) + rows[f"death_cumhaz_h{h:g}y"] = cumulative_hazard + rows[f"death_risk_h{h:g}y"] = -np.expm1(-cumulative_hazard) + rows[f"death_observed_h{h:g}y"] = ( + (rows["death_time"].to_numpy(dtype=np.float64) > rows["landmark_age"].to_numpy(dtype=np.float64)) + & (rows["death_time"].to_numpy(dtype=np.float64) <= rows["landmark_age"].to_numpy(dtype=np.float64) + h) + ).astype(np.int8) + + output_dir = Path(args.output_path) if args.output_path else run_path / "weibull_death_parameter_stats_test" + output_dir.mkdir(parents=True, exist_ok=True) + + rows.to_csv(output_dir / "death_weibull_parameters_by_landmark.csv", index=False) + value_cols = ["death_eta", "death_rate", "death_rho"] + for horizon in horizons.tolist(): + h = float(horizon) + value_cols.extend([f"death_cumhaz_h{h:g}y", f"death_risk_h{h:g}y"]) + + quantile_summary(rows, [], value_cols).to_csv(output_dir / "death_weibull_parameter_summary_overall.csv", index=False) + quantile_summary(rows, ["landmark_age"], value_cols).to_csv(output_dir / "death_weibull_parameter_summary_by_landmark_age.csv", index=False) + quantile_summary(rows, ["sex_label"], value_cols).to_csv(output_dir / "death_weibull_parameter_summary_by_sex.csv", index=False) + quantile_summary(rows, ["sex_label", "landmark_age"], value_cols).to_csv(output_dir / "death_weibull_parameter_summary_by_sex_landmark_age.csv", index=False) + + metadata = { + "run_path": str(run_path), + "config_path": str(config_path), + "checkpoint_path": str(ckpt_path), + "eval_split": str(args.eval_split), + "model_target_mode": model_target_mode, + "time_mode": str(cfg.get("time_mode")), + "dist_mode_config": str(cfg.get("dist_mode")), + "dist_mode_resolved": dist_mode, + "extra_info_types": cfg.get("extra_info_types"), + "death_token_id": death_idx, + "death_label_code": dataset.label_id_to_code.get(death_idx, "Death"), + "n_selected_patients": int(len(subset_indices)), + "n_landmark_rows": int(len(rows)), + "landmark_ages": [float(x) for x in landmark_ages.tolist()], + "horizons": [float(x) for x in horizons.tolist()], + } + with (output_dir / "metadata.json").open("w", encoding="utf-8") as f: + json.dump(metadata, f, indent=2) + + if args.include_all_token_rho_summary and dist_mode == "weibull": + export_all_token_rho_summary( + model=model, + hidden_all=hidden_all, + dataset=dataset, + device=device, + output_dir=output_dir, + token_chunk_size=int(args.token_chunk_size), + row_batch_size=int(args.row_batch_size), + use_amp=use_amp, + horizons=horizons, + ) + elif dist_mode == "mixed": + pd.DataFrame([ + { + "dist_mode": dist_mode, + "disease_shape_available": False, + "death_shape_available": True, + "note": ( + "The mixed model uses Weibull rho only for Death. " + "Non-death disease hazards are exponential, equivalent to fixed rho=1." + ), + } + ]).to_csv(output_dir / "disease_shape_not_available_for_mixed_model.csv", index=False) + + print(f"Wrote Weibull shape parameter statistics to: {output_dir}") + + +if __name__ == "__main__": + main() diff --git a/export_weibull_shape_parameter_stats.py b/export_weibull_shape_parameter_stats.py new file mode 100644 index 0000000..67cc437 --- /dev/null +++ b/export_weibull_shape_parameter_stats.py @@ -0,0 +1,7 @@ +"""Compatibility entry point for Weibull shape-parameter export.""" + +from export_weibull_death_parameter_stats import main + + +if __name__ == "__main__": + main() diff --git a/plot_next_token_to_all_future_auc.R b/plot_next_token_to_all_future_auc.R new file mode 100644 index 0000000..50dc665 --- /dev/null +++ b/plot_next_token_to_all_future_auc.R @@ -0,0 +1,553 @@ +#!/usr/bin/env Rscript + +# Paper-grade single-panel figures supporting the conclusion that fixed-landmark +# horizon evaluation favors all_future over next_token. +# +# Outputs are written as separate panel files. This script intentionally does not +# combine panels with plot_grid(). + +suppressPackageStartupMessages({ + library(cowplot) + library(dplyr) + library(ggplot2) + library(jsonlite) + library(readr) + library(stringr) + library(tibble) + library(tidyr) +}) + +root_dir <- "." +runs_dir <- file.path(root_dir, "runs") +out_dir <- file.path(root_dir, "figures_next_token_to_all_future_absolute_smoking_alcohol_bmi") +dir.create(out_dir, showWarnings = FALSE, recursive = TRUE) + +required_time_mode <- "absolute" +required_extra_info_types <- c(11L, 66L, 67L) +required_extra_info_signature <- paste(sort(required_extra_info_types), collapse = ",") + +theme_set( + theme_cowplot(font_size = 9) + + theme( + plot.background = element_rect(fill = "white", color = NA), + panel.background = element_rect(fill = "white", color = NA), + legend.background = element_rect(fill = "white", color = NA), + legend.key = element_rect(fill = "white", color = NA) + ) +) + +target_cols <- c( + "next_token" = "#B54A3A", + "all_future" = "#2C7FB8" +) + +dist_shapes <- c( + "exponential" = 16, + "weibull" = 17, + "mixed" = 15 +) + +read_run_config <- function(run_path) { + cfg_path <- file.path(run_path, "train_config.json") + if (!file.exists(cfg_path)) return(NULL) + cfg <- jsonlite::read_json(cfg_path, simplifyVector = TRUE) + extra_info_types <- cfg$extra_info_types %||% integer(0) + extra_info_signature <- paste(sort(as.integer(extra_info_types)), collapse = ",") + tibble( + run = basename(run_path), + model_target_mode = as.character(cfg$model_target_mode %||% NA_character_), + target_mode = as.character(cfg$target_mode %||% NA_character_), + dist_mode = as.character(cfg$dist_mode %||% NA_character_), + time_mode = as.character(cfg$time_mode %||% NA_character_), + readout_name = as.character(cfg$readout_name %||% NA_character_), + attn_mask_mode = as.character(cfg$attn_mask_mode %||% NA_character_), + extra_info_signature = extra_info_signature + ) +} + +`%||%` <- function(x, y) { + if (is.null(x) || length(x) == 0) y else x +} + +load_one_result <- function(run_path, file_name, eval_family) { + cfg <- read_run_config(run_path) + if (is.null(cfg)) return(NULL) + fp <- file.path(run_path, file_name) + if (!file.exists(fp)) return(NULL) + + df <- suppressMessages(readr::read_csv(fp, show_col_types = FALSE)) + if (!("auc" %in% names(df)) || nrow(df) == 0) return(NULL) + + out <- df %>% + mutate( + run = basename(run_path), + eval_family = eval_family, + auc = as.numeric(auc) + ) %>% + left_join(cfg, by = "run", suffix = c("", "_cfg")) + + coalesce_joined <- function(data, col) { + cfg_col <- paste0(col, "_cfg") + if (col %in% names(data) && cfg_col %in% names(data)) { + dplyr::coalesce(data[[col]], data[[cfg_col]]) + } else if (col %in% names(data)) { + data[[col]] + } else if (cfg_col %in% names(data)) { + data[[cfg_col]] + } else { + rep(NA_character_, nrow(data)) + } + } + + for (col in c("model_target_mode", "target_mode", "dist_mode", "time_mode", "readout_name", "attn_mask_mode")) { + out[[col]] <- coalesce_joined(out, col) + } + + out %>% + select(-any_of(c( + "model_target_mode_cfg", "target_mode_cfg", "dist_mode_cfg", + "time_mode_cfg", "readout_name_cfg", "attn_mask_mode_cfg" + ))) +} + +run_paths <- list.dirs(runs_dir, recursive = FALSE, full.names = TRUE) + +landmark_auc <- bind_rows(lapply( + run_paths, + load_one_result, + file_name = "df_auc_landmark.csv", + eval_family = "Fixed landmark + horizon" +)) %>% + filter(time_mode == "absolute") + +token_auc <- bind_rows(lapply( + run_paths, + load_one_result, + file_name = "df_both.csv", + eval_family = "Delphi2M-style token" +)) %>% + filter(time_mode == "absolute") + +if (nrow(landmark_auc) == 0) { + stop("No landmark AUC files found under runs/*/df_auc_landmark.csv") +} +if (nrow(token_auc) == 0) { + stop("No token AUC files found under runs/*/df_both.csv") +} + +landmark_auc <- landmark_auc %>% + filter( + time_mode == required_time_mode, + extra_info_signature == required_extra_info_signature + ) + +token_auc <- token_auc %>% + filter( + time_mode == required_time_mode, + extra_info_signature == required_extra_info_signature + ) + +if (nrow(landmark_auc) == 0 || nrow(token_auc) == 0) { + stop( + "No AUC rows remain after filtering for time_mode='", + required_time_mode, + "' and extra_info_types='", + required_extra_info_signature, + "'." + ) +} + +message( + "Using runs with time_mode='", required_time_mode, + "' and extra_info_types='", required_extra_info_signature, "':" +) +print(sort(unique(landmark_auc$run))) + +classify_endpoint <- function(data) { + data %>% + mutate( + endpoint_type = if_else( + str_to_lower(as.character(label_code)) == "death", + "Death", + "Non-death disease" + ), + endpoint_type = factor(endpoint_type, levels = c("Non-death disease", "Death")) + ) +} + +landmark_auc <- classify_endpoint(landmark_auc) +token_auc <- classify_endpoint(token_auc) + +landmark_auc_disease <- landmark_auc %>% filter(endpoint_type == "Non-death disease") +token_auc_disease <- token_auc %>% filter(endpoint_type == "Non-death disease") +landmark_auc_death <- landmark_auc %>% filter(endpoint_type == "Death") +token_auc_death <- token_auc %>% filter(endpoint_type == "Death") + +if (nrow(landmark_auc_death) == 0 || nrow(token_auc_death) == 0) { + warning("Death rows were not found in one or both AUC tables.") +} + +auc_all <- bind_rows( + landmark_auc_disease %>% mutate(horizon = as.numeric(horizon), offset = NA_real_), + token_auc_disease %>% mutate(horizon = NA_real_, offset = as.numeric(offset)) +) %>% + mutate( + model_target_mode = factor(model_target_mode, levels = c("next_token", "all_future")), + eval_family = factor(eval_family, levels = c("Delphi2M-style token", "Fixed landmark + horizon")), + dist_mode = factor(dist_mode, levels = c("exponential", "weibull", "mixed")), + model_label = recode( + as.character(model_target_mode), + "next_token" = "next-token objective", + "all_future" = "all-future objective" + ) + ) + +mean_ci <- function(x) { + x <- x[is.finite(x)] + n <- length(x) + m <- mean(x) + se <- sd(x) / sqrt(n) + tibble(mean = m, ymin = m - 1.96 * se, ymax = m + 1.96 * se, n = n) +} + +save_panel <- function(plot, name, width = 3.6, height = 3.0) { + pdf_path <- file.path(out_dir, paste0(name, ".pdf")) + png_path <- file.path(out_dir, paste0(name, ".png")) + cowplot::save_plot(pdf_path, plot, base_width = width, base_height = height, bg = "white") + cowplot::save_plot(png_path, plot, base_width = width, base_height = height, dpi = 600, bg = "white") + message("Wrote: ", pdf_path) + message("Wrote: ", png_path) +} + +# Panel 1: run-level mean AUC under the clinically aligned landmark/horizon task. +# Death is excluded here and plotted separately below. +landmark_run <- landmark_auc_disease %>% + mutate(model_target_mode = factor(model_target_mode, levels = c("next_token", "all_future"))) %>% + group_by(run, model_target_mode, dist_mode, time_mode, target_mode) %>% + summarise(mean_auc = mean(auc, na.rm = TRUE), median_auc = median(auc, na.rm = TRUE), .groups = "drop") + +landmark_summary <- landmark_run %>% + group_by(model_target_mode) %>% + summarise(mean_ci(mean_auc), .groups = "drop") + +p1 <- ggplot(landmark_run, aes(x = model_target_mode, y = mean_auc)) + + geom_point( + aes(color = model_target_mode, shape = dist_mode), + position = position_jitter(width = 0.09, height = 0, seed = 1), + size = 2.2, + alpha = 0.88 + ) + + geom_errorbar( + data = landmark_summary, + aes(x = model_target_mode, y = mean, ymin = ymin, ymax = ymax, color = model_target_mode), + width = 0.12, + linewidth = 0.55, + inherit.aes = FALSE + ) + + geom_point( + data = landmark_summary, + aes(x = model_target_mode, y = mean, color = model_target_mode), + size = 3.4, + inherit.aes = FALSE + ) + + scale_color_manual(values = target_cols, guide = "none") + + scale_shape_manual(values = dist_shapes, na.translate = FALSE) + + scale_x_discrete(labels = c("next_token", "all_future")) + + coord_cartesian(ylim = c(0.58, 0.78)) + + labs( + x = NULL, + y = "Mean AUC per run", + shape = "Risk head", + title = "Non-death landmark AUC (absolute time)" + ) + + theme( + plot.title = element_text(face = "bold", size = 10), + axis.text.x = element_text(size = 9), + legend.position = c(0.72, 0.20), + legend.background = element_blank() + ) + +save_panel(p1, "panel_01_landmark_overall") + +# Panel 2: landmark AUC by prediction horizon. +landmark_horizon_run <- landmark_auc_disease %>% + mutate( + horizon = as.numeric(horizon), + model_target_mode = factor(model_target_mode, levels = c("next_token", "all_future")) + ) %>% + group_by(run, model_target_mode, horizon) %>% + summarise(mean_auc = mean(auc, na.rm = TRUE), .groups = "drop") + +landmark_horizon_summary <- landmark_horizon_run %>% + group_by(model_target_mode, horizon) %>% + summarise(mean_ci(mean_auc), .groups = "drop") + +p2 <- ggplot(landmark_horizon_run, aes(x = horizon, y = mean_auc, color = model_target_mode)) + + geom_line(aes(group = run), alpha = 0.18, linewidth = 0.35) + + geom_point(alpha = 0.32, size = 1.1) + + geom_ribbon( + data = landmark_horizon_summary, + aes(x = horizon, y = mean, ymin = ymin, ymax = ymax, fill = model_target_mode, group = model_target_mode), + alpha = 0.13, + color = NA, + inherit.aes = FALSE + ) + + geom_line(data = landmark_horizon_summary, aes(y = mean), linewidth = 0.85) + + geom_point(data = landmark_horizon_summary, aes(y = mean), size = 2.0) + + scale_color_manual( + values = target_cols, + labels = c("next_token", "all_future"), + name = NULL + ) + + scale_fill_manual(values = target_cols, guide = "none") + + scale_x_continuous(breaks = c(1, 5, 10)) + + coord_cartesian(ylim = c(0.58, 0.78)) + + labs( + x = "Prediction horizon, years", + y = "Mean AUC per run", + title = "Non-death landmark AUC across horizons" + ) + + theme( + plot.title = element_text(face = "bold", size = 10), + legend.position = c(0.31, 0.20), + legend.background = element_blank() + ) + +save_panel(p2, "panel_02_landmark_by_horizon", width = 3.8, height = 3.0) + +# Panel 3: Delphi2M-style token AUC by offset. This documents why the old +# evaluation can make next_token look competitive, especially near the event. +token_offset_run <- token_auc_disease %>% + mutate( + offset = as.numeric(offset), + model_target_mode = factor(model_target_mode, levels = c("next_token", "all_future")) + ) %>% + group_by(run, model_target_mode, offset) %>% + summarise(mean_auc = mean(auc, na.rm = TRUE), .groups = "drop") + +token_offset_summary <- token_offset_run %>% + group_by(model_target_mode, offset) %>% + summarise(mean_ci(mean_auc), .groups = "drop") + +p3 <- ggplot(token_offset_run, aes(x = offset, y = mean_auc, color = model_target_mode)) + + geom_line(aes(group = run), alpha = 0.18, linewidth = 0.35) + + geom_point(alpha = 0.32, size = 1.1) + + geom_ribbon( + data = token_offset_summary, + aes(x = offset, y = mean, ymin = ymin, ymax = ymax, fill = model_target_mode, group = model_target_mode), + alpha = 0.13, + color = NA, + inherit.aes = FALSE + ) + + geom_line(data = token_offset_summary, aes(y = mean), linewidth = 0.85) + + geom_point(data = token_offset_summary, aes(y = mean), size = 2.0) + + scale_color_manual( + values = target_cols, + labels = c("next_token", "all_future"), + name = NULL + ) + + scale_fill_manual(values = target_cols, guide = "none") + + scale_x_continuous(breaks = c(0.1, 1, 5, 10), trans = "log10") + + coord_cartesian(ylim = c(0.55, 0.82)) + + labs( + x = "Minimum offset before event, years", + y = "Mean AUC per run", + title = "Non-death token AUC by offset" + ) + + theme( + plot.title = element_text(face = "bold", size = 10), + legend.position = c(0.31, 0.20), + legend.background = element_blank() + ) + +save_panel(p3, "panel_03_token_auc_by_offset", width = 3.8, height = 3.0) + +# Panel 4: within-run contrast between old token evaluation and landmark +# evaluation. Each run contributes one point per evaluation family. +run_eval_contrast <- auc_all %>% + group_by(run, model_target_mode, dist_mode, eval_family) %>% + summarise(mean_auc = mean(auc, na.rm = TRUE), .groups = "drop") + +p4 <- ggplot(run_eval_contrast, aes(x = eval_family, y = mean_auc, color = model_target_mode)) + + geom_line(aes(group = run), alpha = 0.34, linewidth = 0.45) + + geom_point(aes(shape = dist_mode), size = 2.0, alpha = 0.84) + + stat_summary( + aes(group = model_target_mode), + fun = mean, + geom = "point", + size = 3.3, + shape = 18, + position = position_dodge(width = 0.16) + ) + + scale_color_manual( + values = target_cols, + labels = c("next_token", "all_future"), + name = NULL + ) + + scale_shape_manual(values = dist_shapes, na.translate = FALSE, name = "Risk head") + + coord_cartesian(ylim = c(0.58, 0.78)) + + labs( + x = NULL, + y = "Mean AUC per run", + title = "Evaluation choice changes the conclusion (absolute time)" + ) + + theme( + plot.title = element_text(face = "bold", size = 10), + axis.text.x = element_text(angle = 18, hjust = 1), + legend.position = "right" + ) + +save_panel(p4, "panel_04_evaluation_contrast", width = 4.3, height = 3.1) + +# Panel 5: disease-level distribution for the landmark task, pooled over +# horizons and runs. This shows the shift without hiding heterogeneity. +landmark_density <- landmark_auc_disease %>% + mutate(model_target_mode = factor(model_target_mode, levels = c("next_token", "all_future"))) %>% + filter(is.finite(auc)) + +p5 <- ggplot(landmark_density, aes(x = auc, fill = model_target_mode, color = model_target_mode)) + + geom_density(alpha = 0.20, linewidth = 0.65, adjust = 1.1) + + geom_vline( + data = landmark_density %>% + group_by(model_target_mode) %>% + summarise(mean_auc = mean(auc), .groups = "drop"), + aes(xintercept = mean_auc, color = model_target_mode), + linewidth = 0.75, + linetype = "22" + ) + + scale_color_manual(values = target_cols, labels = c("next_token", "all_future"), name = NULL) + + scale_fill_manual(values = target_cols, labels = c("next_token", "all_future"), name = NULL) + + coord_cartesian(xlim = c(0.35, 1.0)) + + labs( + x = "AUC", + y = "Density", + title = "Non-death landmark AUC distribution" + ) + + theme( + plot.title = element_text(face = "bold", size = 10), + legend.position = c(0.24, 0.82), + legend.background = element_blank() + ) + +save_panel(p5, "panel_05_landmark_auc_distribution", width = 3.8, height = 3.0) + +# Panel 6: death-only fixed landmark + horizon AUC. Death has one endpoint token, +# so each line is a run trajectory across horizons. +death_landmark_run <- landmark_auc_death %>% + mutate( + horizon = as.numeric(horizon), + model_target_mode = factor(model_target_mode, levels = c("next_token", "all_future")), + dist_mode = factor(dist_mode, levels = c("exponential", "weibull", "mixed")) + ) %>% + group_by(run, model_target_mode, dist_mode, horizon) %>% + summarise(mean_auc = mean(auc, na.rm = TRUE), .groups = "drop") + +death_landmark_summary <- death_landmark_run %>% + group_by(model_target_mode, horizon) %>% + summarise(mean_ci(mean_auc), .groups = "drop") + +p6 <- ggplot(death_landmark_run, aes(x = horizon, y = mean_auc, color = model_target_mode)) + + geom_line(aes(group = run), alpha = 0.42, linewidth = 0.45) + + geom_point(aes(shape = dist_mode), alpha = 0.9, size = 2.0) + + geom_line(data = death_landmark_summary, aes(y = mean, group = model_target_mode), linewidth = 0.9) + + geom_point(data = death_landmark_summary, aes(y = mean), size = 2.2) + + scale_color_manual(values = target_cols, labels = c("next_token", "all_future"), name = NULL) + + scale_shape_manual(values = dist_shapes, na.translate = FALSE, name = "Risk head") + + scale_x_continuous(breaks = c(1, 5, 10)) + + coord_cartesian(ylim = c(0.58, 0.95)) + + labs( + x = "Prediction horizon, years", + y = "AUC", + title = "Death-only landmark AUC" + ) + + theme( + plot.title = element_text(face = "bold", size = 10), + legend.position = "right" + ) + +save_panel(p6, "panel_06_death_landmark_by_horizon", width = 3.9, height = 3.0) + +# Panel 7: death-only Delphi2M-style token AUC by offset. +death_token_run <- token_auc_death %>% + mutate( + offset = as.numeric(offset), + model_target_mode = factor(model_target_mode, levels = c("next_token", "all_future")), + dist_mode = factor(dist_mode, levels = c("exponential", "weibull", "mixed")) + ) %>% + group_by(run, model_target_mode, dist_mode, offset) %>% + summarise(mean_auc = mean(auc, na.rm = TRUE), .groups = "drop") + +death_token_summary <- death_token_run %>% + group_by(model_target_mode, offset) %>% + summarise(mean_ci(mean_auc), .groups = "drop") + +p7 <- ggplot(death_token_run, aes(x = offset, y = mean_auc, color = model_target_mode)) + + geom_line(aes(group = run), alpha = 0.42, linewidth = 0.45) + + geom_point(aes(shape = dist_mode), alpha = 0.9, size = 2.0) + + geom_line(data = death_token_summary, aes(y = mean, group = model_target_mode), linewidth = 0.9) + + geom_point(data = death_token_summary, aes(y = mean), size = 2.2) + + scale_color_manual(values = target_cols, labels = c("next_token", "all_future"), name = NULL) + + scale_shape_manual(values = dist_shapes, na.translate = FALSE, name = "Risk head") + + scale_x_continuous(breaks = c(0.1, 1, 5, 10), trans = "log10") + + coord_cartesian(ylim = c(0.58, 0.95)) + + labs( + x = "Minimum offset before event, years", + y = "AUC", + title = "Death-only token AUC" + ) + + theme( + plot.title = element_text(face = "bold", size = 10), + legend.position = "right" + ) + +save_panel(p7, "panel_07_death_token_auc_by_offset", width = 3.9, height = 3.0) + +# Panel 8: death-only contrast between the two evaluation families. +death_eval_contrast <- bind_rows( + landmark_auc_death %>% mutate(horizon = as.numeric(horizon), offset = NA_real_), + token_auc_death %>% mutate(horizon = NA_real_, offset = as.numeric(offset)) +) %>% + mutate( + model_target_mode = factor(model_target_mode, levels = c("next_token", "all_future")), + eval_family = factor(eval_family, levels = c("Delphi2M-style token", "Fixed landmark + horizon")), + dist_mode = factor(dist_mode, levels = c("exponential", "weibull", "mixed")) + ) %>% + group_by(run, model_target_mode, dist_mode, eval_family) %>% + summarise(mean_auc = mean(auc, na.rm = TRUE), .groups = "drop") + +p8 <- ggplot(death_eval_contrast, aes(x = eval_family, y = mean_auc, color = model_target_mode)) + + geom_line(aes(group = run), alpha = 0.38, linewidth = 0.5) + + geom_point(aes(shape = dist_mode), size = 2.2, alpha = 0.9) + + stat_summary( + aes(group = model_target_mode), + fun = mean, + geom = "point", + size = 3.4, + shape = 18, + position = position_dodge(width = 0.16) + ) + + scale_color_manual(values = target_cols, labels = c("next_token", "all_future"), name = NULL) + + scale_shape_manual(values = dist_shapes, na.translate = FALSE, name = "Risk head") + + coord_cartesian(ylim = c(0.58, 0.95)) + + labs( + x = NULL, + y = "Mean AUC per run", + title = "Death endpoint evaluated separately" + ) + + theme( + plot.title = element_text(face = "bold", size = 10), + axis.text.x = element_text(angle = 18, hjust = 1), + legend.position = "right" + ) + +save_panel(p8, "panel_08_death_evaluation_contrast", width = 4.3, height = 3.1) + +# Export the exact run-level summaries used by the figures. +readr::write_csv(landmark_run, file.path(out_dir, "landmark_run_summary.csv")) +readr::write_csv(token_offset_run, file.path(out_dir, "token_offset_run_summary.csv")) +readr::write_csv(run_eval_contrast, file.path(out_dir, "run_evaluation_contrast.csv")) +readr::write_csv(death_landmark_run, file.path(out_dir, "death_landmark_run_summary.csv")) +readr::write_csv(death_token_run, file.path(out_dir, "death_token_offset_run_summary.csv")) +readr::write_csv(death_eval_contrast, file.path(out_dir, "death_evaluation_contrast.csv")) + +message("Done. Panels are in: ", normalizePath(out_dir, winslash = "/")) diff --git a/run_weibull_shape_exports.sh b/run_weibull_shape_exports.sh new file mode 100755 index 0000000..3bd46e5 --- /dev/null +++ b/run_weibull_shape_exports.sh @@ -0,0 +1,76 @@ +#!/usr/bin/env bash + +# Export Weibull shape-parameter summaries for the all_future models trained +# with smoking/alcohol/BMI extra information. +# +# Bash 4.2 compatible. Run from the DeepHealth repository root on the Linux +# server, for example: +# +# bash run_weibull_shape_exports.sh +# +# Optional overrides: +# PYTHON=python3 DEVICE=cuda BATCH_SIZE=128 NUM_WORKERS=4 ROW_BATCH_SIZE=512 \ +# bash run_weibull_shape_exports.sh + +set -euo pipefail + +PYTHON="${PYTHON:-python}" +DEVICE="${DEVICE:-cuda}" +BATCH_SIZE="${BATCH_SIZE:-128}" +NUM_WORKERS="${NUM_WORKERS:-4}" +ROW_BATCH_SIZE="${ROW_BATCH_SIZE:-512}" +LANDMARK_START="${LANDMARK_START:-40}" +LANDMARK_STOP="${LANDMARK_STOP:-80}" +LANDMARK_STEP="${LANDMARK_STEP:-5}" +HORIZONS="${HORIZONS:-1,5,10}" + +RUNS=( + "runs/relative_weibull_all_future_pure_disease_20260620_095229" + "runs/relative_mixed_all_future_pure_disease_20260620_132415" + "runs/absolute_weibull_all_future_pure_disease_20260620_114816" + "runs/absolute_mixed_all_future_pure_disease_20260620_161804" +) + +echo "Python: ${PYTHON}" +echo "Device: ${DEVICE}" +echo "Batch size: ${BATCH_SIZE}" +echo "Workers: ${NUM_WORKERS}" +echo "Row batch size: ${ROW_BATCH_SIZE}" +echo "Horizons: ${HORIZONS}" +echo + +for run_path in "${RUNS[@]}"; do + if [[ ! -d "${run_path}" ]]; then + echo "[ERROR] Missing run directory: ${run_path}" >&2 + exit 1 + fi + if [[ ! -f "${run_path}/best_model.pt" ]]; then + echo "[ERROR] Missing checkpoint: ${run_path}/best_model.pt" >&2 + exit 1 + fi + if [[ ! -f "${run_path}/train_config.json" ]]; then + echo "[ERROR] Missing config: ${run_path}/train_config.json" >&2 + exit 1 + fi + + output_path="${run_path}/weibull_shape_parameter_stats_test" + echo "=== Exporting Weibull shape stats: ${run_path} ===" + "${PYTHON}" export_weibull_shape_parameter_stats.py \ + --run_path "${run_path}" \ + --output_path "${output_path}" \ + --eval_split test \ + --device "${DEVICE}" \ + --batch_size "${BATCH_SIZE}" \ + --num_workers "${NUM_WORKERS}" \ + --row_batch_size "${ROW_BATCH_SIZE}" \ + --hidden_cache_dtype float32 \ + --landmark_start "${LANDMARK_START}" \ + --landmark_stop "${LANDMARK_STOP}" \ + --landmark_step "${LANDMARK_STEP}" \ + --horizons "${HORIZONS}" \ + --include_all_token_rho_summary + echo "Wrote: ${output_path}" + echo +done + +echo "All Weibull shape exports completed."