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

Remove obsolete evaluation and batch scripts

Browse Source
This commit is contained in:
Jiarui Li
2026-07-25 13:14:42 +08:00
parent b13db5e407
commit 15ace878f4
13 changed files with 0 additions and 4958 deletions
Download Patch File Download Diff File Expand all files Collapse all files

816
evaluate_event_free_survival.py
View File

@@ -1,816 +0,0 @@
"""Compute landmark future death and incident system-disease risks.
For each selected patient and landmark age, this script computes:
* future death risk within tau years;
* future incident disease risk for each ICD-10 chapter-derived system;
* model attribution of each historical organ/system disease set to predicted
mortality risk, computed by deleting that system's historical disease tokens
and re-querying the model;
* historical modeled-disease count;
* historical modeled-disease count within each ICD-10 chapter-derived system.
Death is always token vocab_size - 1. Disease groups are read from
icd10_chapter_organ_mapping.csv.
"""
from __future__ import annotations
import argparse
import json
from pathlib import Path
from typing import Any, Dict, List, Optional, Sequence
import numpy as np
import pandas as pd
import torch
from torch.nn.utils.rnn import pad_sequence
from torch.utils.data import DataLoader, Dataset
from tqdm.auto import tqdm
from dataset import HealthDataset
from eval_data import load_sequence_eval_dataset
from evaluate_auc_v2 import (
LandmarkDataset,
build_model_from_dataset,
cfg_get,
load_checkpoint_state_dict,
load_json_config,
load_model_state,
make_eval_indices,
resolve_dist_mode_for_checkpoint,
resolve_eval_device,
validate_dataset_metadata,
)
from future_risk import (
death_risk_from_probabilities,
new_disease_risk_from_probabilities,
probabilities_from_logits,
)
from models import DeepHealth
from readouts import build_readout
from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX
from train_util import load_eid_file, load_extra_info_types_file
SPECIAL_TOKENS = {PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX}
def parse_int_list(value: Any) -> Optional[List[int]]:
if value is None:
return None
if isinstance(value, (list, tuple, np.ndarray)):
return [int(x) for x in value]
text = str(value).strip()
if text == "":
return None
if text.startswith("["):
values = json.loads(text)
if not isinstance(values, list):
raise ValueError(
f"Expected a JSON list, got {type(values).__name__}")
return [int(x) for x in values]
return [int(x.strip()) for x in text.split(",") if x.strip()]
def load_extra_info_types(value: Any) -> Optional[List[int]]:
if value is None:
return None
text = str(value)
path = Path(text)
if path.exists():
return load_extra_info_types_file(text)
return parse_int_list(value)
def make_landmark_ages(start: float, stop: float, step: float) -> np.ndarray:
if step <= 0:
raise ValueError("landmark_step must be positive")
if stop < start:
raise ValueError("landmark_stop must be >= landmark_start")
# Include stop when it lands on the grid, e.g. 40,45,...,80.
return np.arange(start, stop + step * 0.5, step, dtype=np.float32)
def build_first_occurrence_maps_for_landmarks(
dataset: HealthDataset,
subset_indices: np.ndarray,
) -> Dict[int, tuple[np.ndarray, np.ndarray]]:
first_lists: Dict[int, list[tuple[int, float]]] = {}
for patient_id, dataset_index in enumerate(np.asarray(subset_indices, dtype=np.int64).tolist()):
s = dataset.samples[int(dataset_index)]
seq_event = np.asarray(s["event_seq"], dtype=np.int64)
seq_time = np.asarray(s["time_seq"], dtype=np.float32)
tgt_event = np.asarray(s["target_event_seq"], dtype=np.int64)
tgt_time = np.asarray(s["target_time_seq"], dtype=np.float32)
if seq_event.size == 0 or tgt_event.size == 0:
continue
full_event = np.concatenate([seq_event, tgt_event[-1:]])
full_time = np.concatenate([seq_time, tgt_time[-1:]])
uniq_tokens, first_idx = np.unique(full_event, return_index=True)
for token, idx in zip(uniq_tokens.tolist(), first_idx.tolist()):
token = int(token)
if token in SPECIAL_TOKENS:
continue
first_lists.setdefault(token, []).append(
(patient_id, float(full_time[int(idx)])))
return {
int(token): (
np.asarray([p for p, _ in pairs], dtype=np.int32),
np.asarray([t for _, t in pairs], dtype=np.float32),
)
for token, pairs in first_lists.items()
if pairs
}
def normalize_eval_split(args: argparse.Namespace, cfg: Dict[str, Any]) -> str:
eval_split = str(cfg_get(args, cfg, "eval_split", "test")).lower()
if eval_split in {"valid", "validation"}:
return "val"
if eval_split not in {"train", "val", "test", "all"}:
raise ValueError(f"Unsupported eval_split={eval_split!r}")
return eval_split
def load_eval_sequence_dataset(
args: argparse.Namespace,
cfg: Dict[str, Any],
) -> tuple[Any, np.ndarray, str, str]:
eval_split = normalize_eval_split(args, cfg)
model_target_mode = str(cfg.get("model_target_mode", "next_token")).lower()
data_prefix = str(cfg.get("data_prefix", "ukb"))
labels_file = str(cfg.get("labels_file", "labels.csv"))
no_event_interval_years = float(cfg.get("no_event_interval_years", 5.0))
include_no_event_in_uts_target = bool(
cfg.get("include_no_event_in_uts_target", False))
extra_info_types = load_extra_info_types(args.extra_info_types)
if extra_info_types is None:
extra_info_types = parse_int_list(cfg.get("extra_info_types", None))
print("Loading one sequence eval dataset...")
dataset = load_sequence_eval_dataset(
model_target_mode=model_target_mode,
data_prefix=data_prefix,
labels_file=labels_file,
no_event_interval_years=no_event_interval_years,
include_no_event_in_uts_target=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=extra_info_types,
)
train_eid_file = cfg_get(args, cfg, "train_eid_file", "ukb_train_eid.csv")
val_eid_file = cfg_get(args, cfg, "val_eid_file", "ukb_val_eid.csv")
test_eid_file = cfg_get(args, cfg, "test_eid_file", "ukb_test_eid.csv")
split_files_exist = all(
Path(str(path)).exists()
for path in (train_eid_file, val_eid_file, test_eid_file)
)
if eval_split != "all" and split_files_exist:
split_files = {
"train": train_eid_file,
"val": val_eid_file,
"test": test_eid_file,
}
selected_eids = load_eid_file(split_files[eval_split])
out = np.asarray(
[
idx
for idx, sample in enumerate(dataset.samples)
if int(sample["eid"]) in selected_eids
],
dtype=np.int64,
)
if out.size == 0:
raise ValueError(
f"No samples found for eval_split={eval_split!r} using {split_files[eval_split]}"
)
split_source = "eid_files"
else:
if eval_split == "all":
out = np.arange(len(dataset.samples), dtype=np.int64)
split_source = "all"
else:
out = make_eval_indices(dataset, args, cfg)
split_source = "ratio_split"
subset_size = cfg_get(args, cfg, "dataset_subset_size", None)
if subset_size is not None and int(subset_size) > 0:
out = out[: int(subset_size)]
return dataset, np.asarray(out, dtype=np.int64), eval_split, split_source
def load_organ_groups(
path: Path,
*,
vocab_size: int,
) -> tuple[dict[str, list[int]], dict[str, str], dict[int, str]]:
table = pd.read_csv(path)
required = {"token_id", "organ_system", "organ_system_label", "is_death"}
missing = required - set(table.columns)
if missing:
raise ValueError(f"{path} is missing columns: {sorted(missing)}")
death_idx = int(vocab_size) - 1
groups: dict[str, list[int]] = {}
labels: dict[str, str] = {}
token_to_group: dict[int, str] = {}
for row in table.itertuples(index=False):
token = int(getattr(row, "token_id"))
if token in SPECIAL_TOKENS or token == death_idx:
continue
if token < 0 or token >= int(vocab_size):
continue
if int(getattr(row, "is_death")) == 1:
continue
group = str(getattr(row, "organ_system"))
label = str(getattr(row, "organ_system_label"))
groups.setdefault(group, []).append(token)
labels[group] = label
token_to_group[token] = group
groups = {k: sorted(set(v)) for k, v in groups.items() if v}
return groups, labels, token_to_group
class IndexedLandmarkDataset(Dataset):
def __init__(self, base: LandmarkDataset) -> None:
self.base = base
def __len__(self) -> int:
return len(self.base)
def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
item = dict(self.base[idx])
item["row_idx"] = torch.tensor(int(idx), dtype=torch.long)
return item
def collate_indexed_landmark_fn(batch: List[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]:
event_seq = pad_sequence(
[x["event_seq"] for x in batch], batch_first=True, padding_value=PAD_IDX
)
time_seq = pad_sequence(
[x["time_seq"] for x in batch], batch_first=True, padding_value=0.0
)
readout_mask = pad_sequence(
[x["readout_mask"] for x in batch], batch_first=True, padding_value=False
)
other_type = pad_sequence(
[x["other_type"] for x in batch], batch_first=True, padding_value=0
)
other_value = pad_sequence(
[x["other_value"] for x in batch], batch_first=True, padding_value=0.0
)
other_value_kind = pad_sequence(
[x["other_value_kind"] for x in batch], batch_first=True, padding_value=0
)
other_time = pad_sequence(
[x["other_time"] for x in batch], batch_first=True, padding_value=0.0
)
return {
"event_seq": event_seq,
"time_seq": time_seq,
"padding_mask": event_seq > PAD_IDX,
"readout_mask": readout_mask,
"sex": torch.stack([x["sex"] for x in batch]),
"other_type": other_type,
"other_value": other_value,
"other_value_kind": other_value_kind,
"other_time": other_time,
"landmark_pos": torch.stack([x["landmark_pos"] for x in batch]),
"t_query": torch.stack([x["t_query"] for x in batch]),
"patient_id": torch.stack([x["patient_id"] for x in batch]),
"landmark_age": torch.stack([x["landmark_age"] for x in batch]),
"followup_end_time": torch.stack([x["followup_end_time"] for x in batch]),
"death_time": torch.stack([x["death_time"] for x in batch]),
"row_idx": torch.stack([x["row_idx"] for x in batch]),
}
def build_group_ablated_slice(
batch: Dict[str, torch.Tensor],
token_ids: Sequence[int],
row_indices: torch.Tensor,
) -> Dict[str, torch.Tensor]:
"""Build one fixed-width ablated slice without rebuilding variable-length rows."""
event_seq = batch["event_seq"]
out: Dict[str, torch.Tensor] = {}
out["event_seq"] = event_seq[row_indices].clone()
out["time_seq"] = batch["time_seq"][row_indices]
out["readout_mask"] = batch["readout_mask"][row_indices].clone()
out["padding_mask"] = batch["padding_mask"][row_indices].bool().clone()
out["landmark_pos"] = batch["landmark_pos"][row_indices].clone()
seq_len = int(event_seq.shape[1])
positions = torch.arange(seq_len, device=event_seq.device)[None, :]
ids = torch.as_tensor(token_ids, dtype=event_seq.dtype,
device=event_seq.device)
remove = torch.isin(out["event_seq"], ids) & out["padding_mask"]
out["event_seq"] = torch.where(
remove,
torch.full_like(out["event_seq"], PAD_IDX),
out["event_seq"],
)
out["padding_mask"] &= ~remove
out["readout_mask"] &= ~remove
has_valid = out["padding_mask"].any(dim=1)
if not bool(has_valid.all().item()):
empty_rows = torch.nonzero(~has_valid, as_tuple=False).flatten()
out["event_seq"][empty_rows, 0] = CHECKUP_IDX
out["time_seq"][empty_rows, 0] = batch["t_query"][row_indices[empty_rows]].to(
dtype=out["time_seq"].dtype
)
out["padding_mask"][empty_rows, 0] = True
out["readout_mask"][empty_rows, 0] = True
out["landmark_pos"][empty_rows] = 0
has_readout = out["readout_mask"].any(dim=1)
if not bool(has_readout.all().item()):
rows = torch.nonzero(~has_readout, as_tuple=False).flatten()
local_valid = out["padding_mask"][rows]
last_pos = torch.where(
local_valid,
positions.expand(local_valid.shape[0], -1),
torch.zeros_like(positions.expand(local_valid.shape[0], -1)),
).amax(dim=1)
out["readout_mask"][rows] = False
out["readout_mask"][rows, last_pos] = True
out["landmark_pos"][rows] = last_pos.to(dtype=out["landmark_pos"].dtype)
repeated_keys = (
"sex",
"other_type",
"other_value",
"other_value_kind",
"other_time",
"t_query",
"patient_id",
"landmark_age",
"followup_end_time",
"death_time",
"row_idx",
)
for key in repeated_keys:
out[key] = batch[key][row_indices]
return out
def concat_tensor_batches(chunks: Sequence[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]:
return {
key: torch.cat([chunk[key] for chunk in chunks], dim=0)
for key in chunks[0]
}
def iter_group_ablated_batches(
batch: Dict[str, torch.Tensor],
group_names: Sequence[str],
organ_groups: dict[str, list[int]],
occurred: torch.Tensor,
max_batch_size: int,
):
"""Yield ablated chunks as soon as enough rows are available for a forward pass."""
pending_batches: list[Dict[str, torch.Tensor]] = []
pending_groups: list[str] = []
pending_rows: list[int] = []
pending_n = 0
for group in group_names:
ids = torch.as_tensor(
organ_groups[group], dtype=torch.long, device=occurred.device)
if ids.numel() == 0:
continue
active_rows = torch.nonzero(
occurred[:, ids].any(dim=1), as_tuple=False).flatten()
if active_rows.numel() == 0:
continue
row_offset = 0
while row_offset < int(active_rows.numel()):
capacity = int(max_batch_size) - pending_n
row_stop = min(int(active_rows.numel()), row_offset + capacity)
row_indices = active_rows[row_offset:row_stop].to(
device=batch["event_seq"].device)
chunk = build_group_ablated_slice(
batch=batch,
token_ids=organ_groups[group],
row_indices=row_indices,
)
chunk_n = int(row_indices.numel())
pending_batches.append(chunk)
pending_groups.extend([group] * chunk_n)
pending_rows.extend(int(x)
for x in row_indices.detach().cpu().tolist())
pending_n += chunk_n
row_offset = row_stop
if pending_n >= int(max_batch_size):
yield concat_tensor_batches(pending_batches), pending_groups, pending_rows
pending_batches = []
pending_groups = []
pending_rows = []
pending_n = 0
if pending_batches:
yield concat_tensor_batches(pending_batches), pending_groups, pending_rows
@torch.no_grad()
def infer_landmark_hidden(
*,
model: DeepHealth,
batch: Dict[str, torch.Tensor],
device: torch.device,
model_target_mode: str,
readout_name: str,
readout_reduce: str,
) -> torch.Tensor:
batch_dev = {
k: (v.to(device, non_blocking=True) if isinstance(v, torch.Tensor) else v)
for k, v in batch.items()
}
if model_target_mode == "all_future":
return model(
event_seq=batch_dev["event_seq"].long(),
time_seq=batch_dev["time_seq"].float(),
sex=batch_dev["sex"].long(),
padding_mask=batch_dev["padding_mask"].bool(),
t_query=batch_dev["t_query"].float(),
other_type=batch_dev["other_type"].long(),
other_value=batch_dev["other_value"].float(),
other_value_kind=batch_dev["other_value_kind"].long(),
other_time=batch_dev["other_time"].float(),
target_mode="all_future",
)
hidden = model(
event_seq=batch_dev["event_seq"].long(),
time_seq=batch_dev["time_seq"].float(),
sex=batch_dev["sex"].long(),
padding_mask=batch_dev["padding_mask"].bool(),
other_type=batch_dev["other_type"].long(),
other_value=batch_dev["other_value"].float(),
other_value_kind=batch_dev["other_value_kind"].long(),
other_time=batch_dev["other_time"].float(),
target_mode="next_token",
)
readout = build_readout(readout_name, reduce=readout_reduce)
readout_out = readout(
hidden=hidden,
time_seq=batch_dev["time_seq"].float(),
padding_mask=batch_dev["padding_mask"].bool(),
readout_mask=batch_dev["readout_mask"].bool(),
)
return readout_out.hidden.gather(
1,
batch_dev["landmark_pos"].long()[:, None, None].expand(
-1, 1, readout_out.hidden.shape[-1]
),
).squeeze(1)
def make_occurred_mask(
event_seq: torch.Tensor,
*,
vocab_size: int,
device: torch.device,
) -> torch.Tensor:
occurred = torch.zeros(event_seq.shape[0], int(
vocab_size), dtype=torch.bool, device=device)
valid = (event_seq >= 0) & (event_seq < int(vocab_size))
safe = event_seq.clamp(min=0, max=int(vocab_size) - 1).to(device)
occurred.scatter_(1, safe, valid.to(device))
return occurred
def mortality_hazard_from_risk(risk: torch.Tensor, eps: float = 1e-7) -> torch.Tensor:
return -torch.log1p(-risk.clamp(0.0, 1.0 - float(eps)))
def death_risk_for_batch(
*,
model: DeepHealth,
batch: Dict[str, torch.Tensor],
device: torch.device,
model_target_mode: str,
readout_name: str,
readout_reduce: str,
dist_mode: str,
tau: float,
) -> torch.Tensor:
hidden = infer_landmark_hidden(
model=model,
batch=batch,
device=device,
model_target_mode=model_target_mode,
readout_name=readout_name,
readout_reduce=readout_reduce,
)
logits = model.calc_risk(hidden)
rho = model.calc_weibull_rho(hidden) if dist_mode == "weibull" else None
death_rho = model.calc_death_rho(hidden) if dist_mode == "mixed" else None
probabilities = probabilities_from_logits(
logits,
tau,
dist_mode=dist_mode,
rho=rho,
death_rho=death_rho,
)
return death_risk_from_probabilities(probabilities)
def historical_counts_by_group(
tokens: np.ndarray,
*,
death_idx: int,
token_to_group: dict[int, str],
group_names: Sequence[str],
) -> tuple[int, dict[str, int]]:
unique_tokens = {
int(token)
for token in np.asarray(tokens, dtype=np.int64).tolist()
if int(token) not in SPECIAL_TOKENS and int(token) != int(death_idx)
}
total = len(unique_tokens)
out = {group: 0 for group in group_names}
for token in unique_tokens:
group = token_to_group.get(token)
if group in out:
out[group] += 1
return total, out
def output_name_for_run(run_path: Path, eval_split: str, tau: float) -> Path:
return run_path / f"future_risk_{eval_split}_tau{tau:g}y.csv"
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(
description="Compute landmark death and incident system-disease risks."
)
parser.add_argument("--run_path", type=str, required=True)
parser.add_argument("--output_path", type=str, default=None)
parser.add_argument("--organ_mapping_path", type=str,
default="icd10_chapter_organ_mapping.csv")
parser.add_argument("--eval_split", type=str, default=None)
parser.add_argument("--dataset_subset_size", type=int, default=None)
parser.add_argument("--train_eid_file", type=str, default=None)
parser.add_argument("--val_eid_file", type=str, default=None)
parser.add_argument("--test_eid_file", type=str, default=None)
parser.add_argument("--landmark_start", type=float, default=40.0)
parser.add_argument("--landmark_stop", type=float, default=80.0)
parser.add_argument("--landmark_step", type=float, default=5.0)
parser.add_argument("--tau", type=float, default=5.0)
parser.add_argument("--min_history_events", type=int, default=None)
parser.add_argument("--batch_size", type=int, default=None)
parser.add_argument(
"--attribution_batch_size",
type=int,
default=None,
help="Forward batch size for expanded organ/system ablation queries.",
)
parser.add_argument("--num_workers", type=int, default=None)
parser.add_argument("--device", type=str, default=None)
parser.add_argument("--extra_info_types", type=str, default=None)
return parser.parse_args()
def main() -> None:
args = parse_args()
run_path = Path(args.run_path)
config_path = run_path / "train_config.json"
checkpoint_path = run_path / "best_model.pt"
if not config_path.exists():
raise FileNotFoundError(f"train_config.json not found: {config_path}")
if not checkpoint_path.exists():
raise FileNotFoundError(f"best_model.pt not found: {checkpoint_path}")
cfg = load_json_config(config_path)
model_target_mode = str(cfg.get("model_target_mode", "next_token")).lower()
if model_target_mode not in {"next_token", "all_future"}:
raise ValueError(
f"Unsupported model_target_mode: {model_target_mode!r}")
target_mode = str(cfg.get("target_mode", "uts"))
attn_mask_mode = str(
cfg.get("attn_mask_mode", "non_strict_time" if target_mode ==
"uts" else "target_aware")
)
readout_name = str(cfg.get(
"readout_name", "same_time_group_end" if target_mode == "uts" else "token"))
readout_reduce = str(cfg.get("readout_reduce", "mean"))
dataset, subset_indices, eval_split, split_source = load_eval_sequence_dataset(
args,
cfg,
)
validate_dataset_metadata(dataset, cfg)
landmark_ages = make_landmark_ages(
float(args.landmark_start),
float(args.landmark_stop),
float(args.landmark_step),
)
tau = float(args.tau)
if tau < 0:
raise ValueError("tau must be non-negative")
first_occurrence_by_token = build_first_occurrence_maps_for_landmarks(
dataset,
subset_indices,
)
death_idx = int(dataset.vocab_size) - 1
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_idx],
)
organ_groups, organ_labels, token_to_group = load_organ_groups(
Path(args.organ_mapping_path),
vocab_size=int(dataset.vocab_size),
)
group_names = sorted(organ_groups)
state_dict = load_checkpoint_state_dict(checkpoint_path, map_location="cpu")
dist_mode = resolve_dist_mode_for_checkpoint(
str(cfg.get("dist_mode", "exponential")), state_dict)
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, state_dict=state_dict
).to(device)
load_model_state(model, state_dict)
model.eval()
batch_size = int(cfg_get(args, cfg, "batch_size", 128))
attribution_batch_size = int(
cfg_get(args, cfg, "attribution_batch_size",
max(batch_size * 4, batch_size))
)
if attribution_batch_size <= 0:
raise ValueError("attribution_batch_size must be positive")
num_workers = int(cfg_get(args, cfg, "num_workers", 4))
loader = DataLoader(
IndexedLandmarkDataset(landmark_dataset),
batch_size=batch_size,
shuffle=False,
collate_fn=collate_indexed_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,
)
output_path = Path(args.output_path) if args.output_path else output_name_for_run(
run_path, eval_split, tau)
output_path.parent.mkdir(parents=True, exist_ok=True)
print(f"Eval split: {eval_split}")
print(f"Split source: {split_source}")
print(f"Selected patients: {len(subset_indices)}")
print(f"Landmark ages: {landmark_ages.tolist()}")
print(f"Tau: {tau:g} years")
print(f"Dist mode: {dist_mode}")
print(f"Device: {device}")
print(f"Death token: {death_idx}")
print(f"Organ/system groups: {len(group_names)}")
print(f"Landmark rows: {len(landmark_dataset)}")
print(f"Attribution batch size: {attribution_batch_size}")
print(f"Output: {output_path}")
rows: list[dict[str, Any]] = []
for batch in tqdm(loader, desc="Future risks", dynamic_ncols=True):
hidden = infer_landmark_hidden(
model=model,
batch=batch,
device=device,
model_target_mode=model_target_mode,
readout_name=readout_name,
readout_reduce=readout_reduce,
)
logits = model.calc_risk(hidden)
rho = model.calc_weibull_rho(hidden) if dist_mode == "weibull" else None
death_rho = model.calc_death_rho(
hidden) if dist_mode == "mixed" else None
probabilities = probabilities_from_logits(
logits,
tau,
dist_mode=dist_mode,
rho=rho,
death_rho=death_rho,
)
occurred = make_occurred_mask(
batch["event_seq"].to(device),
vocab_size=int(dataset.vocab_size),
device=device,
)
death_risk_tensor = death_risk_from_probabilities(probabilities)
death_hazard_tensor = mortality_hazard_from_risk(death_risk_tensor)
death_risk = death_risk_tensor.detach().cpu().numpy()
group_risk: dict[str, np.ndarray] = {}
for group in group_names:
group_risk[group] = new_disease_risk_from_probabilities(
probabilities,
occurred,
organ_groups[group],
).detach().cpu().numpy()
group_mortality_attr_prob: dict[str, np.ndarray] = {}
group_mortality_attr_hazard: dict[str, np.ndarray] = {}
batch_n = int(batch["event_seq"].shape[0])
zeros = np.zeros(batch_n, dtype=np.float32)
for group in group_names:
group_mortality_attr_prob[group] = zeros.copy()
group_mortality_attr_hazard[group] = zeros.copy()
for ablated_chunk, chunk_groups, chunk_rows in iter_group_ablated_batches(
batch=batch,
group_names=group_names,
organ_groups=organ_groups,
occurred=occurred,
max_batch_size=attribution_batch_size,
):
ablated_death_risk = death_risk_for_batch(
model=model,
batch=ablated_chunk,
device=device,
model_target_mode=model_target_mode,
readout_name=readout_name,
readout_reduce=readout_reduce,
dist_mode=dist_mode,
tau=tau,
)
row_tensor = torch.as_tensor(
chunk_rows, dtype=torch.long, device=device)
ablated_death_hazard = mortality_hazard_from_risk(
ablated_death_risk)
attr_prob = (
death_risk_tensor[row_tensor] - ablated_death_risk
).detach().cpu().numpy()
attr_hazard = (
death_hazard_tensor[row_tensor] - ablated_death_hazard
).detach().cpu().numpy()
for local_idx, (group, row_idx) in enumerate(zip(chunk_groups, chunk_rows)):
group_mortality_attr_prob[group][row_idx] = attr_prob[local_idx]
group_mortality_attr_hazard[group][row_idx] = attr_hazard[local_idx]
row_indices = batch["row_idx"].cpu().numpy().astype(np.int64)
for j, row_idx in enumerate(row_indices.tolist()):
meta = landmark_dataset.rows[int(row_idx)]
dataset_index = int(meta["dataset_index"])
sample = dataset.samples[dataset_index]
hist_tokens = np.asarray(meta["event_seq"], dtype=np.int64)
total_count, group_counts = historical_counts_by_group(
hist_tokens,
death_idx=death_idx,
token_to_group=token_to_group,
group_names=group_names,
)
out: dict[str, Any] = {
"patient_id": int(meta["patient_id"]),
"dataset_index": dataset_index,
"eid": int(sample.get("eid", -1)),
"sex": int(meta["sex"]),
"landmark_age": float(meta["landmark_age"]),
"tau": tau,
"followup_end_time": float(meta["followup_end_time"]),
"history_disease_count": int(total_count),
"death_risk": float(death_risk[j]),
}
for group in group_names:
out[f"history_count__{group}"] = int(group_counts[group])
out[f"new_disease_risk__{group}"] = float(group_risk[group][j])
if int(group_counts[group]) == 0:
group_mortality_attr_prob[group][j] = 0.0
group_mortality_attr_hazard[group][j] = 0.0
out[f"mortality_attribution_probability__{group}"] = float(
group_mortality_attr_prob[group][j]
)
out[f"mortality_attribution_hazard__{group}"] = float(
group_mortality_attr_hazard[group][j]
)
rows.append(out)
df = pd.DataFrame(rows)
df.to_csv(output_path, index=False)
print(f"Wrote {len(df)} rows to {output_path}")
if __name__ == "__main__":
main()

