From 93450ab06b5f1651a1a6c36f8e45724db77de8b0 Mon Sep 17 00:00:00 2001 From: Jiarui Li Date: Wed, 1 Jul 2026 09:19:42 +0800 Subject: [PATCH] Update missing evaluation runner --- .gitattributes | 1 + evaluate_doa_auc.py | 773 ------------------------------- evaluate_event_free_survival.py | 795 -------------------------------- run_missing_evaluations.sh | 152 ++++++ 4 files changed, 153 insertions(+), 1568 deletions(-) create mode 100644 .gitattributes delete mode 100644 evaluate_doa_auc.py delete mode 100644 evaluate_event_free_survival.py create mode 100644 run_missing_evaluations.sh diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000..dfdb8b7 --- /dev/null +++ b/.gitattributes @@ -0,0 +1 @@ +*.sh text eol=lf diff --git a/evaluate_doa_auc.py b/evaluate_doa_auc.py deleted file mode 100644 index 36d4441..0000000 --- a/evaluate_doa_auc.py +++ /dev/null @@ -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 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() diff --git a/evaluate_event_free_survival.py b/evaluate_event_free_survival.py deleted file mode 100644 index d1f33a7..0000000 --- a/evaluate_event_free_survival.py +++ /dev/null @@ -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() diff --git a/run_missing_evaluations.sh b/run_missing_evaluations.sh new file mode 100644 index 0000000..069e4b3 --- /dev/null +++ b/run_missing_evaluations.sh @@ -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."