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, _observed_formed_burden, 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 _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"]), "t_query": float(row["t_query"]), "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["t_query"])) 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": np.float32(float(spec["landmark_age"])), "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 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) class ReadoutJobIterableDataset(IterableDataset): def __init__( self, *, rows: list[dict[str, Any]], formed_mode: str, horizon: float, ) -> None: super().__init__() self.rows = rows self.formed_mode = str(formed_mode) self.horizon = float(horizon) if self.formed_mode not in {"observed", "model_weighted"}: raise ValueError(f"Unknown formed_mode={self.formed_mode!r}") 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] if self.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()): yield _make_readout_job(row, row_idx, "formed", query_time, delta) if self.horizon > 0: yield _make_readout_job( row, row_idx, "future", float(row["t_query"]), self.horizon, ) def _make_readout_job( row: dict[str, Any], row_idx: int, kind: str, 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), "kind": torch.tensor(0 if kind == "formed" else 1, 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 ) time_seq = pad_sequence( [x["time_seq"] for x in batch], batch_first=True, padding_value=0.0 ) 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 ) return { "event_seq": event_seq, "time_seq": time_seq, "padding_mask": event_seq > PAD_IDX, "sex": torch.stack([x["sex"] for x in batch]), "other_type": other_type, "other_value": other_value, "other_value_kind": other_value_kind, "other_time": other_time, "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]), "kind": torch.stack([x["kind"] for x in batch]), } @torch.inference_mode() def _query_hidden_readout_batch( *, ctx: Any, batch: dict[str, torch.Tensor], ) -> torch.Tensor: event = batch["event_seq"].long().to(ctx.device, non_blocking=True) return 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", ) @torch.inference_mode() def _probabilities_from_hidden_torch( *, ctx: Any, hidden: torch.Tensor, disease_ids: np.ndarray, deltas: np.ndarray, ) -> torch.Tensor: if hidden.ndim != 2: raise ValueError(f"hidden must have shape (N, H), got {tuple(hidden.shape)}") if deltas.ndim != 1 or deltas.size != hidden.shape[0]: raise ValueError( "deltas must be 1D with the same length as hidden rows, got " f"{deltas.shape} vs {tuple(hidden.shape)}" ) ids = torch.as_tensor(disease_ids, dtype=torch.long, device=ctx.device) logits = ctx.model.calc_risk(hidden)[:, ids] rate = torch.nn.functional.softplus(logits).clamp_min(1e-8) delta_t = torch.as_tensor(deltas, dtype=rate.dtype, device=ctx.device).clamp_min(0) if ctx.dist_mode == "weibull": rho = ctx.model.calc_weibull_rho(hidden)[:, ids] exposure = torch.pow(delta_t[:, None], rho) elif ctx.dist_mode == "mixed": exposure = delta_t[:, None].expand_as(rate) death_idx = int(getattr(ctx.model, "death_idx", getattr(ctx.model, "vocab_size", 0) - 1)) death_cols = [j for j, token in enumerate(disease_ids.tolist()) if int(token) == death_idx] if death_cols: death_rho = ctx.model.calc_death_rho(hidden) for col in death_cols: exposure[:, int(col)] = torch.pow(delta_t, death_rho) else: exposure = delta_t[:, None].expand_as(rate) return -torch.expm1(-rate * exposure) @torch.inference_mode() def _readout_probabilities_from_batch( *, ctx: Any, batch: dict[str, torch.Tensor], union_disease_ids: np.ndarray, ) -> torch.Tensor: hidden = _query_hidden_readout_batch(ctx=ctx, batch=batch) deltas = batch["delta"].detach().cpu().numpy().astype(np.float32, copy=False) return _probabilities_from_hidden_torch( ctx=ctx, hidden=hidden, disease_ids=union_disease_ids, deltas=deltas, ).to(dtype=torch.float32) 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 _apply_readout_batch_to_accumulators( *, batch: dict[str, torch.Tensor], readout_prob: torch.Tensor, survival_by_row: torch.Tensor | None, future_prob_by_row: torch.Tensor, ctx: Any, ) -> None: if readout_prob.numel() == 0: return row_indices = batch["row_idx"].long().to(ctx.device, non_blocking=True) kinds = batch["kind"].long().to(ctx.device, non_blocking=True) kind_is_formed = kinds == 0 kind_is_future = kinds == 1 if survival_by_row is not None and bool(kind_is_formed.any().item()): formed_rows = row_indices[kind_is_formed] formed_survival = 1.0 - readout_prob[kind_is_formed].clamp(0.0, 1.0) if hasattr(survival_by_row, "scatter_reduce_"): survival_by_row.scatter_reduce_( dim=0, index=formed_rows[:, None].expand_as(formed_survival), src=formed_survival, reduce="prod", include_self=True, ) else: for job_idx in torch.nonzero(kind_is_formed, as_tuple=False).flatten().tolist(): survival_by_row[int(row_indices[job_idx].item())] *= ( 1.0 - readout_prob[job_idx].clamp(0.0, 1.0) ) if bool(kind_is_future.any().item()): future_rows = row_indices[kind_is_future] future_prob_by_row[future_rows] = readout_prob[kind_is_future] def _project_bi_rows( *, rows: list[dict[str, Any]], horizon: float, matrices: list[dict[str, Any]], formed_mode: str, formed_by_row: torch.Tensor, future_prob_by_row: torch.Tensor, ctx: Any, ) -> list[dict[str, Any]]: disease_future_by_row = (1.0 - formed_by_row) * future_prob_by_row disease_total_by_row = formed_by_row + disease_future_by_row projected: list[dict[str, Any]] = [] for matrix in matrices: A = torch.as_tensor(matrix["A_union"], dtype=formed_by_row.dtype, device=ctx.device) projected.append( { "matrix": matrix, "historical": torch.matmul(formed_by_row, A.T).detach().cpu().numpy(), "future": torch.matmul(disease_future_by_row, A.T).detach().cpu().numpy(), "total": torch.matmul(disease_total_by_row, A.T).detach().cpu().numpy(), } ) out: list[dict[str, Any]] = [] for row_idx, row in enumerate(rows): for item in projected: matrix = item["matrix"] historical = item["historical"][row_idx] future = item["future"][row_idx] total = item["total"][row_idx] 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 _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 _compute_bi_from_streamed_readouts( *, rows: list[dict[str, Any]], horizon: float, matrices: list[dict[str, Any]], union_disease_ids: np.ndarray, formed_mode: str, readout_batch_size: int, num_workers: int, ctx: Any, log_prefix: str | None = None, ) -> tuple[list[dict[str, Any]], int, dict[str, float]]: horizon = float(horizon) if horizon < 0: raise ValueError(f"horizon must be non-negative, got {horizon}") dtype = torch.float32 if formed_mode == "observed": formed_by_row = torch.as_tensor( _observed_formed_for_rows( rows=rows, union_disease_ids=union_disease_ids, ), dtype=dtype, device=ctx.device, ) survival_by_row = None elif formed_mode == "model_weighted": survival_by_row = torch.ones( (len(rows), union_disease_ids.size), dtype=dtype, device=ctx.device, ) formed_by_row = None else: raise ValueError(f"Unknown formed_mode={formed_mode!r}") future_prob_by_row = torch.zeros( (len(rows), union_disease_ids.size), dtype=dtype, device=ctx.device, ) readout_jobs = 0 n_batches = 0 build_readout_sec = 0.0 forward_sec = 0.0 reduce_sec = 0.0 t_loader0 = time.perf_counter() readout_dataset = ReadoutJobIterableDataset( rows=rows, formed_mode=formed_mode, horizon=horizon, ) readout_loader = DataLoader( readout_dataset, 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, ) build_readout_sec += time.perf_counter() - t_loader0 for batch in readout_loader: t_forward0 = time.perf_counter() readout_prob = _readout_probabilities_from_batch( ctx=ctx, batch=batch, union_disease_ids=union_disease_ids, ) if ctx.device.type == "cuda": torch.cuda.synchronize(ctx.device) forward_sec += time.perf_counter() - t_forward0 t_reduce0 = time.perf_counter() _apply_readout_batch_to_accumulators( batch=batch, readout_prob=readout_prob, survival_by_row=survival_by_row, future_prob_by_row=future_prob_by_row, ctx=ctx, ) if ctx.device.type == "cuda": torch.cuda.synchronize(ctx.device) reduce_sec += time.perf_counter() - t_reduce0 n_batches += 1 readout_jobs += int(batch["row_idx"].numel()) if log_prefix and (n_batches == 1 or n_batches % 50 == 0): print( f"{log_prefix} processed {readout_jobs} readout jobs " f"in {n_batches} batches", flush=True, ) if survival_by_row is not None: formed_by_row = 1.0 - survival_by_row assert formed_by_row is not None t_project0 = time.perf_counter() rows_out = _project_bi_rows( rows=rows, horizon=horizon, matrices=matrices, formed_mode=formed_mode, formed_by_row=formed_by_row, future_prob_by_row=future_prob_by_row, ctx=ctx, ) if ctx.device.type == "cuda": torch.cuda.synchronize(ctx.device) reduce_sec += time.perf_counter() - t_project0 timings = { "build_readout_sec": build_readout_sec, "forward_sec": forward_sec, "reduce_sec": reduce_sec, } return rows_out, readout_jobs, timings def _compute_chunk_worker(payload: dict[str, Any]) -> dict[str, Any]: device = payload["device"] run_path = Path(payload["run_path"]) shard_path = Path(payload["shard_path"]) print( f"[BI worker {device}] starting with {len(payload['row_specs'])} rows", flush=True, ) load_start = time.perf_counter() ctx = load_burden_context(run_path, device=device) print( f"[BI worker {device}] context loaded in {time.perf_counter() - load_start:.2f}s", flush=True, ) materialize_start = time.perf_counter() rows = _materialize_worker_rows(ctx.dataset, payload["row_specs"]) print( f"[BI worker {device}] rows materialized in " f"{time.perf_counter() - materialize_start:.2f}s", flush=True, ) out, readout_jobs, timings = _compute_bi_from_streamed_readouts( rows=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"]), num_workers=int(payload["num_workers"]), ctx=ctx, log_prefix=f"[BI worker {device}]", ) print( f"[BI worker {device}] done: readout_jobs={readout_jobs}, " f"build={timings['build_readout_sec']:.2f}s, " f"forward={timings['forward_sec']:.2f}s, " f"reduce={timings['reduce_sec']:.2f}s", flush=True, ) write_start = time.perf_counter() row_count = _write_rows_csv(out, shard_path) timings["write_csv_sec"] = time.perf_counter() - write_start print( f"[BI worker {device}] wrote {row_count} rows to {shard_path} " f"in {timings['write_csv_sec']:.2f}s", flush=True, ) return { "shard_path": str(shard_path), "row_count": row_count, "readout_jobs": readout_jobs, "timings": timings, } 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], *, formed_mode: str, horizon: float, ) -> 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) weighted_rows = sorted( rows, key=lambda row: _estimate_readout_jobs_for_row( row, formed_mode=formed_mode, horizon=horizon, ), reverse=True, ) for row in weighted_rows: bucket_idx = int(np.argmin(loads)) buckets[bucket_idx].append(row) loads[bucket_idx] += _estimate_readout_jobs_for_row( row, formed_mode=formed_mode, horizon=horizon, ) chunks: list[tuple[str | None, list[dict[str, Any]]]] = [] for device, bucket in zip(devices, buckets): if not bucket: continue chunks.append((device, bucket)) return chunks def _estimate_readout_jobs_for_row( row: dict[str, Any], *, formed_mode: str, horizon: float, ) -> int: n_jobs = 0 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)]) n_jobs += int(np.sum(np.maximum(end_times - grid, 0.0) > 0)) if horizon > 0: n_jobs += 1 return max(n_jobs, 1) 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( "--num_workers", type=int, default=4, help="DataLoader workers per GPU process for readout job generation.", ) 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( "--mp_start_method", type=str, default="fork", choices=["fork", "forkserver"], help="Multiprocessing start method for Linux multi-GPU runs.", ) 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) total_readout_jobs = 0 output_written = False shard_paths: list[Path] = [] total_timings = { "build_readout_sec": 0.0, "forward_sec": 0.0, "reduce_sec": 0.0, "write_csv_sec": 0.0, } row_chunks = _split_rows_for_devices( rows, devices, formed_mode=args.formed_mode, horizon=horizon, ) for device, chunk_rows in row_chunks: estimated_jobs = sum( _estimate_readout_jobs_for_row( row, formed_mode=args.formed_mode, horizon=horizon, ) for row in chunk_rows ) print( f"Assigned {len(chunk_rows)} rows / ~{estimated_jobs} readout jobs to {device}", flush=True, ) if len(row_chunks) == 1: all_rows, total_readout_jobs, timings = _compute_bi_from_streamed_readouts( 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), num_workers=int(args.num_workers), ctx=ctx, log_prefix="[BI main]", ) for key, value in timings.items(): total_timings[key] += float(value) write_start = time.perf_counter() _write_rows_csv(all_rows, output_path) total_timings["write_csv_sec"] = time.perf_counter() - write_start output_written = True 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), "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_rows], "horizon": horizon, "matrices": matrices, "union_disease_ids": union_disease_ids, "readout_batch_size": int(args.readout_batch_size), "num_workers": int(args.num_workers), "formed_mode": args.formed_mode, } for part_idx, (device, chunk_rows) in enumerate(row_chunks) ] with ProcessPoolExecutor( max_workers=len(payloads), mp_context=mp.get_context(args.mp_start_method), ) 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() shard_paths.append(Path(result["shard_path"])) total_readout_jobs += int(result["readout_jobs"]) for key, value in result["timings"].items(): total_timings[key] += float(value) write_start = time.perf_counter() _concat_csv_shards(sorted(shard_paths), output_path) total_timings["write_csv_sec"] += time.perf_counter() - write_start output_written = True if not output_written: raise RuntimeError("BI output was not written.") 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"DataLoader workers per GPU process: {int(args.num_workers)}") print(f"Multiprocessing start method: {args.mp_start_method}") print( "Timing seconds: " f"build_readout={total_timings['build_readout_sec']:.2f}, " f"forward={total_timings['forward_sec']:.2f}, " f"reduce={total_timings['reduce_sec']:.2f}, " f"write_csv={total_timings['write_csv_sec']:.2f}" ) 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()