JiaruiLi/DeepHealth
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Implement parallel processing for DOA AUC evaluation and enhance task management

Browse Source
This commit is contained in:
Jiarui Li
2026-06-13 15:45:07 +08:00
parent 76787d2fb2
commit 58f253d9b6
1 changed files with 174 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files

209
evaluate_doa_auc.py
View File

@@ -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(
dataset: DOAStatusDataset,
hidden_all: np.ndarray,
row_arrays: Dict[str, np.ndarray],
model,
disease_ids: Sequence[int],
_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,
min_cases: int,
device: torch.device,
logit_batch_size: int,
use_amp: bool,
) -> 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,
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": {},
}
)
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]):
first_time = first_time_array(dataset.first_occurrence_by_token, token, patient_count)[patient_ids]
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
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],
rho=rho_col,
score_mode=score_mode,
logits=logits,
rho=rho,
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)
else:
sex_mask = sex == sex_value
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()) < 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()