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 import torch from tqdm.auto import tqdm from burden_index import ( _build_readout_grid, _observed_formed_burden, _probabilities_from_hidden, 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_horizon(value: Any) -> float: horizon = float(value) if horizon < 0: raise ValueError(f"horizon must be non-negative, got {horizon}") return horizon def _format_horizon_for_filename(horizon: float) -> str: text = f"{float(horizon):g}".replace("-", "m").replace(".", "p") return f"h{text}" 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 _iter_readout_batches( n: int, batch_size: int, ) -> Iterable[slice]: batch_size = max(1, int(batch_size)) for start in range(0, n, batch_size): yield slice(start, min(start + batch_size, n)) @torch.inference_mode() def _query_hidden_jobs( *, ctx: Any, jobs: list[tuple[dict[str, Any], float]], ) -> torch.Tensor: if not jobs: return torch.empty(0, ctx.model.n_embd, device=ctx.device) batch_size = len(jobs) max_event_len = max(int(np.asarray(row["event_seq"]).size) for row, _ in jobs) max_other_len = max(int(np.asarray(row["other_type"]).size) for row, _ in jobs) event = np.full((batch_size, max_event_len), PAD_IDX, dtype=np.int64) time = np.zeros((batch_size, max_event_len), dtype=np.float32) other_type = np.zeros((batch_size, max_other_len), dtype=np.int64) other_value = np.zeros((batch_size, max_other_len), dtype=np.float32) other_value_kind = np.zeros((batch_size, max_other_len), dtype=np.int64) other_time = np.zeros((batch_size, max_other_len), dtype=np.float32) sex = np.zeros(batch_size, dtype=np.int64) query_times = np.zeros(batch_size, dtype=np.float32) for i, (row, query_time) in enumerate(jobs): event_seq = np.asarray(row["event_seq"], dtype=np.int64) time_seq = np.asarray(row["time_seq"], dtype=np.float32) other_type_seq = np.asarray(row["other_type"], dtype=np.int64) other_value_seq = np.asarray(row["other_value"], dtype=np.float32) other_value_kind_seq = np.asarray(row["other_value_kind"], dtype=np.int64) other_time_seq = np.asarray(row["other_time"], dtype=np.float32) event[i, : event_seq.size] = event_seq time[i, : time_seq.size] = time_seq other_type[i, : other_type_seq.size] = other_type_seq other_value[i, : other_value_seq.size] = other_value_seq other_value_kind[i, : other_value_kind_seq.size] = other_value_kind_seq other_time[i, : other_time_seq.size] = other_time_seq sex[i] = int(row["sex"]) query_times[i] = np.float32(query_time) event_t = torch.from_numpy(event).long().to(ctx.device) return ctx.model( event_seq=event_t, time_seq=torch.from_numpy(time).float().to(ctx.device), sex=torch.from_numpy(sex).long().to(ctx.device), padding_mask=event_t > PAD_IDX, t_query=torch.from_numpy(query_times).float().to(ctx.device), other_type=torch.from_numpy(other_type).long().to(ctx.device), other_value=torch.from_numpy(other_value).float().to(ctx.device), other_value_kind=torch.from_numpy(other_value_kind).long().to(ctx.device), other_time=torch.from_numpy(other_time).float().to(ctx.device), target_mode="all_future", ) def _build_readout_table( *, rows: list[dict[str, Any]], formed_mode: str, horizon: float, ) -> dict[str, Any]: jobs: list[tuple[dict[str, Any], float]] = [] row_indices: list[int] = [] kinds: list[str] = [] deltas: list[float] = [] if formed_mode not in {"observed", "model_weighted"}: raise ValueError(f"Unknown formed_mode={formed_mode!r}") for row_idx, row in enumerate(rows): if formed_mode == "model_weighted": 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: end_times = np.concatenate( [grid[1:], np.asarray([row["t_query"]], dtype=np.float32)] ) row_deltas = np.maximum(end_times - grid, 0.0).astype(np.float32) valid = row_deltas > 0 for query_time, delta in zip(grid[valid].tolist(), row_deltas[valid].tolist()): jobs.append((row, float(query_time))) row_indices.append(row_idx) kinds.append("formed") deltas.append(float(delta)) if horizon > 0: jobs.append((row, float(row["t_query"]))) row_indices.append(row_idx) kinds.append("future") deltas.append(float(horizon)) return { "jobs": jobs, "row_indices": np.asarray(row_indices, dtype=np.int64), "kinds": np.asarray(kinds, dtype=object), "deltas": np.asarray(deltas, dtype=np.float32), } @torch.inference_mode() def _readout_probabilities( *, ctx: Any, readout_table: dict[str, Any], union_disease_ids: np.ndarray, readout_batch_size: int, ) -> np.ndarray: jobs = readout_table["jobs"] if not jobs: return np.zeros((0, union_disease_ids.size), dtype=np.float64) out = np.empty((len(jobs), union_disease_ids.size), dtype=np.float64) deltas = np.asarray(readout_table["deltas"], dtype=np.float32) for slc in _iter_readout_batches(len(jobs), readout_batch_size): hidden = _query_hidden_jobs(ctx=ctx, jobs=jobs[slc]) out[slc] = _probabilities_from_hidden( ctx=ctx, hidden=hidden, disease_ids=union_disease_ids, deltas=deltas[slc], ) return out def _observed_formed_for_rows( *, rows: list[dict[str, Any]], union_disease_ids: np.ndarray, ) -> np.ndarray: formed = np.zeros((len(rows), union_disease_ids.size), dtype=np.float64) for row_idx, row in enumerate(rows): formed[row_idx] = _observed_formed_burden( disease_ids=union_disease_ids, event_seq=row["event_seq"], time_seq=row["time_seq"], t_query=float(row["t_query"]), ) return formed def _reduce_readout_table_to_bi_rows( *, rows: list[dict[str, Any]], horizon: float, matrices: list[dict[str, Any]], union_disease_ids: np.ndarray, formed_mode: str, readout_table: dict[str, Any], readout_prob: np.ndarray, ) -> list[dict[str, Any]]: if formed_mode == "observed": formed_by_row = _observed_formed_for_rows( rows=rows, union_disease_ids=union_disease_ids, ) elif formed_mode == "model_weighted": formed_by_row = np.zeros((len(rows), union_disease_ids.size), dtype=np.float64) else: raise ValueError(f"Unknown formed_mode={formed_mode!r}") future_prob_by_row = np.zeros((len(rows), union_disease_ids.size), dtype=np.float64) row_indices = np.asarray(readout_table["row_indices"], dtype=np.int64) kinds = np.asarray(readout_table["kinds"], dtype=object) if formed_mode == "model_weighted": survival_by_row = np.ones((len(rows), union_disease_ids.size), dtype=np.float64) else: survival_by_row = None for job_idx, row_idx in enumerate(row_indices.tolist()): kind = str(kinds[job_idx]) if kind == "formed" and survival_by_row is not None: survival_by_row[int(row_idx)] *= 1.0 - np.clip(readout_prob[job_idx], 0.0, 1.0) elif kind == "future": future_prob_by_row[int(row_idx)] = readout_prob[job_idx] if survival_by_row is not None: formed_by_row = 1.0 - survival_by_row out: list[dict[str, Any]] = [] for row_idx, row in enumerate(rows): formed = formed_by_row[row_idx] historical_by_matrix = { matrix["burden_type"]: matrix["A_union"] @ formed for matrix in matrices } disease_future = (1.0 - formed) * future_prob_by_row[row_idx] disease_total = formed + disease_future for matrix in matrices: historical = historical_by_matrix[matrix["burden_type"]] future = matrix["A_union"] @ disease_future total = matrix["A_union"] @ disease_total for dim_idx, meta in matrix["category_meta"].iterrows(): out.append( { "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"]), "horizon": float(horizon), "formed_mode": formed_mode, "burden_type": matrix["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(historical[int(dim_idx)]), "bi_future": float(future[int(dim_idx)]), "bi_total": float(total[int(dim_idx)]), } ) return out def _compute_bi_from_readout_table( *, rows: list[dict[str, Any]], horizon: float, matrices: list[dict[str, Any]], union_disease_ids: np.ndarray, formed_mode: str, readout_batch_size: int, ctx: Any, ) -> tuple[list[dict[str, Any]], int]: horizon = float(horizon) if horizon < 0: raise ValueError(f"horizon must be non-negative, got {horizon}") readout_table = _build_readout_table( rows=rows, formed_mode=formed_mode, horizon=horizon, ) readout_prob = _readout_probabilities( ctx=ctx, readout_table=readout_table, union_disease_ids=union_disease_ids, readout_batch_size=readout_batch_size, ) rows_out = _reduce_readout_table_to_bi_rows( rows=rows, horizon=horizon, matrices=matrices, union_disease_ids=union_disease_ids, formed_mode=formed_mode, readout_table=readout_table, readout_prob=readout_prob, ) return rows_out, len(readout_table["jobs"]) def _compute_chunk_worker(payload: dict[str, Any]) -> dict[str, Any]: device = payload["device"] run_path = Path(payload["run_path"]) ctx = load_burden_context(run_path, device=device) out, readout_jobs = _compute_bi_from_readout_table( rows=payload["rows"], horizon=payload["horizon"], matrices=payload["matrices"], union_disease_ids=payload["union_disease_ids"], formed_mode=payload["formed_mode"], readout_batch_size=int(payload["readout_batch_size"]), ctx=ctx, ) return {"rows": out, "readout_jobs": readout_jobs} def _attach_union_projection( matrices: list[dict[str, Any]], ) -> tuple[np.ndarray, list[dict[str, Any]]]: union_disease_ids = np.asarray( sorted( { int(token) for matrix in matrices for token in np.asarray(matrix["disease_ids"], dtype=np.int64).tolist() } ), dtype=np.int64, ) if union_disease_ids.size == 0: raise ValueError("No disease tokens are covered by the requested burden matrices.") union_pos = {int(token): i for i, token in enumerate(union_disease_ids.tolist())} projected: list[dict[str, Any]] = [] for matrix in matrices: disease_ids = np.asarray(matrix["disease_ids"], dtype=np.int64) A = np.asarray(matrix["A"], dtype=np.float64) A_union = np.zeros((A.shape[0], union_disease_ids.size), dtype=np.float64) for local_col, token in enumerate(disease_ids.tolist()): A_union[:, union_pos[int(token)]] += A[:, int(local_col)] item = dict(matrix) item["A_union"] = A_union projected.append(item) return union_disease_ids, projected 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( "--horizon", type=float, required=True, help=( "Future horizon in years. Use 0 to compute historical burden only " "(bi_future=0 and bi_total=bi_historical)." ), ) 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, help=( "Number of readout points forwarded together inside each worker. " "Increase this to improve GPU utilization if memory allows." ), ) 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, } ) union_disease_ids, matrices = _attach_union_projection(matrices) landmark_ages = _parse_landmark_ages(args) horizon = _parse_horizon(args.horizon) 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}_{_format_horizon_for_filename(horizon)}.csv" ) output_path.parent.mkdir(parents=True, exist_ok=True) all_rows: list[dict[str, Any]] = [] total_readout_jobs = 0 row_chunks = _split_rows_for_devices(rows, devices) if len(row_chunks) == 1: all_rows, total_readout_jobs = _compute_bi_from_readout_table( rows=rows, horizon=horizon, matrices=matrices, union_disease_ids=union_disease_ids, formed_mode=args.formed_mode, readout_batch_size=int(args.readout_batch_size), ctx=ctx, ) 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, "horizon": horizon, "matrices": matrices, "union_disease_ids": union_disease_ids, "readout_batch_size": int(args.readout_batch_size), "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, ): result = future.result() all_rows.extend(result["rows"]) total_readout_jobs += int(result["readout_jobs"]) 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"Horizon: {horizon:g}") print(f"Landmark rows: {len(rows)}") print(f"Readout jobs: {total_readout_jobs}") print(f"Readout batch size per worker: {int(args.readout_batch_size)}") 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"Union disease tokens evaluated once per sample: {union_disease_ids.size}") print(f"Output: {output_path}") if __name__ == "__main__": main()