952
evaluate_extra_info_attribution.py
View File

@@ -1,952 +0,0 @@
"""Evaluate extra-info attribution to death and disease distribution parameters.
For each landmark query, this script scans selected extra-info types that are
available at or before the query age. For each such type it re-runs the model
with that extra-info type removed and summarizes:
* death distribution parameters before and after ablation;
* disease distribution parameters before and after ablation, by ICD-10
chapter-derived organ/system groups.
Death is always token vocab_size - 1.
"""
from __future__ import annotations
import argparse
import json
import re
from concurrent.futures import ProcessPoolExecutor, as_completed
from pathlib import Path
from typing import Any, Sequence
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 evaluate_auc_v2 import (
build_model_from_dataset,
cfg_get,
load_checkpoint_state_dict,
load_json_config,
load_model_state,
resolve_dist_mode_for_checkpoint,
resolve_eval_device,
validate_dataset_metadata,
)
from landmark_eval_utils import (
IndexedLandmarkDataset,
LandmarkDataset,
build_first_occurrence_maps_for_landmarks,
collate_indexed_landmark_fn,
infer_landmark_hidden,
load_eval_sequence_dataset,
load_organ_groups,
make_landmark_ages,
)
EXTRA_KEY_COLUMNS = [
"selected_extra_info_type_id",
"selected_extra_info_var_name",
"selected_extra_info_full_name",
"landmark_age",
"sex",
]
DEATH_PARAMETER_COLUMNS = [
"original_death_lambda",
"ablated_death_lambda",
"original_death_scale",
"ablated_death_scale",
"original_death_shape",
"ablated_death_shape",
]
DISEASE_PARAMETER_KEY_COLUMNS = [
*EXTRA_KEY_COLUMNS,
"target_group",
"target_group_label",
]
DISEASE_PARAMETER_COLUMNS = [
"original_disease_lambda",
"ablated_disease_lambda",
"original_disease_scale",
"ablated_disease_scale",
"original_disease_shape",
"ablated_disease_shape",
]
def parse_int_list(value: Any) -> list[int] | None:
if value is None:
return None
if isinstance(value, (list, tuple, np.ndarray)):
return [int(x) for x in value]
text = str(value).strip()
if text == "":
return None
if text.startswith("["):
raw = json.loads(text)
if not isinstance(raw, list):
raise ValueError("Expected JSON list for integer list")
return [int(x) for x in raw]
return [int(x.strip()) for x in re.split(r"[,;\s]+", text) if x.strip()]
def load_extra_info_metadata(
*,
dataset_extra_info_types: Sequence[int],
search_root: Path = Path("."),
) -> dict[int, dict[str, Any]]:
metadata: dict[int, dict[str, Any]] = {
int(type_id): {
"type_id": int(type_id),
"var_name": f"extra_info_{int(type_id)}",
"full_name": f"extra-info type {int(type_id)}",
}
for type_id in dataset_extra_info_types
}
line_re = re.compile(r"^\s*(\d+)\s*#\s*([^|#]+?)(?:\s*\|\s*(.*?))?\s*$")
for path in sorted(search_root.glob("extra_info_types*.txt")):
for line in path.read_text(encoding="utf-8").splitlines():
match = line_re.match(line)
if not match:
continue
type_id = int(match.group(1))
if type_id not in metadata:
continue
var_name = match.group(2).strip()
full_name = (match.group(3) or var_name).strip()
metadata[type_id] = {
"type_id": type_id,
"var_name": var_name,
"full_name": full_name,
}
return metadata
def resolve_extra_info_types(
value: str | None,
*,
dataset_extra_info_types: Sequence[int],
metadata: dict[int, dict[str, Any]],
) -> list[int]:
available = [int(x) for x in dataset_extra_info_types]
if value is None or str(value).strip() == "":
return available
out: list[int] = []
seen: set[int] = set()
by_var = {
str(meta.get("var_name", "")).lower(): int(type_id)
for type_id, meta in metadata.items()
}
by_full = {
str(meta.get("full_name", "")).lower(): int(type_id)
for type_id, meta in metadata.items()
}
for part in re.split(r"[,;]+", str(value)):
text = part.strip()
if not text:
continue
if text.isdigit() or (text.startswith("-") and text[1:].isdigit()):
type_id = int(text)
else:
lower = text.lower()
if lower in by_var:
type_id = by_var[lower]
elif lower in by_full:
type_id = by_full[lower]
else:
matches = [
int(t)
for t, meta in metadata.items()
if lower in str(meta.get("var_name", "")).lower()
or lower in str(meta.get("full_name", "")).lower()
]
if len(matches) != 1:
raise ValueError(
f"--extra_info={text!r} matched {len(matches)} types; "
"use a type id or exact variable name."
)
type_id = matches[0]
if type_id not in available:
raise ValueError(
f"extra-info type {type_id} is not available in this dataset/run"
)
if type_id not in seen:
out.append(type_id)
seen.add(type_id)
return out
def death_distribution_parameters(
model,
hidden: torch.Tensor,
*,
dist_mode: str,
eps: float = 1e-8,
) -> tuple[str, torch.Tensor]:
logits = model.calc_risk(hidden)
death_idx = int(logits.shape[1]) - 1
death_lambda = F.softplus(logits[:, death_idx]) + float(eps)
if dist_mode == "exponential":
nan = torch.full_like(death_lambda, float("nan"))
return "exponential", torch.stack([death_lambda, nan, nan], dim=1)
if dist_mode == "weibull":
rho = model.calc_weibull_rho(hidden)[:, death_idx].to(dtype=death_lambda.dtype)
elif dist_mode == "mixed":
rho = model.calc_death_rho(hidden).to(dtype=death_lambda.dtype)
else:
raise ValueError(f"Unsupported dist_mode={dist_mode!r}")
shape = rho.clamp_min(float(eps))
scale = torch.pow(death_lambda.clamp_min(float(eps)), -1.0 / shape)
nan = torch.full_like(death_lambda, float("nan"))
return "weibull", torch.stack([nan, scale, shape], dim=1)
def parameter_pair_block(original: torch.Tensor, ablated: torch.Tensor) -> torch.Tensor:
return torch.stack(
[
original[:, 0],
ablated[:, 0],
original[:, 1],
ablated[:, 1],
original[:, 2],
ablated[:, 2],
],
dim=1,
)
def all_disease_parameter_pair_block(
*,
original_logits: torch.Tensor,
ablated_logits: torch.Tensor,
dist_mode: str,
original_rho: torch.Tensor | None = None,
ablated_rho: torch.Tensor | None = None,
eps: float = 1e-8,
) -> torch.Tensor:
original_lambda = F.softplus(original_logits) + float(eps)
ablated_lambda = F.softplus(ablated_logits) + float(eps)
if dist_mode in {"exponential", "mixed"}:
nan = torch.full_like(original_lambda, float("nan"))
return torch.stack(
[
original_lambda,
ablated_lambda,
nan,
nan,
nan,
nan,
],
dim=2,
)
if dist_mode == "weibull":
if original_rho is None or ablated_rho is None:
raise ValueError("rho tensors are required for weibull disease parameters")
original_shape = original_rho.to(dtype=original_lambda.dtype).clamp_min(float(eps))
ablated_shape = ablated_rho.to(dtype=ablated_lambda.dtype).clamp_min(float(eps))
original_scale = torch.pow(original_lambda.clamp_min(float(eps)), -1.0 / original_shape)
ablated_scale = torch.pow(ablated_lambda.clamp_min(float(eps)), -1.0 / ablated_shape)
nan = torch.full_like(original_lambda, float("nan"))
return torch.stack(
[
nan,
nan,
original_scale,
ablated_scale,
original_shape,
ablated_shape,
],
dim=2,
)
raise ValueError(f"Unsupported dist_mode={dist_mode!r}")
def grouped_parameter_stats(
values: torch.Tensor,
group_token_mask: torch.Tensor,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
finite = torch.isfinite(values)
values64 = values.to(dtype=torch.float64)
safe_values = torch.where(finite, values64, torch.zeros_like(values64))
mask = group_token_mask.to(device=values.device, dtype=torch.float64)
sums = torch.einsum("nvc,gv->ngc", safe_values, mask)
sumsq = torch.einsum("nvc,gv->ngc", safe_values * safe_values, mask)
counts = torch.einsum("nvc,gv->ngc", finite.to(dtype=torch.float64), mask)
return (
sums.detach().cpu().numpy().astype(np.float64, copy=False),
sumsq.detach().cpu().numpy().astype(np.float64, copy=False),
counts.detach().cpu().numpy().astype(np.float64, copy=False),
)
def build_extra_info_ablated_slice(
batch: dict[str, torch.Tensor],
*,
row_indices: torch.Tensor,
extra_info_type_id: int,
) -> dict[str, torch.Tensor]:
out: dict[str, torch.Tensor] = {}
repeated_keys = (
"event_seq",
"time_seq",
"padding_mask",
"readout_mask",
"sex",
"landmark_pos",
"t_query",
"patient_id",
"landmark_age",
"followup_end_time",
"death_time",
"row_idx",
)
for key in repeated_keys:
out[key] = batch[key][row_indices]
out["other_type"] = batch["other_type"][row_indices].clone()
out["other_value"] = batch["other_value"][row_indices].clone()
out["other_value_kind"] = batch["other_value_kind"][row_indices].clone()
out["other_time"] = batch["other_time"][row_indices].clone()
remove = out["other_type"] == int(extra_info_type_id)
out["other_type"][remove] = 0
out["other_value"][remove] = 0
out["other_value_kind"][remove] = 0
out["other_time"][remove] = 0
return out
def concat_tensor_batches(chunks: Sequence[dict[str, torch.Tensor]]) -> dict[str, torch.Tensor]:
return {key: torch.cat([chunk[key] for chunk in chunks], dim=0) for key in chunks[0]}
def iter_extra_info_ablated_batches(
batch: dict[str, torch.Tensor],
*,
selected_extra_info_types: Sequence[int],
max_batch_size: int,
):
pending_batches: list[dict[str, torch.Tensor]] = []
pending_types: list[int] = []
pending_rows: list[int] = []
pending_n = 0
other_type = batch["other_type"]
visible = other_type > 0
visible &= batch["other_time"] <= batch["t_query"][:, None].to(batch["other_time"].dtype)
for type_id in selected_extra_info_types:
active_rows = torch.nonzero(
((other_type == int(type_id)) & visible).any(dim=1),
as_tuple=False,
).flatten()
if active_rows.numel() == 0:
continue
row_offset = 0
while row_offset < int(active_rows.numel()):
capacity = int(max_batch_size) - pending_n
row_stop = min(int(active_rows.numel()), row_offset + capacity)
row_indices = active_rows[row_offset:row_stop].to(device=batch["event_seq"].device)
chunk = build_extra_info_ablated_slice(
batch,
row_indices=row_indices,
extra_info_type_id=int(type_id),
)
chunk_n = int(row_indices.numel())
pending_batches.append(chunk)
pending_types.extend([int(type_id)] * chunk_n)
pending_rows.extend(int(x) for x in row_indices.detach().cpu().tolist())
pending_n += chunk_n
row_offset = row_stop
if pending_n >= int(max_batch_size):
yield concat_tensor_batches(pending_batches), pending_types, pending_rows
pending_batches = []
pending_types = []
pending_rows = []
pending_n = 0
if pending_batches:
yield concat_tensor_batches(pending_batches), pending_types, pending_rows
def finite_float64(values: Any) -> np.ndarray:
arr = np.asarray(values, dtype=np.float64)
return arr[np.isfinite(arr)]
def update_death_summary(
summary: dict[tuple[Any, ...], dict[str, float]],
*,
key_rows: pd.DataFrame,
values: np.ndarray,
) -> None:
if key_rows.empty:
return
table = key_rows.copy()
for idx, column in enumerate(DEATH_PARAMETER_COLUMNS):
table[column] = values[:, idx]
for key, group in table.groupby(EXTRA_KEY_COLUMNS, dropna=False, sort=False):
if not isinstance(key, tuple):
key = (key,)
acc = summary.setdefault(
key,
{
"n": 0.0,
**{f"count__{col}": 0.0 for col in DEATH_PARAMETER_COLUMNS},
**{f"sum__{col}": 0.0 for col in DEATH_PARAMETER_COLUMNS},
**{f"sumsq__{col}": 0.0 for col in DEATH_PARAMETER_COLUMNS},
},
)
acc["n"] += float(len(group))
for column in DEATH_PARAMETER_COLUMNS:
vals = finite_float64(pd.to_numeric(group[column], errors="coerce"))
acc[f"count__{column}"] += float(vals.size)
acc[f"sum__{column}"] += float(vals.sum())
acc[f"sumsq__{column}"] += float(np.square(vals).sum())
def update_disease_parameter_summary_from_group_stats(
summary: dict[tuple[Any, ...], dict[str, float]],
*,
key_rows: pd.DataFrame,
group_names: Sequence[str],
group_labels: Sequence[str],
sums: np.ndarray,
sumsq: np.ndarray,
counts: np.ndarray,
) -> None:
if key_rows.empty or sums.size == 0:
return
rows = key_rows.reset_index(drop=True)
for row_idx, row in rows.iterrows():
base_key = tuple(row[column] for column in EXTRA_KEY_COLUMNS)
for group_idx, (group, label) in enumerate(zip(group_names, group_labels)):
count_row = counts[int(row_idx), int(group_idx)]
n_add = float(np.nanmax(count_row)) if count_row.size else 0.0
if n_add <= 0:
continue
full_key = (*base_key, str(group), str(label))
acc = summary.setdefault(
full_key,
{
"n": 0.0,
**{f"count__{col}": 0.0 for col in DISEASE_PARAMETER_COLUMNS},
**{f"sum__{col}": 0.0 for col in DISEASE_PARAMETER_COLUMNS},
**{f"sumsq__{col}": 0.0 for col in DISEASE_PARAMETER_COLUMNS},
},
)
acc["n"] += n_add
for col_idx, column in enumerate(DISEASE_PARAMETER_COLUMNS):
count = float(counts[int(row_idx), int(group_idx), int(col_idx)])
if count <= 0:
continue
acc[f"count__{column}"] += count
acc[f"sum__{column}"] += float(sums[int(row_idx), int(group_idx), int(col_idx)])
acc[f"sumsq__{column}"] += float(sumsq[int(row_idx), int(group_idx), int(col_idx)])
def merge_summary_dict(
dst: dict[tuple[Any, ...], dict[str, float]],
src: dict[tuple[Any, ...], dict[str, float]],
) -> None:
for key, src_acc in src.items():
dst_acc = dst.setdefault(key, {name: 0.0 for name in src_acc})
for name, value in src_acc.items():
dst_acc[name] = dst_acc.get(name, 0.0) + float(value)
def reduce_attribution_chunk_bundle(
payload: tuple[
list[tuple[pd.DataFrame, np.ndarray]],
list[tuple[pd.DataFrame, np.ndarray, np.ndarray, np.ndarray]],
list[str],
list[str],
],
) -> tuple[dict[tuple[Any, ...], dict[str, float]], dict[tuple[Any, ...], dict[str, float]]]:
death_items, disease_items, group_names, group_labels = payload
death_summary: dict[tuple[Any, ...], dict[str, float]] = {}
disease_summary: dict[tuple[Any, ...], dict[str, float]] = {}
for key_rows, values in death_items:
update_death_summary(
death_summary,
key_rows=key_rows,
values=values,
)
for key_rows, sums, sumsq, counts in disease_items:
update_disease_parameter_summary_from_group_stats(
disease_summary,
key_rows=key_rows,
group_names=group_names,
group_labels=group_labels,
sums=sums,
sumsq=sumsq,
counts=counts,
)
return death_summary, disease_summary
def reduce_attribution_chunks(
*,
death_key_chunks: list[pd.DataFrame],
death_value_chunks: list[np.ndarray],
disease_stat_chunks: list[tuple[pd.DataFrame, np.ndarray, np.ndarray, np.ndarray]],
group_names: list[str],
group_labels: list[str],
cpu_reduce_workers: int,
) -> tuple[dict[tuple[Any, ...], dict[str, float]], dict[tuple[Any, ...], dict[str, float]]]:
n_chunks = max(len(death_key_chunks), len(disease_stat_chunks))
if n_chunks == 0:
return {}, {}
worker_count = max(1, min(int(cpu_reduce_workers), n_chunks))
if worker_count == 1:
return reduce_attribution_chunk_bundle(
(
list(zip(death_key_chunks, death_value_chunks)),
disease_stat_chunks,
group_names,
group_labels,
)
)
bundles = []
for worker_idx in range(worker_count):
start = worker_idx * n_chunks // worker_count
stop = (worker_idx + 1) * n_chunks // worker_count
if start >= stop:
continue
death_items = [
(death_key_chunks[i], death_value_chunks[i])
for i in range(start, min(stop, len(death_key_chunks)))
]
disease_items = disease_stat_chunks[start:min(stop, len(disease_stat_chunks))]
bundles.append((death_items, disease_items, group_names, group_labels))
merged_death: dict[tuple[Any, ...], dict[str, float]] = {}
merged_disease: dict[tuple[Any, ...], dict[str, float]] = {}
with ProcessPoolExecutor(max_workers=len(bundles)) as executor:
futures = [executor.submit(reduce_attribution_chunk_bundle, bundle) for bundle in bundles]
for future in tqdm(as_completed(futures), total=len(futures), desc="CPU summary reduction", dynamic_ncols=True):
death_part, disease_part = future.result()
merge_summary_dict(merged_death, death_part)
merge_summary_dict(merged_disease, disease_part)
return merged_death, merged_disease
def write_death_summary_csv(
path: Path,
summary: dict[tuple[Any, ...], dict[str, float]],
*,
death_distribution: str,
) -> int:
rows: list[dict[str, Any]] = []
for key, acc in summary.items():
n = int(acc["n"])
row = {column: value for column, value in zip(EXTRA_KEY_COLUMNS, key)}
row["n"] = n
row["death_distribution"] = death_distribution
for column in DEATH_PARAMETER_COLUMNS:
count = int(acc[f"count__{column}"])
mean = acc[f"sum__{column}"] / count if count > 0 else np.nan
second = acc[f"sumsq__{column}"] / count if count > 0 else np.nan
row[f"mean__{column}"] = mean
row[f"var__{column}"] = second - mean * mean if count > 0 else np.nan
rows.append(row)
columns = [
*EXTRA_KEY_COLUMNS,
"n",
"death_distribution",
*[
name
for column in DEATH_PARAMETER_COLUMNS
for name in (f"mean__{column}", f"var__{column}")
],
]
pd.DataFrame(rows, columns=columns).sort_values(
["selected_extra_info_type_id", "landmark_age", "sex"],
kind="mergesort",
).to_csv(path, index=False)
return len(rows)
def write_disease_parameter_summary_csv(
path: Path,
summary: dict[tuple[Any, ...], dict[str, float]],
) -> int:
rows: list[dict[str, Any]] = []
for key, acc in summary.items():
n = int(acc["n"])
row = {column: value for column, value in zip(DISEASE_PARAMETER_KEY_COLUMNS, key)}
row["n"] = n
for column in DISEASE_PARAMETER_COLUMNS:
count = int(acc[f"count__{column}"])
mean = acc[f"sum__{column}"] / count if count > 0 else np.nan
second = acc[f"sumsq__{column}"] / count if count > 0 else np.nan
row[f"mean__{column}"] = mean
row[f"var__{column}"] = second - mean * mean if count > 0 else np.nan
rows.append(row)
columns = [
*DISEASE_PARAMETER_KEY_COLUMNS,
"n",
*[
name
for column in DISEASE_PARAMETER_COLUMNS
for name in (f"mean__{column}", f"var__{column}")
],
]
pd.DataFrame(rows, columns=columns).sort_values(
["selected_extra_info_type_id", "target_group", "landmark_age", "sex"],
kind="mergesort",
).to_csv(path, index=False)
return len(rows)
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(
description="Compute extra-info ablation attribution for death and disease distribution parameters."
)
parser.add_argument("--run_path", type=str, required=True)
parser.add_argument(
"--extra_info",
type=str,
default=None,
help=(
"Optional type id, variable name, exact full name, or comma-separated list. "
"If omitted, scan all extra-info types available in the run."
),
)
parser.add_argument("--output_dir", type=str, default=None)
parser.add_argument("--organ_mapping_path", type=str, default="icd10_chapter_organ_mapping.csv")
parser.add_argument("--eval_split", type=str, default=None)
parser.add_argument("--dataset_subset_size", type=int, default=None)
parser.add_argument("--train_eid_file", type=str, default=None)
parser.add_argument("--val_eid_file", type=str, default=None)
parser.add_argument("--test_eid_file", type=str, default=None)
parser.add_argument("--landmark_start", type=float, default=40.0)
parser.add_argument("--landmark_stop", type=float, default=80.0)
parser.add_argument("--landmark_step", type=float, default=5.0)
parser.add_argument("--min_history_events", type=int, default=None)
parser.add_argument("--batch_size", type=int, default=None)
parser.add_argument(
"--attribution_batch_size",
type=int,
default=None,
help="Forward batch size for expanded extra-info ablation queries.",
)
parser.add_argument("--num_workers", type=int, default=None)
parser.add_argument(
"--cpu_reduce_workers",
type=int,
default=None,
help="Worker processes for CPU-side summary reduction. Defaults to --num_workers.",
)
parser.add_argument("--device", type=str, default=None)
return parser.parse_args()
def main() -> None:
args = parse_args()
# Dataset extra-info types must reproduce the checkpoint training config.
# --extra_info only filters which already-trained types are ablated.
args.extra_info_types = None
run_path = Path(args.run_path)
config_path = run_path / "train_config.json"
checkpoint_path = run_path / "best_model.pt"
if not config_path.exists():
raise FileNotFoundError(f"train_config.json not found: {config_path}")
if not checkpoint_path.exists():
raise FileNotFoundError(f"best_model.pt not found: {checkpoint_path}")
cfg = load_json_config(config_path)
model_target_mode = str(cfg.get("model_target_mode", "next_token")).lower()
if model_target_mode not in {"next_token", "all_future"}:
raise ValueError(f"Unsupported model_target_mode: {model_target_mode!r}")
target_mode = str(cfg.get("target_mode", "uts"))
attn_mask_mode = str(
cfg.get("attn_mask_mode", "non_strict_time" if target_mode == "uts" else "target_aware")
)
readout_name = str(
cfg.get("readout_name", "same_time_group_end" if target_mode == "uts" else "token")
)
readout_reduce = str(cfg.get("readout_reduce", "mean"))
dataset, subset_indices, eval_split, split_source = load_eval_sequence_dataset(args, cfg)
validate_dataset_metadata(dataset, cfg)
extra_metadata = load_extra_info_metadata(
dataset_extra_info_types=dataset.extra_info_types,
search_root=Path("."),
)
selected_extra_info_types = resolve_extra_info_types(
args.extra_info,
dataset_extra_info_types=dataset.extra_info_types,
metadata=extra_metadata,
)
if not selected_extra_info_types:
raise ValueError("No extra-info types selected for attribution")
landmark_ages = make_landmark_ages(
float(args.landmark_start),
float(args.landmark_stop),
float(args.landmark_step),
)
first_occurrence_by_token = build_first_occurrence_maps_for_landmarks(
dataset,
subset_indices,
)
death_idx = int(dataset.vocab_size) - 1
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_idx],
)
organ_groups, organ_labels, _token_to_group = load_organ_groups(
Path(args.organ_mapping_path),
vocab_size=int(dataset.vocab_size),
)
all_disease_tokens = sorted(
{
int(token)
for tokens in organ_groups.values()
for token in tokens
if int(token) != death_idx
}
)
risk_groups = {
"all_modeled_diseases": all_disease_tokens,
**{group: tokens for group, tokens in sorted(organ_groups.items())},
}
risk_group_labels = {
"all_modeled_diseases": "All modeled diseases",
**organ_labels,
}
group_names = list(risk_groups.keys())
group_labels = [str(risk_group_labels[group]) for group in group_names]
state_dict = load_checkpoint_state_dict(checkpoint_path, map_location="cpu")
dist_mode = resolve_dist_mode_for_checkpoint(str(cfg.get("dist_mode", "exponential")), state_dict)
death_distribution_name = "exponential" if dist_mode == "exponential" else "weibull"
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, state_dict=state_dict
).to(device)
load_model_state(model, state_dict)
model.eval()
group_token_mask = torch.zeros(
(len(group_names), int(dataset.vocab_size)),
dtype=torch.float32,
device=device,
)
for group_idx, group in enumerate(group_names):
valid_tokens = [
int(token)
for token in risk_groups[group]
if 0 <= int(token) < int(dataset.vocab_size) and int(token) != death_idx
]
if valid_tokens:
group_token_mask[group_idx, torch.as_tensor(valid_tokens, dtype=torch.long, device=device)] = 1.0
batch_size = int(cfg_get(args, cfg, "batch_size", 128))
attribution_batch_size = int(
cfg_get(args, cfg, "attribution_batch_size", max(batch_size * 32, 4096))
)
if attribution_batch_size <= 0:
raise ValueError("attribution_batch_size must be positive")
num_workers = int(cfg_get(args, cfg, "num_workers", 4))
cpu_reduce_workers = int(
args.cpu_reduce_workers
if args.cpu_reduce_workers is not None
else max(1, num_workers)
)
if cpu_reduce_workers <= 0:
raise ValueError("--cpu_reduce_workers must be positive")
loader = DataLoader(
IndexedLandmarkDataset(landmark_dataset),
batch_size=batch_size,
shuffle=False,
collate_fn=collate_indexed_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,
)
output_dir = (
Path(args.output_dir)
if args.output_dir
else run_path / f"extra_info_attribution_{eval_split}"
)
output_dir.mkdir(parents=True, exist_ok=True)
print(f"Eval split: {eval_split}")
print(f"Split source: {split_source}")
print(f"Selected patients: {len(subset_indices)}")
print(f"Landmark ages: {landmark_ages.tolist()}")
print(f"Dist mode: {dist_mode}")
print(f"Device: {device}")
print(f"Death token: {death_idx}")
print(f"Extra-info types: {selected_extra_info_types}")
print(f"Landmark rows: {len(landmark_dataset)}")
print(f"Attribution batch size: {attribution_batch_size}")
print(f"CPU reduce workers: {cpu_reduce_workers}")
print(f"Output directory: {output_dir}")
death_key_chunks: list[pd.DataFrame] = []
death_value_chunks: list[np.ndarray] = []
disease_stat_chunks: list[tuple[pd.DataFrame, np.ndarray, np.ndarray, np.ndarray]] = []
for batch in tqdm(loader, desc="Extra-info attribution", 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()
}
with torch.no_grad():
hidden = infer_landmark_hidden(
model=model,
batch=batch_dev,
device=device,
model_target_mode=model_target_mode,
readout_name=readout_name,
readout_reduce=readout_reduce,
)
_death_distribution, original_death_params = death_distribution_parameters(
model,
hidden,
dist_mode=dist_mode,
)
original_logits = model.calc_risk(hidden)
original_rho = model.calc_weibull_rho(hidden) if dist_mode == "weibull" else None
for ablated_batch, type_ids, local_rows in iter_extra_info_ablated_batches(
batch_dev,
selected_extra_info_types=selected_extra_info_types,
max_batch_size=attribution_batch_size,
):
row_tensor = torch.as_tensor(local_rows, dtype=torch.long, device=device)
with torch.no_grad():
ablated_hidden = infer_landmark_hidden(
model=model,
batch=ablated_batch,
device=device,
model_target_mode=model_target_mode,
readout_name=readout_name,
readout_reduce=readout_reduce,
)
_ablated_distribution, ablated_death_params = death_distribution_parameters(
model,
ablated_hidden,
dist_mode=dist_mode,
)
ablated_logits = model.calc_risk(ablated_hidden)
ablated_rho = model.calc_weibull_rho(ablated_hidden) if dist_mode == "weibull" else None
key_rows = []
for type_id, local_row in zip(type_ids, local_rows):
meta = extra_metadata[int(type_id)]
key_rows.append(
{
"selected_extra_info_type_id": int(type_id),
"selected_extra_info_var_name": str(meta.get("var_name", "")),
"selected_extra_info_full_name": str(meta.get("full_name", "")),
"landmark_age": float(batch["landmark_age"][int(local_row)].item()),
"sex": int(batch["sex"][int(local_row)].item()),
}
)
key_table = pd.DataFrame(key_rows, columns=EXTRA_KEY_COLUMNS)
value_block = parameter_pair_block(
original_death_params[row_tensor],
ablated_death_params,
).detach().cpu().numpy()
death_key_chunks.append(key_table)
death_value_chunks.append(value_block)
disease_values = all_disease_parameter_pair_block(
original_logits=original_logits[row_tensor],
ablated_logits=ablated_logits,
dist_mode=dist_mode,
original_rho=None if original_rho is None else original_rho[row_tensor],
ablated_rho=ablated_rho,
)
sums, sumsq, counts = grouped_parameter_stats(
disease_values,
group_token_mask,
)
disease_stat_chunks.append((key_table, sums, sumsq, counts))
death_summary, disease_parameter_summary = reduce_attribution_chunks(
death_key_chunks=death_key_chunks,
death_value_chunks=death_value_chunks,
disease_stat_chunks=disease_stat_chunks,
group_names=group_names,
group_labels=group_labels,
cpu_reduce_workers=cpu_reduce_workers,
)
death_summary_path = output_dir / "summary_extra_info_death_parameters.csv"
disease_summary_path = output_dir / "summary_extra_info_disease_parameters.csv"
death_rows = write_death_summary_csv(
death_summary_path,
death_summary,
death_distribution=death_distribution_name,
)
disease_rows = write_disease_parameter_summary_csv(
disease_summary_path,
disease_parameter_summary,
)
manifest = {
"death_summary_file": death_summary_path.name,
"disease_parameter_summary_file": disease_summary_path.name,
"death_summary_rows": int(death_rows),
"disease_parameter_summary_rows": int(disease_rows),
"eval_split": eval_split,
"split_source": split_source,
"dist_mode": dist_mode,
"landmark_start": float(args.landmark_start),
"landmark_stop": float(args.landmark_stop),
"landmark_step": float(args.landmark_step),
"selected_extra_info_types": [
extra_metadata[int(type_id)] for type_id in selected_extra_info_types
],
}
with (output_dir / "manifest.json").open("w", encoding="utf-8") as f:
json.dump(manifest, f, ensure_ascii=False, indent=2)
print(f"Wrote {death_rows} death summary rows to {death_summary_path}")
print(f"Wrote {disease_rows} disease-parameter summary rows to {disease_summary_path}")
if __name__ == "__main__":
main()

