JiaruiLi/DeepHealth
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add Weibull shape export scripts

Browse Source
This commit is contained in:
Jiarui Li
2026-07-01 15:31:31 +08:00
parent f417a91a74
commit d08e5b34f4
4 changed files with 1138 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

553
plot_next_token_to_all_future_auc.R Normal file
View File

@@ -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 = "/"))