270 lines
9.3 KiB
Python
270 lines
9.3 KiB
Python
from __future__ import annotations
|
|
|
|
from collections.abc import Sequence
|
|
from typing import Any, overload
|
|
|
|
import numpy as np
|
|
|
|
try:
|
|
import torch
|
|
import torch.nn.functional as F
|
|
except ModuleNotFoundError: # pragma: no cover - optional for numpy-only use
|
|
torch = None
|
|
F = None
|
|
|
|
|
|
ArrayLike = Any
|
|
|
|
|
|
def _death_token(vocab_size: int) -> int:
|
|
if int(vocab_size) <= 0:
|
|
raise ValueError(f"vocab_size must be positive, got {vocab_size}")
|
|
return int(vocab_size) - 1
|
|
|
|
|
|
def _infer_vocab_size(x: ArrayLike, vocab_size: int | None) -> int:
|
|
if x.ndim != 2:
|
|
raise ValueError(f"Expected a 2D array/tensor with shape (N, V), got {tuple(x.shape)}")
|
|
inferred = int(x.shape[1])
|
|
if vocab_size is None:
|
|
return inferred
|
|
if int(vocab_size) != inferred:
|
|
raise ValueError(f"vocab_size={vocab_size} does not match input width {inferred}")
|
|
return int(vocab_size)
|
|
|
|
|
|
def _normalize_disease_ids(
|
|
disease_ids: Sequence[int] | np.ndarray | torch.Tensor | None,
|
|
*,
|
|
vocab_size: int,
|
|
excluded_token_ids: Sequence[int],
|
|
) -> list[int]:
|
|
death_idx = _death_token(vocab_size)
|
|
excluded = {
|
|
int(idx)
|
|
for idx in excluded_token_ids
|
|
if 0 <= int(idx) < vocab_size
|
|
}
|
|
excluded.add(death_idx)
|
|
if disease_ids is None:
|
|
return [idx for idx in range(vocab_size) if idx not in excluded]
|
|
|
|
if torch is not None and isinstance(disease_ids, torch.Tensor):
|
|
raw = disease_ids.detach().cpu().reshape(-1).tolist()
|
|
else:
|
|
raw = np.asarray(disease_ids).reshape(-1).tolist()
|
|
|
|
out: list[int] = []
|
|
seen: set[int] = set()
|
|
for value in raw:
|
|
idx = int(value)
|
|
if idx < 0 or idx >= vocab_size:
|
|
raise ValueError(f"disease id {idx} is outside [0, {vocab_size})")
|
|
if idx in excluded:
|
|
continue
|
|
if idx not in seen:
|
|
seen.add(idx)
|
|
out.append(idx)
|
|
return out
|
|
|
|
|
|
@overload
|
|
def future_event_free_survival_from_probabilities(
|
|
probabilities: torch.Tensor,
|
|
occurred: torch.Tensor,
|
|
disease_ids: Sequence[int] | np.ndarray | torch.Tensor | None = None,
|
|
*,
|
|
vocab_size: int | None = None,
|
|
excluded_token_ids: Sequence[int] = (0, 1, 2),
|
|
eps: float = 1e-7,
|
|
) -> torch.Tensor:
|
|
...
|
|
|
|
|
|
@overload
|
|
def future_event_free_survival_from_probabilities(
|
|
probabilities: np.ndarray,
|
|
occurred: np.ndarray,
|
|
disease_ids: Sequence[int] | np.ndarray | torch.Tensor | None = None,
|
|
*,
|
|
vocab_size: int | None = None,
|
|
excluded_token_ids: Sequence[int] = (0, 1, 2),
|
|
eps: float = 1e-7,
|
|
) -> np.ndarray:
|
|
...
|
|
|
|
|
|
def future_event_free_survival_from_probabilities(
|
|
probabilities: ArrayLike,
|
|
occurred: ArrayLike,
|
|
disease_ids: Sequence[int] | np.ndarray | torch.Tensor | None = None,
|
|
*,
|
|
vocab_size: int | None = None,
|
|
excluded_token_ids: Sequence[int] = (0, 1, 2),
|
|
eps: float = 1e-7,
|
|
) -> ArrayLike:
|
|
"""
|
|
Compute P(alive and no new selected disease in the next tau years).
|
|
|
|
Parameters
|
|
----------
|
|
probabilities:
|
|
Matrix with shape (N, V). Entry (i, d) is p_d(t, tau), the model's
|
|
future first-occurrence probability for token d over the chosen tau.
|
|
The death probability is always read from token V - 1.
|
|
|
|
occurred:
|
|
Boolean matrix with shape (N, V). Entry (i, d) is True if disease d has
|
|
already occurred at or before query time t. Already occurred diseases do
|
|
not contribute to "new disease" risk.
|
|
|
|
disease_ids:
|
|
Optional subset of disease tokens. If None, all non-death tokens are
|
|
included except excluded_token_ids. If provided, death and excluded
|
|
tokens are ignored here and death is still handled separately as
|
|
survival.
|
|
|
|
vocab_size:
|
|
Optional vocabulary size. If omitted, inferred from probabilities.
|
|
|
|
excluded_token_ids:
|
|
Technical tokens to exclude from "new disease" calculations. Defaults
|
|
to (0, 1, 2), matching PAD, CHECKUP, and NO_EVENT.
|
|
|
|
Returns
|
|
-------
|
|
Array/tensor with shape (N,):
|
|
Approximate probability of being alive and having no newly occurring
|
|
disease among the selected disease tokens over the same tau horizon.
|
|
"""
|
|
vocab_size = _infer_vocab_size(probabilities, vocab_size)
|
|
death_idx = _death_token(vocab_size)
|
|
selected = _normalize_disease_ids(
|
|
disease_ids,
|
|
vocab_size=vocab_size,
|
|
excluded_token_ids=excluded_token_ids,
|
|
)
|
|
|
|
if tuple(occurred.shape) != tuple(probabilities.shape):
|
|
raise ValueError(
|
|
"occurred must have the same shape as probabilities, got "
|
|
f"{tuple(occurred.shape)} vs {tuple(probabilities.shape)}"
|
|
)
|
|
|
|
if torch is not None and isinstance(probabilities, torch.Tensor):
|
|
probs = probabilities.clamp(min=0.0, max=1.0 - float(eps))
|
|
occurred_bool = occurred.to(device=probs.device, dtype=torch.bool)
|
|
log_survival = torch.log1p(-probs[:, death_idx])
|
|
if selected:
|
|
ids = torch.as_tensor(selected, dtype=torch.long, device=probs.device)
|
|
new_mask = ~occurred_bool[:, ids]
|
|
log_no_new = torch.log1p(-probs[:, ids]) * new_mask.to(probs.dtype)
|
|
log_survival = log_survival + log_no_new.sum(dim=1)
|
|
return torch.exp(log_survival)
|
|
|
|
probs_np = np.clip(np.asarray(probabilities), 0.0, 1.0 - float(eps))
|
|
occurred_bool_np = np.asarray(occurred, dtype=bool)
|
|
log_survival_np = np.log1p(-probs_np[:, death_idx])
|
|
if selected:
|
|
selected_arr = np.asarray(selected, dtype=np.int64)
|
|
new_mask_np = ~occurred_bool_np[:, selected_arr]
|
|
log_no_new_np = np.log1p(-probs_np[:, selected_arr]) * new_mask_np
|
|
log_survival_np = log_survival_np + log_no_new_np.sum(axis=1)
|
|
return np.exp(log_survival_np)
|
|
|
|
|
|
def probabilities_from_logits(
|
|
logits: torch.Tensor,
|
|
tau_years: float | torch.Tensor,
|
|
*,
|
|
dist_mode: str = "exponential",
|
|
rho: torch.Tensor | None = None,
|
|
death_rho: torch.Tensor | None = None,
|
|
eps: float = 1e-8,
|
|
) -> torch.Tensor:
|
|
"""
|
|
Convert all-future logits to tau-year event probabilities.
|
|
|
|
The death token is always treated as vocab_size - 1. For dist_mode="mixed",
|
|
non-death tokens use exponential hazards and the death token uses
|
|
death_rho. For dist_mode="weibull", rho must have the same shape as logits.
|
|
"""
|
|
if torch is None or F is None:
|
|
raise ImportError("probabilities_from_logits requires PyTorch.")
|
|
if logits.ndim != 2:
|
|
raise ValueError(f"logits must have shape (N, V), got {tuple(logits.shape)}")
|
|
if float(torch.as_tensor(tau_years).detach().min().cpu()) < 0:
|
|
raise ValueError("tau_years must be non-negative")
|
|
|
|
mode = str(dist_mode).lower()
|
|
if mode not in {"exponential", "weibull", "mixed"}:
|
|
raise ValueError("dist_mode must be one of: exponential, weibull, mixed")
|
|
|
|
rate = F.softplus(logits) + float(eps)
|
|
tau = torch.as_tensor(tau_years, dtype=rate.dtype, device=rate.device)
|
|
if tau.ndim == 0:
|
|
tau = tau.expand(logits.shape[0])
|
|
if tau.ndim != 1 or tau.shape[0] != logits.shape[0]:
|
|
raise ValueError(
|
|
"tau_years must be a scalar or a 1D tensor with length N, got "
|
|
f"{tuple(tau.shape)} for N={logits.shape[0]}"
|
|
)
|
|
|
|
if mode == "exponential":
|
|
exposure = tau[:, None].expand_as(rate)
|
|
elif mode == "weibull":
|
|
if rho is None or rho.shape != logits.shape:
|
|
raise ValueError("rho must have the same shape as logits for dist_mode='weibull'")
|
|
exposure = torch.pow(tau[:, None].clamp_min(float(eps)), rho.to(rate.dtype))
|
|
else:
|
|
exposure = tau[:, None].expand_as(rate).clone()
|
|
if death_rho is None:
|
|
raise ValueError("death_rho is required for dist_mode='mixed'")
|
|
death_idx = _death_token(logits.shape[1])
|
|
death_shape = tuple(death_rho.shape)
|
|
death_rho = death_rho.to(device=rate.device, dtype=rate.dtype)
|
|
if death_rho.ndim == 2 and death_rho.shape[1] == 1:
|
|
death_rho = death_rho.squeeze(1)
|
|
if death_rho.ndim != 1 or death_rho.shape[0] != logits.shape[0]:
|
|
raise ValueError(
|
|
"death_rho must have shape (N,) or (N, 1), got "
|
|
f"{death_shape} for N={logits.shape[0]}"
|
|
)
|
|
exposure[:, death_idx] = torch.pow(tau.clamp_min(float(eps)), death_rho)
|
|
|
|
return -torch.expm1(-rate * exposure)
|
|
|
|
|
|
def future_event_free_survival_from_logits(
|
|
logits: torch.Tensor,
|
|
occurred: torch.Tensor,
|
|
tau_years: float | torch.Tensor,
|
|
disease_ids: Sequence[int] | np.ndarray | torch.Tensor | None = None,
|
|
*,
|
|
dist_mode: str = "exponential",
|
|
rho: torch.Tensor | None = None,
|
|
death_rho: torch.Tensor | None = None,
|
|
eps: float = 1e-8,
|
|
) -> torch.Tensor:
|
|
"""
|
|
Convenience wrapper for computing future event-free survival from logits.
|
|
|
|
Returns P(alive and no new selected disease in the next tau years), with
|
|
death fixed to token vocab_size - 1.
|
|
"""
|
|
probabilities = probabilities_from_logits(
|
|
logits=logits,
|
|
tau_years=tau_years,
|
|
dist_mode=dist_mode,
|
|
rho=rho,
|
|
death_rho=death_rho,
|
|
eps=eps,
|
|
)
|
|
return future_event_free_survival_from_probabilities(
|
|
probabilities=probabilities,
|
|
occurred=occurred,
|
|
disease_ids=disease_ids,
|
|
vocab_size=logits.shape[1],
|
|
excluded_token_ids=excluded_token_ids,
|
|
)
|