Update missing evaluation runner
This commit is contained in:
1
.gitattributes
vendored
Normal file
1
.gitattributes
vendored
Normal file
@@ -0,0 +1 @@
|
|||||||
|
*.sh text eol=lf
|
||||||
@@ -1,773 +0,0 @@
|
|||||||
"""Evaluate disease AUC at date of assessment (DOA).
|
|
||||||
|
|
||||||
Cases are patients whose first occurrence of a disease is after DOA and within
|
|
||||||
the requested horizon. Controls are patients who never have that disease in the
|
|
||||||
full observed record. Patients prevalent at/before DOA or incident after the
|
|
||||||
horizon are not used for that disease-horizon AUC.
|
|
||||||
|
|
||||||
The script adapts automatically to checkpoint target mode:
|
|
||||||
- next_token: use the CHECKUP token position at DOA;
|
|
||||||
- all_future: query the model directly with t_query=DOA. The history includes
|
|
||||||
the CHECKUP token at DOA.
|
|
||||||
"""
|
|
||||||
from __future__ import annotations
|
|
||||||
|
|
||||||
import argparse
|
|
||||||
import contextlib
|
|
||||||
import json
|
|
||||||
import os
|
|
||||||
from concurrent.futures import ProcessPoolExecutor
|
|
||||||
from pathlib import Path
|
|
||||||
from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple
|
|
||||||
|
|
||||||
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, Subset
|
|
||||||
from tqdm.auto import tqdm
|
|
||||||
|
|
||||||
from dataset import _ExpoBaseDataset
|
|
||||||
from evaluate_auc_v2 import (
|
|
||||||
build_metadata_for_merge,
|
|
||||||
build_model_from_dataset,
|
|
||||||
get_auc_delong_var,
|
|
||||||
load_checkpoint_state_dict,
|
|
||||||
load_json_config,
|
|
||||||
load_model_state,
|
|
||||||
parse_float_list,
|
|
||||||
parse_int_list,
|
|
||||||
project_distribution_chunk,
|
|
||||||
resolve_dist_mode_for_checkpoint,
|
|
||||||
select_disease_tokens,
|
|
||||||
validate_dataset_metadata,
|
|
||||||
_score_to_probability,
|
|
||||||
)
|
|
||||||
from readouts import build_readout
|
|
||||||
from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX, DAYS_PER_YEAR
|
|
||||||
|
|
||||||
|
|
||||||
SPECIAL_TOKENS = {PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX}
|
|
||||||
|
|
||||||
|
|
||||||
def cfg_get(args: argparse.Namespace, cfg: Dict[str, Any], name: str, default: Any) -> Any:
|
|
||||||
value = getattr(args, name, None)
|
|
||||||
if value is not None:
|
|
||||||
return value
|
|
||||||
return cfg.get(name, default)
|
|
||||||
|
|
||||||
|
|
||||||
def split_indices(
|
|
||||||
n: int,
|
|
||||||
train_ratio: float,
|
|
||||||
val_ratio: float,
|
|
||||||
test_ratio: float,
|
|
||||||
seed: int,
|
|
||||||
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
|
|
||||||
total = float(train_ratio) + float(val_ratio) + float(test_ratio)
|
|
||||||
if not np.isclose(total, 1.0, atol=1e-6):
|
|
||||||
raise ValueError(f"train/val/test ratios must sum to 1.0, got {total}")
|
|
||||||
indices = np.random.RandomState(int(seed)).permutation(int(n))
|
|
||||||
n_train = int(n * train_ratio)
|
|
||||||
n_val = int(n * val_ratio)
|
|
||||||
return indices[:n_train], indices[n_train:n_train + n_val], indices[n_train + n_val:]
|
|
||||||
|
|
||||||
|
|
||||||
def make_eval_indices(
|
|
||||||
dataset: Dataset,
|
|
||||||
args: argparse.Namespace,
|
|
||||||
cfg: Dict[str, Any],
|
|
||||||
) -> np.ndarray:
|
|
||||||
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))
|
|
||||||
eval_split = str(cfg_get(args, cfg, "eval_split", "test")).lower()
|
|
||||||
if eval_split in {"valid", "validation"}:
|
|
||||||
eval_split = "val"
|
|
||||||
|
|
||||||
train_idx, val_idx, test_idx = split_indices(
|
|
||||||
len(dataset), train_ratio, val_ratio, test_ratio, seed
|
|
||||||
)
|
|
||||||
split_map = {
|
|
||||||
"train": train_idx,
|
|
||||||
"val": val_idx,
|
|
||||||
"test": test_idx,
|
|
||||||
"all": np.arange(len(dataset), dtype=np.int64),
|
|
||||||
}
|
|
||||||
if eval_split not in split_map:
|
|
||||||
raise ValueError(f"Unsupported eval_split={eval_split!r}")
|
|
||||||
|
|
||||||
indices = np.asarray(split_map[eval_split], dtype=np.int64)
|
|
||||||
subset_size = cfg_get(args, cfg, "dataset_subset_size", None)
|
|
||||||
if subset_size is not None and int(subset_size) > 0:
|
|
||||||
indices = indices[: int(subset_size)]
|
|
||||||
return indices
|
|
||||||
|
|
||||||
|
|
||||||
def subset_first_occurrence_map(
|
|
||||||
first_occurrence_by_token: Dict[int, Tuple[np.ndarray, np.ndarray]],
|
|
||||||
selected_patient_ids: np.ndarray,
|
|
||||||
) -> Dict[int, Tuple[np.ndarray, np.ndarray]]:
|
|
||||||
selected = set(int(x) for x in np.asarray(selected_patient_ids, dtype=np.int64).tolist())
|
|
||||||
out: Dict[int, Tuple[np.ndarray, np.ndarray]] = {}
|
|
||||||
for token, pairs in first_occurrence_by_token.items():
|
|
||||||
p, t = pairs
|
|
||||||
keep = np.array([int(x) in selected for x in p], dtype=bool)
|
|
||||||
if np.any(keep):
|
|
||||||
out[int(token)] = (
|
|
||||||
np.asarray(p, dtype=np.int32)[keep],
|
|
||||||
np.asarray(t, dtype=np.float32)[keep],
|
|
||||||
)
|
|
||||||
return out
|
|
||||||
|
|
||||||
|
|
||||||
class DOAStatusDataset(_ExpoBaseDataset):
|
|
||||||
def __init__(
|
|
||||||
self,
|
|
||||||
data_prefix: str,
|
|
||||||
labels_file: str,
|
|
||||||
model_target_mode: str,
|
|
||||||
extra_info_types: Iterable[int] | None = None,
|
|
||||||
) -> None:
|
|
||||||
super().__init__(
|
|
||||||
data_prefix=data_prefix,
|
|
||||||
labels_file=labels_file,
|
|
||||||
no_event_interval_years=5.0,
|
|
||||||
include_no_event_in_uts_target=False,
|
|
||||||
extra_info_types=extra_info_types,
|
|
||||||
)
|
|
||||||
self.model_target_mode = str(model_target_mode).lower()
|
|
||||||
if self.model_target_mode not in {"next_token", "all_future"}:
|
|
||||||
raise ValueError(f"Unknown model_target_mode: {model_target_mode!r}")
|
|
||||||
|
|
||||||
self.records: List[Dict[str, Any]] = []
|
|
||||||
self.first_occurrence_by_token: Dict[int, Tuple[np.ndarray, np.ndarray]] = {}
|
|
||||||
|
|
||||||
unique_eids, starts = np.unique(self.event_data[:, 0], return_index=True)
|
|
||||||
ends = np.concatenate([starts[1:], [len(self.event_data)]])
|
|
||||||
first_lists: Dict[int, List[Tuple[int, float]]] = {}
|
|
||||||
|
|
||||||
for eid_raw, start, end in zip(unique_eids, starts, ends):
|
|
||||||
eid = int(eid_raw)
|
|
||||||
rows = self.event_data[start:end]
|
|
||||||
checkup_rows = rows[rows[:, 2].astype(np.int64) == CHECKUP_IDX]
|
|
||||||
if len(checkup_rows) == 0:
|
|
||||||
continue
|
|
||||||
|
|
||||||
features = self._split_features(eid)
|
|
||||||
if features is None:
|
|
||||||
continue
|
|
||||||
|
|
||||||
doa_days = float(np.min(checkup_rows[:, 1].astype(np.float32)))
|
|
||||||
doa_years = np.float32(doa_days / DAYS_PER_YEAR)
|
|
||||||
|
|
||||||
raw_times = rows[:, 1].astype(np.float32) / DAYS_PER_YEAR
|
|
||||||
raw_labels = rows[:, 2].astype(np.int64)
|
|
||||||
shifted_labels = np.where(
|
|
||||||
raw_labels >= NO_EVENT_IDX,
|
|
||||||
raw_labels + 1,
|
|
||||||
raw_labels,
|
|
||||||
).astype(np.int64)
|
|
||||||
order = np.lexsort((shifted_labels, raw_times))
|
|
||||||
event_times = raw_times[order].astype(np.float32)
|
|
||||||
event_labels = shifted_labels[order].astype(np.int64)
|
|
||||||
|
|
||||||
disease_mask = event_labels != CHECKUP_IDX
|
|
||||||
disease_times = event_times[disease_mask]
|
|
||||||
disease_labels = event_labels[disease_mask]
|
|
||||||
|
|
||||||
patient_id = len(self.records)
|
|
||||||
for token in np.unique(disease_labels).tolist():
|
|
||||||
token = int(token)
|
|
||||||
if token in SPECIAL_TOKENS:
|
|
||||||
continue
|
|
||||||
hit = np.where(disease_labels == token)[0]
|
|
||||||
if hit.size:
|
|
||||||
first_lists.setdefault(token, []).append(
|
|
||||||
(patient_id, float(disease_times[int(hit[0])]))
|
|
||||||
)
|
|
||||||
|
|
||||||
hist = event_times <= doa_years
|
|
||||||
hist_events = event_labels[hist]
|
|
||||||
hist_times = event_times[hist]
|
|
||||||
|
|
||||||
if self.model_target_mode == "next_token":
|
|
||||||
checkup_at_doa = (
|
|
||||||
(hist_events == CHECKUP_IDX)
|
|
||||||
& np.isclose(hist_times, doa_years, rtol=0.0, atol=1e-6)
|
|
||||||
)
|
|
||||||
if not np.any(checkup_at_doa):
|
|
||||||
raise RuntimeError(f"Missing CHECKUP token at DOA for eid={eid}")
|
|
||||||
event_seq = hist_events
|
|
||||||
time_seq = hist_times
|
|
||||||
readout_pos = int(np.where(checkup_at_doa)[0][-1])
|
|
||||||
else:
|
|
||||||
event_seq = hist_events
|
|
||||||
time_seq = hist_times
|
|
||||||
readout_pos = -1
|
|
||||||
|
|
||||||
self.records.append(
|
|
||||||
{
|
|
||||||
"patient_id": patient_id,
|
|
||||||
"eid": eid,
|
|
||||||
"doa": doa_years,
|
|
||||||
"event_seq": event_seq.astype(np.int64),
|
|
||||||
"time_seq": time_seq.astype(np.float32),
|
|
||||||
"readout_pos": readout_pos,
|
|
||||||
"full_events": disease_labels,
|
|
||||||
"full_times": disease_times,
|
|
||||||
"sex": int(features["sex"]),
|
|
||||||
"other_type": features["other_type"],
|
|
||||||
"other_value": features["other_value"],
|
|
||||||
"other_value_kind": features["other_value_kind"],
|
|
||||||
"other_time": features["other_time"],
|
|
||||||
}
|
|
||||||
)
|
|
||||||
|
|
||||||
for token, pairs in first_lists.items():
|
|
||||||
self.first_occurrence_by_token[int(token)] = (
|
|
||||||
np.asarray([p for p, _ in pairs], dtype=np.int32),
|
|
||||||
np.asarray([t for _, t in pairs], dtype=np.float32),
|
|
||||||
)
|
|
||||||
|
|
||||||
if not self.records:
|
|
||||||
raise RuntimeError("No DOA records were built from checkup events.")
|
|
||||||
|
|
||||||
def __len__(self) -> int:
|
|
||||||
return len(self.records)
|
|
||||||
|
|
||||||
def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
|
|
||||||
s = self.records[idx]
|
|
||||||
return {
|
|
||||||
"event_seq": torch.from_numpy(s["event_seq"]).long(),
|
|
||||||
"time_seq": torch.from_numpy(s["time_seq"]).float(),
|
|
||||||
"readout_pos": torch.tensor(s["readout_pos"], dtype=torch.long),
|
|
||||||
"t_query": torch.tensor(float(s["doa"]), dtype=torch.float32),
|
|
||||||
"patient_id": torch.tensor(s["patient_id"], dtype=torch.long),
|
|
||||||
"sex": torch.tensor(s["sex"], dtype=torch.long),
|
|
||||||
"other_type": torch.from_numpy(s["other_type"]).long(),
|
|
||||||
"other_value": torch.from_numpy(s["other_value"]).float(),
|
|
||||||
"other_value_kind": torch.from_numpy(s["other_value_kind"]).long(),
|
|
||||||
"other_time": torch.from_numpy(s["other_time"]).float(),
|
|
||||||
}
|
|
||||||
|
|
||||||
|
|
||||||
def collate_doa_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
|
|
||||||
)
|
|
||||||
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
|
|
||||||
)
|
|
||||||
|
|
||||||
readout_mask = torch.zeros_like(event_seq, dtype=torch.bool)
|
|
||||||
readout_pos = torch.stack([x["readout_pos"] for x in batch])
|
|
||||||
for i, pos in enumerate(readout_pos.tolist()):
|
|
||||||
if pos >= 0:
|
|
||||||
readout_mask[i, int(pos)] = True
|
|
||||||
|
|
||||||
return {
|
|
||||||
"event_seq": event_seq,
|
|
||||||
"time_seq": time_seq,
|
|
||||||
"padding_mask": event_seq > PAD_IDX,
|
|
||||||
"readout_mask": readout_mask,
|
|
||||||
"readout_pos": readout_pos,
|
|
||||||
"t_query": torch.stack([x["t_query"] for x in batch]),
|
|
||||||
"patient_id": torch.stack([x["patient_id"] for x in batch]),
|
|
||||||
"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,
|
|
||||||
}
|
|
||||||
|
|
||||||
|
|
||||||
@torch.inference_mode()
|
|
||||||
def infer_doa_hidden(
|
|
||||||
model,
|
|
||||||
loader: DataLoader,
|
|
||||||
device: torch.device,
|
|
||||||
model_target_mode: str,
|
|
||||||
readout_name: str,
|
|
||||||
readout_reduce: str,
|
|
||||||
use_amp: bool,
|
|
||||||
) -> Tuple[np.ndarray, Dict[str, np.ndarray]]:
|
|
||||||
model_target_mode = str(model_target_mode).lower()
|
|
||||||
readout = None
|
|
||||||
if model_target_mode == "next_token":
|
|
||||||
if readout_name == "same_time_group_end":
|
|
||||||
readout = build_readout("same_time_group_end", reduce=readout_reduce).to(device)
|
|
||||||
else:
|
|
||||||
readout = build_readout(readout_name).to(device)
|
|
||||||
readout.eval()
|
|
||||||
|
|
||||||
hidden_parts: List[np.ndarray] = []
|
|
||||||
patient_parts: List[np.ndarray] = []
|
|
||||||
sex_parts: List[np.ndarray] = []
|
|
||||||
autocast_enabled = bool(use_amp and device.type == "cuda")
|
|
||||||
|
|
||||||
for batch in tqdm(loader, desc="DOA inference", leave=False, 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_context = (
|
|
||||||
torch.autocast(device_type=device.type, dtype=torch.float16)
|
|
||||||
if autocast_enabled
|
|
||||||
else contextlib.nullcontext()
|
|
||||||
)
|
|
||||||
with amp_context:
|
|
||||||
if model_target_mode == "all_future":
|
|
||||||
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",
|
|
||||||
)
|
|
||||||
else:
|
|
||||||
hidden_raw = model(
|
|
||||||
event_seq=batch_dev["event_seq"],
|
|
||||||
time_seq=batch_dev["time_seq"],
|
|
||||||
sex=batch_dev["sex"],
|
|
||||||
padding_mask=batch_dev["padding_mask"],
|
|
||||||
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="next_token",
|
|
||||||
)
|
|
||||||
ro = readout(
|
|
||||||
hidden=hidden_raw,
|
|
||||||
time_seq=batch_dev["time_seq"],
|
|
||||||
padding_mask=batch_dev["padding_mask"],
|
|
||||||
readout_mask=batch_dev["readout_mask"],
|
|
||||||
)
|
|
||||||
if ro.hidden.dim() == 2:
|
|
||||||
hidden = ro.hidden
|
|
||||||
else:
|
|
||||||
hidden = ro.hidden[batch_dev["readout_mask"]]
|
|
||||||
|
|
||||||
hidden_parts.append(hidden.detach().float().cpu().numpy().astype(np.float32, copy=False))
|
|
||||||
patient_parts.append(batch["patient_id"].cpu().numpy().astype(np.int32, copy=False))
|
|
||||||
sex_parts.append(batch["sex"].cpu().numpy().astype(np.int8, copy=False))
|
|
||||||
|
|
||||||
return (
|
|
||||||
np.concatenate(hidden_parts, axis=0),
|
|
||||||
{
|
|
||||||
"patient_id": np.concatenate(patient_parts, axis=0),
|
|
||||||
"sex": np.concatenate(sex_parts, axis=0),
|
|
||||||
},
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
def first_time_array(
|
|
||||||
first_occurrence_by_token: Dict[int, Tuple[np.ndarray, np.ndarray]],
|
|
||||||
token: int,
|
|
||||||
patient_count: int,
|
|
||||||
) -> np.ndarray:
|
|
||||||
out = np.full(patient_count, np.inf, dtype=np.float32)
|
|
||||||
pairs = first_occurrence_by_token.get(int(token))
|
|
||||||
if pairs is not None:
|
|
||||||
p, t = pairs
|
|
||||||
out[np.asarray(p, dtype=np.int64)] = np.asarray(t, dtype=np.float32)
|
|
||||||
return out
|
|
||||||
|
|
||||||
|
|
||||||
_DOA_WORKER: Dict[str, Any] = {}
|
|
||||||
|
|
||||||
|
|
||||||
def _init_doa_worker(
|
|
||||||
disease_ids: np.ndarray,
|
|
||||||
logits_all: np.ndarray,
|
|
||||||
rho_all: Optional[np.ndarray],
|
|
||||||
row_patient_id: np.ndarray,
|
|
||||||
row_sex: np.ndarray,
|
|
||||||
row_doa: np.ndarray,
|
|
||||||
first_occurrence_by_token: Dict[int, Tuple[np.ndarray, np.ndarray]],
|
|
||||||
patient_count: int,
|
|
||||||
horizons: np.ndarray,
|
|
||||||
min_cases: int,
|
|
||||||
dist_mode: str,
|
|
||||||
score_mode: str,
|
|
||||||
death_idx: int,
|
|
||||||
) -> None:
|
|
||||||
os.environ.setdefault("OMP_NUM_THREADS", "1")
|
|
||||||
os.environ.setdefault("MKL_NUM_THREADS", "1")
|
|
||||||
os.environ.setdefault("OPENBLAS_NUM_THREADS", "1")
|
|
||||||
os.environ.setdefault("NUMEXPR_NUM_THREADS", "1")
|
|
||||||
_DOA_WORKER.clear()
|
|
||||||
_DOA_WORKER.update(
|
|
||||||
{
|
|
||||||
"disease_ids": np.asarray(disease_ids, dtype=np.int64),
|
|
||||||
"logits_all": np.asarray(logits_all, dtype=np.float32),
|
|
||||||
"rho_all": None if rho_all is None else np.asarray(rho_all, dtype=np.float32),
|
|
||||||
"row_patient_id": np.asarray(row_patient_id, dtype=np.int32),
|
|
||||||
"row_sex": np.asarray(row_sex, dtype=np.int8),
|
|
||||||
"row_doa": np.asarray(row_doa, dtype=np.float32),
|
|
||||||
"first_occurrence_by_token": first_occurrence_by_token,
|
|
||||||
"patient_count": int(patient_count),
|
|
||||||
"horizons": np.asarray(horizons, dtype=np.float32),
|
|
||||||
"min_cases": int(min_cases),
|
|
||||||
"dist_mode": str(dist_mode).lower(),
|
|
||||||
"score_mode": str(score_mode).lower(),
|
|
||||||
"death_idx": int(death_idx),
|
|
||||||
"first_time_cache": {},
|
|
||||||
}
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
def _doa_first_time_by_patient(token: int) -> np.ndarray:
|
|
||||||
cache = _DOA_WORKER["first_time_cache"]
|
|
||||||
if int(token) in cache:
|
|
||||||
return cache[int(token)]
|
|
||||||
|
|
||||||
out = np.full(int(_DOA_WORKER["patient_count"]), np.inf, dtype=np.float32)
|
|
||||||
pairs = _DOA_WORKER["first_occurrence_by_token"].get(int(token))
|
|
||||||
if pairs is not None:
|
|
||||||
p, t = pairs
|
|
||||||
out[np.asarray(p, dtype=np.int64)] = np.asarray(t, dtype=np.float32)
|
|
||||||
cache[int(token)] = out
|
|
||||||
return out
|
|
||||||
|
|
||||||
|
|
||||||
def _eval_doa_token(task: Tuple[int, int]) -> List[Dict[str, Any]]:
|
|
||||||
col, token = task
|
|
||||||
col = int(col)
|
|
||||||
token = int(token)
|
|
||||||
|
|
||||||
patient_ids = _DOA_WORKER["row_patient_id"]
|
|
||||||
sex = _DOA_WORKER["row_sex"]
|
|
||||||
doa = _DOA_WORKER["row_doa"]
|
|
||||||
logits = _DOA_WORKER["logits_all"][:, col]
|
|
||||||
rho_all = _DOA_WORKER["rho_all"]
|
|
||||||
rho = None if rho_all is None else rho_all[:, col]
|
|
||||||
first_time = _doa_first_time_by_patient(token)[patient_ids]
|
|
||||||
never = np.isinf(first_time)
|
|
||||||
incident_after_doa = first_time > doa
|
|
||||||
|
|
||||||
rows: List[Dict[str, Any]] = []
|
|
||||||
for horizon in _DOA_WORKER["horizons"].tolist():
|
|
||||||
horizon = float(horizon)
|
|
||||||
case_mask = incident_after_doa & (first_time <= doa + np.float32(horizon))
|
|
||||||
control_mask = never
|
|
||||||
if int(case_mask.sum()) < int(_DOA_WORKER["min_cases"]) or int(control_mask.sum()) < int(_DOA_WORKER["min_cases"]):
|
|
||||||
continue
|
|
||||||
|
|
||||||
scores = _score_to_probability(
|
|
||||||
logits=logits,
|
|
||||||
rho=rho,
|
|
||||||
score_mode=_DOA_WORKER["score_mode"],
|
|
||||||
horizon=horizon,
|
|
||||||
dist_mode=_DOA_WORKER["dist_mode"],
|
|
||||||
token=token,
|
|
||||||
death_idx=int(_DOA_WORKER["death_idx"]),
|
|
||||||
)
|
|
||||||
|
|
||||||
for sex_value, sex_name in [(0, "female"), (1, "male"), (-1, "all")]:
|
|
||||||
sex_mask = np.ones_like(case_mask, dtype=bool) if sex_value == -1 else sex == sex_value
|
|
||||||
cm = case_mask & sex_mask
|
|
||||||
nm = control_mask & sex_mask
|
|
||||||
if int(cm.sum()) < int(_DOA_WORKER["min_cases"]) or int(nm.sum()) < int(_DOA_WORKER["min_cases"]):
|
|
||||||
continue
|
|
||||||
auc, var = get_auc_delong_var(scores[cm], scores[nm])
|
|
||||||
rows.append(
|
|
||||||
{
|
|
||||||
"token": token,
|
|
||||||
"horizon": horizon,
|
|
||||||
"sex": sex_name,
|
|
||||||
"n_case": int(cm.sum()),
|
|
||||||
"n_control": int(nm.sum()),
|
|
||||||
"auc": auc,
|
|
||||||
"auc_var": var,
|
|
||||||
"auc_se": float(np.sqrt(max(var, 0.0))) if np.isfinite(var) else np.nan,
|
|
||||||
}
|
|
||||||
)
|
|
||||||
return rows
|
|
||||||
|
|
||||||
|
|
||||||
def _doa_task_block(tasks: Sequence[Tuple[int, int]]) -> List[Dict[str, Any]]:
|
|
||||||
rows: List[Dict[str, Any]] = []
|
|
||||||
for task in tasks:
|
|
||||||
rows.extend(_eval_doa_token(task))
|
|
||||||
return rows
|
|
||||||
|
|
||||||
|
|
||||||
def _split_tasks(tasks: Sequence[Tuple[int, int]], chunk_size: int) -> List[List[Tuple[int, int]]]:
|
|
||||||
if not tasks:
|
|
||||||
return []
|
|
||||||
if chunk_size <= 0:
|
|
||||||
chunk_size = max(1, int(np.ceil(len(tasks) / 8)))
|
|
||||||
return [list(tasks[i:i + chunk_size]) for i in range(0, len(tasks), chunk_size)]
|
|
||||||
|
|
||||||
|
|
||||||
def evaluate_doa_auc_chunk(
|
|
||||||
dataset: DOAStatusDataset,
|
|
||||||
hidden_all: np.ndarray,
|
|
||||||
row_arrays: Dict[str, np.ndarray],
|
|
||||||
model,
|
|
||||||
disease_ids: Sequence[int],
|
|
||||||
horizons: np.ndarray,
|
|
||||||
dist_mode: str,
|
|
||||||
score_mode: str,
|
|
||||||
min_cases: int,
|
|
||||||
device: torch.device,
|
|
||||||
logit_batch_size: int,
|
|
||||||
use_amp: bool,
|
|
||||||
num_workers_auc: int,
|
|
||||||
auc_task_chunk_size: int,
|
|
||||||
) -> pd.DataFrame:
|
|
||||||
logits_all, rho_all = project_distribution_chunk(
|
|
||||||
model=model,
|
|
||||||
hidden_all=hidden_all,
|
|
||||||
disease_ids=disease_ids,
|
|
||||||
dist_mode=dist_mode,
|
|
||||||
device=device,
|
|
||||||
logit_batch_size=logit_batch_size,
|
|
||||||
use_amp=use_amp,
|
|
||||||
)
|
|
||||||
patient_ids = row_arrays["patient_id"].astype(np.int32)
|
|
||||||
doa = np.asarray([r["doa"] for r in dataset.records], dtype=np.float32)[patient_ids]
|
|
||||||
patient_count = len(dataset.records)
|
|
||||||
death_idx = int(getattr(model, "death_idx", getattr(model, "vocab_size", 0) - 1))
|
|
||||||
disease_ids_arr = np.asarray([int(x) for x in disease_ids], dtype=np.int64)
|
|
||||||
tasks = [(j, int(token)) for j, token in enumerate(disease_ids_arr.tolist())]
|
|
||||||
|
|
||||||
init_args = (
|
|
||||||
disease_ids_arr,
|
|
||||||
logits_all,
|
|
||||||
rho_all,
|
|
||||||
patient_ids,
|
|
||||||
row_arrays["sex"].astype(np.int8),
|
|
||||||
doa,
|
|
||||||
dataset.first_occurrence_by_token,
|
|
||||||
patient_count,
|
|
||||||
horizons,
|
|
||||||
min_cases,
|
|
||||||
dist_mode,
|
|
||||||
score_mode,
|
|
||||||
death_idx,
|
|
||||||
)
|
|
||||||
|
|
||||||
if int(num_workers_auc) <= 1 or len(tasks) <= 1:
|
|
||||||
_init_doa_worker(*init_args)
|
|
||||||
rows = _doa_task_block(tasks)
|
|
||||||
return pd.DataFrame(rows)
|
|
||||||
|
|
||||||
rows: List[Dict[str, Any]] = []
|
|
||||||
task_blocks = _split_tasks(tasks, int(auc_task_chunk_size))
|
|
||||||
with ProcessPoolExecutor(
|
|
||||||
max_workers=int(num_workers_auc),
|
|
||||||
initializer=_init_doa_worker,
|
|
||||||
initargs=init_args,
|
|
||||||
) as pool:
|
|
||||||
futures = [pool.submit(_doa_task_block, block) for block in task_blocks]
|
|
||||||
for fut in tqdm(futures, desc="DOA AUC workers", leave=False, dynamic_ncols=True):
|
|
||||||
rows.extend(fut.result())
|
|
||||||
return pd.DataFrame(rows)
|
|
||||||
|
|
||||||
|
|
||||||
def iter_chunks(values: Sequence[int], chunk_size: int) -> Iterable[List[int]]:
|
|
||||||
values = [int(x) for x in values]
|
|
||||||
if chunk_size <= 0:
|
|
||||||
yield values
|
|
||||||
return
|
|
||||||
for start in range(0, len(values), chunk_size):
|
|
||||||
yield values[start:start + chunk_size]
|
|
||||||
|
|
||||||
|
|
||||||
def main() -> None:
|
|
||||||
parser = argparse.ArgumentParser(description="Evaluate DOA fixed-horizon disease AUC")
|
|
||||||
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=None,
|
|
||||||
choices=["train", "val", "valid", "validation", "test", "all"])
|
|
||||||
parser.add_argument("--dataset_subset_size", 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("--num_workers_auc", type=int, default=None)
|
|
||||||
parser.add_argument("--auc_task_chunk_size", type=int, default=None)
|
|
||||||
parser.add_argument("--logit_batch_size", type=int, default=None)
|
|
||||||
parser.add_argument("--disease_chunk_size", type=int, default=None)
|
|
||||||
parser.add_argument("--horizons", type=str, default=None)
|
|
||||||
parser.add_argument("--score_mode", type=str, choices=["risk", "eta"], default=None)
|
|
||||||
parser.add_argument("--filter_min_total", type=int, default=None)
|
|
||||||
parser.add_argument("--min_cases", type=int, default=None)
|
|
||||||
parser.add_argument("--labels_meta_path", type=str, default=None)
|
|
||||||
parser.add_argument("--use_amp", action=argparse.BooleanOptionalAction, default=None)
|
|
||||||
args = parser.parse_args()
|
|
||||||
|
|
||||||
run_path = Path(args.run_path)
|
|
||||||
cfg = load_json_config(run_path / "train_config.json")
|
|
||||||
ckpt_path = run_path / "best_model.pt"
|
|
||||||
if not ckpt_path.exists():
|
|
||||||
raise FileNotFoundError(f"best_model.pt not found in {run_path}")
|
|
||||||
|
|
||||||
output_path = Path(args.output_path or run_path)
|
|
||||||
output_path.mkdir(parents=True, exist_ok=True)
|
|
||||||
|
|
||||||
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}")
|
|
||||||
|
|
||||||
labels_meta_path = cfg_get(args, cfg, "labels_meta_path", None)
|
|
||||||
if labels_meta_path is None:
|
|
||||||
labels_meta_path = cfg.get("labels_meta_path", "delphi_labels_chapters_colours_icd.csv")
|
|
||||||
labels_meta = pd.read_csv(labels_meta_path) if labels_meta_path and Path(labels_meta_path).exists() else None
|
|
||||||
|
|
||||||
dataset = DOAStatusDataset(
|
|
||||||
data_prefix=cfg.get("data_prefix", "ukb"),
|
|
||||||
labels_file=cfg.get("labels_file", "labels.csv"),
|
|
||||||
model_target_mode=model_target_mode,
|
|
||||||
extra_info_types=parse_int_list(cfg.get("extra_info_types", None)),
|
|
||||||
)
|
|
||||||
validate_dataset_metadata(dataset, cfg)
|
|
||||||
eval_indices = make_eval_indices(dataset, args, cfg)
|
|
||||||
eval_patient_ids = np.asarray(
|
|
||||||
[dataset.records[int(i)]["patient_id"] for i in eval_indices],
|
|
||||||
dtype=np.int32,
|
|
||||||
)
|
|
||||||
eval_first_occurrence = subset_first_occurrence_map(
|
|
||||||
dataset.first_occurrence_by_token,
|
|
||||||
eval_patient_ids,
|
|
||||||
)
|
|
||||||
|
|
||||||
disease_requested = parse_int_list(cfg_get(args, cfg, "diseases_of_interest", None))
|
|
||||||
disease_ids = select_disease_tokens(
|
|
||||||
dataset=dataset,
|
|
||||||
labels_meta=labels_meta,
|
|
||||||
requested_tokens=disease_requested,
|
|
||||||
filter_min_total=int(cfg_get(args, cfg, "filter_min_total", 0)),
|
|
||||||
first_occurrence_by_token=eval_first_occurrence,
|
|
||||||
)
|
|
||||||
if not disease_ids:
|
|
||||||
raise RuntimeError("No disease tokens selected after filtering.")
|
|
||||||
|
|
||||||
horizons = np.asarray(
|
|
||||||
parse_float_list(cfg_get(args, cfg, "horizons", "1,5,10")) or [1.0, 5.0, 10.0],
|
|
||||||
dtype=np.float32,
|
|
||||||
)
|
|
||||||
score_mode = str(cfg_get(args, cfg, "score_mode", "risk")).lower()
|
|
||||||
min_cases = int(cfg_get(args, cfg, "min_cases", 2))
|
|
||||||
|
|
||||||
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)
|
|
||||||
cfg_model = dict(cfg)
|
|
||||||
cfg_model["dist_mode"] = dist_mode
|
|
||||||
|
|
||||||
device = torch.device(cfg.get("device", "cuda") if torch.cuda.is_available() else "cpu")
|
|
||||||
model = build_model_from_dataset(args, cfg_model, dataset).to(device)
|
|
||||||
load_model_state(model, state_dict)
|
|
||||||
model.eval()
|
|
||||||
|
|
||||||
if model_target_mode == "next_token" and (
|
|
||||||
model.token_embedding.num_embeddings <= CHECKUP_IDX
|
|
||||||
or model.risk_head.out_features <= CHECKUP_IDX
|
|
||||||
):
|
|
||||||
raise RuntimeError("Next-token DOA evaluation requires <CHECKUP> in the model vocabulary.")
|
|
||||||
|
|
||||||
eval_dataset = Subset(dataset, eval_indices)
|
|
||||||
loader = DataLoader(
|
|
||||||
eval_dataset,
|
|
||||||
batch_size=int(cfg_get(args, cfg, "batch_size", 128)),
|
|
||||||
shuffle=False,
|
|
||||||
collate_fn=collate_doa_fn,
|
|
||||||
num_workers=int(cfg_get(args, cfg, "num_workers", 4)),
|
|
||||||
pin_memory=device.type == "cuda",
|
|
||||||
persistent_workers=int(cfg_get(args, cfg, "num_workers", 4)) > 0,
|
|
||||||
prefetch_factor=2 if int(cfg_get(args, cfg, "num_workers", 4)) > 0 else None,
|
|
||||||
)
|
|
||||||
|
|
||||||
target_mode = cfg.get("target_mode", "uts")
|
|
||||||
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"))
|
|
||||||
|
|
||||||
eval_split = str(cfg_get(args, cfg, "eval_split", "test")).lower()
|
|
||||||
if eval_split in {"valid", "validation"}:
|
|
||||||
eval_split = "val"
|
|
||||||
print(f"DOA records: total={len(dataset)}, eval_{eval_split}={len(eval_dataset)}")
|
|
||||||
print(f"Model target mode: {model_target_mode}")
|
|
||||||
print(f"Dist mode: {dist_mode}")
|
|
||||||
print(f"Score mode: {score_mode}")
|
|
||||||
print(f"Horizons: {horizons.tolist()}")
|
|
||||||
print(f"Disease tokens: {len(disease_ids)}")
|
|
||||||
|
|
||||||
hidden_all, row_arrays = infer_doa_hidden(
|
|
||||||
model=model,
|
|
||||||
loader=loader,
|
|
||||||
device=device,
|
|
||||||
model_target_mode=model_target_mode,
|
|
||||||
readout_name=readout_name,
|
|
||||||
readout_reduce=readout_reduce,
|
|
||||||
use_amp=bool(cfg_get(args, cfg, "use_amp", False)),
|
|
||||||
)
|
|
||||||
chunk_size = int(cfg_get(args, cfg, "disease_chunk_size", 256))
|
|
||||||
num_workers_auc = int(cfg_get(args, cfg, "num_workers_auc", max(1, (os.cpu_count() or 2) - 1)))
|
|
||||||
auc_task_chunk_size = int(cfg_get(args, cfg, "auc_task_chunk_size", 0))
|
|
||||||
print(f"Disease chunk size: {chunk_size}")
|
|
||||||
print(f"AUC workers: {num_workers_auc}")
|
|
||||||
result_parts = []
|
|
||||||
for disease_chunk in tqdm(
|
|
||||||
iter_chunks(disease_ids, chunk_size),
|
|
||||||
desc="Disease chunks",
|
|
||||||
leave=True,
|
|
||||||
dynamic_ncols=True,
|
|
||||||
):
|
|
||||||
result_parts.append(
|
|
||||||
evaluate_doa_auc_chunk(
|
|
||||||
dataset=dataset,
|
|
||||||
hidden_all=hidden_all,
|
|
||||||
row_arrays=row_arrays,
|
|
||||||
model=model,
|
|
||||||
disease_ids=disease_chunk,
|
|
||||||
horizons=horizons,
|
|
||||||
dist_mode=dist_mode,
|
|
||||||
score_mode=score_mode,
|
|
||||||
min_cases=min_cases,
|
|
||||||
device=device,
|
|
||||||
logit_batch_size=int(cfg_get(args, cfg, "logit_batch_size", cfg_get(args, cfg, "batch_size", 128))),
|
|
||||||
use_amp=bool(cfg_get(args, cfg, "use_amp", False)),
|
|
||||||
num_workers_auc=num_workers_auc,
|
|
||||||
auc_task_chunk_size=auc_task_chunk_size,
|
|
||||||
)
|
|
||||||
)
|
|
||||||
result = pd.concat(result_parts, ignore_index=True) if result_parts else pd.DataFrame()
|
|
||||||
if result.empty:
|
|
||||||
raise RuntimeError("No DOA AUC rows produced. Check disease selection and min_cases.")
|
|
||||||
|
|
||||||
meta = build_metadata_for_merge(dataset, labels_meta)
|
|
||||||
result = result.merge(meta, on="token", how="left")
|
|
||||||
result.insert(0, "eval_split", eval_split)
|
|
||||||
out_file = output_path / f"doa_auc_{eval_split}.csv"
|
|
||||||
result.to_csv(out_file, index=False)
|
|
||||||
|
|
||||||
summary = result.groupby(["token", "label_code", "horizon"], dropna=False, as_index=False).agg(
|
|
||||||
auc_mean=("auc", "mean"),
|
|
||||||
n_case=("n_case", "sum"),
|
|
||||||
n_control=("n_control", "sum"),
|
|
||||||
)
|
|
||||||
summary.insert(0, "eval_split", eval_split)
|
|
||||||
summary.to_csv(output_path / f"doa_auc_{eval_split}_summary.csv", index=False)
|
|
||||||
print(f"Wrote {out_file}")
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
|
||||||
main()
|
|
||||||
@@ -1,795 +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).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()
|
|
||||||
152
run_missing_evaluations.sh
Normal file
152
run_missing_evaluations.sh
Normal file
@@ -0,0 +1,152 @@
|
|||||||
|
#!/usr/bin/env bash
|
||||||
|
set -euo pipefail
|
||||||
|
|
||||||
|
# Run all non-wrapper evaluation scripts for every completed experiment under
|
||||||
|
# runs/. The script is written for Linux servers with bash 4.2.
|
||||||
|
|
||||||
|
cd "$(dirname "${BASH_SOURCE[0]}")"
|
||||||
|
|
||||||
|
PYTHON_BIN="${PYTHON_BIN:-python}"
|
||||||
|
DEVICE="${DEVICE:-cuda}"
|
||||||
|
EVAL_SPLIT="${EVAL_SPLIT:-test}"
|
||||||
|
TAU="${TAU:-5}"
|
||||||
|
NUM_WORKERS="${NUM_WORKERS:-4}"
|
||||||
|
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}"
|
||||||
|
TAU_LABEL="$("${PYTHON_BIN}" -c 'import sys; print(f"{float(sys.argv[1]):g}")' "${TAU}")"
|
||||||
|
|
||||||
|
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
|
||||||
|
}
|
||||||
|
|
||||||
|
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_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
|
||||||
|
}
|
||||||
|
|
||||||
|
for run_path in runs/*; do
|
||||||
|
[[ -d "${run_path}" ]] || continue
|
||||||
|
|
||||||
|
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
|
||||||
|
|
||||||
|
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)
|
||||||
|
|
||||||
|
run_dir_result_if_missing \
|
||||||
|
"evaluate_auc.py" \
|
||||||
|
"${run_path}" \
|
||||||
|
"df_both.csv" \
|
||||||
|
"df_auc_unpooled.csv" \
|
||||||
|
"${PYTHON_BIN}" evaluate_auc.py "${common[@]}" "${auc_extra[@]}"
|
||||||
|
|
||||||
|
run_dir_result_if_missing \
|
||||||
|
"evaluate_auc_v2.py" \
|
||||||
|
"${run_path}" \
|
||||||
|
"df_auc_landmark.csv" \
|
||||||
|
"df_auc_landmark_unpooled.csv" \
|
||||||
|
"${PYTHON_BIN}" evaluate_auc_v2.py "${common[@]}" "${auc_extra[@]}"
|
||||||
|
|
||||||
|
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}_tau${TAU_LABEL}y" \
|
||||||
|
"manifest.json" \
|
||||||
|
"summary_extra_info_future_disease_risk.csv" \
|
||||||
|
"${PYTHON_BIN}" evaluate_extra_info_attribution.py "${common[@]}" --tau "${TAU}"
|
||||||
|
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."
|
||||||
Reference in New Issue
Block a user