709 lines
27 KiB
Python
709 lines
27 KiB
Python
"""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()
|