Batch burden index readout computation

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}")