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

Add event-free survival evaluation

Browse Source
This commit is contained in:
Jiarui Li
2026-06-27 11:48:47 +08:00
parent df6b45b95a
commit 3cd1109249
6 changed files with 2568 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files

332
build_icd10_chapter_organ_mapping.py Normal file
View File

@@ -0,0 +1,332 @@
from __future__ import annotations
import csv
import re
from dataclasses import dataclass
from pathlib import Path
LABEL_OFFSET = 3
@dataclass(frozen=True)
class ChapterRule:
chapter: str
start: str
end: str
title: str
organ_system: str
organ_system_label: str
CHAPTER_RULES = [
ChapterRule(
"I",
"A00",
"B99",
"Certain infectious and parasitic diseases",
"infectious_systemic",
"Infectious / systemic",
),
ChapterRule(
"II",
"C00",
"D48",
"Neoplasms",
"neoplasm",
"Neoplasm / oncology",
),
ChapterRule(
"III",
"D50",
"D89",
"Diseases of the blood and blood-forming organs and certain immune disorders",
"hematologic_immune",
"Blood / immune",
),
ChapterRule(
"IV",
"E00",
"E90",
"Endocrine, nutritional and metabolic diseases",
"endocrine_metabolic",
"Endocrine / metabolic",
),
ChapterRule(
"V",
"F00",
"F99",
"Mental and behavioural disorders",
"mental_behavioral",
"Mental / behavioral",
),
ChapterRule(
"VI",
"G00",
"G99",
"Diseases of the nervous system",
"nervous_system",
"Nervous system",
),
ChapterRule(
"VII",
"H00",
"H59",
"Diseases of the eye and adnexa",
"eye",
"Eye",
),
ChapterRule(
"VIII",
"H60",
"H95",
"Diseases of the ear and mastoid process",
"ear",
"Ear",
),
ChapterRule(
"IX",
"I00",
"I99",
"Diseases of the circulatory system",
"circulatory",
"Circulatory system",
),
ChapterRule(
"X",
"J00",
"J99",
"Diseases of the respiratory system",
"respiratory",
"Respiratory system",
),
ChapterRule(
"XI",
"K00",
"K93",
"Diseases of the digestive system",
"digestive",
"Digestive system",
),
ChapterRule(
"XII",
"L00",
"L99",
"Diseases of the skin and subcutaneous tissue",
"skin",
"Skin / subcutaneous tissue",
),
ChapterRule(
"XIII",
"M00",
"M99",
"Diseases of the musculoskeletal system and connective tissue",
"musculoskeletal",
"Musculoskeletal / connective tissue",
),
ChapterRule(
"XIV",
"N00",
"N99",
"Diseases of the genitourinary system",
"genitourinary",
"Genitourinary system",
),
ChapterRule(
"XV",
"O00",
"O99",
"Pregnancy, childbirth and the puerperium",
"pregnancy_childbirth",
"Pregnancy / childbirth",
),
ChapterRule(
"XVI",
"P00",
"P96",
"Certain conditions originating in the perinatal period",
"perinatal",
"Perinatal conditions",
),
ChapterRule(
"XVII",
"Q00",
"Q99",
"Congenital malformations, deformations and chromosomal abnormalities",
"congenital_chromosomal",
"Congenital / chromosomal",
),
ChapterRule(
"XVIII",
"R00",
"R99",
"Symptoms, signs and abnormal clinical and laboratory findings",
"symptoms_findings",
"Symptoms / findings",
),
ChapterRule(
"XIX",
"S00",
"T98",
"Injury, poisoning and certain other consequences of external causes",
"injury_poisoning",
"Injury / poisoning",
),
ChapterRule(
"XX",
"V01",
"Y98",
"External causes of morbidity and mortality",
"external_causes",
"External causes",
),
ChapterRule(
"XXI",
"Z00",
"Z99",
"Factors influencing health status and contact with health services",
"health_services_factors",
"Health status / services factors",
),
ChapterRule(
"XXII",
"U00",
"U99",
"Codes for special purposes",
"special_purposes",
"Special purposes",
),
]
CODE_RE = re.compile(r"^([A-Z][0-9]{2})(?:\.[0-9A-Z]+)?\b")
UNKNOWN_CANCER_RE = re.compile(r"^(CXX)\b\s*(.*)$", re.IGNORECASE)
def code_key(code: str) -> tuple[str, int]:
code = code.upper().strip()
return code[0], int(code[1:3])
def in_range(code: str, start: str, end: str) -> bool:
letter, number = code_key(code)
start_letter, start_number = code_key(start)
end_letter, end_number = code_key(end)
return (start_letter, start_number) <= (letter, number) <= (end_letter, end_number)
def parse_label(line: str) -> tuple[str, str]:
text = line.strip()
unknown_cancer = UNKNOWN_CANCER_RE.match(text)
if unknown_cancer:
return unknown_cancer.group(1).upper(), unknown_cancer.group(2).strip()
match = CODE_RE.match(text)
if not match:
return text, ""
code = match.group(1).upper()
name = text[len(code):].strip()
if name.startswith("(") and name.endswith(")"):
name = name[1:-1]
return code, name
def assign_chapter(code: str) -> ChapterRule | None:
if code.upper().strip() == "CXX":
return CHAPTER_RULES[1]
if re.fullmatch(r"[A-Z][0-9]{2}", code.upper().strip()) is None:
return None
for rule in CHAPTER_RULES:
if in_range(code, rule.start, rule.end):
return rule
return None
def build_mapping(labels_path: Path, output_path: Path) -> None:
rows = []
with labels_path.open("r", encoding="utf-8") as f:
for label_index, raw in enumerate(f):
text = raw.strip()
if not text:
continue
token_id = LABEL_OFFSET + label_index
code, name = parse_label(text)
if code.lower() == "death":
rows.append(
{
"label_index": label_index,
"token_id": token_id,
"code": "Death",
"name": "Death",
"icd10_chapter": "Death",
"icd10_range": "",
"icd10_chapter_title": "Death endpoint",
"organ_system": "death",
"organ_system_label": "Death",
"is_death": 1,
}
)
continue
rule = assign_chapter(code)
if rule is None:
rows.append(
{
"label_index": label_index,
"token_id": token_id,
"code": code,
"name": name,
"icd10_chapter": "Unmapped",
"icd10_range": "",
"icd10_chapter_title": "Unmapped",
"organ_system": "unmapped",
"organ_system_label": "Unmapped",
"is_death": 0,
}
)
continue
rows.append(
{
"label_index": label_index,
"token_id": token_id,
"code": code,
"name": name,
"icd10_chapter": rule.chapter,
"icd10_range": f"{rule.start}-{rule.end}",
"icd10_chapter_title": rule.title,
"organ_system": rule.organ_system,
"organ_system_label": rule.organ_system_label,
"is_death": 0,
}
)
fieldnames = [
"label_index",
"token_id",
"code",
"name",
"icd10_chapter",
"icd10_range",
"icd10_chapter_title",
"organ_system",
"organ_system_label",
"is_death",
]
with output_path.open("w", encoding="utf-8", newline="") as f:
writer = csv.DictWriter(f, fieldnames=fieldnames)
writer.writeheader()
writer.writerows(rows)
def main() -> None:
build_mapping(
labels_path=Path("labels.csv"),
output_path=Path("icd10_chapter_organ_mapping.csv"),
)
if __name__ == "__main__":
main()

