Implement parallel processing for DOA AUC evaluation and enhance task management

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,67 +395,101 @@ def first_time_array(
     return out
 
 
-def evaluate_doa_auc(
+_DOA_WORKER: Dict[str, Any] = {}
-    dataset: DOAStatusDataset,
+
-    hidden_all: np.ndarray,
+
-    row_arrays: Dict[str, np.ndarray],
+def _init_doa_worker(
-    model,
+    disease_ids: np.ndarray,
-    disease_ids: Sequence[int],
+    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,
-    min_cases: int,
+    death_idx: int,
-    device: torch.device,
+) -> None:
-    logit_batch_size: int,
+    os.environ.setdefault("OMP_NUM_THREADS", "1")
-    use_amp: bool,
+    os.environ.setdefault("MKL_NUM_THREADS", "1")
-) -> pd.DataFrame:
+    os.environ.setdefault("OPENBLAS_NUM_THREADS", "1")
-    logits_all, rho_all = project_distribution_chunk(
+    os.environ.setdefault("NUMEXPR_NUM_THREADS", "1")
-        model=model,
+    _DOA_WORKER.clear()
-        hidden_all=hidden_all,
+    _DOA_WORKER.update(
-        disease_ids=disease_ids,
+        {
-        dist_mode=dist_mode,
+            "disease_ids": np.asarray(disease_ids, dtype=np.int64),
-        device=device,
+            "logits_all": np.asarray(logits_all, dtype=np.float32),
-        logit_batch_size=logit_batch_size,
+            "rho_all": None if rho_all is None else np.asarray(rho_all, dtype=np.float32),
-        use_amp=use_amp,
+            "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": {},
+        }
     )
-    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))
 
-    rows: List[Dict[str, Any]] = []
+
-    for col, token in enumerate([int(x) for x in disease_ids]):
+def _doa_first_time_by_patient(token: int) -> np.ndarray:
-        first_time = first_time_array(dataset.first_occurrence_by_token, token, patient_count)[patient_ids]
+    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
 
-        for horizon in horizons.tolist():
+    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()) < min_cases or int(control_mask.sum()) < min_cases:
+        if int(case_mask.sum()) < int(_DOA_WORKER["min_cases"]) or int(control_mask.sum()) < int(_DOA_WORKER["min_cases"]):
             continue
 
-            rho_col = None if rho_all is None else rho_all[:, col]
         scores = _score_to_probability(
-                logits=logits_all[:, col],
+            logits=logits,
-                rho=rho_col,
+            rho=rho,
-                score_mode=score_mode,
+            score_mode=_DOA_WORKER["score_mode"],
             horizon=horizon,
-                dist_mode=dist_mode,
+            dist_mode=_DOA_WORKER["dist_mode"],
             token=token,
-                death_idx=death_idx,
+            death_idx=int(_DOA_WORKER["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) if sex_value == -1 else sex == sex_value
-                    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:
+            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(
@@ -468,6 +504,87 @@ def evaluate_doa_auc(
                     "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)
 
 
@@ -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()