diff --git a/all_future_model_interpretation.pdf b/all_future_model_interpretation.pdf index 5f9cb8f..e776721 100644 Binary files a/all_future_model_interpretation.pdf and b/all_future_model_interpretation.pdf differ diff --git a/all_future_model_interpretation.tex b/all_future_model_interpretation.tex index 970b73d..7d7c980 100644 --- a/all_future_model_interpretation.tex +++ b/all_future_model_interpretation.tex @@ -37,12 +37,12 @@ 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 \[ - h_i(t) = f_\theta(\mathcal{H}_i(t)), + h_i(t)=f_\theta(\mathcal{H}_i(t)), \] and, for each disease token \(d\) in the modeled disease vocabulary \(\mathcal{D}\), estimates a future first-occurrence risk over a horizon @@ -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} diff --git a/evaluate_event_free_survival.py b/evaluate_event_free_survival.py index a03f055..5e9af6e 100644 --- a/evaluate_event_free_survival.py +++ b/evaluate_event_free_survival.py @@ -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( - probabilities, - occurred, - disease_ids=None, - vocab_size=int(dataset.vocab_size), - ).detach().cpu().numpy() + 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_survival: dict[str, np.ndarray] = {} + group_risk: dict[str, np.ndarray] = {} for group in group_names: - group_survival[group] = future_event_free_survival_from_probabilities( + group_risk[group] = new_disease_risk_from_probabilities( probabilities, occurred, - disease_ids=organ_groups[group], - vocab_size=int(dataset.vocab_size), + organ_groups[group], + ).detach().cpu().numpy() + + group_mortality_attr_prob: dict[str, np.ndarray] = {} + group_mortality_attr_hazard: dict[str, np.ndarray] = {} + for group in group_names: + 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) diff --git a/future_event_free_survival.py b/future_event_free_survival.py deleted file mode 100644 index bc51aca..0000000 --- a/future_event_free_survival.py +++ /dev/null @@ -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, - ) diff --git a/future_risk.py b/future_risk.py new file mode 100644 index 0000000..8480b1a --- /dev/null +++ b/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))