14
evaluate_auc.py
View File

@@ -200,19 +200,7 @@ def _build_first_occurrence_maps(
def _get_death_token_ids(dataset: HealthDataset) -> List[int]:
ids: List[int] = []
exact_codes = {"death", "<death>", "dth", "deceased", "mortality"}
for token, code in dataset.label_id_to_code.items():
token = int(token)
if token in SPECIAL_TOKENS:
continue
text = str(code).strip().lower()
if text in exact_codes or ("death" in text) or ("mortality" in text):
ids.append(token)
death_ids = sorted(set(int(x)
for x in ids if int(x) not in SPECIAL_TOKENS))
death_ids = [int(dataset.vocab_size) - 1]
print(f"[INFO] death token ids: {death_ids}")
return death_ids

35
evaluate_auc_v2.py
View File

@@ -395,40 +395,7 @@ def _metadata_count_map(dataset: HealthDataset, labels_meta: Optional[pd.DataFra
def _get_death_token_ids(dataset: HealthDataset, labels_meta: Optional[pd.DataFrame]) -> List[int]:
ids: List[int] = []
if labels_meta is not None and not labels_meta.empty:
meta = labels_meta.copy()
if "ICD-10 Chapter (short)" in meta.columns:
death_rows = meta[meta["ICD-10 Chapter (short)"].astype(
str) == "Death"]
code_col = _first_existing_column(
death_rows, ["Name", "code", "ICD10", "icd10", "label", "token", "disease_code"])
if code_col is not None:
for raw in death_rows[code_col].astype(str).tolist():
code = raw.split()[0].strip()
if code in dataset.label_code_to_id:
ids.append(int(dataset.label_code_to_id[code]))
elif "index" in death_rows.columns:
idx = pd.to_numeric(death_rows["index"], errors="coerce")
has_no_event = (
NO_EVENT_IDX in dataset.label_id_to_code
and dataset.label_id_to_code.get(NO_EVENT_IDX) == "<NO_EVENT>"
)
if has_no_event:
idx = idx.where(idx < NO_EVENT_IDX, idx + 1)
ids.extend(int(x) for x in idx.dropna().astype(int).tolist())
exact_codes = {"death", "<death>", "dth", "deceased", "mortality"}
for token, code in dataset.label_id_to_code.items():
token = int(token)
if token in SPECIAL_TOKENS:
continue
text = str(code).strip().lower()
if text in exact_codes or ("death" in text) or ("mortality" in text):
ids.append(token)
return sorted(set(int(x) for x in ids if int(x) not in SPECIAL_TOKENS))
return [int(dataset.vocab_size) - 1]
def _build_first_occurrence_maps(

708
evaluate_event_free_survival.py Normal file
View File

@@ -0,0 +1,708 @@
"""Compute landmark future event-free survival summaries for DeepHealth.
For each selected patient and landmark age, this script computes:
* P(alive and no new modeled disease within tau years);
* P(alive and no new disease in each ICD-10 chapter-derived system);
* 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, Iterable, 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 AllFutureHealthDataset, HealthDataset
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 future_event_free_survival import (
future_event_free_survival_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_extra_info_types_file,
split_all_future_datasets,
split_all_future_datasets_by_eid_files,
split_dataset,
split_dataset_by_eid_files,
)
SPECIAL_TOKENS = {PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX}
class AllFutureSelectedSequenceDataset:
"""Sequence-view dataset built from selected AllFutureHealthDataset patients."""
def __init__(
self,
base: AllFutureHealthDataset,
patient_indices: Iterable[int],
) -> None:
self.base = base
self.label_code_to_id = base.label_code_to_id
self.label_id_to_code = base.label_id_to_code
self.vocab_size = base.vocab_size
self.n_types = base.n_types
self.n_cont_types = base.n_cont_types
self.n_categories = base.n_categories
self.cont_type_ids = base.cont_type_ids
self.extra_info_types = base.extra_info_types
seen: set[int] = set()
self.samples: List[Dict[str, Any]] = []
for pidx in patient_indices:
pidx = int(pidx)
if pidx in seen:
continue
seen.add(pidx)
patient = base.patients[pidx]
labels = np.asarray(patient["labels"], dtype=np.int64)
times = np.asarray(patient["times"], dtype=np.float32)
if labels.size < 2:
continue
input_len = int(labels.size - 1)
self.samples.append(
{
"eid": int(patient["eid"]),
"event_seq": labels[:-1],
"time_seq": times[:-1],
"target_event_seq": labels[1:],
"target_time_seq": times[1:],
"readout_mask": np.ones(input_len, dtype=bool),
"sex": int(patient["sex"]),
"other_type": np.asarray(patient["other_type"], dtype=np.int64),
"other_value": np.asarray(patient["other_value"], dtype=np.float32),
"other_value_kind": np.asarray(patient["other_value_kind"], dtype=np.int64),
"other_time": np.asarray(patient["other_time"], dtype=np.float32),
}
)
def __len__(self) -> int:
return len(self.samples)
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 _subset_indices(subset: Any) -> np.ndarray:
if not hasattr(subset, "indices"):
raise TypeError(f"Expected a torch Subset-like object, got {type(subset).__name__}")
return np.asarray(subset.indices, dtype=np.int64)
def _patient_indices_from_all_future_subset(
dataset: AllFutureHealthDataset,
subset: Any,
) -> np.ndarray:
indices = _subset_indices(subset)
if dataset.split == "train":
return indices
patient_indices = [
int(dataset.valid_queries[int(query_idx)][0])
for query_idx in indices.tolist()
]
return np.asarray(sorted(set(patient_indices)), dtype=np.int64)
def load_training_style_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))
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 model_target_mode == "all_future":
print("Loading AllFutureHealthDataset objects using the training path...")
train_dataset = AllFutureHealthDataset(
data_prefix=data_prefix,
labels_file=labels_file,
split="train",
min_history_events=int(cfg.get("all_future_min_history_events", 1)),
min_future_events=int(cfg.get("all_future_min_future_events", 1)),
validation_query_seed=int(cfg.get("all_future_validation_query_seed", cfg.get("seed", 42))),
extra_info_types=extra_info_types,
)
val_dataset = AllFutureHealthDataset(
data_prefix=data_prefix,
labels_file=labels_file,
split="valid",
min_history_events=int(cfg.get("all_future_min_history_events", 1)),
min_future_events=int(cfg.get("all_future_min_future_events", 1)),
validation_query_seed=int(cfg.get("all_future_validation_query_seed", cfg.get("seed", 42))),
extra_info_types=extra_info_types,
)
test_dataset = AllFutureHealthDataset(
data_prefix=data_prefix,
labels_file=labels_file,
split="test",
min_history_events=int(cfg.get("all_future_min_history_events", 1)),
min_future_events=int(cfg.get("all_future_min_future_events", 1)),
validation_query_seed=int(cfg.get("all_future_validation_query_seed", cfg.get("seed", 42))),
extra_info_types=extra_info_types,
)
if split_files_exist:
train_subset, val_subset, test_subset = split_all_future_datasets_by_eid_files(
train_dataset=train_dataset,
val_dataset=val_dataset,
test_dataset=test_dataset,
train_eid_file=train_eid_file,
val_eid_file=val_eid_file,
test_eid_file=test_eid_file,
)
split_source = "eid_files"
else:
train_subset, val_subset, test_subset = split_all_future_datasets(
train_dataset=train_dataset,
val_dataset=val_dataset,
test_dataset=test_dataset,
train_ratio=float(cfg_get(args, cfg, "train_ratio", 0.7)),
val_ratio=float(cfg_get(args, cfg, "val_ratio", 0.15)),
test_ratio=float(cfg_get(args, cfg, "test_ratio", 0.15)),
seed=int(cfg_get(args, cfg, "seed", 42)),
)
split_source = "ratio_split"
split_map = {
"train": (train_dataset, train_subset),
"val": (val_dataset, val_subset),
"test": (test_dataset, test_subset),
}
if eval_split == "all":
patient_indices = np.arange(len(train_dataset.patients), dtype=np.int64)
dataset = AllFutureSelectedSequenceDataset(train_dataset, patient_indices)
else:
source_dataset, subset = split_map[eval_split]
patient_indices = _patient_indices_from_all_future_subset(source_dataset, subset)
dataset = AllFutureSelectedSequenceDataset(source_dataset, patient_indices)
out = np.arange(len(dataset.samples), dtype=np.int64)
else:
print("Loading HealthDataset using the training path...")
dataset = HealthDataset(
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,
extra_info_types=extra_info_types,
)
if split_files_exist:
train_subset, val_subset, test_subset = split_dataset_by_eid_files(
dataset=dataset,
train_eid_file=train_eid_file,
val_eid_file=val_eid_file,
test_eid_file=test_eid_file,
)
split_source = "eid_files"
else:
train_subset, val_subset, test_subset = split_dataset(
dataset=dataset,
train_ratio=float(cfg_get(args, cfg, "train_ratio", 0.7)),
val_ratio=float(cfg_get(args, cfg, "val_ratio", 0.15)),
test_ratio=float(cfg_get(args, cfg, "test_ratio", 0.15)),
seed=int(cfg_get(args, cfg, "seed", 42)),
)
split_source = "ratio_split"
split_map = {
"train": _subset_indices(train_subset),
"val": _subset_indices(val_subset),
"test": _subset_indices(test_subset),
"all": np.arange(len(dataset.samples), dtype=np.int64),
}
out = split_map[eval_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]),
}
@torch.inference_mode()
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 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"event_free_survival_{eval_split}_tau{tau:g}y.csv"
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(
description="Compute landmark event-free survival summaries."
)
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("--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_training_style_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).to(device)
load_model_state(model, state_dict)
model.eval()
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,
)
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"Death token: {death_idx}")
print(f"Organ/system groups: {len(group_names)}")
print(f"Landmark rows: {len(landmark_dataset)}")
print(f"Output: {output_path}")
rows: list[dict[str, Any]] = []
for batch in tqdm(loader, desc="Event-free survival", 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,
)
all_survival = future_event_free_survival_from_probabilities(
probabilities,
occurred,
disease_ids=None,
vocab_size=int(dataset.vocab_size),
).detach().cpu().numpy()
group_survival: dict[str, np.ndarray] = {}
for group in group_names:
group_survival[group] = future_event_free_survival_from_probabilities(
probabilities,
occurred,
disease_ids=organ_groups[group],
vocab_size=int(dataset.vocab_size),
).detach().cpu().numpy()
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),
"event_free_survival_all": float(all_survival[j]),
}
for group in group_names:
out[f"history_count__{group}"] = int(group_counts[group])
out[f"event_free_survival__{group}"] = float(group_survival[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()

269
future_event_free_survival.py Normal file
View File

@@ -0,0 +1,269 @@
from __future__ import annotations
from collections.abc import Sequence
from typing import Any, overload
import numpy as np
try:
import torch
import torch.nn.functional as F
except ModuleNotFoundError: # pragma: no cover - optional for numpy-only use
torch = None
F = None
ArrayLike = Any
def _death_token(vocab_size: int) -> int:
if int(vocab_size) <= 0:
raise ValueError(f"vocab_size must be positive, got {vocab_size}")
return int(vocab_size) - 1
def _infer_vocab_size(x: ArrayLike, vocab_size: int | None) -> int:
if x.ndim != 2:
raise ValueError(f"Expected a 2D array/tensor with shape (N, V), got {tuple(x.shape)}")
inferred = int(x.shape[1])
if vocab_size is None:
return inferred
if int(vocab_size) != inferred:
raise ValueError(f"vocab_size={vocab_size} does not match input width {inferred}")
return int(vocab_size)
def _normalize_disease_ids(
disease_ids: Sequence[int] | np.ndarray | torch.Tensor | None,
*,
vocab_size: int,
excluded_token_ids: Sequence[int],
) -> list[int]:
death_idx = _death_token(vocab_size)
excluded = {
int(idx)
for idx in excluded_token_ids
if 0 <= int(idx) < vocab_size
}
excluded.add(death_idx)
if disease_ids is None:
return [idx for idx in range(vocab_size) if idx not in excluded]
if torch is not None and isinstance(disease_ids, torch.Tensor):
raw = disease_ids.detach().cpu().reshape(-1).tolist()
else:
raw = np.asarray(disease_ids).reshape(-1).tolist()
out: list[int] = []
seen: set[int] = set()
for value in raw:
idx = int(value)
if idx < 0 or idx >= vocab_size:
raise ValueError(f"disease id {idx} is outside [0, {vocab_size})")
if idx in excluded:
continue
if idx not in seen:
seen.add(idx)
out.append(idx)
return out
@overload
def future_event_free_survival_from_probabilities(
probabilities: torch.Tensor,
occurred: torch.Tensor,
disease_ids: Sequence[int] | np.ndarray | torch.Tensor | None = None,
*,
vocab_size: int | None = None,
excluded_token_ids: Sequence[int] = (0, 1, 2),
eps: float = 1e-7,
) -> torch.Tensor:
...
@overload
def future_event_free_survival_from_probabilities(
probabilities: np.ndarray,
occurred: np.ndarray,
disease_ids: Sequence[int] | np.ndarray | torch.Tensor | None = None,
*,
vocab_size: int | None = None,
excluded_token_ids: Sequence[int] = (0, 1, 2),
eps: float = 1e-7,
) -> np.ndarray:
...
def future_event_free_survival_from_probabilities(
probabilities: ArrayLike,
occurred: ArrayLike,
disease_ids: Sequence[int] | np.ndarray | torch.Tensor | None = None,
*,
vocab_size: int | None = None,
excluded_token_ids: Sequence[int] = (0, 1, 2),
eps: float = 1e-7,
) -> ArrayLike:
"""
Compute P(alive and no new selected disease in the next tau years).
Parameters
----------
probabilities:
Matrix with shape (N, V). Entry (i, d) is p_d(t, tau), the model's
future first-occurrence probability for token d over the chosen tau.
The death probability is always read from token V - 1.
occurred:
Boolean matrix with shape (N, V). Entry (i, d) is True if disease d has
already occurred at or before query time t. Already occurred diseases do
not contribute to "new disease" risk.
disease_ids:
Optional subset of disease tokens. If None, all non-death tokens are
included except excluded_token_ids. If provided, death and excluded
tokens are ignored here and death is still handled separately as
survival.
vocab_size:
Optional vocabulary size. If omitted, inferred from probabilities.
excluded_token_ids:
Technical tokens to exclude from "new disease" calculations. Defaults
to (0, 1, 2), matching PAD, CHECKUP, and NO_EVENT.
Returns
-------
Array/tensor with shape (N,):
Approximate probability of being alive and having no newly occurring
disease among the selected disease tokens over the same tau horizon.
"""
vocab_size = _infer_vocab_size(probabilities, vocab_size)
death_idx = _death_token(vocab_size)
selected = _normalize_disease_ids(
disease_ids,
vocab_size=vocab_size,
excluded_token_ids=excluded_token_ids,
)
if tuple(occurred.shape) != tuple(probabilities.shape):
raise ValueError(
"occurred must have the same shape as probabilities, got "
f"{tuple(occurred.shape)} vs {tuple(probabilities.shape)}"
)
if torch is not None and isinstance(probabilities, torch.Tensor):
probs = probabilities.clamp(min=0.0, max=1.0 - float(eps))
occurred_bool = occurred.to(device=probs.device, dtype=torch.bool)
log_survival = torch.log1p(-probs[:, death_idx])
if selected:
ids = torch.as_tensor(selected, dtype=torch.long, device=probs.device)
new_mask = ~occurred_bool[:, ids]
log_no_new = torch.log1p(-probs[:, ids]) * new_mask.to(probs.dtype)
log_survival = log_survival + log_no_new.sum(dim=1)
return torch.exp(log_survival)
probs_np = np.clip(np.asarray(probabilities), 0.0, 1.0 - float(eps))
occurred_bool_np = np.asarray(occurred, dtype=bool)
log_survival_np = np.log1p(-probs_np[:, death_idx])
if selected:
selected_arr = np.asarray(selected, dtype=np.int64)
new_mask_np = ~occurred_bool_np[:, selected_arr]
log_no_new_np = np.log1p(-probs_np[:, selected_arr]) * new_mask_np
log_survival_np = log_survival_np + log_no_new_np.sum(axis=1)
return np.exp(log_survival_np)
def probabilities_from_logits(
logits: torch.Tensor,
tau_years: float | torch.Tensor,
*,
dist_mode: str = "exponential",
rho: torch.Tensor | None = None,
death_rho: torch.Tensor | None = None,
eps: float = 1e-8,
) -> torch.Tensor:
"""
Convert all-future logits to tau-year event probabilities.
The death token is always treated as vocab_size - 1. For dist_mode="mixed",
non-death tokens use exponential hazards and the death token uses
death_rho. For dist_mode="weibull", rho must have the same shape as logits.
"""
if torch is None or F is None:
raise ImportError("probabilities_from_logits requires PyTorch.")
if logits.ndim != 2:
raise ValueError(f"logits must have shape (N, V), got {tuple(logits.shape)}")
if float(torch.as_tensor(tau_years).detach().min().cpu()) < 0:
raise ValueError("tau_years must be non-negative")
mode = str(dist_mode).lower()
if mode not in {"exponential", "weibull", "mixed"}:
raise ValueError("dist_mode must be one of: exponential, weibull, mixed")
rate = F.softplus(logits) + float(eps)
tau = torch.as_tensor(tau_years, dtype=rate.dtype, device=rate.device)
if tau.ndim == 0:
tau = tau.expand(logits.shape[0])
if tau.ndim != 1 or tau.shape[0] != logits.shape[0]:
raise ValueError(
"tau_years must be a scalar or a 1D tensor with length N, got "
f"{tuple(tau.shape)} for N={logits.shape[0]}"
)
if mode == "exponential":
exposure = tau[:, None].expand_as(rate)
elif mode == "weibull":
if rho is None or rho.shape != logits.shape:
raise ValueError("rho must have the same shape as logits for dist_mode='weibull'")
exposure = torch.pow(tau[:, None].clamp_min(float(eps)), rho.to(rate.dtype))
else:
exposure = tau[:, None].expand_as(rate).clone()
if death_rho is None:
raise ValueError("death_rho is required for dist_mode='mixed'")
death_idx = _death_token(logits.shape[1])
death_shape = tuple(death_rho.shape)
death_rho = death_rho.to(device=rate.device, dtype=rate.dtype)
if death_rho.ndim == 2 and death_rho.shape[1] == 1:
death_rho = death_rho.squeeze(1)
if death_rho.ndim != 1 or death_rho.shape[0] != logits.shape[0]:
raise ValueError(
"death_rho must have shape (N,) or (N, 1), got "
f"{death_shape} for N={logits.shape[0]}"
)
exposure[:, death_idx] = torch.pow(tau.clamp_min(float(eps)), death_rho)
return -torch.expm1(-rate * exposure)
def future_event_free_survival_from_logits(
logits: torch.Tensor,
occurred: torch.Tensor,
tau_years: float | torch.Tensor,
disease_ids: Sequence[int] | np.ndarray | torch.Tensor | None = None,
*,
dist_mode: str = "exponential",
rho: torch.Tensor | None = None,
death_rho: torch.Tensor | None = None,
eps: float = 1e-8,
) -> torch.Tensor:
"""
Convenience wrapper for computing future event-free survival from logits.
Returns P(alive and no new selected disease in the next tau years), with
death fixed to token vocab_size - 1.
"""
probabilities = probabilities_from_logits(
logits=logits,
tau_years=tau_years,
dist_mode=dist_mode,
rho=rho,
death_rho=death_rho,
eps=eps,
)
return future_event_free_survival_from_probabilities(
probabilities=probabilities,
occurred=occurred,
disease_ids=disease_ids,
vocab_size=logits.shape[1],
excluded_token_ids=excluded_token_ids,
)

1257
icd10_chapter_organ_mapping.csv Normal file
View File

File diff suppressed because it is too large Load Diff