116 lines
4.4 KiB
Python
116 lines
4.4 KiB
Python
from __future__ import annotations
|
|
|
|
from collections.abc import Sequence
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
|
|
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 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.
|
|
|
|
Death is always treated as token vocab_size - 1. For dist_mode="mixed",
|
|
non-death tokens use exponential hazards and death uses death_rho.
|
|
"""
|
|
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 death_risk_from_probabilities(probabilities: torch.Tensor) -> torch.Tensor:
|
|
"""Return p_death(t, tau), with death fixed to token vocab_size - 1."""
|
|
if probabilities.ndim != 2:
|
|
raise ValueError(
|
|
f"probabilities must have shape (N, V), got {tuple(probabilities.shape)}"
|
|
)
|
|
return probabilities[:, death_token(probabilities.shape[1])]
|
|
|
|
|
|
def new_disease_risk_from_probabilities(
|
|
probabilities: torch.Tensor,
|
|
occurred: torch.Tensor,
|
|
disease_ids: Sequence[int],
|
|
) -> torch.Tensor:
|
|
"""
|
|
Compute P(at least one selected disease newly occurs within tau years).
|
|
|
|
Already occurred diseases are masked out. Death is not included here and
|
|
should be reported separately with death_risk_from_probabilities.
|
|
"""
|
|
if probabilities.ndim != 2 or occurred.shape != probabilities.shape:
|
|
raise ValueError(
|
|
"probabilities and occurred must both have shape (N, V), got "
|
|
f"{tuple(probabilities.shape)} and {tuple(occurred.shape)}"
|
|
)
|
|
if not disease_ids:
|
|
return probabilities.new_zeros(probabilities.shape[0])
|
|
|
|
death_idx = death_token(probabilities.shape[1])
|
|
ids = [
|
|
idx
|
|
for idx in dict.fromkeys(int(x) for x in disease_ids)
|
|
if 0 <= idx < probabilities.shape[1] and idx != death_idx
|
|
]
|
|
if not ids:
|
|
return probabilities.new_zeros(probabilities.shape[0])
|
|
|
|
idx_tensor = torch.as_tensor(ids, dtype=torch.long, device=probabilities.device)
|
|
p = probabilities[:, idx_tensor].clamp(0.0, 1.0 - 1e-7)
|
|
new_mask = ~occurred[:, idx_tensor].to(dtype=torch.bool)
|
|
log_no_new = torch.log1p(-p) * new_mask.to(dtype=p.dtype)
|
|
return -torch.expm1(log_no_new.sum(dim=1))
|