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 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) 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 _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( *, ctx: Any, readout_table: dict[str, Any], union_disease_ids: np.ndarray, readout_batch_size: int, ) -> torch.Tensor: jobs = readout_table["jobs"] if not jobs: return torch.empty((0, union_disease_ids.size), dtype=torch.float32, device=ctx.device) out = torch.empty((len(jobs), union_disease_ids.size), dtype=torch.float32, device=ctx.device) 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_torch( ctx=ctx, hidden=hidden, disease_ids=union_disease_ids, deltas=deltas[slc], ).to(dtype=out.dtype) 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: torch.Tensor, ctx: Any, ) -> list[dict[str, Any]]: if formed_mode == "observed": formed_by_row = torch.as_tensor( _observed_formed_for_rows( rows=rows, union_disease_ids=union_disease_ids, ), dtype=readout_prob.dtype, device=ctx.device, ) elif formed_mode == "model_weighted": formed_by_row = torch.zeros( (len(rows), union_disease_ids.size), dtype=readout_prob.dtype, device=ctx.device, ) else: raise ValueError(f"Unknown formed_mode={formed_mode!r}") future_prob_by_row = torch.zeros( (len(rows), union_disease_ids.size), dtype=readout_prob.dtype, device=ctx.device, ) row_indices = torch.as_tensor( np.asarray(readout_table["row_indices"], dtype=np.int64), dtype=torch.long, device=ctx.device, ) kinds = np.asarray(readout_table["kinds"], dtype=object) if formed_mode == "model_weighted": survival_by_row = torch.ones( (len(rows), union_disease_ids.size), dtype=readout_prob.dtype, device=ctx.device, ) else: survival_by_row = None if readout_prob.numel() > 0: kind_is_formed = torch.as_tensor( np.asarray(kinds == "formed", dtype=np.bool_), dtype=torch.bool, device=ctx.device, ) kind_is_future = torch.as_tensor( np.asarray(kinds == "future", dtype=np.bool_), dtype=torch.bool, device=ctx.device, ) 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] if survival_by_row is not None: formed_by_row = 1.0 - survival_by_row 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=readout_prob.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 _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, dict[str, float]]: horizon = float(horizon) if horizon < 0: raise ValueError(f"horizon must be non-negative, got {horizon}") t0 = time.perf_counter() readout_table = _build_readout_table( rows=rows, formed_mode=formed_mode, horizon=horizon, ) t1 = time.perf_counter() readout_prob = _readout_probabilities( ctx=ctx, readout_table=readout_table, union_disease_ids=union_disease_ids, readout_batch_size=readout_batch_size, ) if ctx.device.type == "cuda": torch.cuda.synchronize(ctx.device) t2 = time.perf_counter() 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, ctx=ctx, ) if ctx.device.type == "cuda": torch.cuda.synchronize(ctx.device) t3 = time.perf_counter() timings = { "build_readout_sec": t1 - t0, "forward_sec": t2 - t1, "reduce_sec": t3 - t2, } return rows_out, len(readout_table["jobs"]), timings def _compute_chunk_worker(payload: dict[str, Any]) -> dict[str, Any]: device = payload["device"] run_path = Path(payload["run_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_readout_table( 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"]), ctx=ctx, ) 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, ) return {"rows": out, "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("--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) all_rows: list[dict[str, Any]] = [] total_readout_jobs = 0 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_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, ) for key, value in timings.items(): total_timings[key] += float(value) 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), "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), "formed_mode": args.formed_mode, } for device, chunk_rows in 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() all_rows.extend(result["rows"]) total_readout_jobs += int(result["readout_jobs"]) for key, value in result["timings"].items(): total_timings[key] += float(value) write_start = time.perf_counter() out_df = pd.DataFrame(all_rows) out_df.to_csv(output_path, index=False) total_timings["write_csv_sec"] = time.perf_counter() - write_start 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"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()