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,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]):
+    task_blocks = _split_tasks(tasks, int(auc_task_chunk_size))
-        first_time = first_time_array(dataset.first_occurrence_by_token, token, patient_count)[patient_ids]
+    with ProcessPoolExecutor(
-        never = np.isinf(first_time)
+        max_workers=int(num_workers_auc),
-        incident_after_doa = first_time > doa
+        initializer=_init_doa_worker,
-
+        initargs=init_args,
-        for horizon in horizons.tolist():
+    ) as pool:
-            horizon = float(horizon)
+        futures = [pool.submit(_doa_task_block, block) for block in task_blocks]
-            case_mask = incident_after_doa & (first_time <= doa + np.float32(horizon))
+        for fut in tqdm(futures, desc="DOA AUC workers", leave=False, dynamic_ncols=True):
-            control_mask = never
+            rows.extend(fut.result())
-            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,
-                    }
-                )
     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()