7
evaluate_landmark_auc.py
View File

@@ -1,7 +0,0 @@
from __future__ import annotations
from evaluate_auc_v2 import main
if __name__ == "__main__":
main()

819
evaluate_single_disease_mortality_attribution.py
View File

@@ -1,819 +0,0 @@
"""Compute per-disease attribution to predicted mortality distribution parameters.
For each selected patient and landmark age, this script keeps only rows where
each scanned disease token has already occurred in the history. It then deletes
that historical disease token, re-queries the model, and reports the original
and ablated fitted death distribution parameters. If --disease is omitted, all
disease tokens in the mapping are scanned.
Death is always token vocab_size - 1.
"""
from __future__ import annotations
import argparse
import json
from pathlib import Path
from typing import Any, Dict
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 evaluate_auc_v2 import (
build_model_from_dataset,
cfg_get,
load_checkpoint_state_dict,
load_json_config,
load_model_state,
resolve_dist_mode_for_checkpoint,
resolve_eval_device,
validate_dataset_metadata,
)
from landmark_eval_utils import (
IndexedLandmarkDataset,
LandmarkDataset,
build_first_occurrence_maps_for_landmarks,
collate_indexed_landmark_fn,
historical_counts_by_group,
infer_landmark_hidden,
load_eval_sequence_dataset,
load_organ_groups,
make_landmark_ages,
)
from targets import CHECKUP_IDX, PAD_IDX
OUTPUT_COLUMNS = [
"patient_id",
"dataset_index",
"eid",
"sex",
"landmark_age",
"followup_end_time",
"history_disease_count",
"selected_disease_history_count",
"selected_disease_token_id",
"selected_disease_code",
"selected_disease_name",
"selected_disease_organ_system",
"selected_disease_organ_system_label",
"history_count__selected_organ_system",
"death_distribution",
"original_death_lambda",
"ablated_death_lambda",
"original_death_scale",
"ablated_death_scale",
"original_death_shape",
"ablated_death_shape",
]
SUMMARY_KEY_COLUMNS = [
"selected_disease_token_id",
"selected_disease_code",
"selected_disease_name",
"selected_disease_organ_system",
"selected_disease_organ_system_label",
"landmark_age",
"sex",
]
SUMMARY_MEAN_COLUMNS = [
"history_disease_count",
"selected_disease_history_count",
"history_count__selected_organ_system",
]
SUMMARY_PARAMETER_COLUMNS = [
"original_death_lambda",
"ablated_death_lambda",
"original_death_scale",
"ablated_death_scale",
"original_death_shape",
"ablated_death_shape",
]
def write_compressed_npz_table(path: Path, table: pd.DataFrame) -> int:
table = table.reindex(columns=OUTPUT_COLUMNS)
arrays: dict[str, np.ndarray] = {
"__columns__": np.asarray(OUTPUT_COLUMNS, dtype="U"),
}
for column in OUTPUT_COLUMNS:
values = table[column] if column in table else pd.Series([], dtype=object)
if values.dtype == object:
arrays[column] = values.fillna("").astype(str).to_numpy(dtype="U")
else:
arrays[column] = values.to_numpy()
np.savez_compressed(path, **arrays)
return int(len(table))
def normalize_output_dir(path: Path) -> Path:
if path.suffix:
return path.with_suffix(path.suffix + "_shards")
return path
def write_manifest(
output_dir: Path,
*,
rows: int,
shards: list[dict[str, Any]],
summary_file: str,
scanned_diseases: list[dict[str, Any]],
eval_split: str,
dist_mode: str,
landmark_start: float,
landmark_stop: float,
landmark_step: float,
) -> None:
payload = {
"format": "compressed_npz_shards",
"columns": OUTPUT_COLUMNS,
"rows": int(rows),
"shards": shards,
"summary_file": summary_file,
"scanned_diseases": scanned_diseases,
"eval_split": eval_split,
"dist_mode": str(dist_mode),
"landmark_start": float(landmark_start),
"landmark_stop": float(landmark_stop),
"landmark_step": float(landmark_step),
}
with (output_dir / "manifest.json").open("w", encoding="utf-8") as f:
json.dump(payload, f, ensure_ascii=False, indent=2)
def update_summary_accumulator(
summary: dict[tuple[Any, ...], dict[str, float]],
table: pd.DataFrame,
) -> None:
if table.empty:
return
grouped = table.groupby(SUMMARY_KEY_COLUMNS, dropna=False, sort=False)
for key, group in grouped:
if not isinstance(key, tuple):
key = (key,)
acc = summary.setdefault(
key,
{
"n": 0.0,
**{column: 0.0 for column in SUMMARY_MEAN_COLUMNS},
**{f"count__{column}": 0.0 for column in SUMMARY_PARAMETER_COLUMNS},
**{f"sum__{column}": 0.0 for column in SUMMARY_PARAMETER_COLUMNS},
**{f"sumsq__{column}": 0.0 for column in SUMMARY_PARAMETER_COLUMNS},
},
)
n = int(len(group))
acc["n"] += float(n)
for column in SUMMARY_MEAN_COLUMNS:
acc[column] += float(pd.to_numeric(group[column], errors="coerce").sum())
for column in SUMMARY_PARAMETER_COLUMNS:
values = pd.to_numeric(group[column], errors="coerce").dropna()
acc[f"count__{column}"] += float(len(values))
acc[f"sum__{column}"] += float(values.sum())
acc[f"sumsq__{column}"] += float((values * values).sum())
def write_summary_csv(
path: Path,
summary: dict[tuple[Any, ...], dict[str, float]],
) -> int:
rows: list[dict[str, Any]] = []
for key, acc in summary.items():
n = int(acc["n"])
out = {column: value for column, value in zip(SUMMARY_KEY_COLUMNS, key)}
out["n"] = n
for column in SUMMARY_MEAN_COLUMNS:
out[f"mean__{column}"] = acc[column] / n if n > 0 else np.nan
for column in SUMMARY_PARAMETER_COLUMNS:
count = int(acc[f"count__{column}"])
mean = acc[f"sum__{column}"] / count if count > 0 else np.nan
second = acc[f"sumsq__{column}"] / count if count > 0 else np.nan
out[f"mean__{column}"] = mean
out[f"var__{column}"] = second - mean * mean if count > 0 else np.nan
rows.append(out)
columns = [
*SUMMARY_KEY_COLUMNS,
"n",
*[f"mean__{column}" for column in SUMMARY_MEAN_COLUMNS],
*[
name
for column in SUMMARY_PARAMETER_COLUMNS
for name in (f"mean__{column}", f"var__{column}")
],
]
pd.DataFrame(rows, columns=columns).sort_values(
["selected_disease_token_id", "landmark_age", "sex"],
kind="mergesort",
).to_csv(path, index=False)
return len(rows)
def build_disease_ablated_slice(
batch: Dict[str, torch.Tensor],
row_indices: torch.Tensor,
token_ids: torch.Tensor,
) -> Dict[str, torch.Tensor]:
"""Build an ablated slice for aligned (row, disease_token) pairs."""
event_seq = batch["event_seq"]
row_indices = row_indices.to(device=event_seq.device, dtype=torch.long)
token_ids = token_ids.to(device=event_seq.device, dtype=event_seq.dtype)
out: Dict[str, torch.Tensor] = {}
out["event_seq"] = event_seq[row_indices].clone()
out["time_seq"] = batch["time_seq"][row_indices]
out["readout_mask"] = batch["readout_mask"][row_indices].clone()
out["padding_mask"] = batch["padding_mask"][row_indices].bool().clone()
out["landmark_pos"] = batch["landmark_pos"][row_indices].clone()
seq_len = int(event_seq.shape[1])
positions = torch.arange(seq_len, device=event_seq.device)[None, :]
remove = (out["event_seq"] == token_ids[:, None]) & out["padding_mask"]
out["event_seq"] = torch.where(
remove,
torch.full_like(out["event_seq"], PAD_IDX),
out["event_seq"],
)
out["padding_mask"] &= ~remove
out["readout_mask"] &= ~remove
has_valid = out["padding_mask"].any(dim=1)
empty_rows = ~has_valid
out["event_seq"][empty_rows, 0] = CHECKUP_IDX
out["time_seq"][empty_rows, 0] = batch["t_query"][row_indices[empty_rows]].to(
dtype=out["time_seq"].dtype
)
out["padding_mask"][empty_rows, 0] = True
out["readout_mask"][empty_rows, 0] = True
out["landmark_pos"][empty_rows] = 0
has_readout = out["readout_mask"].any(dim=1)
missing_readout = ~has_readout
local_valid = out["padding_mask"]
last_pos = torch.where(
local_valid,
positions.expand(local_valid.shape[0], -1),
torch.zeros_like(positions.expand(local_valid.shape[0], -1)),
).amax(dim=1)
out["readout_mask"][missing_readout] = False
out["readout_mask"][missing_readout, last_pos[missing_readout]] = True
out["landmark_pos"][missing_readout] = last_pos[missing_readout].to(
dtype=out["landmark_pos"].dtype
)
repeated_keys = (
"sex",
"other_type",
"other_value",
"other_value_kind",
"other_time",
"t_query",
"patient_id",
"landmark_age",
"followup_end_time",
"death_time",
"row_idx",
)
for key in repeated_keys:
out[key] = batch[key][row_indices]
return out
def load_disease_metadata(
mapping_path: Path,
*,
vocab_size: int,
) -> dict[int, dict[str, Any]]:
if not mapping_path.exists():
raise FileNotFoundError(f"Disease mapping file not found: {mapping_path}")
table = pd.read_csv(mapping_path)
required = {"token_id", "code", "name", "is_death"}
missing = required - set(table.columns)
if missing:
raise ValueError(f"{mapping_path} is missing columns: {sorted(missing)}")
death_idx = int(vocab_size) - 1
out: dict[int, dict[str, Any]] = {}
for row in table.itertuples(index=False):
token = int(getattr(row, "token_id"))
if token < 0 or token >= int(vocab_size) or token == death_idx:
continue
if int(getattr(row, "is_death")) == 1:
continue
meta = {
"token_id": token,
"code": str(getattr(row, "code")),
"name": str(getattr(row, "name")),
}
for column in (
"icd10_chapter",
"icd10_chapter_title",
"organ_system",
"organ_system_label",
):
if hasattr(row, column):
meta[column] = str(getattr(row, column))
out[token] = meta
return out
def resolve_disease_token(
value: str,
metadata: dict[int, dict[str, Any]],
) -> tuple[int, dict[str, Any]]:
text = str(value).strip()
if text == "":
raise ValueError("--disease must not be empty")
if text.isdigit() or (text.startswith("-") and text[1:].isdigit()):
token = int(text)
if token not in metadata:
raise ValueError(f"Disease token_id {token} was not found in the mapping")
return token, metadata[token]
lower = text.lower()
exact = [
(token, meta)
for token, meta in metadata.items()
if str(meta.get("code", "")).lower() == lower
or str(meta.get("name", "")).lower() == lower
]
if len(exact) == 1:
return exact[0]
if len(exact) > 1:
raise ValueError(f"--disease={value!r} matched multiple diseases exactly")
contains = [
(token, meta)
for token, meta in metadata.items()
if lower in str(meta.get("code", "")).lower()
or lower in str(meta.get("name", "")).lower()
]
if len(contains) == 1:
return contains[0]
if not contains:
raise ValueError(f"--disease={value!r} did not match any disease token")
preview = ", ".join(
f"{token}:{meta.get('code')} ({meta.get('name')})"
for token, meta in contains[:10]
)
raise ValueError(
f"--disease={value!r} matched {len(contains)} diseases; use token_id or code. "
f"First matches: {preview}"
)
def resolve_disease_tokens(
value: str | None,
metadata: dict[int, dict[str, Any]],
) -> list[tuple[int, dict[str, Any]]]:
if value is None or str(value).strip() == "":
return [(token, metadata[token]) for token in sorted(metadata)]
out: list[tuple[int, dict[str, Any]]] = []
seen: set[int] = set()
for part in str(value).split(","):
token, meta = resolve_disease_token(part, metadata)
if token not in seen:
out.append((token, meta))
seen.add(token)
return out
def death_distribution_parameters(
model,
hidden: torch.Tensor,
*,
dist_mode: str,
eps: float = 1e-8,
) -> tuple[str, torch.Tensor]:
"""Return death distribution parameters with columns matching PARAMETER_VALUE_COLUMNS."""
logits = model.calc_risk(hidden)
death_idx = int(logits.shape[1]) - 1
death_lambda = F.softplus(logits[:, death_idx]) + float(eps)
if dist_mode == "exponential":
nan = torch.full_like(death_lambda, float("nan"))
return "exponential", torch.stack([death_lambda, nan, nan], dim=1)
if dist_mode == "weibull":
rho = model.calc_weibull_rho(hidden)[:, death_idx].to(dtype=death_lambda.dtype)
elif dist_mode == "mixed":
rho = model.calc_death_rho(hidden).to(dtype=death_lambda.dtype)
else:
raise ValueError(f"Unsupported dist_mode={dist_mode!r}")
shape = rho.clamp_min(float(eps))
scale = torch.pow(death_lambda.clamp_min(float(eps)), -1.0 / shape)
nan = torch.full_like(death_lambda, float("nan"))
return "weibull", torch.stack([nan, scale, shape], dim=1)
def parameter_pair_block(original: torch.Tensor, ablated: torch.Tensor) -> torch.Tensor:
return torch.stack(
[
original[:, 0],
ablated[:, 0],
original[:, 1],
ablated[:, 1],
original[:, 2],
ablated[:, 2],
],
dim=1,
)
def output_name_for_run(run_path: Path, eval_split: str, *, all_diseases: bool) -> Path:
scope = "all_diseases" if all_diseases else "selected_diseases"
return run_path / f"single_disease_mortality_parameters_{eval_split}_{scope}"
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(
description="Compute per-disease model attribution to mortality distribution parameters."
)
parser.add_argument("--run_path", type=str, required=True)
parser.add_argument(
"--disease",
type=str,
default=None,
help=(
"Optional disease token_id, ICD-10 code, exact name, unambiguous name "
"substring, or comma-separated list. If omitted, scan all disease tokens."
),
)
parser.add_argument(
"--output_path",
type=str,
default=None,
help="Output directory for compressed .npz shards.",
)
parser.add_argument("--organ_mapping_path", type=str, default="icd10_chapter_organ_mapping.csv")
parser.add_argument("--eval_split", type=str, default=None)
parser.add_argument("--dataset_subset_size", type=int, default=None)
parser.add_argument("--train_eid_file", type=str, default=None)
parser.add_argument("--val_eid_file", type=str, default=None)
parser.add_argument("--test_eid_file", type=str, default=None)
parser.add_argument("--landmark_start", type=float, default=40.0)
parser.add_argument("--landmark_stop", type=float, default=80.0)
parser.add_argument("--landmark_step", type=float, default=5.0)
parser.add_argument("--min_history_events", type=int, default=None)
parser.add_argument("--batch_size", type=int, default=None)
parser.add_argument(
"--attribution_batch_size",
type=int,
default=None,
help="Forward batch size for disease-token ablation queries.",
)
parser.add_argument("--num_workers", type=int, default=None)
parser.add_argument("--device", type=str, default=None)
parser.add_argument("--extra_info_types", type=str, default=None)
parser.add_argument(
"--shard_rows",
type=int,
default=200_000,
help="Approximate number of detailed rows to buffer before writing one .npz shard.",
)
return parser.parse_args()
def main() -> None:
args = parse_args()
run_path = Path(args.run_path)
config_path = run_path / "train_config.json"
checkpoint_path = run_path / "best_model.pt"
if not config_path.exists():
raise FileNotFoundError(f"train_config.json not found: {config_path}")
if not checkpoint_path.exists():
raise FileNotFoundError(f"best_model.pt not found: {checkpoint_path}")
cfg = load_json_config(config_path)
model_target_mode = str(cfg.get("model_target_mode", "next_token")).lower()
if model_target_mode not in {"next_token", "all_future"}:
raise ValueError(f"Unsupported model_target_mode: {model_target_mode!r}")
target_mode = str(cfg.get("target_mode", "uts"))
attn_mask_mode = str(
cfg.get("attn_mask_mode", "non_strict_time" if target_mode == "uts" else "target_aware")
)
readout_name = str(
cfg.get("readout_name", "same_time_group_end" if target_mode == "uts" else "token")
)
readout_reduce = str(cfg.get("readout_reduce", "mean"))
dataset, subset_indices, eval_split, split_source = load_eval_sequence_dataset(args, cfg)
validate_dataset_metadata(dataset, cfg)
metadata = load_disease_metadata(
Path(args.organ_mapping_path),
vocab_size=int(dataset.vocab_size),
)
scanned_disease_items = resolve_disease_tokens(args.disease, metadata)
if not scanned_disease_items:
raise ValueError("No diseases selected for attribution")
scanned_disease_tokens = [token for token, _meta in scanned_disease_items]
landmark_ages = make_landmark_ages(
float(args.landmark_start),
float(args.landmark_stop),
float(args.landmark_step),
)
first_occurrence_by_token = build_first_occurrence_maps_for_landmarks(
dataset,
subset_indices,
)
death_idx = int(dataset.vocab_size) - 1
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_idx],
)
organ_groups, _organ_labels, token_to_group = load_organ_groups(
Path(args.organ_mapping_path),
vocab_size=int(dataset.vocab_size),
)
group_names = sorted(organ_groups)
state_dict = load_checkpoint_state_dict(checkpoint_path, map_location="cpu")
dist_mode = resolve_dist_mode_for_checkpoint(str(cfg.get("dist_mode", "exponential")), state_dict)
death_distribution_name = "exponential" if dist_mode == "exponential" else "weibull"
cfg_model = dict(cfg)
cfg_model["dist_mode"] = dist_mode
device = resolve_eval_device(args.device)
selected_token_mask = np.zeros(int(dataset.vocab_size), dtype=bool)
selected_token_mask[np.asarray(scanned_disease_tokens, dtype=np.int64)] = True
model = build_model_from_dataset(
args, cfg_model, dataset, state_dict=state_dict
).to(device)
load_model_state(model, state_dict)
model.eval()
batch_size = int(cfg_get(args, cfg, "batch_size", 128))
attribution_batch_size = int(
cfg_get(args, cfg, "attribution_batch_size", max(batch_size * 32, 4096))
)
if attribution_batch_size <= 0:
raise ValueError("attribution_batch_size must be positive")
if int(args.shard_rows) <= 0:
raise ValueError("--shard_rows must be positive")
num_workers = int(cfg_get(args, cfg, "num_workers", 4))
loader = DataLoader(
IndexedLandmarkDataset(landmark_dataset),
batch_size=batch_size,
shuffle=False,
collate_fn=collate_indexed_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,
)
output_path = (
Path(args.output_path)
if args.output_path
else output_name_for_run(
run_path,
eval_split,
all_diseases=args.disease is None or str(args.disease).strip() == "",
)
)
output_dir = normalize_output_dir(output_path)
output_dir.mkdir(parents=True, exist_ok=True)
print(f"Eval split: {eval_split}")
print(f"Split source: {split_source}")
print(f"Selected patients: {len(subset_indices)}")
print(f"Landmark ages: {landmark_ages.tolist()}")
print(f"Dist mode: {dist_mode}")
print(f"Device: {device}")
print(f"Death token: {death_idx}")
if len(scanned_disease_items) == len(metadata):
print(f"Diseases: all mapped diseases ({len(scanned_disease_items)})")
else:
preview = ", ".join(
f"{token}:{meta.get('code')}" for token, meta in scanned_disease_items[:10]
)
print(f"Diseases: {len(scanned_disease_items)} selected ({preview})")
print(f"Landmark rows: {len(landmark_dataset)}")
print(f"Attribution batch size: {attribution_batch_size}")
print(f"Output directory: {output_dir}")
written_rows = 0
shard_index = 0
shards: list[dict[str, Any]] = []
row_base_cache: dict[int, dict[str, Any]] = {}
result_row_idx_chunks: list[np.ndarray] = []
result_disease_token_chunks: list[np.ndarray] = []
result_value_chunks: list[np.ndarray] = []
def get_row_base(row_idx: int) -> dict[str, Any]:
cached = row_base_cache.get(row_idx)
if cached is not None:
return cached
meta = landmark_dataset.rows[int(row_idx)]
dataset_index = int(meta["dataset_index"])
sample = dataset.samples[dataset_index]
hist_tokens = np.asarray(meta["event_seq"], dtype=np.int64)
unique_tokens, token_counts = np.unique(hist_tokens, return_counts=True)
total_count, group_counts = historical_counts_by_group(
hist_tokens,
death_idx=death_idx,
token_to_group=token_to_group,
group_names=group_names,
)
cached = {
"patient_id": int(meta["patient_id"]),
"dataset_index": dataset_index,
"eid": int(sample.get("eid", -1)),
"sex": int(meta["sex"]),
"landmark_age": float(meta["landmark_age"]),
"followup_end_time": float(meta["followup_end_time"]),
"history_disease_count": int(total_count),
"_hist_tokens": hist_tokens,
"_token_counts": {
int(token): int(count)
for token, count in zip(unique_tokens.tolist(), token_counts.tolist())
},
"_group_counts": group_counts,
}
row_base_cache[row_idx] = cached
return cached
for batch in tqdm(loader, desc="Per-disease mortality attribution", 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()
}
hidden = infer_landmark_hidden(
model=model,
batch=batch_dev,
device=device,
model_target_mode=model_target_mode,
readout_name=readout_name,
readout_reduce=readout_reduce,
)
with torch.no_grad():
_death_distribution, original_params = death_distribution_parameters(
model,
hidden,
dist_mode=dist_mode,
)
event_np = batch["event_seq"].numpy()
valid_event = (event_np >= 0) & (event_np < int(dataset.vocab_size))
selected_event = np.zeros_like(valid_event, dtype=bool)
selected_event[valid_event] = selected_token_mask[event_np[valid_event]]
pair_row_np, pair_pos_np = np.nonzero(selected_event)
if pair_row_np.size == 0:
continue
pair_disease_np = event_np[pair_row_np, pair_pos_np].astype(np.int64, copy=False)
pair_offset = 0
while pair_offset < int(pair_row_np.shape[0]):
pair_stop = min(int(pair_row_np.shape[0]), pair_offset + int(attribution_batch_size))
local_rows_np = pair_row_np[pair_offset:pair_stop].astype(np.int64, copy=False)
disease_tokens_np = pair_disease_np[pair_offset:pair_stop]
local_rows = torch.as_tensor(local_rows_np, dtype=torch.long, device=device)
disease_token_ids = torch.as_tensor(disease_tokens_np, dtype=torch.long, device=device)
ablated_chunk = build_disease_ablated_slice(
batch=batch_dev,
row_indices=local_rows,
token_ids=disease_token_ids,
)
with torch.no_grad():
ablated_hidden = infer_landmark_hidden(
model=model,
batch=ablated_chunk,
device=device,
model_target_mode=model_target_mode,
readout_name=readout_name,
readout_reduce=readout_reduce,
)
_ablated_distribution, ablated_params = death_distribution_parameters(
model,
ablated_hidden,
dist_mode=dist_mode,
)
value_block = parameter_pair_block(
original_params[local_rows],
ablated_params,
).detach().cpu().numpy()
row_ids = batch["row_idx"][local_rows_np].numpy().astype(np.int64, copy=False)
disease_tokens_list = disease_tokens_np
result_row_idx_chunks.append(row_ids)
result_disease_token_chunks.append(disease_tokens_list)
result_value_chunks.append(value_block)
pair_offset = pair_stop
if result_value_chunks:
all_row_ids = np.concatenate(result_row_idx_chunks).astype(np.int64, copy=False)
all_disease_tokens = np.concatenate(result_disease_token_chunks).astype(
np.int64,
copy=False,
)
all_values = np.concatenate(result_value_chunks, axis=0)
rows: list[dict[str, Any]] = []
for i, (row_idx, disease_token) in enumerate(
zip(all_row_ids.tolist(), all_disease_tokens.tolist())
):
disease_token = int(disease_token)
disease_meta = metadata[disease_token]
row_base = get_row_base(int(row_idx))
group_counts = row_base["_group_counts"]
disease_history_count = int(row_base["_token_counts"].get(disease_token, 0))
if disease_history_count <= 0:
raise RuntimeError(
"Internal mismatch: occurred mask selected disease "
f"{disease_token} for row {row_idx}, but cached history has count 0"
)
rows.append(
{
"patient_id": row_base["patient_id"],
"dataset_index": row_base["dataset_index"],
"eid": row_base["eid"],
"sex": row_base["sex"],
"landmark_age": row_base["landmark_age"],
"followup_end_time": row_base["followup_end_time"],
"history_disease_count": row_base["history_disease_count"],
"selected_disease_history_count": disease_history_count,
"selected_disease_token_id": int(disease_token),
"selected_disease_code": str(disease_meta.get("code", "")),
"selected_disease_name": str(disease_meta.get("name", "")),
"selected_disease_organ_system": str(disease_meta.get("organ_system", "")),
"selected_disease_organ_system_label": str(
disease_meta.get("organ_system_label", "")
),
"history_count__selected_organ_system": int(
group_counts.get(str(disease_meta.get("organ_system", "")), 0)
),
"death_distribution": death_distribution_name,
"original_death_lambda": float(all_values[i, 0]),
"ablated_death_lambda": float(all_values[i, 1]),
"original_death_scale": float(all_values[i, 2]),
"ablated_death_scale": float(all_values[i, 3]),
"original_death_shape": float(all_values[i, 4]),
"ablated_death_shape": float(all_values[i, 5]),
}
)
result_table = pd.DataFrame(rows).reindex(columns=OUTPUT_COLUMNS)
written_rows = int(len(result_table))
summary_accumulator: dict[tuple[Any, ...], dict[str, float]] = {}
update_summary_accumulator(summary_accumulator, result_table)
for start in range(0, written_rows, int(args.shard_rows)):
stop = min(written_rows, start + int(args.shard_rows))
shard_name = f"part-{shard_index:06d}.npz"
shard_path = output_dir / shard_name
shard_rows = write_compressed_npz_table(
shard_path,
result_table.iloc[start:stop],
)
shards.append({"file": shard_name, "rows": int(shard_rows)})
shard_index += 1
else:
result_table = pd.DataFrame(columns=OUTPUT_COLUMNS)
summary_accumulator = {}
if not shards:
empty_path = output_dir / "part-000000.npz"
write_compressed_npz_table(empty_path, pd.DataFrame(columns=OUTPUT_COLUMNS))
shards.append({"file": empty_path.name, "rows": 0})
summary_path = output_dir / "summary_by_disease_age_sex.csv"
summary_rows = write_summary_csv(summary_path, summary_accumulator)
write_manifest(
output_dir,
rows=written_rows,
shards=shards,
summary_file=summary_path.name,
scanned_diseases=[
{"token_id": int(token), **{k: v for k, v in meta.items() if k != "token_id"}}
for token, meta in scanned_disease_items
],
eval_split=eval_split,
dist_mode=dist_mode,
landmark_start=float(args.landmark_start),
landmark_stop=float(args.landmark_stop),
landmark_step=float(args.landmark_step),
)
print(f"Wrote {written_rows} rows in {len(shards)} shard(s) to {output_dir}")
print(f"Wrote {summary_rows} summary rows to {summary_path}")
if __name__ == "__main__":
main()

