Add mortality attribution evaluation

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.}
+\paragraph{Future incident disease risk in a specified set.}
-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.}
 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]
+    1-\prod_{d\in G}\left[1-p^{\mathrm{new}}_{i,d}(t,\tau)\right].
-    \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
+This quantity answers: ``What is the model-estimated probability that at least
-an organ system, the disease-to-organ grouping is external to the model and
+one not-yet-observed disease in \(G\) first occurs within the next \(\tau\)
-must not be described as a learned organ score.
+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{未来无新病且存活的概率。}
+\paragraph{指定疾病集合内的未来新发风险。}
-利用疾病层面风险和死亡风险，可以汇总未来 \(\tau\) 年内无任何模型内新病且存活的概率：
+对于任意由分析者预先指定的疾病 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]
+    1-\prod_{d\in G}\left[1-p^{\mathrm{new}}_{i,d}(t,\tau)\right].
-    \prod_{d\in\mathcal{D}}
-    \left[1-p^{\mathrm{new}}_{i,d}(t,\tau)\right].
 \]
-如果用累计 hazard 表示，令
+这个量回答的是：模型估计该个体在未来 \(\tau\) 年内，至少新发生一个 \(G\) 中尚未发生
-\(\Lambda_{i,d}(t,\tau)=-\log[1-p_{i,d}(t,\tau)]\)，则
+疾病的概率是多少。它不包含死亡；死亡风险应单独报告为
+\(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(
+    p_{i,\mathrm{death}}^{\mathrm{orig}}(t,\tau)
-        -\Lambda_{i,\mathrm{death}}(t,\tau)
+    -
-        -\sum_{d\in\mathcal{D}}
+    p_{i,\mathrm{death}}^{(-G)}(t,\tau).
-        [1-m_{i,d}(t)]\Lambda_{i,d}(t,\tau)
+\]
-    \right).
+更推荐的主尺度是累计 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 层面严重程度标签，例如住院强度、治疗升级或经过
+如果复发 episode 还带有可靠的 episode 层面严重程度标签，例如住院强度、治疗升级或
-验证的严重程度分级，那么还可以进一步训练有监督的严重程度相关预后模型。但这些扩展都
+经过验证的严重程度分级，那么还可以进一步训练有监督的严重程度相关预后模型。但这些
-需要新的数据和新的训练目标，并不是当前 all-future first-occurrence 模型已经具备的能力。
+扩展都需要新的数据和新的训练目标，并不是当前 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
+risk model. Its outputs support separate reporting of future mortality risk and
-modeled disease vocabulary, but do not directly quantify current disease
+future incident disease risk over analyst-specified disease sets. Historical
-burden, organ damage, frailty, or disease severity.
+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}

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);
+* future death risk within tau years;
-* P(alive and no new disease in each ICD-10 chapter-derived system);
+* 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 (
+from future_risk import (
-    future_event_free_survival_from_probabilities,
+    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)
-            probabilities,
+        death_hazard_tensor = mortality_hazard_from_risk(death_risk_tensor)
-            occurred,
+        death_risk = death_risk_tensor.detach().cpu().numpy()
-            disease_ids=None,
-            vocab_size=int(dataset.vocab_size),
-        ).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],
+                organ_groups[group],
-                vocab_size=int(dataset.vocab_size),
+            ).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)

future_event_free_survival.py

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

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