From 63df61e1dd29e0ac5d046f757585fe012cf1539b Mon Sep 17 00:00:00 2001 From: Jiarui Li Date: Fri, 26 Jun 2026 11:15:15 +0800 Subject: [PATCH] Batch burden index readout computation --- compute_burden_index_landmarks.py | 449 +++++++++++++++++++++++------- 1 file changed, 356 insertions(+), 93 deletions(-) diff --git a/compute_burden_index_landmarks.py b/compute_burden_index_landmarks.py index 189b974..4599e77 100644 --- a/compute_burden_index_landmarks.py +++ b/compute_burden_index_landmarks.py @@ -8,9 +8,15 @@ from typing import Any, Iterable import numpy as np import pandas as pd +import torch from tqdm.auto import tqdm -from burden_index import compute_burden_index, load_burden_context +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, @@ -34,13 +40,16 @@ def _parse_landmark_ages(args: argparse.Namespace) -> np.ndarray: 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_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]: @@ -241,80 +250,312 @@ def _config_split_indices( return make_eval_indices(_Sized(), args, cfg) -def _result_rows_for_sample( +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( *, - 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)]), - } + 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 _compute_chunk_worker(payload: dict[str, Any]) -> list[dict[str, Any]]: +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: 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 + 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( @@ -348,13 +589,30 @@ def main() -> None: 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( + "--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", @@ -393,8 +651,9 @@ def main() -> None: "category_meta": category_meta, } ) + union_disease_ids, matrices = _attach_union_projection(matrices) landmark_ages = _parse_landmark_ages(args) - horizons = _parse_horizons(args.horizons) + horizon = _parse_horizon(args.horizon) eval_split = str(args.eval_split).lower() if eval_split in {"valid", "validation"}: eval_split = "val" @@ -416,28 +675,24 @@ def main() -> None: ) output_path = Path(args.output_path) if args.output_path else ( - run_path / f"burden_index_{eval_split}_{args.formed_mode}.csv" + 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: - 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, - ) - ) + 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. @@ -447,8 +702,10 @@ def main() -> None: "device": device, "run_path": str(run_path), "rows": chunk_rows, - "horizons": horizons.tolist(), + "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 @@ -464,19 +721,25 @@ def main() -> None: desc="Computing BI chunks", dynamic_ncols=True, ): - all_rows.extend(future.result()) + 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}")