7
evaluate_token_auc.py
View File

@@ -1,7 +0,0 @@
from __future__ import annotations
from evaluate_auc import main
if __name__ == "__main__":
main()

327
export_tquery_logits_hidden.py
View File

@@ -1,327 +0,0 @@
"""Export landmark risk logits and hidden states for t_query ages.
This script follows evaluate_event_free_survival.py's data loading,
landmark construction, checkpoint loading, and readout logic, but only exports:
* all token/disease risk logits from ``model.calc_risk(hidden)``;
* the corresponding landmark hidden state.
The two large arrays are saved separately as .npy files. Row metadata is saved
as a CSV with matching row order.
"""
from __future__ import annotations
import argparse
import json
from pathlib import Path
from typing import Any, Optional
import numpy as np
import pandas as pd
import torch
from torch.utils.data import DataLoader
from tqdm.auto import tqdm
from evaluate_auc_v2 import (
LandmarkDataset,
build_model_from_dataset,
cfg_get,
load_checkpoint_state_dict,
load_json_config,
load_model_state,
resolve_dist_mode_for_checkpoint,
resolve_eval_device,
validate_dataset_metadata,
)
from evaluate_event_free_survival import (
IndexedLandmarkDataset,
build_first_occurrence_maps_for_landmarks,
collate_indexed_landmark_fn,
infer_landmark_hidden,
load_eval_sequence_dataset,
make_landmark_ages,
)
def numpy_float_dtype(name: str) -> np.dtype:
key = str(name).lower()
if key in {"float16", "fp16", "half"}:
return np.dtype(np.float16)
if key in {"float32", "fp32", "single"}:
return np.dtype(np.float32)
raise ValueError(f"dtype must be float16 or float32, got {name!r}")
def output_paths_for_run(
run_path: Path,
eval_split: str,
landmark_start: float,
landmark_stop: float,
landmark_step: float,
output_dir: Optional[str],
) -> tuple[Path, Path, Path, Path]:
suffix = f"{eval_split}_t{landmark_start:g}-{landmark_stop:g}_step{landmark_step:g}"
base_dir = Path(output_dir) if output_dir else run_path
return (
base_dir / f"tquery_logits_{suffix}.npy",
base_dir / f"tquery_hidden_{suffix}.npy",
base_dir / f"tquery_metadata_{suffix}.csv",
base_dir / f"tquery_export_config_{suffix}.json",
)
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(
description="Export landmark risk logits and hidden states for t_query ages."
)
parser.add_argument("--run_path", type=str, required=True)
parser.add_argument(
"--output_dir",
type=str,
default=None,
help="Directory for output files. Defaults to run_path.",
)
parser.add_argument("--logits_path", type=str, default=None)
parser.add_argument("--hidden_path", type=str, default=None)
parser.add_argument("--metadata_path", type=str, default=None)
parser.add_argument("--export_config_path", type=str, default=None)
parser.add_argument("--eval_split", type=str, default=None)
parser.add_argument("--dataset_subset_size", type=int, default=None)
parser.add_argument("--train_eid_file", type=str, default=None)
parser.add_argument("--val_eid_file", type=str, default=None)
parser.add_argument("--test_eid_file", type=str, default=None)
parser.add_argument("--landmark_start", type=float, default=40.0)
parser.add_argument("--landmark_stop", type=float, default=80.0)
parser.add_argument(
"--landmark_step",
type=float,
default=1.0,
help="t_query grid step in years. Default exports every integer age 40..80.",
)
parser.add_argument("--min_history_events", type=int, 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("--extra_info_types", type=str, default=None)
parser.add_argument(
"--logits_dtype",
type=str,
default="float32",
choices=["float16", "float32"],
)
parser.add_argument(
"--hidden_dtype",
type=str,
default="float32",
choices=["float16", "float32"],
)
return parser.parse_args()
def main() -> None:
args = parse_args()
run_path = Path(args.run_path)
config_path = run_path / "train_config.json"
checkpoint_path = run_path / "best_model.pt"
if not config_path.exists():
raise FileNotFoundError(f"train_config.json not found: {config_path}")
if not checkpoint_path.exists():
raise FileNotFoundError(f"best_model.pt not found: {checkpoint_path}")
cfg = load_json_config(config_path)
model_target_mode = str(cfg.get("model_target_mode", "next_token")).lower()
if model_target_mode not in {"next_token", "all_future"}:
raise ValueError(f"Unsupported model_target_mode: {model_target_mode!r}")
target_mode = str(cfg.get("target_mode", "uts"))
attn_mask_mode = str(
cfg.get(
"attn_mask_mode",
"non_strict_time" if target_mode == "uts" else "target_aware",
)
)
readout_name = str(
cfg.get("readout_name", "same_time_group_end" if target_mode == "uts" else "token")
)
readout_reduce = str(cfg.get("readout_reduce", "mean"))
dataset, subset_indices, eval_split, split_source = load_eval_sequence_dataset(
args,
cfg,
)
validate_dataset_metadata(dataset, cfg)
landmark_ages = make_landmark_ages(
float(args.landmark_start),
float(args.landmark_stop),
float(args.landmark_step),
)
first_occurrence_by_token = build_first_occurrence_maps_for_landmarks(
dataset,
subset_indices,
)
death_idx = int(dataset.vocab_size) - 1
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_idx],
)
state_dict = load_checkpoint_state_dict(checkpoint_path, map_location="cpu")
dist_mode = resolve_dist_mode_for_checkpoint(
str(cfg.get("dist_mode", "exponential")),
state_dict,
)
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, state_dict=state_dict
).to(device)
load_model_state(model, state_dict)
model.eval()
default_logits_path, default_hidden_path, default_metadata_path, default_config_path = (
output_paths_for_run(
run_path=run_path,
eval_split=eval_split,
landmark_start=float(args.landmark_start),
landmark_stop=float(args.landmark_stop),
landmark_step=float(args.landmark_step),
output_dir=args.output_dir,
)
)
logits_path = Path(args.logits_path) if args.logits_path else default_logits_path
hidden_path = Path(args.hidden_path) if args.hidden_path else default_hidden_path
metadata_path = Path(args.metadata_path) if args.metadata_path else default_metadata_path
export_config_path = (
Path(args.export_config_path) if args.export_config_path else default_config_path
)
for path in (logits_path, hidden_path, metadata_path, export_config_path):
path.parent.mkdir(parents=True, exist_ok=True)
n_rows = len(landmark_dataset)
vocab_size = int(dataset.vocab_size)
hidden_dim = int(getattr(model, "n_embd", cfg_get(args, cfg_model, "n_embd", 120)))
logits_dtype = numpy_float_dtype(args.logits_dtype)
hidden_dtype = numpy_float_dtype(args.hidden_dtype)
logits_memmap = np.lib.format.open_memmap(
logits_path,
mode="w+",
dtype=logits_dtype,
shape=(n_rows, vocab_size),
)
hidden_memmap = np.lib.format.open_memmap(
hidden_path,
mode="w+",
dtype=hidden_dtype,
shape=(n_rows, hidden_dim),
)
batch_size = int(cfg_get(args, cfg, "batch_size", 128))
num_workers = int(cfg_get(args, cfg, "num_workers", 4))
loader = DataLoader(
IndexedLandmarkDataset(landmark_dataset),
batch_size=batch_size,
shuffle=False,
collate_fn=collate_indexed_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,
)
print(f"Eval split: {eval_split}")
print(f"Split source: {split_source}")
print(f"Selected patients: {len(subset_indices)}")
print(f"t_query ages: {landmark_ages.tolist()}")
print(f"Dist mode: {dist_mode}")
print(f"Device: {device}")
print(f"Landmark rows: {n_rows}")
print(f"Logits: {logits_path} shape={(n_rows, vocab_size)} dtype={logits_dtype}")
print(f"Hidden: {hidden_path} shape={(n_rows, hidden_dim)} dtype={hidden_dtype}")
print(f"Metadata: {metadata_path}")
meta_rows: list[dict[str, Any]] = []
written = 0
with torch.no_grad():
for batch in tqdm(loader, desc="Export logits/hidden", dynamic_ncols=True):
hidden = infer_landmark_hidden(
model=model,
batch=batch,
device=device,
model_target_mode=model_target_mode,
readout_name=readout_name,
readout_reduce=readout_reduce,
)
logits = model.calc_risk(hidden)
row_indices = batch["row_idx"].detach().cpu().numpy().astype(np.int64)
if not np.array_equal(row_indices, np.arange(written, written + len(row_indices))):
raise RuntimeError("DataLoader row order changed; export requires shuffle=False.")
batch_n = int(logits.shape[0])
logits_memmap[written : written + batch_n] = (
logits.detach().cpu().numpy().astype(logits_dtype, copy=False)
)
hidden_memmap[written : written + batch_n] = (
hidden.detach().cpu().numpy().astype(hidden_dtype, copy=False)
)
for row_idx in row_indices.tolist():
meta = landmark_dataset.rows[int(row_idx)]
sample = dataset.samples[int(meta["dataset_index"])]
meta_rows.append(
{
"row_index": int(row_idx),
"patient_id": int(meta["patient_id"]),
"dataset_index": int(meta["dataset_index"]),
"eid": int(sample.get("eid", -1)),
"sex": int(meta["sex"]),
"t_query": float(meta["t_query"]),
"landmark_age": float(meta["landmark_age"]),
"followup_end_time": float(meta["followup_end_time"]),
"death_time": float(meta["death_time"]),
}
)
written += batch_n
logits_memmap.flush()
hidden_memmap.flush()
pd.DataFrame(meta_rows).to_csv(metadata_path, index=False)
export_config = {
"run_path": str(run_path),
"eval_split": eval_split,
"split_source": split_source,
"model_target_mode": model_target_mode,
"target_mode": target_mode,
"attn_mask_mode": attn_mask_mode,
"readout_name": readout_name,
"readout_reduce": readout_reduce,
"dist_mode": dist_mode,
"landmark_ages": [float(x) for x in landmark_ages.tolist()],
"n_rows": int(n_rows),
"vocab_size": int(vocab_size),
"hidden_dim": int(hidden_dim),
"death_token": int(death_idx),
"logits_path": str(logits_path),
"hidden_path": str(hidden_path),
"metadata_path": str(metadata_path),
"logits_dtype": str(logits_dtype),
"hidden_dtype": str(hidden_dtype),
}
with export_config_path.open("w", encoding="utf-8") as f:
json.dump(export_config, f, indent=2)
print(f"Wrote {written} rows.")
print(f"Wrote export config: {export_config_path}")
if __name__ == "__main__":
main()

