Add mortality attribution evaluation

future_risk.py

@@ -0,0 +1,115 @@
+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))