485 lines
18 KiB
Python
485 lines
18 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import multiprocessing as mp
|
|
from concurrent.futures import ProcessPoolExecutor, as_completed
|
|
from pathlib import Path
|
|
from typing import Any, Iterable
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
from tqdm.auto import tqdm
|
|
|
|
from burden_index import compute_burden_index, load_burden_context
|
|
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_horizons(value: Any) -> np.ndarray:
|
|
horizons = np.asarray(parse_float_list(value) or [5.0], dtype=np.float32)
|
|
if horizons.size == 0:
|
|
raise ValueError("horizons must contain at least one value.")
|
|
if np.any(horizons < 0):
|
|
raise ValueError(f"horizons must be non-negative, got {horizons}")
|
|
return horizons
|
|
|
|
|
|
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 valid devices were parsed.")
|
|
return devices
|
|
return [args.device]
|
|
|
|
|
|
def _build_burden_matrix_from_mapping(
|
|
mapping_csv: Path,
|
|
*,
|
|
token_col: str,
|
|
category_id_col: str,
|
|
category_col: str,
|
|
key_area_col: str,
|
|
weight_col: str,
|
|
) -> tuple[np.ndarray, np.ndarray, pd.DataFrame]:
|
|
df = pd.read_csv(mapping_csv)
|
|
required = {token_col, category_id_col, category_col, weight_col}
|
|
missing = sorted(required - set(df.columns))
|
|
if missing:
|
|
raise ValueError(f"{mapping_csv} is missing required columns: {missing}")
|
|
|
|
df = df.copy()
|
|
df[token_col] = pd.to_numeric(df[token_col], errors="raise").astype(int)
|
|
raw_category = df[category_id_col]
|
|
numeric_category = pd.to_numeric(raw_category, errors="coerce")
|
|
if numeric_category.notna().all():
|
|
df["_burden_category_key"] = numeric_category.astype(int)
|
|
else:
|
|
df["_burden_category_key"] = raw_category.astype(str)
|
|
df[weight_col] = pd.to_numeric(df[weight_col], errors="raise").astype(float)
|
|
df = df[df[weight_col] != 0].copy()
|
|
if df.empty:
|
|
raise ValueError(f"{mapping_csv} has no non-zero burden weights.")
|
|
|
|
disease_ids = np.asarray(sorted(df[token_col].unique().tolist()), dtype=np.int64)
|
|
category_keys = sorted(df["_burden_category_key"].unique().tolist())
|
|
|
|
disease_pos = {int(token): j for j, token in enumerate(disease_ids.tolist())}
|
|
category_pos = {cat: i for i, cat in enumerate(category_keys)}
|
|
A = np.zeros((len(category_keys), disease_ids.size), dtype=np.float64)
|
|
for _, row in df.iterrows():
|
|
token = int(row[token_col])
|
|
cat = row["_burden_category_key"]
|
|
weight = float(row[weight_col])
|
|
A[category_pos[cat], disease_pos[token]] += weight
|
|
|
|
meta_cols = list(dict.fromkeys(["_burden_category_key", category_id_col, category_col]))
|
|
if key_area_col in df.columns:
|
|
meta_cols.append(key_area_col)
|
|
category_meta = (
|
|
df[meta_cols]
|
|
.drop_duplicates(subset=["_burden_category_key"])
|
|
.sort_values("_burden_category_key")
|
|
.reset_index(drop=True)
|
|
)
|
|
category_meta = category_meta.rename(
|
|
columns={
|
|
"_burden_category_key": "burden_dimension_id",
|
|
category_col: "burden_dimension",
|
|
key_area_col: "burden_key_area",
|
|
}
|
|
)
|
|
if category_id_col in category_meta.columns:
|
|
category_meta = category_meta.drop(columns=[category_id_col])
|
|
if "burden_key_area" not in category_meta.columns:
|
|
category_meta["burden_key_area"] = ""
|
|
|
|
return A, disease_ids, category_meta
|
|
|
|
|
|
def _parse_burden_types(value: str) -> list[str]:
|
|
out = [x.strip().lower() for x in str(value).split(",") if x.strip()]
|
|
if not out:
|
|
raise ValueError("burden_types must contain at least one value.")
|
|
valid = {"functional", "organ"}
|
|
unknown = sorted(set(out) - valid)
|
|
if unknown:
|
|
raise ValueError(f"Unknown burden_types {unknown}; valid values are {sorted(valid)}")
|
|
return list(dict.fromkeys(out))
|
|
|
|
|
|
def _load_mapping_specs(args: argparse.Namespace) -> list[dict[str, Any]]:
|
|
specs: list[dict[str, Any]] = []
|
|
for burden_type in _parse_burden_types(args.burden_types):
|
|
if burden_type == "functional":
|
|
path = Path(args.functional_mapping_csv)
|
|
specs.append(
|
|
{
|
|
"burden_type": "functional",
|
|
"mapping_csv": path,
|
|
"token_col": "token_id",
|
|
"category_id_col": "hfrm_category_id",
|
|
"category_col": "hfrm_category",
|
|
"key_area_col": "hfrm_key_area",
|
|
"weight_col": args.functional_weight_col,
|
|
}
|
|
)
|
|
elif burden_type == "organ":
|
|
path = Path(args.organ_mapping_csv)
|
|
specs.append(
|
|
{
|
|
"burden_type": "organ",
|
|
"mapping_csv": path,
|
|
"token_col": "token_id",
|
|
"category_id_col": "organ_system",
|
|
"category_col": "organ_system",
|
|
"key_area_col": "icd10_chapter_name",
|
|
"weight_col": "organ_weight",
|
|
}
|
|
)
|
|
for spec in specs:
|
|
if not spec["mapping_csv"].exists():
|
|
raise FileNotFoundError(
|
|
f"{spec['burden_type']} mapping_csv not found: {spec['mapping_csv']}"
|
|
)
|
|
return specs
|
|
|
|
|
|
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():
|
|
landmark_age = float(landmark_age)
|
|
if not (followup_end > landmark_age):
|
|
continue
|
|
prefix_mask = full_time <= np.float32(landmark_age)
|
|
if not np.any(prefix_mask):
|
|
continue
|
|
prefix_events = full_event[prefix_mask].astype(np.int64, copy=False)
|
|
prefix_times = full_time[prefix_mask].astype(np.float32, 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": np.float32(landmark_age),
|
|
"t_query": np.float32(landmark_age),
|
|
"followup_end_time": np.float32(followup_end),
|
|
"event_seq": prefix_events,
|
|
"time_seq": prefix_times,
|
|
"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 _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,
|
|
)
|
|
# make_eval_indices reads len(dataset), so a small shim is enough.
|
|
class _Sized:
|
|
def __len__(self) -> int:
|
|
return n
|
|
|
|
return make_eval_indices(_Sized(), args, cfg)
|
|
|
|
|
|
def _result_rows_for_sample(
|
|
*,
|
|
sample_row: dict[str, Any],
|
|
horizons: Iterable[float],
|
|
A: np.ndarray,
|
|
disease_ids: np.ndarray,
|
|
category_meta: pd.DataFrame,
|
|
burden_type: str,
|
|
formed_mode: str,
|
|
ctx: Any,
|
|
run_path: Path,
|
|
) -> list[dict[str, Any]]:
|
|
out: list[dict[str, Any]] = []
|
|
for horizon in horizons:
|
|
result = compute_burden_index(
|
|
run_path=run_path,
|
|
burden_matrix=A,
|
|
disease_ids=disease_ids,
|
|
event_seq=sample_row["event_seq"],
|
|
time_seq=sample_row["time_seq"],
|
|
sex=int(sample_row["sex"]),
|
|
other_type=sample_row["other_type"],
|
|
other_value=sample_row["other_value"],
|
|
other_value_kind=sample_row["other_value_kind"],
|
|
other_time=sample_row["other_time"],
|
|
t_query=float(sample_row["t_query"]),
|
|
horizon=float(horizon),
|
|
formed_mode=formed_mode,
|
|
context=ctx,
|
|
)
|
|
for dim_idx, meta in category_meta.iterrows():
|
|
out.append(
|
|
{
|
|
"patient_id": sample_row["patient_id"],
|
|
"dataset_index": sample_row["dataset_index"],
|
|
"sex": sample_row["sex"],
|
|
"landmark_age": float(sample_row["landmark_age"]),
|
|
"t_query": float(sample_row["t_query"]),
|
|
"followup_end_time": float(sample_row["followup_end_time"]),
|
|
"horizon": float(horizon),
|
|
"formed_mode": formed_mode,
|
|
"burden_type": burden_type,
|
|
"burden_dimension_id": meta["burden_dimension_id"],
|
|
"burden_dimension": str(meta["burden_dimension"]),
|
|
"burden_key_area": str(meta.get("burden_key_area", "")),
|
|
"bi_historical": float(result.historical[int(dim_idx)]),
|
|
"bi_future": float(result.future[int(dim_idx)]),
|
|
"bi_total": float(result.total[int(dim_idx)]),
|
|
}
|
|
)
|
|
return out
|
|
|
|
|
|
def _compute_chunk_worker(payload: dict[str, Any]) -> list[dict[str, Any]]:
|
|
device = payload["device"]
|
|
run_path = Path(payload["run_path"])
|
|
ctx = load_burden_context(run_path, device=device)
|
|
out: list[dict[str, Any]] = []
|
|
for row in payload["rows"]:
|
|
for matrix in payload["matrices"]:
|
|
out.extend(
|
|
_result_rows_for_sample(
|
|
sample_row=row,
|
|
horizons=payload["horizons"],
|
|
A=matrix["A"],
|
|
disease_ids=matrix["disease_ids"],
|
|
category_meta=matrix["category_meta"],
|
|
burden_type=matrix["burden_type"],
|
|
formed_mode=payload["formed_mode"],
|
|
ctx=ctx,
|
|
run_path=run_path,
|
|
)
|
|
)
|
|
return out
|
|
|
|
|
|
def _split_rows_for_devices(
|
|
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)]
|
|
index_chunks = np.array_split(np.arange(len(rows)), len(devices))
|
|
chunks: list[tuple[str | None, list[dict[str, Any]]]] = []
|
|
for device, idx in zip(devices, index_chunks):
|
|
if idx.size == 0:
|
|
continue
|
|
chunks.append((device, [rows[int(i)] for i in idx.tolist()]))
|
|
return chunks
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(
|
|
description="Compute DeepHealth Burden Indices at landmark ages."
|
|
)
|
|
parser.add_argument("--run_path", type=str, required=True)
|
|
parser.add_argument("--burden_types", type=str, default="functional,organ",
|
|
help="Comma-separated burden types to compute: functional,organ.")
|
|
parser.add_argument("--functional_mapping_csv", type=str,
|
|
default="cihi_hfrm_label_mapping.csv")
|
|
parser.add_argument("--organ_mapping_csv", type=str,
|
|
default="icd10_organ_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("--formed_mode", type=str, default="model_weighted",
|
|
choices=["observed", "model_weighted"])
|
|
parser.add_argument("--horizons", type=str, default="5")
|
|
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("--device", type=str, default=None)
|
|
parser.add_argument(
|
|
"--devices",
|
|
type=str,
|
|
default=None,
|
|
help=(
|
|
"Comma-separated devices for data-parallel BI computation, e.g. "
|
|
"'cuda:0,cuda:1'. Overrides --device when provided."
|
|
),
|
|
)
|
|
parser.add_argument("--functional_weight_col", type=str,
|
|
default="hfrm_normalized_weight")
|
|
args = parser.parse_args()
|
|
|
|
run_path = Path(args.run_path)
|
|
mapping_specs = _load_mapping_specs(args)
|
|
devices = _parse_devices(args)
|
|
|
|
initial_device = "cpu" if len(devices) > 1 else devices[0]
|
|
ctx = load_burden_context(run_path, device=initial_device)
|
|
matrices = []
|
|
for spec in mapping_specs:
|
|
A, disease_ids, category_meta = _build_burden_matrix_from_mapping(
|
|
spec["mapping_csv"],
|
|
token_col=spec["token_col"],
|
|
category_id_col=spec["category_id_col"],
|
|
category_col=spec["category_col"],
|
|
key_area_col=spec["key_area_col"],
|
|
weight_col=spec["weight_col"],
|
|
)
|
|
matrices.append(
|
|
{
|
|
"burden_type": spec["burden_type"],
|
|
"A": A,
|
|
"disease_ids": disease_ids,
|
|
"category_meta": category_meta,
|
|
}
|
|
)
|
|
landmark_ages = _parse_landmark_ages(args)
|
|
horizons = _parse_horizons(args.horizons)
|
|
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. Check eval split, landmark ages, and min_history_events."
|
|
)
|
|
|
|
output_path = Path(args.output_path) if args.output_path else (
|
|
run_path / f"burden_index_{eval_split}_{args.formed_mode}.csv"
|
|
)
|
|
output_path.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
all_rows: list[dict[str, Any]] = []
|
|
row_chunks = _split_rows_for_devices(rows, devices)
|
|
if len(row_chunks) == 1:
|
|
for row in tqdm(rows, desc="Computing BI", dynamic_ncols=True):
|
|
for matrix in matrices:
|
|
all_rows.extend(
|
|
_result_rows_for_sample(
|
|
sample_row=row,
|
|
horizons=horizons.tolist(),
|
|
A=matrix["A"],
|
|
disease_ids=matrix["disease_ids"],
|
|
category_meta=matrix["category_meta"],
|
|
burden_type=matrix["burden_type"],
|
|
formed_mode=args.formed_mode,
|
|
ctx=ctx,
|
|
run_path=run_path,
|
|
)
|
|
)
|
|
else:
|
|
# The main-process context is only needed to build the dataset and rows.
|
|
# Workers load their own model copy on the assigned device.
|
|
del ctx
|
|
payloads = [
|
|
{
|
|
"device": device,
|
|
"run_path": str(run_path),
|
|
"rows": chunk_rows,
|
|
"horizons": horizons.tolist(),
|
|
"matrices": matrices,
|
|
"formed_mode": args.formed_mode,
|
|
}
|
|
for device, chunk_rows in row_chunks
|
|
]
|
|
with ProcessPoolExecutor(
|
|
max_workers=len(payloads),
|
|
mp_context=mp.get_context("spawn"),
|
|
) as executor:
|
|
futures = [executor.submit(_compute_chunk_worker, p) for p in payloads]
|
|
for future in tqdm(
|
|
as_completed(futures),
|
|
total=len(futures),
|
|
desc="Computing BI chunks",
|
|
dynamic_ncols=True,
|
|
):
|
|
all_rows.extend(future.result())
|
|
|
|
out_df = pd.DataFrame(all_rows)
|
|
out_df.to_csv(output_path, index=False)
|
|
print(f"Run path: {run_path}")
|
|
print(f"Eval split: {eval_split}")
|
|
print(f"Landmark rows: {len(rows)}")
|
|
print(f"Devices: {', '.join(str(d) for d, _ in row_chunks)}")
|
|
for matrix in matrices:
|
|
print(
|
|
f"{matrix['burden_type']} dimensions: {matrix['A'].shape[0]}, "
|
|
f"mapped disease tokens: {matrix['A'].shape[1]}"
|
|
)
|
|
print(f"Output: {output_path}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|