519
export_weibull_death_parameter_stats.py
View File

@@ -1,519 +0,0 @@
"""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
import torch.multiprocessing as torch_mp
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,
)
try:
torch_mp.set_sharing_strategy("file_system")
except RuntimeError:
pass
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=0,
help=(
"DataLoader workers. Default 0 avoids Linux multiprocessing "
"'received 0 items of ancdata' failures on shared filesystems."
),
)
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, state_dict=state_dict
).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", 0))
loader_kwargs = {
"batch_size": batch_size,
"shuffle": False,
"collate_fn": collate_landmark_fn,
"num_workers": num_workers,
"pin_memory": device.type == "cuda",
}
if num_workers > 0:
loader_kwargs["persistent_workers"] = True
loader_kwargs["prefetch_factor"] = 2
loader = DataLoader(landmark_dataset, **loader_kwargs)
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()

7
export_weibull_shape_parameter_stats.py
View File

@@ -1,7 +0,0 @@
"""Compatibility entry point for Weibull shape-parameter export."""
from export_weibull_death_parameter_stats import main
if __name__ == "__main__":
main()

511
landmark_eval_utils.py
View File

@@ -1,511 +0,0 @@
"""Shared landmark evaluation helpers for attribution scripts."""
from __future__ import annotations
import argparse
import json
from pathlib import Path
from typing import Any, Dict, List, Optional, Sequence
import numpy as np
import pandas as pd
import torch
from torch.nn.utils.rnn import pad_sequence
from torch.utils.data import Dataset
from dataset import HealthDataset
from eval_data import load_sequence_eval_dataset
from evaluate_auc_v2 import (
LandmarkDataset,
build_model_from_dataset,
cfg_get,
make_eval_indices,
)
from models import DeepHealth
from readouts import build_readout
from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX
from train_util import load_eid_file, load_extra_info_types_file
SPECIAL_TOKENS = {PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX}
def parse_int_list(value: Any) -> Optional[List[int]]:
if value is None:
return None
if isinstance(value, (list, tuple, np.ndarray)):
return [int(x) for x in value]
text = str(value).strip()
if text == "":
return None
if text.startswith("["):
values = json.loads(text)
if not isinstance(values, list):
raise ValueError(f"Expected a JSON list, got {type(values).__name__}")
return [int(x) for x in values]
return [int(x.strip()) for x in text.split(",") if x.strip()]
def load_extra_info_types(value: Any) -> Optional[List[int]]:
if value is None:
return None
text = str(value)
path = Path(text)
if path.exists():
return load_extra_info_types_file(text)
return parse_int_list(value)
def make_landmark_ages(start: float, stop: float, step: float) -> np.ndarray:
if step <= 0:
raise ValueError("landmark_step must be positive")
if stop < start:
raise ValueError("landmark_stop must be >= landmark_start")
# Include stop when it lands on the grid, e.g. 40,45,...,80.
return np.arange(start, stop + step * 0.5, step, dtype=np.float32)
def build_first_occurrence_maps_for_landmarks(
dataset: HealthDataset,
subset_indices: np.ndarray,
) -> Dict[int, tuple[np.ndarray, np.ndarray]]:
first_lists: Dict[int, list[tuple[int, float]]] = {}
for patient_id, dataset_index in enumerate(np.asarray(subset_indices, dtype=np.int64).tolist()):
s = dataset.samples[int(dataset_index)]
seq_event = np.asarray(s["event_seq"], dtype=np.int64)
seq_time = np.asarray(s["time_seq"], dtype=np.float32)
tgt_event = np.asarray(s["target_event_seq"], dtype=np.int64)
tgt_time = np.asarray(s["target_time_seq"], dtype=np.float32)
if seq_event.size == 0 or tgt_event.size == 0:
continue
full_event = np.concatenate([seq_event, tgt_event[-1:]])
full_time = np.concatenate([seq_time, tgt_time[-1:]])
uniq_tokens, first_idx = np.unique(full_event, return_index=True)
for token, idx in zip(uniq_tokens.tolist(), first_idx.tolist()):
token = int(token)
if token in SPECIAL_TOKENS:
continue
first_lists.setdefault(token, []).append((patient_id, float(full_time[int(idx)])))
return {
int(token): (
np.asarray([p for p, _ in pairs], dtype=np.int32),
np.asarray([t for _, t in pairs], dtype=np.float32),
)
for token, pairs in first_lists.items()
if pairs
}
def normalize_eval_split(args: argparse.Namespace, cfg: Dict[str, Any]) -> str:
eval_split = str(cfg_get(args, cfg, "eval_split", "test")).lower()
if eval_split in {"valid", "validation"}:
return "val"
if eval_split not in {"train", "val", "test", "all"}:
raise ValueError(f"Unsupported eval_split={eval_split!r}")
return eval_split
def load_eval_sequence_dataset(
args: argparse.Namespace,
cfg: Dict[str, Any],
) -> tuple[Any, np.ndarray, str, str]:
eval_split = normalize_eval_split(args, cfg)
model_target_mode = str(cfg.get("model_target_mode", "next_token")).lower()
data_prefix = str(cfg.get("data_prefix", "ukb"))
labels_file = str(cfg.get("labels_file", "labels.csv"))
no_event_interval_years = float(cfg.get("no_event_interval_years", 5.0))
include_no_event_in_uts_target = bool(cfg.get("include_no_event_in_uts_target", False))
extra_info_types = load_extra_info_types(args.extra_info_types)
if extra_info_types is None:
extra_info_types = parse_int_list(cfg.get("extra_info_types", None))
print("Loading one sequence eval dataset...")
dataset = load_sequence_eval_dataset(
model_target_mode=model_target_mode,
data_prefix=data_prefix,
labels_file=labels_file,
no_event_interval_years=no_event_interval_years,
include_no_event_in_uts_target=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=extra_info_types,
)
train_eid_file = cfg_get(args, cfg, "train_eid_file", "ukb_train_eid.csv")
val_eid_file = cfg_get(args, cfg, "val_eid_file", "ukb_val_eid.csv")
test_eid_file = cfg_get(args, cfg, "test_eid_file", "ukb_test_eid.csv")
split_files_exist = all(
Path(str(path)).exists()
for path in (train_eid_file, val_eid_file, test_eid_file)
)
if eval_split != "all" and split_files_exist:
split_files = {
"train": train_eid_file,
"val": val_eid_file,
"test": test_eid_file,
}
selected_eids = load_eid_file(split_files[eval_split])
out = np.asarray(
[
idx
for idx, sample in enumerate(dataset.samples)
if int(sample["eid"]) in selected_eids
],
dtype=np.int64,
)
if out.size == 0:
raise ValueError(
f"No samples found for eval_split={eval_split!r} using {split_files[eval_split]}"
)
split_source = "eid_files"
else:
if eval_split == "all":
out = np.arange(len(dataset.samples), dtype=np.int64)
split_source = "all"
else:
out = make_eval_indices(dataset, args, cfg)
split_source = "ratio_split"
subset_size = cfg_get(args, cfg, "dataset_subset_size", None)
if subset_size is not None and int(subset_size) > 0:
out = out[: int(subset_size)]
return dataset, np.asarray(out, dtype=np.int64), eval_split, split_source
def load_organ_groups(
path: Path,
*,
vocab_size: int,
) -> tuple[dict[str, list[int]], dict[str, str], dict[int, str]]:
table = pd.read_csv(path)
required = {"token_id", "organ_system", "organ_system_label", "is_death"}
missing = required - set(table.columns)
if missing:
raise ValueError(f"{path} is missing columns: {sorted(missing)}")
death_idx = int(vocab_size) - 1
groups: dict[str, list[int]] = {}
labels: dict[str, str] = {}
token_to_group: dict[int, str] = {}
for row in table.itertuples(index=False):
token = int(getattr(row, "token_id"))
if token in SPECIAL_TOKENS or token == death_idx:
continue
if token < 0 or token >= int(vocab_size):
continue
if int(getattr(row, "is_death")) == 1:
continue
group = str(getattr(row, "organ_system"))
label = str(getattr(row, "organ_system_label"))
groups.setdefault(group, []).append(token)
labels[group] = label
token_to_group[token] = group
groups = {k: sorted(set(v)) for k, v in groups.items() if v}
return groups, labels, token_to_group
class IndexedLandmarkDataset(Dataset):
def __init__(self, base: LandmarkDataset) -> None:
self.base = base
def __len__(self) -> int:
return len(self.base)
def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
item = dict(self.base[idx])
item["row_idx"] = torch.tensor(int(idx), dtype=torch.long)
return item
def collate_indexed_landmark_fn(batch: List[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]:
event_seq = pad_sequence(
[x["event_seq"] for x in batch], batch_first=True, padding_value=PAD_IDX
)
time_seq = pad_sequence(
[x["time_seq"] for x in batch], batch_first=True, padding_value=0.0
)
readout_mask = pad_sequence(
[x["readout_mask"] for x in batch], batch_first=True, padding_value=False
)
other_type = pad_sequence(
[x["other_type"] for x in batch], batch_first=True, padding_value=0
)
other_value = pad_sequence(
[x["other_value"] for x in batch], batch_first=True, padding_value=0.0
)
other_value_kind = pad_sequence(
[x["other_value_kind"] for x in batch], batch_first=True, padding_value=0
)
other_time = pad_sequence(
[x["other_time"] for x in batch], batch_first=True, padding_value=0.0
)
return {
"event_seq": event_seq,
"time_seq": time_seq,
"padding_mask": event_seq > PAD_IDX,
"readout_mask": readout_mask,
"sex": torch.stack([x["sex"] for x in batch]),
"other_type": other_type,
"other_value": other_value,
"other_value_kind": other_value_kind,
"other_time": other_time,
"landmark_pos": torch.stack([x["landmark_pos"] for x in batch]),
"t_query": torch.stack([x["t_query"] for x in batch]),
"patient_id": torch.stack([x["patient_id"] for x in batch]),
"landmark_age": torch.stack([x["landmark_age"] for x in batch]),
"followup_end_time": torch.stack([x["followup_end_time"] for x in batch]),
"death_time": torch.stack([x["death_time"] for x in batch]),
"row_idx": torch.stack([x["row_idx"] for x in batch]),
}
def build_group_ablated_slice(
batch: Dict[str, torch.Tensor],
token_ids: Sequence[int],
row_indices: torch.Tensor,
) -> Dict[str, torch.Tensor]:
"""Build one fixed-width ablated slice without rebuilding variable-length rows."""
event_seq = batch["event_seq"]
out: Dict[str, torch.Tensor] = {}
out["event_seq"] = event_seq[row_indices].clone()
out["time_seq"] = batch["time_seq"][row_indices]
out["readout_mask"] = batch["readout_mask"][row_indices].clone()
out["padding_mask"] = batch["padding_mask"][row_indices].bool().clone()
out["landmark_pos"] = batch["landmark_pos"][row_indices].clone()
seq_len = int(event_seq.shape[1])
positions = torch.arange(seq_len, device=event_seq.device)[None, :]
ids = torch.as_tensor(token_ids, dtype=event_seq.dtype, device=event_seq.device)
remove = torch.isin(out["event_seq"], ids) & out["padding_mask"]
out["event_seq"] = torch.where(
remove,
torch.full_like(out["event_seq"], PAD_IDX),
out["event_seq"],
)
out["padding_mask"] &= ~remove
out["readout_mask"] &= ~remove
has_valid = out["padding_mask"].any(dim=1)
if not bool(has_valid.all().item()):
empty_rows = torch.nonzero(~has_valid, as_tuple=False).flatten()
out["event_seq"][empty_rows, 0] = CHECKUP_IDX
out["time_seq"][empty_rows, 0] = batch["t_query"][row_indices[empty_rows]].to(
dtype=out["time_seq"].dtype
)
out["padding_mask"][empty_rows, 0] = True
out["readout_mask"][empty_rows, 0] = True
out["landmark_pos"][empty_rows] = 0
has_readout = out["readout_mask"].any(dim=1)
if not bool(has_readout.all().item()):
rows = torch.nonzero(~has_readout, as_tuple=False).flatten()
local_valid = out["padding_mask"][rows]
last_pos = torch.where(
local_valid,
positions.expand(local_valid.shape[0], -1),
torch.zeros_like(positions.expand(local_valid.shape[0], -1)),
).amax(dim=1)
out["readout_mask"][rows] = False
out["readout_mask"][rows, last_pos] = True
out["landmark_pos"][rows] = last_pos.to(dtype=out["landmark_pos"].dtype)
repeated_keys = (
"sex",
"other_type",
"other_value",
"other_value_kind",
"other_time",
"t_query",
"patient_id",
"landmark_age",
"followup_end_time",
"death_time",
"row_idx",
)
for key in repeated_keys:
out[key] = batch[key][row_indices]
return out
def concat_tensor_batches(chunks: Sequence[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]:
return {
key: torch.cat([chunk[key] for chunk in chunks], dim=0)
for key in chunks[0]
}
def iter_group_ablated_batches(
batch: Dict[str, torch.Tensor],
group_names: Sequence[str],
organ_groups: dict[str, list[int]],
occurred: torch.Tensor,
max_batch_size: int,
):
"""Yield ablated chunks as soon as enough rows are available for a forward pass."""
pending_batches: list[Dict[str, torch.Tensor]] = []
pending_groups: list[str] = []
pending_rows: list[int] = []
pending_n = 0
for group in group_names:
ids = torch.as_tensor(organ_groups[group], dtype=torch.long, device=occurred.device)
if ids.numel() == 0:
continue
active_rows = torch.nonzero(occurred[:, ids].any(dim=1), as_tuple=False).flatten()
if active_rows.numel() == 0:
continue
row_offset = 0
while row_offset < int(active_rows.numel()):
capacity = int(max_batch_size) - pending_n
row_stop = min(int(active_rows.numel()), row_offset + capacity)
row_indices = active_rows[row_offset:row_stop].to(device=batch["event_seq"].device)
chunk = build_group_ablated_slice(
batch=batch,
token_ids=organ_groups[group],
row_indices=row_indices,
)
chunk_n = int(row_indices.numel())
pending_batches.append(chunk)
pending_groups.extend([group] * chunk_n)
pending_rows.extend(int(x) for x in row_indices.detach().cpu().tolist())
pending_n += chunk_n
row_offset = row_stop
if pending_n >= int(max_batch_size):
yield concat_tensor_batches(pending_batches), pending_groups, pending_rows
pending_batches = []
pending_groups = []
pending_rows = []
pending_n = 0
if pending_batches:
yield concat_tensor_batches(pending_batches), pending_groups, pending_rows
@torch.no_grad()
def infer_landmark_hidden(
*,
model: DeepHealth,
batch: Dict[str, torch.Tensor],
device: torch.device,
model_target_mode: str,
readout_name: str,
readout_reduce: str,
) -> torch.Tensor:
batch_dev = {
k: (v.to(device, non_blocking=True) if isinstance(v, torch.Tensor) else v)
for k, v in batch.items()
}
if model_target_mode == "all_future":
return model(
event_seq=batch_dev["event_seq"].long(),
time_seq=batch_dev["time_seq"].float(),
sex=batch_dev["sex"].long(),
padding_mask=batch_dev["padding_mask"].bool(),
t_query=batch_dev["t_query"].float(),
other_type=batch_dev["other_type"].long(),
other_value=batch_dev["other_value"].float(),
other_value_kind=batch_dev["other_value_kind"].long(),
other_time=batch_dev["other_time"].float(),
target_mode="all_future",
)
hidden = model(
event_seq=batch_dev["event_seq"].long(),
time_seq=batch_dev["time_seq"].float(),
sex=batch_dev["sex"].long(),
padding_mask=batch_dev["padding_mask"].bool(),
other_type=batch_dev["other_type"].long(),
other_value=batch_dev["other_value"].float(),
other_value_kind=batch_dev["other_value_kind"].long(),
other_time=batch_dev["other_time"].float(),
target_mode="next_token",
)
readout = build_readout(readout_name, reduce=readout_reduce)
readout_out = readout(
hidden=hidden,
time_seq=batch_dev["time_seq"].float(),
padding_mask=batch_dev["padding_mask"].bool(),
readout_mask=batch_dev["readout_mask"].bool(),
)
return readout_out.hidden.gather(
1,
batch_dev["landmark_pos"].long()[:, None, None].expand(
-1, 1, readout_out.hidden.shape[-1]
),
).squeeze(1)
def make_occurred_mask(
event_seq: torch.Tensor,
*,
vocab_size: int,
device: torch.device,
) -> torch.Tensor:
occurred = torch.zeros(event_seq.shape[0], int(vocab_size), dtype=torch.bool, device=device)
valid = (event_seq >= 0) & (event_seq < int(vocab_size))
safe = event_seq.clamp(min=0, max=int(vocab_size) - 1).to(device)
occurred.scatter_(1, safe, valid.to(device))
return occurred
def mortality_hazard_from_risk(risk: torch.Tensor, eps: float = 1e-7) -> torch.Tensor:
return -torch.log1p(-risk.clamp(0.0, 1.0 - float(eps)))
def death_risk_for_batch(
*,
model: DeepHealth,
batch: Dict[str, torch.Tensor],
device: torch.device,
model_target_mode: str,
readout_name: str,
readout_reduce: str,
dist_mode: str,
tau: float,
) -> torch.Tensor:
hidden = infer_landmark_hidden(
model=model,
batch=batch,
device=device,
model_target_mode=model_target_mode,
readout_name=readout_name,
readout_reduce=readout_reduce,
)
logits = model.calc_risk(hidden)
rho = model.calc_weibull_rho(hidden) if dist_mode == "weibull" else None
death_rho = model.calc_death_rho(hidden) if dist_mode == "mixed" else None
probabilities = probabilities_from_logits(
logits,
tau,
dist_mode=dist_mode,
rho=rho,
death_rho=death_rho,
)
return death_risk_from_probabilities(probabilities)
def historical_counts_by_group(
tokens: np.ndarray,
*,
death_idx: int,
token_to_group: dict[int, str],
group_names: Sequence[str],
) -> tuple[int, dict[str, int]]:
unique_tokens = {
int(token)
for token in np.asarray(tokens, dtype=np.int64).tolist()
if int(token) not in SPECIAL_TOKENS and int(token) != int(death_idx)
}
total = len(unique_tokens)
out = {group: 0 for group in group_names}
for token in unique_tokens:
group = token_to_group.get(token)
if group in out:
out[group] += 1
return total, out

553
plot_next_token_to_all_future_auc.R
View File

@@ -1,553 +0,0 @@
#!/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 = "/"))

221
run_missing_evaluations.sh
View File

@@ -1,221 +0,0 @@
#!/usr/bin/env bash
set -euo pipefail
# Run all non-wrapper evaluation scripts for every completed current-format
# experiment under runs/. The script is written for Linux servers with bash 4.2.
cd "$(dirname "${BASH_SOURCE[0]}")"
shopt -s globstar nullglob
PYTHON_BIN="${PYTHON_BIN:-python}"
DEVICE="${DEVICE:-cuda}"
EVAL_SPLIT="${EVAL_SPLIT:-test}"
NUM_WORKERS="${NUM_WORKERS:-4}"
CPU_REDUCE_WORKERS="${CPU_REDUCE_WORKERS:-}"
NUM_WORKERS_AUC="${NUM_WORKERS_AUC:-}"
BATCH_SIZE="${BATCH_SIZE:-}"
DATASET_SUBSET_SIZE="${DATASET_SUBSET_SIZE:-}"
DRY_RUN="${DRY_RUN:-0}"
# These attribution jobs can be expensive, but they are part of the evaluation
# surface in this repository. Set either variable to 0 to leave that family out.
RUN_EXTRA_INFO_ATTRIBUTION="${RUN_EXTRA_INFO_ATTRIBUTION:-1}"
RUN_SINGLE_DISEASE_MORTALITY_ATTRIBUTION="${RUN_SINGLE_DISEASE_MORTALITY_ATTRIBUTION:-1}"
common_args_base() {
printf '%s\n' --run_path "$1" --eval_split "${EVAL_SPLIT}" --num_workers "${NUM_WORKERS}"
if [[ -n "${BATCH_SIZE}" ]]; then
printf '%s\n' --batch_size "${BATCH_SIZE}"
fi
if [[ -n "${DATASET_SUBSET_SIZE}" ]]; then
printf '%s\n' --dataset_subset_size "${DATASET_SUBSET_SIZE}"
fi
}
common_args_with_device() {
common_args_base "$1"
printf '%s\n' --device "${DEVICE}"
}
auc_args() {
if [[ -n "${NUM_WORKERS_AUC}" ]]; then
printf '%s\n' --num_workers_auc "${NUM_WORKERS_AUC}"
fi
}
cpu_reduce_args() {
if [[ -n "${CPU_REDUCE_WORKERS}" ]]; then
printf '%s\n' --cpu_reduce_workers "${CPU_REDUCE_WORKERS}"
fi
}
has_completed_dir() {
local dir="$1"
shift
[[ -d "${dir}" ]] || return 1
local required
for required in "$@"; do
[[ -s "${dir}/${required}" ]] || return 1
done
}
run_command() {
echo " run: $*"
if [[ "${DRY_RUN}" == "1" ]]; then
return 0
fi
"$@"
}
run_dir_result_if_missing() {
local label="$1"
local result_dir="$2"
local required_1="$3"
local required_2="$4"
shift 4
if has_completed_dir "${result_dir}" "${required_1}" "${required_2}"; then
echo " skip ${label}: found ${result_dir}"
return 0
fi
run_command "$@"
}
run_file_result_if_missing() {
local label="$1"
local result_dir="$2"
local required="$3"
shift 3
if [[ -s "${result_dir}/${required}" ]]; then
echo " skip ${label}: found ${result_dir}/${required}"
return 0
fi
run_command "$@"
}
run_has_extra_info() {
"${PYTHON_BIN}" - "$1" <<'PY'
import json
import sys
from pathlib import Path
cfg_path = Path(sys.argv[1]) / "train_config.json"
try:
cfg = json.loads(cfg_path.read_text(encoding="utf-8"))
except Exception:
raise SystemExit(1)
extra = cfg.get("extra_info_types", [])
raise SystemExit(0 if isinstance(extra, list) and len(extra) > 0 else 1)
PY
}
run_is_all_future() {
"${PYTHON_BIN}" - "$1" <<'PY'
import json
import sys
from pathlib import Path
cfg_path = Path(sys.argv[1]) / "train_config.json"
try:
cfg = json.loads(cfg_path.read_text(encoding="utf-8"))
except Exception:
raise SystemExit(1)
mode = str(cfg.get("model_target_mode", "next_token")).lower()
raise SystemExit(0 if mode == "all_future" else 1)
PY
}
run_has_current_model_config() {
"${PYTHON_BIN}" - "$1" <<'PY'
import json
import sys
from pathlib import Path
cfg_path = Path(sys.argv[1]) / "train_config.json"
try:
cfg = json.loads(cfg_path.read_text(encoding="utf-8"))
n_layer = int(cfg.get("n_layer", 0))
except Exception:
raise SystemExit(1)
supported = {"transformer_ffn_v1", "traj_mixer_v5"}
raise SystemExit(
0
if cfg.get("model_architecture") in supported and n_layer >= 1
else 1
)
PY
}
for config_path in runs/**/train_config.json; do
run_path="${config_path%/train_config.json}"
echo "==> ${run_path}"
if [[ ! -f "${run_path}/train_config.json" ]]; then
echo " skip run: missing train_config.json"
continue
fi
if [[ ! -s "${run_path}/best_model.pt" ]]; then
echo " skip run: missing best_model.pt"
continue
fi
if ! run_has_current_model_config "${run_path}"; then
echo " skip run: config lacks current model_architecture/n_layer fields"
continue
fi
common=()
while IFS= read -r arg; do common+=("${arg}"); done < <(common_args_with_device "${run_path}")
auc_extra=()
while IFS= read -r arg; do auc_extra+=("${arg}"); done < <(auc_args)
cpu_reduce_extra=()
while IFS= read -r arg; do cpu_reduce_extra+=("${arg}"); done < <(cpu_reduce_args)
run_file_result_if_missing \
"evaluate_auc.py" \
"${run_path}" \
"df_auc_delphi2m_report.csv" \
"${PYTHON_BIN}" evaluate_auc.py "${common[@]}" "${auc_extra[@]}"
run_file_result_if_missing \
"evaluate_auc_v2.py" \
"${run_path}" \
"df_auc_landmark_delphi2m_report.csv" \
"${PYTHON_BIN}" evaluate_auc_v2.py "${common[@]}" "${auc_extra[@]}"
if ! run_is_all_future "${run_path}"; then
echo " skip attribution evaluations: model_target_mode is not all_future"
continue
fi
if [[ "${RUN_EXTRA_INFO_ATTRIBUTION}" == "1" ]]; then
if run_has_extra_info "${run_path}"; then
run_dir_result_if_missing \
"evaluate_extra_info_attribution.py" \
"${run_path}/extra_info_attribution_${EVAL_SPLIT}" \
"manifest.json" \
"summary_extra_info_disease_parameters.csv" \
"${PYTHON_BIN}" evaluate_extra_info_attribution.py "${common[@]}" "${cpu_reduce_extra[@]}"
else
echo " skip evaluate_extra_info_attribution.py: run has no extra-info types"
fi
fi
if [[ "${RUN_SINGLE_DISEASE_MORTALITY_ATTRIBUTION}" == "1" ]]; then
run_dir_result_if_missing \
"evaluate_single_disease_mortality_attribution.py" \
"${run_path}/single_disease_mortality_parameters_${EVAL_SPLIT}_all_diseases" \
"manifest.json" \
"summary_by_disease_age_sex.csv" \
"${PYTHON_BIN}" evaluate_single_disease_mortality_attribution.py "${common[@]}"
fi
done
echo "All missing evaluations are complete."

143
run_missing_training_runs.sh
View File

@@ -1,143 +0,0 @@
#!/usr/bin/env bash
set -euo pipefail
# Linux bash 5.2+ training-only script.
#
# Based on the existing runs, the objective/time/death-distribution checks are
# already covered. The remaining gap for the current proof chain is the
# extra-info ablation under the final candidate model:
#
# all_future + relative time + mixed death/risk head
#
# This script only launches those missing training jobs. It intentionally does
# not call evaluate_*.py and does not add extra random seeds. Set
# MODEL_ARCHITECTURE=traj_mixer_v5 to run the TrajMixer variant.
cd "$(dirname "${BASH_SOURCE[0]}")"
PYTHON_BIN="${PYTHON_BIN:-python}"
DEVICE="${DEVICE:-cuda}"
NUM_WORKERS="${NUM_WORKERS:-4}"
PROGRESS_INTERVAL="${PROGRESS_INTERVAL:-20}"
MODEL_ARCHITECTURE="${MODEL_ARCHITECTURE:-transformer_ffn_v1}"
N_LAYER="${N_LAYER:-12}"
TIME_MODE="relative"
DIST_MODE="mixed"
SEED="42"
VALIDATION_QUERY_SEED="42"
COMMON_ARGS=(
--data_prefix ukb
--labels_file labels.csv
--seed "${SEED}"
--validation_query_seed "${VALIDATION_QUERY_SEED}"
--train_eid_file ukb_train_eid.csv
--val_eid_file ukb_val_eid.csv
--test_eid_file ukb_test_eid.csv
--min_history_events 1
--min_future_events 1
--n_embd 120
--n_head 10
--n_layer "${N_LAYER}"
--model_architecture "${MODEL_ARCHITECTURE}"
--n_bins 16
--extra_pool_reduce mean
--dropout 0.0
--batch_size 256
--base_lr 0.0003
--weight_decay 0.1
--betas 0.9 0.99
--grad_clip 1.0
--max_epochs 200
--warmup_epochs 10
--patience 15
--min_lr_ratio 0.1
--num_workers "${NUM_WORKERS}"
--device "${DEVICE}"
--progress_interval "${PROGRESS_INTERVAL}"
)
already_trained() {
local extra_file="$1"
"${PYTHON_BIN}" - "$TIME_MODE" "$DIST_MODE" "$extra_file" "$SEED" "$VALIDATION_QUERY_SEED" "$MODEL_ARCHITECTURE" "$N_LAYER" <<'PY'
import json
import sys
from pathlib import Path
(
time_mode,
dist_mode,
extra_file,
seed,
validation_query_seed,
model_architecture,
n_layer,
) = sys.argv[1:8]
extra_name = Path(extra_file).name
for config_path in Path("runs").rglob("train_config.json"):
try:
cfg = json.loads(config_path.read_text(encoding="utf-8"))
except Exception:
continue
observed_query_seed = cfg.get(
"all_future_validation_query_seed",
cfg.get("validation_query_seed", -1),
)
if (
cfg.get("model_target_mode") == "all_future"
and cfg.get("model_architecture") == model_architecture
and int(cfg.get("n_layer", -1)) == int(n_layer)
and cfg.get("time_mode") == time_mode
and cfg.get("dist_mode") == dist_mode
and Path(str(cfg.get("extra_info_types_file", ""))).name == extra_name
and int(cfg.get("seed", -1)) == int(seed)
and int(observed_query_seed) == int(validation_query_seed)
):
print(config_path.parent)
raise SystemExit(0)
raise SystemExit(1)
PY
}
train_if_missing() {
local label="$1"
local extra_file="$2"
if [[ ! -f "${extra_file}" ]]; then
echo "ERROR: missing extra-info type file: ${extra_file}" >&2
return 2
fi
echo "==> Checking ${label}: ${MODEL_ARCHITECTURE} n_layer=${N_LAYER} ${TIME_MODE} ${DIST_MODE} all_future with ${extra_file}"
if existing_run="$(already_trained "$extra_file")"; then
echo " skip: already trained at ${existing_run}"
return 0
fi
echo " train: ${label}"
"${PYTHON_BIN}" train_all_future.py \
"${COMMON_ARGS[@]}" \
--time_mode "${TIME_MODE}" \
--dist_mode "${DIST_MODE}" \
--extra_info_types_file "${extra_file}"
}
# Already present in runs/:
# - next-token objective checks under SAB, plus older absolute extra ablations
# - all-future absolute/relative x exponential/weibull/mixed under SAB
#
# Still needed:
# - final all-future relative+mixed extra-info ablations beyond the existing
# SAB baseline. These close the disease-only question without expanding seed
# count or running downstream evaluation.
train_if_missing "true_disease_only" "extra_info_types_none.txt"
train_if_missing "assessment_only_extra" "extra_info_types_assessment_only.txt"
train_if_missing "exposure_only_extra" "extra_info_types_exposure_only.txt"
train_if_missing "all_extra_info" "extra_info_types_all.txt"
echo "All requested training-only missing configurations are done."

76
run_weibull_shape_exports.sh
View File

@@ -1,76 +0,0 @@
#!/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=0 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:-0}"
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."