"""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()