Refactor DeepHealth indices around disease expression
This commit is contained in:
717
compute_deephealth_indices_landmarks.py
Normal file
717
compute_deephealth_indices_landmarks.py
Normal file
@@ -0,0 +1,717 @@
|
||||
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()
|
||||
Reference in New Issue
Block a user