Add burden index computation interfaces

burden_index.py

@@ -0,0 +1,547 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Literal, Sequence
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+from evaluate_auc_v2 import (
+    build_model_from_dataset,
+    load_checkpoint_state_dict,
+    load_json_config,
+    load_model_state,
+    resolve_dist_mode_for_checkpoint,
+    validate_dataset_metadata,
+)
+from eval_data import load_sequence_eval_dataset
+from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX
+
+
+FormedBurdenMode = Literal["observed", "model_weighted"]
+
+
+@dataclass(frozen=True)
+class BurdenContext:
+    model: torch.nn.Module
+    dataset: Any
+    cfg: dict[str, Any]
+    dist_mode: str
+    device: torch.device
+    run_path: Path
+
+
+@dataclass(frozen=True)
+class DiseaseBurdenResult:
+    disease_id: int
+    formed: float
+    future: float
+    total: float
+    formed_mode: str
+    horizon: float
+
+
+@dataclass(frozen=True)
+class BurdenIndexResult:
+    historical: np.ndarray
+    future: np.ndarray
+    total: np.ndarray
+    disease_ids: np.ndarray
+    formed: np.ndarray
+    disease_future: np.ndarray
+    disease_total: np.ndarray
+    formed_mode: str
+    horizon: float
+
+
+def load_burden_context(
+    run_path: str | Path,
+    *,
+    device: str | torch.device | None = None,
+) -> BurdenContext:
+    run_path = Path(run_path)
+    config_path = run_path / "train_config.json"
+    model_ckpt_path = run_path / "best_model.pt"
+    if not config_path.exists():
+        raise FileNotFoundError(f"train_config.json not found in {run_path}")
+    if not model_ckpt_path.exists():
+        raise FileNotFoundError(f"best_model.pt not found in {run_path}")
+
+    cfg = load_json_config(config_path)
+    model_target_mode = str(cfg.get("model_target_mode", "next_token")).lower()
+    if model_target_mode == "next_token":
+        raise RuntimeError(
+            "Burden Index computation requires an all_future checkpoint because "
+            "it uses p_d(h, Delta). The provided run is model_target_mode='next_token'."
+        )
+    if model_target_mode != "all_future":
+        raise ValueError(
+            "train_config.json model_target_mode must be all_future for burden "
+            f"computation, got {model_target_mode!r}."
+        )
+
+    device_obj = torch.device(
+        device if device is not None else ("cuda" if torch.cuda.is_available() else "cpu")
+    )
+    if device_obj.type == "cuda" and not torch.cuda.is_available():
+        raise RuntimeError(f"Requested device {device_obj}, but CUDA is not available.")
+
+    data_prefix = cfg.get("data_prefix", "ukb")
+    labels_file = cfg.get("labels_file", "labels.csv")
+    extra_info_types = cfg.get("extra_info_types", None)
+    dataset = load_sequence_eval_dataset(
+        model_target_mode="all_future",
+        data_prefix=data_prefix,
+        labels_file=labels_file,
+        no_event_interval_years=float(cfg.get("no_event_interval_years", 5.0)),
+        include_no_event_in_uts_target=bool(
+            cfg.get("include_no_event_in_uts_target", False)
+        ),
+        min_history_events=int(cfg.get("all_future_min_history_events", 1)),
+        min_future_events=int(cfg.get("all_future_min_future_events", 1)),
+        extra_info_types=extra_info_types,
+    )
+    validate_dataset_metadata(dataset, cfg)
+
+    state_dict = load_checkpoint_state_dict(model_ckpt_path, map_location="cpu")
+    dist_mode = resolve_dist_mode_for_checkpoint(
+        str(cfg.get("dist_mode", "exponential")),
+        state_dict,
+    )
+    if dist_mode not in {"exponential", "weibull", "mixed"}:
+        raise ValueError(f"Unsupported dist_mode={dist_mode!r}")
+
+    cfg_model = dict(cfg)
+    cfg_model["dist_mode"] = dist_mode
+    args = _ConfigNamespace()
+    model = build_model_from_dataset(args, cfg_model, dataset)
+    load_model_state(model, state_dict)
+    model.eval().to(device_obj)
+
+    return BurdenContext(
+        model=model,
+        dataset=dataset,
+        cfg=cfg,
+        dist_mode=dist_mode,
+        device=device_obj,
+        run_path=run_path,
+    )
+
+
+@torch.inference_mode()
+def compute_disease_burden(
+    *,
+    run_path: str | Path,
+    disease_id: int,
+    event_seq: Sequence[int] | np.ndarray,
+    time_seq: Sequence[float] | np.ndarray,
+    sex: int,
+    other_type: Sequence[int] | np.ndarray,
+    other_value: Sequence[float] | np.ndarray,
+    other_value_kind: Sequence[int] | np.ndarray,
+    other_time: Sequence[float] | np.ndarray,
+    t_query: float,
+    horizon: float,
+    formed_mode: FormedBurdenMode,
+    device: str | torch.device | None = None,
+    context: BurdenContext | None = None,
+) -> DiseaseBurdenResult:
+    ctx = context or load_burden_context(run_path, device=device)
+    disease_ids = np.asarray([int(disease_id)], dtype=np.int64)
+    formed, future_prob = _compute_disease_components(
+        ctx=ctx,
+        disease_ids=disease_ids,
+        event_seq=event_seq,
+        time_seq=time_seq,
+        sex=sex,
+        other_type=other_type,
+        other_value=other_value,
+        other_value_kind=other_value_kind,
+        other_time=other_time,
+        t_query=float(t_query),
+        horizon=float(horizon),
+        formed_mode=formed_mode,
+    )
+    future = (1.0 - formed) * future_prob
+    total = formed + future
+    return DiseaseBurdenResult(
+        disease_id=int(disease_id),
+        formed=float(formed[0]),
+        future=float(future[0]),
+        total=float(total[0]),
+        formed_mode=str(formed_mode),
+        horizon=float(horizon),
+    )
+
+
+@torch.inference_mode()
+def compute_burden_index(
+    *,
+    run_path: str | Path,
+    burden_matrix: np.ndarray,
+    disease_ids: Sequence[int] | np.ndarray,
+    event_seq: Sequence[int] | np.ndarray,
+    time_seq: Sequence[float] | np.ndarray,
+    sex: int,
+    other_type: Sequence[int] | np.ndarray,
+    other_value: Sequence[float] | np.ndarray,
+    other_value_kind: Sequence[int] | np.ndarray,
+    other_time: Sequence[float] | np.ndarray,
+    t_query: float,
+    horizon: float,
+    formed_mode: FormedBurdenMode,
+    device: str | torch.device | None = None,
+    context: BurdenContext | None = None,
+) -> BurdenIndexResult:
+    ctx = context or load_burden_context(run_path, device=device)
+    disease_ids_arr = np.asarray(disease_ids, dtype=np.int64)
+    if disease_ids_arr.ndim != 1 or disease_ids_arr.size == 0:
+        raise ValueError("disease_ids must be a non-empty 1D sequence.")
+
+    A = np.asarray(burden_matrix, dtype=np.float64)
+    if A.ndim != 2:
+        raise ValueError(f"burden_matrix must be 2D, got shape {A.shape}")
+    if A.shape[1] != disease_ids_arr.size:
+        raise ValueError(
+            "burden_matrix columns must match disease_ids length, got "
+            f"{A.shape[1]} columns vs {disease_ids_arr.size} disease ids."
+        )
+
+    formed, future_prob = _compute_disease_components(
+        ctx=ctx,
+        disease_ids=disease_ids_arr,
+        event_seq=event_seq,
+        time_seq=time_seq,
+        sex=sex,
+        other_type=other_type,
+        other_value=other_value,
+        other_value_kind=other_value_kind,
+        other_time=other_time,
+        t_query=float(t_query),
+        horizon=float(horizon),
+        formed_mode=formed_mode,
+    )
+    disease_future = (1.0 - formed) * future_prob
+    disease_total = formed + disease_future
+
+    historical = A @ formed
+    future = A @ disease_future
+    total = A @ disease_total
+    return BurdenIndexResult(
+        historical=historical,
+        future=future,
+        total=total,
+        disease_ids=disease_ids_arr.copy(),
+        formed=formed,
+        disease_future=disease_future,
+        disease_total=disease_total,
+        formed_mode=str(formed_mode),
+        horizon=float(horizon),
+    )
+
+
+class _ConfigNamespace:
+    def __getattr__(self, _name: str) -> None:
+        return None
+
+
+def _compute_disease_components(
+    *,
+    ctx: BurdenContext,
+    disease_ids: np.ndarray,
+    event_seq: Sequence[int] | np.ndarray,
+    time_seq: Sequence[float] | np.ndarray,
+    sex: int,
+    other_type: Sequence[int] | np.ndarray,
+    other_value: Sequence[float] | np.ndarray,
+    other_value_kind: Sequence[int] | np.ndarray,
+    other_time: Sequence[float] | np.ndarray,
+    t_query: float,
+    horizon: float,
+    formed_mode: FormedBurdenMode,
+) -> tuple[np.ndarray, np.ndarray]:
+    formed_mode = _validate_formed_mode(formed_mode)
+    disease_ids = np.asarray(disease_ids, dtype=np.int64)
+    _validate_disease_ids(ctx, disease_ids)
+    event_seq_arr, time_seq_arr = _validate_event_inputs(event_seq, time_seq)
+    other_type_arr, other_value_arr, other_value_kind_arr, other_time_arr = (
+        _validate_other_inputs(other_type, other_value, other_value_kind, other_time)
+    )
+    if horizon < 0:
+        raise ValueError(f"horizon must be non-negative, got {horizon}")
+
+    if formed_mode == "observed":
+        formed = _observed_formed_burden(
+            disease_ids=disease_ids,
+            event_seq=event_seq_arr,
+            time_seq=time_seq_arr,
+            t_query=t_query,
+        )
+    else:
+        formed = _model_weighted_formed_burden(
+            ctx=ctx,
+            disease_ids=disease_ids,
+            event_seq=event_seq_arr,
+            time_seq=time_seq_arr,
+            sex=sex,
+            other_type=other_type_arr,
+            other_value=other_value_arr,
+            other_value_kind=other_value_kind_arr,
+            other_time=other_time_arr,
+            t_query=t_query,
+        )
+
+    hidden_query = _query_hidden(
+        ctx=ctx,
+        event_seq=event_seq_arr,
+        time_seq=time_seq_arr,
+        sex=sex,
+        other_type=other_type_arr,
+        other_value=other_value_arr,
+        other_value_kind=other_value_kind_arr,
+        other_time=other_time_arr,
+        query_times=np.asarray([t_query], dtype=np.float32),
+    )
+    future_prob = _probabilities_from_hidden(
+        ctx=ctx,
+        hidden=hidden_query,
+        disease_ids=disease_ids,
+        deltas=np.asarray([horizon], dtype=np.float32),
+    )[0]
+    return formed.astype(np.float64), future_prob.astype(np.float64)
+
+
+def _model_weighted_formed_burden(
+    *,
+    ctx: BurdenContext,
+    disease_ids: np.ndarray,
+    event_seq: np.ndarray,
+    time_seq: np.ndarray,
+    sex: int,
+    other_type: np.ndarray,
+    other_value: np.ndarray,
+    other_value_kind: np.ndarray,
+    other_time: np.ndarray,
+    t_query: float,
+) -> np.ndarray:
+    grid = _build_readout_grid(
+        event_seq=event_seq,
+        time_seq=time_seq,
+        other_type=other_type,
+        other_time=other_time,
+        t_query=t_query,
+    )
+    if grid.size == 0:
+        return np.zeros(disease_ids.size, dtype=np.float64)
+
+    end_times = np.concatenate([grid[1:], np.asarray([t_query], dtype=np.float32)])
+    deltas = np.maximum(end_times - grid, 0.0).astype(np.float32)
+    if np.all(deltas <= 0):
+        return np.zeros(disease_ids.size, dtype=np.float64)
+
+    hidden = _query_hidden(
+        ctx=ctx,
+        event_seq=event_seq,
+        time_seq=time_seq,
+        sex=sex,
+        other_type=other_type,
+        other_value=other_value,
+        other_value_kind=other_value_kind,
+        other_time=other_time,
+        query_times=grid.astype(np.float32),
+    )
+    interval_prob = _probabilities_from_hidden(
+        ctx=ctx,
+        hidden=hidden,
+        disease_ids=disease_ids,
+        deltas=deltas,
+    ).astype(np.float64)
+    survival = np.prod(1.0 - np.clip(interval_prob, 0.0, 1.0), axis=0)
+    return 1.0 - survival
+
+
+def _observed_formed_burden(
+    *,
+    disease_ids: np.ndarray,
+    event_seq: np.ndarray,
+    time_seq: np.ndarray,
+    t_query: float,
+) -> np.ndarray:
+    valid = (
+        (time_seq <= np.float32(t_query))
+        & (event_seq > PAD_IDX)
+        & (event_seq != CHECKUP_IDX)
+        & (event_seq != NO_EVENT_IDX)
+    )
+    observed = set(int(x) for x in event_seq[valid].tolist())
+    return np.asarray([1.0 if int(d) in observed else 0.0 for d in disease_ids])
+
+
+def _build_readout_grid(
+    *,
+    event_seq: np.ndarray,
+    time_seq: np.ndarray,
+    other_type: np.ndarray,
+    other_time: np.ndarray,
+    t_query: float,
+) -> np.ndarray:
+    event_mask = (event_seq > PAD_IDX) & (time_seq <= np.float32(t_query))
+    other_mask = (other_type > 0) & (other_time <= np.float32(t_query))
+    times = np.concatenate(
+        [
+            time_seq[event_mask].astype(np.float32, copy=False),
+            other_time[other_mask].astype(np.float32, copy=False),
+        ]
+    )
+    if times.size == 0:
+        return np.zeros(0, dtype=np.float32)
+    return np.unique(times)
+
+
+def _query_hidden(
+    *,
+    ctx: BurdenContext,
+    event_seq: np.ndarray,
+    time_seq: np.ndarray,
+    sex: int,
+    other_type: np.ndarray,
+    other_value: np.ndarray,
+    other_value_kind: np.ndarray,
+    other_time: np.ndarray,
+    query_times: np.ndarray,
+) -> torch.Tensor:
+    if query_times.ndim != 1:
+        raise ValueError("query_times must be 1D.")
+    batch_size = int(query_times.size)
+    if batch_size == 0:
+        return torch.empty(0, ctx.model.n_embd, device=ctx.device)
+
+    event = torch.from_numpy(event_seq[None, :].repeat(batch_size, axis=0)).long()
+    times = torch.from_numpy(time_seq[None, :].repeat(batch_size, axis=0)).float()
+    other_t = torch.from_numpy(other_type[None, :].repeat(batch_size, axis=0)).long()
+    other_v = torch.from_numpy(other_value[None, :].repeat(batch_size, axis=0)).float()
+    other_k = torch.from_numpy(
+        other_value_kind[None, :].repeat(batch_size, axis=0)
+    ).long()
+    other_tm = torch.from_numpy(other_time[None, :].repeat(batch_size, axis=0)).float()
+    sex_t = torch.full((batch_size,), int(sex), dtype=torch.long)
+    tq = torch.from_numpy(query_times.astype(np.float32, copy=False)).float()
+
+    event = event.to(ctx.device)
+    times = times.to(ctx.device)
+    other_t = other_t.to(ctx.device)
+    other_v = other_v.to(ctx.device)
+    other_k = other_k.to(ctx.device)
+    other_tm = other_tm.to(ctx.device)
+    sex_t = sex_t.to(ctx.device)
+    tq = tq.to(ctx.device)
+
+    return ctx.model(
+        event_seq=event,
+        time_seq=times,
+        sex=sex_t,
+        padding_mask=event > PAD_IDX,
+        t_query=tq,
+        other_type=other_t,
+        other_value=other_v,
+        other_value_kind=other_k,
+        other_time=other_tm,
+        target_mode="all_future",
+    )
+
+
+def _probabilities_from_hidden(
+    *,
+    ctx: BurdenContext,
+    hidden: torch.Tensor,
+    disease_ids: np.ndarray,
+    deltas: np.ndarray,
+) -> np.ndarray:
+    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 = F.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)
+
+    prob = -torch.expm1(-rate * exposure)
+    return prob.detach().cpu().numpy().astype(np.float64, copy=False)
+
+
+def _validate_formed_mode(mode: str) -> FormedBurdenMode:
+    if mode not in {"observed", "model_weighted"}:
+        raise ValueError(
+            "formed_mode must be either 'observed' or 'model_weighted', "
+            f"got {mode!r}."
+        )
+    return mode  # type: ignore[return-value]
+
+
+def _validate_disease_ids(ctx: BurdenContext, disease_ids: np.ndarray) -> None:
+    if disease_ids.ndim != 1 or disease_ids.size == 0:
+        raise ValueError("disease_ids must be a non-empty 1D array.")
+    vocab_size = int(getattr(ctx.model, "vocab_size", ctx.model.risk_head.out_features))
+    if np.any(disease_ids < 0) or np.any(disease_ids >= vocab_size):
+        raise ValueError(f"disease_ids must be in [0, {vocab_size}), got {disease_ids}")
+
+
+def _validate_event_inputs(
+    event_seq: Sequence[int] | np.ndarray,
+    time_seq: Sequence[float] | np.ndarray,
+) -> tuple[np.ndarray, np.ndarray]:
+    events = np.asarray(event_seq, dtype=np.int64)
+    times = np.asarray(time_seq, dtype=np.float32)
+    if events.ndim != 1 or times.ndim != 1:
+        raise ValueError("event_seq and time_seq must be 1D.")
+    if events.shape != times.shape:
+        raise ValueError(
+            f"event_seq and time_seq must have the same shape, got {events.shape} vs {times.shape}"
+        )
+    if events.size == 0:
+        raise ValueError("event_seq must contain at least one token.")
+    if not np.all(np.isfinite(times)):
+        raise ValueError("time_seq contains non-finite values.")
+    return events, times
+
+
+def _validate_other_inputs(
+    other_type: Sequence[int] | np.ndarray,
+    other_value: Sequence[float] | np.ndarray,
+    other_value_kind: Sequence[int] | np.ndarray,
+    other_time: Sequence[float] | np.ndarray,
+) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+    typ = np.asarray(other_type, dtype=np.int64)
+    val = np.asarray(other_value, dtype=np.float32)
+    kind = np.asarray(other_value_kind, dtype=np.int64)
+    tm = np.asarray(other_time, dtype=np.float32)
+    if not (typ.shape == val.shape == kind.shape == tm.shape):
+        raise ValueError(
+            "other_type, other_value, other_value_kind, and other_time must "
+            f"have the same shape, got {typ.shape}, {val.shape}, {kind.shape}, {tm.shape}."
+        )
+    if typ.ndim != 1:
+        raise ValueError("other_* inputs must be 1D.")
+    if not np.all(np.isfinite(tm)):
+        raise ValueError("other_time contains non-finite values.")
+    return typ, val, kind, tm