JiaruiLi/DeepHealth
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add mortality attribution evaluation

Browse Source
This commit is contained in:
Jiarui Li
2026-06-27 13:34:30 +08:00
parent 00de2ff4a9
commit 0b7c866292
5 changed files with 403 additions and 372 deletions
Download Patch File Download Diff File Expand all files Collapse all files

BIN
all_future_model_interpretation.pdf
View File

Binary file not shown.

209
all_future_model_interpretation.tex
View File

@@ -37,7 +37,7 @@ all-future DeepHealth 模型应被解释为一个基于既往轨迹的未来新
\section{English}
\subsection{Model object}
\subsection{Model Object}
For individual \(i\), let \(\mathcal{H}_i(t)\) denote the observed history up to
query time \(t\). The all-future model produces
@@ -68,7 +68,7 @@ history itself we know
This historical indicator is not learned by the model; it is read directly from
the event history.
\subsection{Masking already occurred diseases}
\subsection{Masking Already Occurred Diseases}
For a disease that has already occurred before or at \(t\), the model output for
that disease should not be interpreted as recurrence risk or current disease
@@ -77,10 +77,10 @@ should be masked:
\[
p^{\mathrm{new}}_{i,d}(t,\tau)
=
\left[1-m_{i,d}(t)\right]p_{i,d}(t,\tau).
[1-m_{i,d}(t)]p_{i,d}(t,\tau).
\]
\subsection{Directly supported model-derived quantities}
\subsection{Directly Supported Quantities}
The current model directly supports the following quantities.
@@ -100,42 +100,64 @@ is the estimated probability of death within the next \(\tau\) years. Death is a
terminal endpoint and should not be treated as an ordinary disease burden
weight.
\paragraph{Probability of being alive with no new modeled disease.}
Using the model's disease-specific and death risks, one can summarize the
probability of no new modeled disease and survival over the next \(\tau\) years:
\[
S^{\mathrm{all}}_i(t,\tau)
=
\left[1-p_{i,\mathrm{death}}(t,\tau)\right]
\prod_{d\in\mathcal{D}}
\left[1-p^{\mathrm{new}}_{i,d}(t,\tau)\right].
\]
Equivalently, if the model is represented through cumulative hazards
\(\Lambda_{i,d}(t,\tau)=-\log[1-p_{i,d}(t,\tau)]\),
\[
S^{\mathrm{all}}_i(t,\tau)
=
\exp\left(
-\Lambda_{i,\mathrm{death}}(t,\tau)
-\sum_{d\in\mathcal{D}}
[1-m_{i,d}(t)]\Lambda_{i,d}(t,\tau)
\right).
\]
\paragraph{Probability of being alive with no new disease in a specified set.}
\paragraph{Future incident disease risk in a specified set.}
For any analyst-specified subset of disease tokens \(G\subseteq\mathcal{D}\),
the model can summarize future incident risk within that set:
\[
S^{G}_i(t,\tau)
R^G_i(t,\tau)
=
\left[1-p_{i,\mathrm{death}}(t,\tau)\right]
\prod_{d\in G}
\left[1-p^{\mathrm{new}}_{i,d}(t,\tau)\right].
1-\prod_{d\in G}\left[1-p^{\mathrm{new}}_{i,d}(t,\tau)\right].
\]
This is a subset-level future disease-free survival summary. If \(G\) is called
an organ system, the disease-to-organ grouping is external to the model and
must not be described as a learned organ score.
This quantity answers: ``What is the model-estimated probability that at least
one not-yet-observed disease in \(G\) first occurs within the next \(\tau\)
years?'' It does not include death. Mortality risk should be reported
separately as \(p_{i,\mathrm{death}}(t,\tau)\).
\subsection{Historical counts are not model-derived burden}
If \(G\) is called an organ system, the disease-to-organ grouping is external to
the model and should not be described as a learned organ score.
\subsection{Model Attribution to Predicted Mortality Risk}
The model can also be queried for a model-internal attribution of historical
disease sets to predicted mortality risk. For a disease set \(G\), define the
original mortality risk as
\[
p_{i,\mathrm{death}}^{\mathrm{orig}}(t,\tau),
\]
and the risk after deleting historical disease tokens in \(G\) from the input
history as
\[
p_{i,\mathrm{death}}^{(-G)}(t,\tau).
\]
On the probability scale, the attribution is
\[
\Delta p^G_i(t,\tau)
=
p_{i,\mathrm{death}}^{\mathrm{orig}}(t,\tau)
-
p_{i,\mathrm{death}}^{(-G)}(t,\tau).
\]
A more stable primary scale is the cumulative-hazard scale:
\[
\Lambda_{i,\mathrm{death}}(t,\tau)
=
-\log\left[1-p_{i,\mathrm{death}}(t,\tau)\right],
\]
\[
\Delta \Lambda^G_i(t,\tau)
=
\Lambda_{i,\mathrm{death}}^{\mathrm{orig}}(t,\tau)
-
\Lambda_{i,\mathrm{death}}^{(-G)}(t,\tau).
\]
This quantity should be described as model attribution to predicted mortality
risk. It is not a causal contribution, not an organ damage score, and not a
clinical disease-burden weight. Because diseases can interact within the
trajectory model, deleting an organ/system as a whole is generally more
interpretable than summing single-disease attributions.
\subsection{Historical Counts Are Not Model-Derived Burden}
One may count observed historical diseases:
\[
@@ -148,7 +170,7 @@ model and should not be presented as a model-derived disease burden score.
Without disease severity labels or disease weights, it treats all disease
tokens equally.
\subsection{What the current model does not estimate}
\subsection{What the Current Model Does Not Estimate}
The current model does not directly estimate:
\begin{itemize}[leftmargin=1.5em]
@@ -164,7 +186,7 @@ These interpretations require additional labels, mappings, weights, or model
training objectives. Using the present all-future model to claim these
quantities would be over-interpretation.
\subsection{Post-onset prognosis for the same new disease}
\subsection{Post-Onset Prognosis for the Same New Disease}
For a disease \(d\) that newly occurs at time \(T_{i,d}\), the model cannot
infer the clinical severity of that disease itself. However, after the disease
@@ -180,7 +202,7 @@ be followed by different future disease and mortality risk profiles in people
with different prior trajectories. It should not be described as direct disease
severity.
\subsection{Future extension with reliable recurrence data}
\subsection{Future Extension with Reliable Recurrence Data}
The above interpretation is constrained by the first-occurrence nature of the
current disease sequence. UK Biobank does not provide a reliable longitudinal
@@ -228,8 +250,8 @@ all-future 模型产生
\[
h_i(t)=f_\theta(\mathcal{H}_i(t)),
\]
并且对模型疾病词表 \(\mathcal{D}\) 中的每个疾病 \(d\)，估计未来 \(\tau\) 年内的首次
发生风险：
并且对模型疾病词表 \(\mathcal{D}\) 中的每个疾病 \(d\)，估计未来 \(\tau\) 年内的
首次发生风险：
\[
p_{i,d}(t,\tau)
=
@@ -248,23 +270,23 @@ all-future 模型产生
\]
这个历史发生指示量不是模型学出来的，而是直接从事件历史中读出的。
\subsection{已经发生疾病的 mask}
\subsection{已经发生疾病的 Mask}
如果某个疾病在 \(t\) 之前或 \(t\) 时已经发生，那么该疾病对应的模型输出不应解释为复发
风险，也不应解释为当前疾病活跃程度。在汇总未来新发疾病风险时，应对已经发生过的疾病
风险，也不应解释为当前疾病活动程度。在汇总未来新发疾病风险时，应对已经发生过的疾病
进行 mask：
\[
p^{\mathrm{new}}_{i,d}(t,\tau)
=
\left[1-m_{i,d}(t)\right]p_{i,d}(t,\tau).
[1-m_{i,d}(t)]p_{i,d}(t,\tau).
\]
\subsection{当前模型直接支持的派生量}
\subsection{当前模型直接支持的量}
当前模型直接支持以下几类量。
\paragraph{疾病层面的未来首次发生风险。}
对每个模型内疾病 \(d\)，
对每一个模型内疾病 \(d\)，
\[
p^{\mathrm{new}}_{i,d}(t,\tau)
\]
@@ -277,38 +299,57 @@ all-future 模型产生
表示未来 \(\tau\) 年内死亡的概率。死亡是终末结局，不应作为普通疾病负担权重加入疾病
负担求和。
\paragraph{未来无新病且存活的概率。}
利用疾病层面风险和死亡风险，可以汇总未来 \(\tau\) 年内无任何模型内新病且存活的概率：
\paragraph{指定疾病集合内的未来新发风险。}
对于任意由分析者预先指定的疾病 token 集合 \(G\subseteq\mathcal{D}\)，模型可以汇总该
集合内的未来新发风险：
\[
S^{\mathrm{all}}_i(t,\tau)
R^G_i(t,\tau)
=
\left[1-p_{i,\mathrm{death}}(t,\tau)\right]
\prod_{d\in\mathcal{D}}
\left[1-p^{\mathrm{new}}_{i,d}(t,\tau)\right].
1-\prod_{d\in G}\left[1-p^{\mathrm{new}}_{i,d}(t,\tau)\right].
\]
如果用累计 hazard 表示，令
\(\Lambda_{i,d}(t,\tau)=-\log[1-p_{i,d}(t,\tau)]\)，则
这个量回答的是：模型估计该个体在未来 \(\tau\) 年内，至少新发生一个 \(G\) 中尚未发生
疾病的概率是多少。它不包含死亡；死亡风险应单独报告为
\(p_{i,\mathrm{death}}(t,\tau)\)。
如果将 \(G\) 称为某个器官系统，那么疾病到器官的分组来自模型外部，不应表述为模型
学到的器官评分。
\subsection{死亡风险预测的模型归因}
模型还可以用于计算历史疾病集合对死亡风险预测的模型内部归因。对于疾病集合 \(G\)，
令原始死亡风险为
\[
S^{\mathrm{all}}_i(t,\tau)
p_{i,\mathrm{death}}^{\mathrm{orig}}(t,\tau),
\]
将输入历史中属于 \(G\) 的历史疾病 token 删除后，再次查询模型，得到
\[
p_{i,\mathrm{death}}^{(-G)}(t,\tau).
\]
在概率尺度上，归因可以写为
\[
\Delta p^G_i(t,\tau)
=
\exp\left(
-\Lambda_{i,\mathrm{death}}(t,\tau)
-\sum_{d\in\mathcal{D}}
[1-m_{i,d}(t)]\Lambda_{i,d}(t,\tau)
\right).
p_{i,\mathrm{death}}^{\mathrm{orig}}(t,\tau)
-
p_{i,\mathrm{death}}^{(-G)}(t,\tau).
\]
更推荐的主尺度是累计 hazard：
\[
\Lambda_{i,\mathrm{death}}(t,\tau)
=
-\log\left[1-p_{i,\mathrm{death}}(t,\tau)\right],
\]
\[
\Delta \Lambda^G_i(t,\tau)
=
\Lambda_{i,\mathrm{death}}^{\mathrm{orig}}(t,\tau)
-
\Lambda_{i,\mathrm{death}}^{(-G)}(t,\tau).
\]
\paragraph{未来无指定疾病集合新发且存活的概率。}
对于任意由分析者预先指定的疾病 token 集合 \(G\subseteq\mathcal{D}\)，可以定义
\[
S^{G}_i(t,\tau)
=
\left[1-p_{i,\mathrm{death}}(t,\tau)\right]
\prod_{d\in G}
\left[1-p^{\mathrm{new}}_{i,d}(t,\tau)\right].
\]
这只是指定疾病集合层面的未来 disease-free survival 汇总。如果将 \(G\) 称为某个器官系统，
那么疾病到器官的分组来自模型外部，不能描述为模型学到的器官评分。
这个量应表述为疾病集合对死亡风险预测的模型归因。它不是因果贡献，不是器官损伤评分，
也不是临床疾病负担权重。由于轨迹模型中疾病之间可能存在交互，整体删除一个器官系统
通常比逐个疾病归因后相加更容易解释。
\subsection{历史计数不是模型派生的疾病负担}
@@ -333,8 +374,8 @@ all-future 模型产生
\item 不同疾病 token 之间的相对临床重要性。
\end{itemize}
这些解释都需要额外标签、映射、权重或新的训练目标。用当前 all-future 模型直接声称这些量，
属于过分解读。
这些解释都需要额外标签、映射、权重或新的训练目标。用当前 all-future 模型直接声称
这些量，属于过分解读。
\subsection{同一新发疾病后的预后差异}
@@ -373,21 +414,25 @@ all-future 模型产生
\item 复发或重复事件风险：疾病已经发生后的未来 episode 风险；
\item 死亡风险：与非致死事件竞争的终末结局风险。
\end{itemize}
如果复发 episode 还带有可靠的 episode 层面严重程度标签，例如住院强度、治疗升级或经过
验证的严重程度分级，那么还可以进一步训练有监督的严重程度相关预后模型。但这些扩展都
需要新的数据和新的训练目标，并不是当前 all-future first-occurrence 模型已经具备的能力。
如果复发 episode 还带有可靠的 episode 层面严重程度标签，例如住院强度、治疗升级或
经过验证的严重程度分级，那么还可以进一步训练有监督的严重程度相关预后模型。但这些
扩展都需要新的数据和新的训练目标，并不是当前 all-future first-occurrence 模型已经
具备的能力。
\section{Recommended wording}
\section{Recommended Wording}
\paragraph{English.}
The all-future model is a history-conditioned incident disease and mortality
risk model. Its outputs support future disease-free survival summaries over the
modeled disease vocabulary, but do not directly quantify current disease
burden, organ damage, frailty, or disease severity.
risk model. Its outputs support separate reporting of future mortality risk and
future incident disease risk over analyst-specified disease sets. Historical
disease sets may be ablated to obtain model attribution to predicted mortality
risk, but this is not causal contribution and does not directly quantify current
disease burden, organ damage, frailty, or disease severity.
\paragraph{中文。}
all-future 模型是基于既往轨迹的未来新发疾病和死亡风险模型。它的输出可以支持模型词表
范围内的未来无新病且存活概率汇总，但不能直接量化当前疾病负担、器官损伤、衰弱程度或
疾病严重度。
all-future 模型是基于既往轨迹的未来新发疾病和死亡风险模型。它的输出可以支持分别
报告未来死亡风险，以及指定疾病集合内的未来新发疾病风险。通过删除历史疾病集合可以
得到其对死亡风险预测的模型归因，但这不是因果贡献，也不能直接量化当前疾病负担、
器官损伤、衰弱程度或疾病严重度。
\end{document}

178
evaluate_event_free_survival.py
View File

@@ -1,9 +1,12 @@
"""Compute landmark future event-free survival summaries for DeepHealth.
"""Compute landmark future death and incident system-disease risks.
For each selected patient and landmark age, this script computes:
* P(alive and no new modeled disease within tau years);
* P(alive and no new disease in each ICD-10 chapter-derived system);
* future death risk within tau years;
* future incident disease risk for each ICD-10 chapter-derived system;
* model attribution of each historical organ/system disease set to predicted
mortality risk, computed by deleting that system's historical disease tokens
and re-querying the model;
* historical modeled-disease count;
* historical modeled-disease count within each ICD-10 chapter-derived system.
@@ -38,8 +41,9 @@ from evaluate_auc_v2 import (
resolve_eval_device,
validate_dataset_metadata,
)
from future_event_free_survival import (
future_event_free_survival_from_probabilities,
from future_risk import (
death_risk_from_probabilities,
new_disease_risk_from_probabilities,
probabilities_from_logits,
)
from models import DeepHealth
@@ -284,6 +288,75 @@ def collate_indexed_landmark_fn(batch: List[Dict[str, torch.Tensor]]) -> Dict[st
}
def ablate_event_history_for_tokens(
batch: Dict[str, torch.Tensor],
token_ids: Sequence[int],
) -> Dict[str, torch.Tensor]:
"""Return a batch with selected disease tokens removed from event history."""
selected = {int(token) for token in token_ids}
if not selected:
return batch
event_rows: list[torch.Tensor] = []
time_rows: list[torch.Tensor] = []
readout_rows: list[torch.Tensor] = []
landmark_positions: list[torch.Tensor] = []
event_seq = batch["event_seq"]
time_seq = batch["time_seq"]
readout_mask = batch["readout_mask"]
padding_mask = batch["padding_mask"].bool()
for i in range(event_seq.shape[0]):
valid = padding_mask[i]
events = event_seq[i, valid]
times = time_seq[i, valid]
reads = readout_mask[i, valid]
keep = torch.ones_like(events, dtype=torch.bool)
for token in selected:
keep &= events != int(token)
kept_events = events[keep]
kept_times = times[keep]
kept_reads = reads[keep]
if kept_events.numel() == 0:
kept_events = torch.tensor(
[CHECKUP_IDX],
dtype=event_seq.dtype,
device=event_seq.device,
)
kept_times = batch["t_query"][i : i + 1].to(
dtype=time_seq.dtype,
device=time_seq.device,
)
kept_reads = torch.ones(1, dtype=torch.bool, device=readout_mask.device)
if bool(kept_reads.any()):
landmark_pos = torch.nonzero(kept_reads, as_tuple=False)[-1, 0]
else:
landmark_pos = torch.tensor(
int(kept_events.numel() - 1),
dtype=batch["landmark_pos"].dtype,
device=batch["landmark_pos"].device,
)
kept_reads = torch.zeros_like(kept_events, dtype=torch.bool)
kept_reads[int(landmark_pos.item())] = True
event_rows.append(kept_events)
time_rows.append(kept_times)
readout_rows.append(kept_reads)
landmark_positions.append(landmark_pos.to(dtype=batch["landmark_pos"].dtype))
out = dict(batch)
out["event_seq"] = pad_sequence(event_rows, batch_first=True, padding_value=PAD_IDX)
out["time_seq"] = pad_sequence(time_rows, batch_first=True, padding_value=0.0)
out["readout_mask"] = pad_sequence(
readout_rows, batch_first=True, padding_value=False
)
out["padding_mask"] = out["event_seq"] > PAD_IDX
out["landmark_pos"] = torch.stack(landmark_positions)
return out
@torch.no_grad()
def infer_landmark_hidden(
*,
@@ -351,6 +424,42 @@ def make_occurred_mask(
return occurred
def mortality_hazard_from_risk(risk: torch.Tensor, eps: float = 1e-7) -> torch.Tensor:
return -torch.log1p(-risk.clamp(0.0, 1.0 - float(eps)))
def death_risk_for_batch(
*,
model: DeepHealth,
batch: Dict[str, torch.Tensor],
device: torch.device,
model_target_mode: str,
readout_name: str,
readout_reduce: str,
dist_mode: str,
tau: float,
) -> torch.Tensor:
hidden = infer_landmark_hidden(
model=model,
batch=batch,
device=device,
model_target_mode=model_target_mode,
readout_name=readout_name,
readout_reduce=readout_reduce,
)
logits = model.calc_risk(hidden)
rho = model.calc_weibull_rho(hidden) if dist_mode == "weibull" else None
death_rho = model.calc_death_rho(hidden) if dist_mode == "mixed" else None
probabilities = probabilities_from_logits(
logits,
tau,
dist_mode=dist_mode,
rho=rho,
death_rho=death_rho,
)
return death_risk_from_probabilities(probabilities)
def historical_counts_by_group(
tokens: np.ndarray,
*,
@@ -373,12 +482,12 @@ def historical_counts_by_group(
def output_name_for_run(run_path: Path, eval_split: str, tau: float) -> Path:
return run_path / f"event_free_survival_{eval_split}_tau{tau:g}y.csv"
return run_path / f"future_risk_{eval_split}_tau{tau:g}y.csv"
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(
description="Compute landmark event-free survival summaries."
description="Compute landmark death and incident system-disease risks."
)
parser.add_argument("--run_path", type=str, required=True)
parser.add_argument("--output_path", type=str, default=None)
@@ -496,7 +605,7 @@ def main() -> None:
print(f"Output: {output_path}")
rows: list[dict[str, Any]] = []
for batch in tqdm(loader, desc="Event-free survival", dynamic_ncols=True):
for batch in tqdm(loader, desc="Future risks", dynamic_ncols=True):
hidden = infer_landmark_hidden(
model=model,
batch=batch,
@@ -521,20 +630,45 @@ def main() -> None:
device=device,
)
all_survival = future_event_free_survival_from_probabilities(
death_risk_tensor = death_risk_from_probabilities(probabilities)
death_hazard_tensor = mortality_hazard_from_risk(death_risk_tensor)
death_risk = death_risk_tensor.detach().cpu().numpy()
group_risk: dict[str, np.ndarray] = {}
for group in group_names:
group_risk[group] = new_disease_risk_from_probabilities(
probabilities,
occurred,
disease_ids=None,
vocab_size=int(dataset.vocab_size),
organ_groups[group],
).detach().cpu().numpy()
group_survival: dict[str, np.ndarray] = {}
group_mortality_attr_prob: dict[str, np.ndarray] = {}
group_mortality_attr_hazard: dict[str, np.ndarray] = {}
for group in group_names:
group_survival[group] = future_event_free_survival_from_probabilities(
probabilities,
occurred,
disease_ids=organ_groups[group],
vocab_size=int(dataset.vocab_size),
ids = torch.as_tensor(organ_groups[group], dtype=torch.long, device=device)
if ids.numel() == 0 or not bool(occurred[:, ids].any().item()):
zeros = np.zeros(batch["event_seq"].shape[0], dtype=np.float32)
group_mortality_attr_prob[group] = zeros
group_mortality_attr_hazard[group] = zeros
continue
ablated_batch = ablate_event_history_for_tokens(batch, organ_groups[group])
ablated_death_risk = death_risk_for_batch(
model=model,
batch=ablated_batch,
device=device,
model_target_mode=model_target_mode,
readout_name=readout_name,
readout_reduce=readout_reduce,
dist_mode=dist_mode,
tau=tau,
)
ablated_death_hazard = mortality_hazard_from_risk(ablated_death_risk)
group_mortality_attr_prob[group] = (
death_risk_tensor - ablated_death_risk
).detach().cpu().numpy()
group_mortality_attr_hazard[group] = (
death_hazard_tensor - ablated_death_hazard
).detach().cpu().numpy()
row_indices = batch["row_idx"].cpu().numpy().astype(np.int64)
@@ -559,11 +693,17 @@ def main() -> None:
"tau": tau,
"followup_end_time": float(meta["followup_end_time"]),
"history_disease_count": int(total_count),
"event_free_survival_all": float(all_survival[j]),
"death_risk": float(death_risk[j]),
}
for group in group_names:
out[f"history_count__{group}"] = int(group_counts[group])
out[f"event_free_survival__{group}"] = float(group_survival[group][j])
out[f"new_disease_risk__{group}"] = float(group_risk[group][j])
out[f"mortality_attribution_probability__{group}"] = float(
group_mortality_attr_prob[group][j]
)
out[f"mortality_attribution_hazard__{group}"] = float(
group_mortality_attr_hazard[group][j]
)
rows.append(out)
df = pd.DataFrame(rows)

269
future_event_free_survival.py
View File

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

115
future_risk.py Normal file
View File

@@ -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))