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