Keep burden reduction on GPU

compute_burden_index_landmarks.py

@@ -2,6 +2,7 @@ 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
@@ -14,7 +15,6 @@ from tqdm.auto import tqdm
 from burden_index import (
     _build_readout_grid,
     _observed_formed_burden,
-    _probabilities_from_hidden,
     load_burden_context,
 )
 from evaluate_auc_v2 import (
@@ -361,6 +361,44 @@ def _build_readout_table(
     }
 
 
+@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(
     *,
@@ -368,21 +406,21 @@ def _readout_probabilities(
     readout_table: dict[str, Any],
     union_disease_ids: np.ndarray,
     readout_batch_size: int,
-) -> np.ndarray:
+) -> torch.Tensor:
     jobs = readout_table["jobs"]
     if not jobs:
-        return np.zeros((0, union_disease_ids.size), dtype=np.float64)
+        return torch.empty((0, union_disease_ids.size), dtype=torch.float32, device=ctx.device)
 
-    out = np.empty((len(jobs), union_disease_ids.size), dtype=np.float64)
+    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(
+        out[slc] = _probabilities_from_hidden_torch(
             ctx=ctx,
             hidden=hidden,
             disease_ids=union_disease_ids,
             deltas=deltas[slc],
-        )
+        ).to(dtype=out.dtype)
     return out
 
 
@@ -410,49 +448,102 @@ def _reduce_readout_table_to_bi_rows(
     union_disease_ids: np.ndarray,
     formed_mode: str,
     readout_table: dict[str, Any],
-    readout_prob: np.ndarray,
+    readout_prob: torch.Tensor,
+    ctx: Any,
 ) -> list[dict[str, Any]]:
     if formed_mode == "observed":
-        formed_by_row = _observed_formed_for_rows(
+        formed_by_row = torch.as_tensor(
-            rows=rows,
+            _observed_formed_for_rows(
-            union_disease_ids=union_disease_ids,
+                rows=rows,
+                union_disease_ids=union_disease_ids,
+            ),
+            dtype=readout_prob.dtype,
+            device=ctx.device,
         )
     elif formed_mode == "model_weighted":
-        formed_by_row = np.zeros((len(rows), union_disease_ids.size), dtype=np.float64)
+        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 = np.zeros((len(rows), union_disease_ids.size), dtype=np.float64)
+    future_prob_by_row = torch.zeros(
-    row_indices = np.asarray(readout_table["row_indices"], dtype=np.int64)
+        (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 = np.ones((len(rows), union_disease_ids.size), dtype=np.float64)
+        survival_by_row = torch.ones(
+            (len(rows), union_disease_ids.size),
+            dtype=readout_prob.dtype,
+            device=ctx.device,
+        )
     else:
         survival_by_row = None
 
-    for job_idx, row_idx in enumerate(row_indices.tolist()):
+    if readout_prob.numel() > 0:
-        kind = str(kinds[job_idx])
+        kind_is_formed = torch.as_tensor(
-        if kind == "formed" and survival_by_row is not None:
+            np.asarray(kinds == "formed", dtype=np.bool_),
-            survival_by_row[int(row_idx)] *= 1.0 - np.clip(readout_prob[job_idx], 0.0, 1.0)
+            dtype=torch.bool,
-        elif kind == "future":
+            device=ctx.device,
-            future_prob_by_row[int(row_idx)] = readout_prob[job_idx]
+        )
+        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):
-        formed = formed_by_row[row_idx]
+        for item in projected:
-        historical_by_matrix = {
+            matrix = item["matrix"]
-            matrix["burden_type"]: matrix["A_union"] @ formed
+            historical = item["historical"][row_idx]
-            for matrix in matrices
+            future = item["future"][row_idx]
-        }
+            total = item["total"][row_idx]
-        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(
                     {
@@ -485,21 +576,26 @@ def _compute_bi_from_readout_table(
     formed_mode: str,
     readout_batch_size: int,
     ctx: Any,
-) -> tuple[list[dict[str, Any]], int]:
+) -> 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,
@@ -508,15 +604,24 @@ def _compute_bi_from_readout_table(
         formed_mode=formed_mode,
         readout_table=readout_table,
         readout_prob=readout_prob,
+        ctx=ctx,
     )
-    return rows_out, len(readout_table["jobs"])
+    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"])
     ctx = load_burden_context(run_path, device=device)
-    out, readout_jobs = _compute_bi_from_readout_table(
+    out, readout_jobs, timings = _compute_bi_from_readout_table(
         rows=payload["rows"],
         horizon=payload["horizon"],
         matrices=payload["matrices"],
@@ -525,7 +630,7 @@ def _compute_chunk_worker(payload: dict[str, Any]) -> dict[str, Any]:
         readout_batch_size=int(payload["readout_batch_size"]),
         ctx=ctx,
     )
-    return {"rows": out, "readout_jobs": readout_jobs}
+    return {"rows": out, "readout_jobs": readout_jobs, "timings": timings}
 
 
 def _attach_union_projection(
@@ -561,18 +666,64 @@ def _attach_union_projection(
 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)]
-    index_chunks = np.array_split(np.arange(len(rows)), len(devices))
+
+    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, idx in zip(devices, index_chunks):
+    for device, bucket in zip(devices, buckets):
-        if idx.size == 0:
+        if not bucket:
             continue
-        chunks.append((device, [rows[int(i)] for i in idx.tolist()]))
+        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."
@@ -682,9 +833,20 @@ def main() -> None:
 
     all_rows: list[dict[str, Any]] = []
     total_readout_jobs = 0
-    row_chunks = _split_rows_for_devices(rows, devices)
+    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,
+    )
     if len(row_chunks) == 1:
-        all_rows, total_readout_jobs = _compute_bi_from_readout_table(
+        all_rows, total_readout_jobs, timings = _compute_bi_from_readout_table(
             rows=rows,
             horizon=horizon,
             matrices=matrices,
@@ -693,6 +855,8 @@ def main() -> None:
             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.
@@ -724,15 +888,26 @@ def main() -> None:
                 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(
+        "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(