diff --git a/evaluate_doa_auc.py b/evaluate_doa_auc.py index aa74076..1a09118 100644 --- a/evaluate_doa_auc.py +++ b/evaluate_doa_auc.py @@ -16,6 +16,8 @@ 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 @@ -393,7 +395,134 @@ def first_time_array( return out -def evaluate_doa_auc( +_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], @@ -406,6 +535,8 @@ def evaluate_doa_auc( 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, @@ -417,57 +548,43 @@ def evaluate_doa_auc( use_amp=use_amp, ) patient_ids = row_arrays["patient_id"].astype(np.int32) - sex = row_arrays["sex"].astype(np.int8) 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]] = [] - for col, token in enumerate([int(x) for x in disease_ids]): - first_time = first_time_array(dataset.first_occurrence_by_token, token, patient_count)[patient_ids] - never = np.isinf(first_time) - incident_after_doa = first_time > doa - - for horizon in horizons.tolist(): - horizon = float(horizon) - case_mask = incident_after_doa & (first_time <= doa + np.float32(horizon)) - control_mask = never - if int(case_mask.sum()) < min_cases or int(control_mask.sum()) < min_cases: - continue - - rho_col = None if rho_all is None else rho_all[:, col] - scores = _score_to_probability( - logits=logits_all[:, col], - rho=rho_col, - score_mode=score_mode, - horizon=horizon, - dist_mode=dist_mode, - token=token, - death_idx=death_idx, - ) - - for sex_value, sex_name in [(0, "female"), (1, "male"), (-1, "all")]: - if sex_value == -1: - sex_mask = np.ones_like(case_mask, dtype=bool) - else: - sex_mask = sex == sex_value - cm = case_mask & sex_mask - nm = control_mask & sex_mask - if int(cm.sum()) < min_cases or int(nm.sum()) < 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, - } - ) + 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) @@ -489,6 +606,8 @@ def main() -> None: 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) @@ -604,15 +723,19 @@ def main() -> None: 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( - list(iter_chunks(disease_ids, chunk_size)), + iter_chunks(disease_ids, chunk_size), desc="Disease chunks", leave=True, dynamic_ncols=True, ): result_parts.append( - evaluate_doa_auc( + evaluate_doa_auc_chunk( dataset=dataset, hidden_all=hidden_all, row_arrays=row_arrays, @@ -625,6 +748,8 @@ def main() -> None: 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()