718 lines
27 KiB
Python
718 lines
27 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import multiprocessing as mp
|
|
import time
|
|
from concurrent.futures import ProcessPoolExecutor, as_completed
|
|
from pathlib import Path
|
|
from typing import Any, Iterable
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
import torch
|
|
from torch.nn.utils.rnn import pad_sequence
|
|
from torch.utils.data import DataLoader, IterableDataset, get_worker_info
|
|
from tqdm.auto import tqdm
|
|
|
|
from burden_index import (
|
|
build_readout_grid,
|
|
load_deephealth_context,
|
|
probabilities_from_hidden,
|
|
)
|
|
from evaluate_auc_v2 import make_eval_indices, parse_float_list
|
|
from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX
|
|
|
|
|
|
def _parse_landmark_ages(args: argparse.Namespace) -> np.ndarray:
|
|
explicit = parse_float_list(args.landmark_ages)
|
|
if explicit:
|
|
ages = np.asarray(explicit, dtype=np.float32)
|
|
else:
|
|
ages = np.arange(
|
|
float(args.landmark_start),
|
|
float(args.landmark_stop) + 1e-6,
|
|
float(args.landmark_step),
|
|
dtype=np.float32,
|
|
)
|
|
if ages.size == 0:
|
|
raise ValueError("No landmark ages were provided.")
|
|
return ages
|
|
|
|
|
|
def _parse_devices(args: argparse.Namespace) -> list[str | None]:
|
|
if args.devices is not None and str(args.devices).strip():
|
|
devices = [x.strip() for x in str(args.devices).split(",") if x.strip()]
|
|
if not devices:
|
|
raise ValueError("--devices was provided but no devices were parsed.")
|
|
return devices
|
|
return [args.device]
|
|
|
|
|
|
def _load_index_matrices(
|
|
*,
|
|
organ_mapping_csv: Path,
|
|
hfrs_mapping_csv: Path,
|
|
) -> tuple[np.ndarray, list[dict[str, Any]], dict[str, Any]]:
|
|
organ_df = pd.read_csv(organ_mapping_csv)
|
|
organ_required = {"token_id", "organ_id", "organ_label", "organ_weight"}
|
|
missing = sorted(organ_required - set(organ_df.columns))
|
|
if missing:
|
|
raise ValueError(f"{organ_mapping_csv} is missing required columns: {missing}")
|
|
organ_df = organ_df.copy()
|
|
organ_df["token_id"] = pd.to_numeric(organ_df["token_id"], errors="raise").astype(int)
|
|
organ_df["organ_weight"] = pd.to_numeric(
|
|
organ_df["organ_weight"], errors="raise"
|
|
).astype(float)
|
|
organ_df = organ_df[(organ_df["organ_id"].astype(str) != "") & (organ_df["organ_weight"] > 0)]
|
|
if organ_df.empty:
|
|
raise ValueError(f"{organ_mapping_csv} has no mapped organ rows.")
|
|
|
|
hfrs_df = pd.read_csv(hfrs_mapping_csv)
|
|
hfrs_required = {"token_id", "hfrs_weight"}
|
|
missing = sorted(hfrs_required - set(hfrs_df.columns))
|
|
if missing:
|
|
raise ValueError(f"{hfrs_mapping_csv} is missing required columns: {missing}")
|
|
hfrs_df = hfrs_df.copy()
|
|
hfrs_df["token_id"] = pd.to_numeric(hfrs_df["token_id"], errors="raise").astype(int)
|
|
hfrs_df["hfrs_weight"] = pd.to_numeric(
|
|
hfrs_df["hfrs_weight"], errors="raise"
|
|
).astype(float)
|
|
hfrs_df = hfrs_df[hfrs_df["hfrs_weight"] > 0]
|
|
if hfrs_df.empty:
|
|
raise ValueError(f"{hfrs_mapping_csv} has no non-zero HFRS weights.")
|
|
|
|
union_disease_ids = np.asarray(
|
|
sorted(
|
|
set(organ_df["token_id"].astype(int).tolist())
|
|
| set(hfrs_df["token_id"].astype(int).tolist())
|
|
),
|
|
dtype=np.int64,
|
|
)
|
|
union_pos = {int(token): i for i, token in enumerate(union_disease_ids.tolist())}
|
|
|
|
organ_ids = sorted(organ_df["organ_id"].astype(str).unique().tolist())
|
|
organ_pos = {organ_id: i for i, organ_id in enumerate(organ_ids)}
|
|
organ_matrix = np.zeros((len(organ_ids), union_disease_ids.size), dtype=np.float32)
|
|
organ_meta_by_id = {}
|
|
for _, row in organ_df.iterrows():
|
|
organ_id = str(row["organ_id"])
|
|
token = int(row["token_id"])
|
|
organ_matrix[organ_pos[organ_id], union_pos[token]] = 1.0
|
|
organ_meta_by_id.setdefault(
|
|
organ_id,
|
|
{
|
|
"index_type": "organ_involvement",
|
|
"index_id": organ_id,
|
|
"index_label": str(row["organ_label"]),
|
|
},
|
|
)
|
|
organ_meta = [organ_meta_by_id[organ_id] for organ_id in organ_ids]
|
|
|
|
hfrs_weights = np.zeros(union_disease_ids.size, dtype=np.float32)
|
|
for _, row in hfrs_df.iterrows():
|
|
hfrs_weights[union_pos[int(row["token_id"])]] = float(row["hfrs_weight"])
|
|
hfrs_meta = {
|
|
"index_type": "frailty_risk",
|
|
"index_id": "deephealth_hfrs",
|
|
"index_label": "DeepHealth-HFRS frailty risk index",
|
|
}
|
|
|
|
matrices = [
|
|
{
|
|
"kind": "organ_involvement",
|
|
"matrix": organ_matrix,
|
|
"meta": organ_meta,
|
|
},
|
|
{
|
|
"kind": "frailty_risk",
|
|
"weights": hfrs_weights,
|
|
"meta": hfrs_meta,
|
|
},
|
|
]
|
|
return union_disease_ids, matrices, {
|
|
"organ_mapped_tokens": int(organ_df["token_id"].nunique()),
|
|
"hfrs_mapped_tokens": int(hfrs_df["token_id"].nunique()),
|
|
}
|
|
|
|
|
|
def _config_split_indices(
|
|
n: int,
|
|
cfg: dict[str, Any],
|
|
eval_split: str,
|
|
subset_size: int,
|
|
) -> np.ndarray:
|
|
args = argparse.Namespace(
|
|
train_ratio=None,
|
|
val_ratio=None,
|
|
test_ratio=None,
|
|
seed=None,
|
|
eval_split=eval_split,
|
|
dataset_subset_size=subset_size if subset_size > 0 else None,
|
|
)
|
|
|
|
class _Sized:
|
|
def __len__(self) -> int:
|
|
return n
|
|
|
|
return make_eval_indices(_Sized(), args, cfg)
|
|
|
|
|
|
def _eligible_landmark_rows(
|
|
dataset: Any,
|
|
subset_indices: np.ndarray,
|
|
landmark_ages: np.ndarray,
|
|
*,
|
|
min_history_events: int,
|
|
) -> list[dict[str, Any]]:
|
|
rows: list[dict[str, Any]] = []
|
|
special = np.asarray([PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX], dtype=np.int64)
|
|
for patient_id, dataset_index in enumerate(subset_indices.tolist()):
|
|
sample = dataset.samples[int(dataset_index)]
|
|
seq_event = np.asarray(sample["event_seq"], dtype=np.int64)
|
|
seq_time = np.asarray(sample["time_seq"], dtype=np.float32)
|
|
tgt_event = np.asarray(sample["target_event_seq"], dtype=np.int64)
|
|
tgt_time = np.asarray(sample["target_time_seq"], dtype=np.float32)
|
|
if seq_event.size == 0 or tgt_event.size == 0:
|
|
continue
|
|
|
|
full_event = np.concatenate([seq_event, tgt_event[-1:]])
|
|
full_time = np.concatenate([seq_time, tgt_time[-1:]])
|
|
followup_end = float(np.max(full_time))
|
|
|
|
for landmark_age in landmark_ages.tolist():
|
|
t_query = np.float32(float(landmark_age))
|
|
if not (followup_end > float(t_query)):
|
|
continue
|
|
prefix_mask = full_time <= t_query
|
|
if not np.any(prefix_mask):
|
|
continue
|
|
prefix_events = full_event[prefix_mask].astype(np.int64, copy=False)
|
|
valid_history = ~np.isin(prefix_events, special)
|
|
if int(valid_history.sum()) < int(min_history_events):
|
|
continue
|
|
rows.append(
|
|
{
|
|
"patient_id": int(patient_id),
|
|
"dataset_index": int(dataset_index),
|
|
"sex": int(sample["sex"]),
|
|
"landmark_age": t_query,
|
|
"t_query": t_query,
|
|
"followup_end_time": np.float32(followup_end),
|
|
"event_seq": prefix_events,
|
|
"time_seq": full_time[prefix_mask].astype(np.float32, copy=False),
|
|
"other_type": np.asarray(sample["other_type"], dtype=np.int64),
|
|
"other_value": np.asarray(sample["other_value"], dtype=np.float32),
|
|
"other_value_kind": np.asarray(sample["other_value_kind"], dtype=np.int64),
|
|
"other_time": np.asarray(sample["other_time"], dtype=np.float32),
|
|
}
|
|
)
|
|
return rows
|
|
|
|
|
|
def _row_to_worker_spec(row: dict[str, Any]) -> dict[str, Any]:
|
|
return {
|
|
"patient_id": int(row["patient_id"]),
|
|
"dataset_index": int(row["dataset_index"]),
|
|
"landmark_age": float(row["landmark_age"]),
|
|
"followup_end_time": float(row["followup_end_time"]),
|
|
}
|
|
|
|
|
|
def _materialize_worker_rows(
|
|
dataset: Any,
|
|
row_specs: list[dict[str, Any]],
|
|
) -> list[dict[str, Any]]:
|
|
rows: list[dict[str, Any]] = []
|
|
for spec in row_specs:
|
|
sample = dataset.samples[int(spec["dataset_index"])]
|
|
seq_event = np.asarray(sample["event_seq"], dtype=np.int64)
|
|
seq_time = np.asarray(sample["time_seq"], dtype=np.float32)
|
|
tgt_event = np.asarray(sample["target_event_seq"], dtype=np.int64)
|
|
tgt_time = np.asarray(sample["target_time_seq"], dtype=np.float32)
|
|
full_event = np.concatenate([seq_event, tgt_event[-1:]])
|
|
full_time = np.concatenate([seq_time, tgt_time[-1:]])
|
|
t_query = np.float32(float(spec["landmark_age"]))
|
|
prefix_mask = full_time <= t_query
|
|
rows.append(
|
|
{
|
|
"patient_id": int(spec["patient_id"]),
|
|
"dataset_index": int(spec["dataset_index"]),
|
|
"sex": int(sample["sex"]),
|
|
"landmark_age": t_query,
|
|
"t_query": t_query,
|
|
"followup_end_time": np.float32(float(spec["followup_end_time"])),
|
|
"event_seq": full_event[prefix_mask].astype(np.int64, copy=False),
|
|
"time_seq": full_time[prefix_mask].astype(np.float32, copy=False),
|
|
"other_type": np.asarray(sample["other_type"], dtype=np.int64),
|
|
"other_value": np.asarray(sample["other_value"], dtype=np.float32),
|
|
"other_value_kind": np.asarray(sample["other_value_kind"], dtype=np.int64),
|
|
"other_time": np.asarray(sample["other_time"], dtype=np.float32),
|
|
}
|
|
)
|
|
return rows
|
|
|
|
|
|
class HistoricalReadoutDataset(IterableDataset):
|
|
def __init__(self, rows: list[dict[str, Any]]) -> None:
|
|
super().__init__()
|
|
self.rows = rows
|
|
|
|
def __iter__(self) -> Iterable[dict[str, torch.Tensor]]:
|
|
worker = get_worker_info()
|
|
if worker is None:
|
|
start, step = 0, 1
|
|
else:
|
|
start, step = int(worker.id), int(worker.num_workers)
|
|
|
|
for row_idx in range(start, len(self.rows), step):
|
|
row = self.rows[row_idx]
|
|
grid = build_readout_grid(
|
|
event_seq=row["event_seq"],
|
|
time_seq=row["time_seq"],
|
|
other_type=row["other_type"],
|
|
other_time=row["other_time"],
|
|
t_query=float(row["t_query"]),
|
|
)
|
|
if grid.size == 0:
|
|
continue
|
|
end_times = np.concatenate([grid[1:], np.asarray([row["t_query"]], dtype=np.float32)])
|
|
deltas = np.maximum(end_times - grid, 0.0).astype(np.float32)
|
|
valid = deltas > 0
|
|
for query_time, delta in zip(grid[valid].tolist(), deltas[valid].tolist()):
|
|
yield _make_readout_job(row, row_idx, query_time, delta)
|
|
|
|
|
|
def _make_readout_job(
|
|
row: dict[str, Any],
|
|
row_idx: int,
|
|
query_time: float,
|
|
delta: float,
|
|
) -> dict[str, torch.Tensor]:
|
|
return {
|
|
"event_seq": torch.from_numpy(np.asarray(row["event_seq"], dtype=np.int64)).long(),
|
|
"time_seq": torch.from_numpy(np.asarray(row["time_seq"], dtype=np.float32)).float(),
|
|
"sex": torch.tensor(int(row["sex"]), dtype=torch.long),
|
|
"other_type": torch.from_numpy(np.asarray(row["other_type"], dtype=np.int64)).long(),
|
|
"other_value": torch.from_numpy(np.asarray(row["other_value"], dtype=np.float32)).float(),
|
|
"other_value_kind": torch.from_numpy(
|
|
np.asarray(row["other_value_kind"], dtype=np.int64)
|
|
).long(),
|
|
"other_time": torch.from_numpy(np.asarray(row["other_time"], dtype=np.float32)).float(),
|
|
"query_time": torch.tensor(float(query_time), dtype=torch.float32),
|
|
"delta": torch.tensor(float(delta), dtype=torch.float32),
|
|
"row_idx": torch.tensor(int(row_idx), dtype=torch.long),
|
|
}
|
|
|
|
|
|
def _collate_readout_jobs(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
|
|
)
|
|
return {
|
|
"event_seq": event_seq,
|
|
"time_seq": pad_sequence(
|
|
[x["time_seq"] for x in batch], batch_first=True, padding_value=0.0
|
|
),
|
|
"padding_mask": event_seq > PAD_IDX,
|
|
"sex": torch.stack([x["sex"] for x in batch]),
|
|
"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
|
|
),
|
|
"query_time": torch.stack([x["query_time"] for x in batch]),
|
|
"delta": torch.stack([x["delta"] for x in batch]),
|
|
"row_idx": torch.stack([x["row_idx"] for x in batch]),
|
|
}
|
|
|
|
|
|
@torch.inference_mode()
|
|
def _readout_probabilities(
|
|
*,
|
|
ctx: Any,
|
|
batch: dict[str, torch.Tensor],
|
|
disease_ids: np.ndarray,
|
|
) -> torch.Tensor:
|
|
event = batch["event_seq"].long().to(ctx.device, non_blocking=True)
|
|
hidden = ctx.model(
|
|
event_seq=event,
|
|
time_seq=batch["time_seq"].float().to(ctx.device, non_blocking=True),
|
|
sex=batch["sex"].long().to(ctx.device, non_blocking=True),
|
|
padding_mask=event > PAD_IDX,
|
|
t_query=batch["query_time"].float().to(ctx.device, non_blocking=True),
|
|
other_type=batch["other_type"].long().to(ctx.device, non_blocking=True),
|
|
other_value=batch["other_value"].float().to(ctx.device, non_blocking=True),
|
|
other_value_kind=batch["other_value_kind"].long().to(ctx.device, non_blocking=True),
|
|
other_time=batch["other_time"].float().to(ctx.device, non_blocking=True),
|
|
target_mode="all_future",
|
|
)
|
|
deltas = batch["delta"].detach().cpu().numpy().astype(np.float32, copy=False)
|
|
prob = probabilities_from_hidden(
|
|
ctx=ctx,
|
|
hidden=hidden,
|
|
disease_ids=disease_ids,
|
|
deltas=deltas,
|
|
)
|
|
return torch.as_tensor(prob, dtype=torch.float32, device=ctx.device)
|
|
|
|
|
|
def _project_rows(
|
|
*,
|
|
rows: list[dict[str, Any]],
|
|
survival_by_row: torch.Tensor,
|
|
matrices: list[dict[str, Any]],
|
|
ctx: Any,
|
|
) -> list[dict[str, Any]]:
|
|
disease_expression = 1.0 - survival_by_row.clamp(0.0, 1.0)
|
|
disease_intensity = -torch.log(survival_by_row.clamp(1e-7, 1.0))
|
|
out: list[dict[str, Any]] = []
|
|
|
|
organ_matrix = torch.as_tensor(
|
|
matrices[0]["matrix"], dtype=torch.float32, device=ctx.device
|
|
)
|
|
organ_values = -torch.expm1(-torch.matmul(disease_intensity, organ_matrix.T))
|
|
organ_values_np = organ_values.detach().cpu().numpy()
|
|
|
|
hfrs_weights = torch.as_tensor(
|
|
matrices[1]["weights"], dtype=torch.float32, device=ctx.device
|
|
)
|
|
hfrs_values = torch.matmul(disease_expression, hfrs_weights)
|
|
hfrs_values_np = hfrs_values.detach().cpu().numpy()
|
|
|
|
for row_idx, row in enumerate(rows):
|
|
base = {
|
|
"patient_id": row["patient_id"],
|
|
"dataset_index": row["dataset_index"],
|
|
"sex": row["sex"],
|
|
"landmark_age": float(row["landmark_age"]),
|
|
"t_query": float(row["t_query"]),
|
|
"followup_end_time": float(row["followup_end_time"]),
|
|
}
|
|
for dim_idx, meta in enumerate(matrices[0]["meta"]):
|
|
out.append(
|
|
{
|
|
**base,
|
|
"index_type": meta["index_type"],
|
|
"index_id": meta["index_id"],
|
|
"index_label": meta["index_label"],
|
|
"index_value": float(organ_values_np[row_idx, dim_idx]),
|
|
}
|
|
)
|
|
out.append(
|
|
{
|
|
**base,
|
|
"index_type": matrices[1]["meta"]["index_type"],
|
|
"index_id": matrices[1]["meta"]["index_id"],
|
|
"index_label": matrices[1]["meta"]["index_label"],
|
|
"index_value": float(hfrs_values_np[row_idx]),
|
|
}
|
|
)
|
|
return out
|
|
|
|
|
|
def _compute_rows(
|
|
*,
|
|
rows: list[dict[str, Any]],
|
|
disease_ids: np.ndarray,
|
|
matrices: list[dict[str, Any]],
|
|
readout_batch_size: int,
|
|
num_workers: int,
|
|
ctx: Any,
|
|
log_prefix: str,
|
|
) -> tuple[list[dict[str, Any]], int, dict[str, float]]:
|
|
survival_by_row = torch.ones(
|
|
(len(rows), disease_ids.size), dtype=torch.float32, device=ctx.device
|
|
)
|
|
loader = DataLoader(
|
|
HistoricalReadoutDataset(rows),
|
|
batch_size=max(1, int(readout_batch_size)),
|
|
collate_fn=_collate_readout_jobs,
|
|
num_workers=max(0, int(num_workers)),
|
|
pin_memory=ctx.device.type == "cuda",
|
|
persistent_workers=int(num_workers) > 0,
|
|
prefetch_factor=2 if int(num_workers) > 0 else None,
|
|
)
|
|
readout_jobs = 0
|
|
n_batches = 0
|
|
forward_sec = 0.0
|
|
reduce_sec = 0.0
|
|
for batch in loader:
|
|
t0 = time.perf_counter()
|
|
prob = _readout_probabilities(ctx=ctx, batch=batch, disease_ids=disease_ids)
|
|
if ctx.device.type == "cuda":
|
|
torch.cuda.synchronize(ctx.device)
|
|
forward_sec += time.perf_counter() - t0
|
|
|
|
t1 = time.perf_counter()
|
|
row_indices = batch["row_idx"].long().to(ctx.device, non_blocking=True)
|
|
interval_survival = 1.0 - prob.clamp(0.0, 1.0)
|
|
if hasattr(survival_by_row, "scatter_reduce_"):
|
|
survival_by_row.scatter_reduce_(
|
|
dim=0,
|
|
index=row_indices[:, None].expand_as(interval_survival),
|
|
src=interval_survival,
|
|
reduce="prod",
|
|
include_self=True,
|
|
)
|
|
else:
|
|
for job_idx in range(interval_survival.shape[0]):
|
|
survival_by_row[int(row_indices[job_idx].item())] *= interval_survival[job_idx]
|
|
if ctx.device.type == "cuda":
|
|
torch.cuda.synchronize(ctx.device)
|
|
reduce_sec += time.perf_counter() - t1
|
|
|
|
n_batches += 1
|
|
readout_jobs += int(batch["row_idx"].numel())
|
|
if n_batches == 1 or n_batches % 50 == 0:
|
|
print(
|
|
f"{log_prefix} processed {readout_jobs} readout jobs in {n_batches} batches",
|
|
flush=True,
|
|
)
|
|
|
|
t2 = time.perf_counter()
|
|
out = _project_rows(
|
|
rows=rows,
|
|
survival_by_row=survival_by_row,
|
|
matrices=matrices,
|
|
ctx=ctx,
|
|
)
|
|
if ctx.device.type == "cuda":
|
|
torch.cuda.synchronize(ctx.device)
|
|
reduce_sec += time.perf_counter() - t2
|
|
return out, readout_jobs, {"forward_sec": forward_sec, "reduce_sec": reduce_sec}
|
|
|
|
|
|
def _write_rows_csv(rows: list[dict[str, Any]], output_path: Path) -> int:
|
|
df = pd.DataFrame(rows)
|
|
df.to_csv(output_path, index=False)
|
|
return int(len(df))
|
|
|
|
|
|
def _concat_csv_shards(shard_paths: list[Path], output_path: Path) -> None:
|
|
wrote_header = False
|
|
with output_path.open("w", encoding="utf-8", newline="") as out_f:
|
|
for shard_path in shard_paths:
|
|
with shard_path.open("r", encoding="utf-8", newline="") as in_f:
|
|
header = in_f.readline()
|
|
if not wrote_header:
|
|
out_f.write(header)
|
|
wrote_header = True
|
|
for line in in_f:
|
|
out_f.write(line)
|
|
shard_path.unlink(missing_ok=True)
|
|
|
|
|
|
def _estimate_jobs(row: dict[str, Any]) -> int:
|
|
grid = build_readout_grid(
|
|
event_seq=row["event_seq"],
|
|
time_seq=row["time_seq"],
|
|
other_type=row["other_type"],
|
|
other_time=row["other_time"],
|
|
t_query=float(row["t_query"]),
|
|
)
|
|
if grid.size == 0:
|
|
return 1
|
|
end_times = np.concatenate([grid[1:], np.asarray([row["t_query"]], dtype=np.float32)])
|
|
return max(int(np.sum(np.maximum(end_times - grid, 0.0) > 0)), 1)
|
|
|
|
|
|
def _split_rows(rows: list[dict[str, Any]], devices: list[str | None]) -> list[tuple[str | None, list[dict[str, Any]]]]:
|
|
if len(devices) <= 1:
|
|
return [(devices[0], rows)]
|
|
buckets: list[list[dict[str, Any]]] = [[] for _ in devices]
|
|
loads = np.zeros(len(devices), dtype=np.int64)
|
|
for row in sorted(rows, key=_estimate_jobs, reverse=True):
|
|
idx = int(np.argmin(loads))
|
|
buckets[idx].append(row)
|
|
loads[idx] += _estimate_jobs(row)
|
|
return [(device, bucket) for device, bucket in zip(devices, buckets) if bucket]
|
|
|
|
|
|
def _worker(payload: dict[str, Any]) -> dict[str, Any]:
|
|
device = payload["device"]
|
|
shard_path = Path(payload["shard_path"])
|
|
print(f"[Index worker {device}] starting with {len(payload['row_specs'])} rows", flush=True)
|
|
ctx = load_deephealth_context(payload["run_path"], device=device)
|
|
rows = _materialize_worker_rows(ctx.dataset, payload["row_specs"])
|
|
out, readout_jobs, timings = _compute_rows(
|
|
rows=rows,
|
|
disease_ids=payload["disease_ids"],
|
|
matrices=payload["matrices"],
|
|
readout_batch_size=int(payload["readout_batch_size"]),
|
|
num_workers=int(payload["num_workers"]),
|
|
ctx=ctx,
|
|
log_prefix=f"[Index worker {device}]",
|
|
)
|
|
t0 = time.perf_counter()
|
|
row_count = _write_rows_csv(out, shard_path)
|
|
timings["write_csv_sec"] = time.perf_counter() - t0
|
|
print(f"[Index worker {device}] wrote {row_count} rows to {shard_path}", flush=True)
|
|
return {
|
|
"shard_path": str(shard_path),
|
|
"row_count": row_count,
|
|
"readout_jobs": readout_jobs,
|
|
"timings": timings,
|
|
}
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(
|
|
description="Compute DeepHealth organ involvement and frailty risk indices."
|
|
)
|
|
parser.add_argument("--run_path", type=str, required=True)
|
|
parser.add_argument(
|
|
"--organ_mapping_csv",
|
|
type=str,
|
|
default="organ_involvement_label_mapping.csv",
|
|
)
|
|
parser.add_argument("--hfrs_mapping_csv", type=str, default="uk_hfrs_label_mapping.csv")
|
|
parser.add_argument("--output_path", type=str, default=None)
|
|
parser.add_argument(
|
|
"--eval_split",
|
|
type=str,
|
|
default="test",
|
|
choices=["train", "val", "valid", "validation", "test", "all"],
|
|
)
|
|
parser.add_argument("--landmark_ages", type=str, default=None)
|
|
parser.add_argument("--landmark_start", type=float, default=40.0)
|
|
parser.add_argument("--landmark_stop", type=float, default=80.0)
|
|
parser.add_argument("--landmark_step", type=float, default=5.0)
|
|
parser.add_argument("--min_history_events", type=int, default=1)
|
|
parser.add_argument("--dataset_subset_size", type=int, default=0)
|
|
parser.add_argument("--readout_batch_size", type=int, default=8192)
|
|
parser.add_argument("--num_workers", type=int, default=4)
|
|
parser.add_argument("--device", type=str, default=None)
|
|
parser.add_argument("--devices", type=str, default=None)
|
|
parser.add_argument(
|
|
"--mp_start_method",
|
|
type=str,
|
|
default="fork",
|
|
choices=["fork", "forkserver"],
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
run_path = Path(args.run_path)
|
|
devices = _parse_devices(args)
|
|
initial_device = "cpu" if len(devices) > 1 else devices[0]
|
|
ctx = load_deephealth_context(run_path, device=initial_device)
|
|
|
|
disease_ids, matrices, mapping_counts = _load_index_matrices(
|
|
organ_mapping_csv=Path(args.organ_mapping_csv),
|
|
hfrs_mapping_csv=Path(args.hfrs_mapping_csv),
|
|
)
|
|
landmark_ages = _parse_landmark_ages(args)
|
|
eval_split = str(args.eval_split).lower()
|
|
if eval_split in {"valid", "validation"}:
|
|
eval_split = "val"
|
|
subset_indices = _config_split_indices(
|
|
len(ctx.dataset),
|
|
ctx.cfg,
|
|
eval_split,
|
|
int(args.dataset_subset_size),
|
|
)
|
|
rows = _eligible_landmark_rows(
|
|
ctx.dataset,
|
|
subset_indices,
|
|
landmark_ages,
|
|
min_history_events=int(args.min_history_events),
|
|
)
|
|
if not rows:
|
|
raise RuntimeError("No eligible landmark rows.")
|
|
|
|
output_path = Path(args.output_path) if args.output_path else (
|
|
run_path / f"deephealth_indices_{eval_split}.csv"
|
|
)
|
|
output_path.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
chunks = _split_rows(rows, devices)
|
|
for device, chunk in chunks:
|
|
print(
|
|
f"Assigned {len(chunk)} rows / ~{sum(_estimate_jobs(r) for r in chunk)} "
|
|
f"readout jobs to {device}",
|
|
flush=True,
|
|
)
|
|
|
|
total_readout_jobs = 0
|
|
timings = {"forward_sec": 0.0, "reduce_sec": 0.0, "write_csv_sec": 0.0}
|
|
if len(chunks) == 1:
|
|
out, total_readout_jobs, chunk_timings = _compute_rows(
|
|
rows=rows,
|
|
disease_ids=disease_ids,
|
|
matrices=matrices,
|
|
readout_batch_size=int(args.readout_batch_size),
|
|
num_workers=int(args.num_workers),
|
|
ctx=ctx,
|
|
log_prefix="[Index main]",
|
|
)
|
|
for key, value in chunk_timings.items():
|
|
timings[key] += float(value)
|
|
t0 = time.perf_counter()
|
|
_write_rows_csv(out, output_path)
|
|
timings["write_csv_sec"] = time.perf_counter() - t0
|
|
else:
|
|
del ctx
|
|
payloads = [
|
|
{
|
|
"device": device,
|
|
"run_path": str(run_path),
|
|
"shard_path": str(
|
|
output_path.with_name(
|
|
f"{output_path.stem}.part{part_idx:03d}{output_path.suffix}"
|
|
)
|
|
),
|
|
"row_specs": [_row_to_worker_spec(row) for row in chunk],
|
|
"disease_ids": disease_ids,
|
|
"matrices": matrices,
|
|
"readout_batch_size": int(args.readout_batch_size),
|
|
"num_workers": int(args.num_workers),
|
|
}
|
|
for part_idx, (device, chunk) in enumerate(chunks)
|
|
]
|
|
shard_paths = []
|
|
with ProcessPoolExecutor(
|
|
max_workers=len(payloads),
|
|
mp_context=mp.get_context(args.mp_start_method),
|
|
) as executor:
|
|
futures = [executor.submit(_worker, payload) for payload in payloads]
|
|
for future in tqdm(
|
|
as_completed(futures),
|
|
total=len(futures),
|
|
desc="Computing DeepHealth index chunks",
|
|
dynamic_ncols=True,
|
|
):
|
|
result = future.result()
|
|
shard_paths.append(Path(result["shard_path"]))
|
|
total_readout_jobs += int(result["readout_jobs"])
|
|
for key, value in result["timings"].items():
|
|
timings[key] += float(value)
|
|
t0 = time.perf_counter()
|
|
_concat_csv_shards(sorted(shard_paths), output_path)
|
|
timings["write_csv_sec"] += time.perf_counter() - t0
|
|
|
|
print(f"Run path: {run_path}")
|
|
print(f"Eval split: {eval_split}")
|
|
print(f"Landmark rows: {len(rows)}")
|
|
print(f"Readout jobs: {total_readout_jobs}")
|
|
print(f"Union disease tokens: {disease_ids.size}")
|
|
print(f"Organ mapped tokens: {mapping_counts['organ_mapped_tokens']}")
|
|
print(f"HFRS mapped tokens: {mapping_counts['hfrs_mapped_tokens']}")
|
|
print(
|
|
"Timing seconds: "
|
|
f"forward={timings['forward_sec']:.2f}, "
|
|
f"reduce={timings['reduce_sec']:.2f}, "
|
|
f"write_csv={timings['write_csv_sec']:.2f}"
|
|
)
|
|
print(f"Devices: {', '.join(str(d) for d, _ in chunks)}")
|
|
print(f"Output: {output_path}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|