Compare commits
24 Commits
6f7b5be405
...
master
| Author | SHA1 | Date | |
|---|---|---|---|
| d728d8585c | |||
| 9ccc6b56ec | |||
| bf5cae8758 | |||
| 75c9f06114 | |||
| d27eca3f3d | |||
| f1ca401783 | |||
| e040707d07 | |||
| 271a678e73 | |||
| 4f102fb271 | |||
| de6f9b75b9 | |||
| dfb22adf2d | |||
| 89dcf4b362 | |||
| efd18f1be2 | |||
| 4580058687 | |||
| e471de030d | |||
| c622ec50f7 | |||
| 31f129a7dc | |||
| 059c756686 | |||
| e016ac501c | |||
| 315f552301 | |||
| 15ace878f4 | |||
| b13db5e407 | |||
| 4526191fe1 | |||
| 3af823f2e1 |
7
AGENTS.md
Normal file
7
AGENTS.md
Normal file
@@ -0,0 +1,7 @@
|
||||
# 项目协作原则
|
||||
|
||||
- 不回答或执行超出用户问题或请求范围的事项。
|
||||
- 如果确有必要超出范围，必须先停止并询问用户；得到明确同意后才能继续。
|
||||
- 回答应简单直接、逻辑清晰。
|
||||
- 永远不要使用 Node.js 读取、处理、分析或转换数据；涉及数据任务时必须使用其他工具。
|
||||
- Python 与数据处理统一使用本机 Miniconda：`C:\ProgramData\miniconda3`；不要使用 Codex 自带或其他 Python/Conda 运行时。
|
||||
@@ -1,389 +0,0 @@
|
||||
# Event–Trajectory Shared Reasoning Backbone
|
||||
|
||||
> 状态：**Frozen implementation baseline**
|
||||
> 架构标识：`event_trajectory_shared_v2`
|
||||
> 固化日期：**2026-07-23**
|
||||
|
||||
## 1. 核心定义
|
||||
|
||||
使用一个共享的 Attention–TrajMixer 推理核心，对固定 Event Memory 进行多轮读取，并持续更新 Trajectory State。
|
||||
|
||||
模型只实例化：
|
||||
|
||||
```python
|
||||
self.reasoning_core = SharedEventTrajectoryCore(...)
|
||||
```
|
||||
|
||||
禁止为不同推理轮创建独立 Transformer blocks。参数只保存一套，计算上顺序运行多轮。
|
||||
|
||||
模型规模固定为五档：
|
||||
|
||||
| model_size | d_model | n_trajectory | trajectory_dim | traj_hidden |
|
||||
|---|---:|---:|---:|---:|
|
||||
| nano | 120 | 6 | 20 | 24 |
|
||||
| tiny | 256 | 8 | 32 | 32 |
|
||||
| small | 512 | 8 | 64 | 32 |
|
||||
| medium | 768 | 12 | 64 | 48 |
|
||||
| huge | 1024 | 16 | 64 | 64 |
|
||||
|
||||
默认使用 `model_size=nano`。`n_reasoning_rounds` 是独立参数，默认值为12，
|
||||
不属于模型规模预设；任意模型规模均可单独指定推理轮数。
|
||||
|
||||
必须满足：
|
||||
|
||||
\[
|
||||
d_{\mathrm{model}}
|
||||
=n_{\mathrm{trajectory}}d_{\mathrm{trajectory}}.
|
||||
\]
|
||||
|
||||
## 2. 固定 Event Memory
|
||||
|
||||
疾病事件与可选协变量首先组成事件序列：
|
||||
|
||||
\[
|
||||
X_E\in\mathbb{R}^{B\times L\times d_{\mathrm{model}}}.
|
||||
\]
|
||||
|
||||
事件特征由以下信息相加：
|
||||
|
||||
```text
|
||||
disease / covariate embedding
|
||||
+ age/time encoding
|
||||
+ sex context
|
||||
```
|
||||
|
||||
然后只编码一次：
|
||||
|
||||
\[
|
||||
E=\operatorname{EventNorm}
|
||||
\left(\operatorname{EventProjection}(X_E)\right).
|
||||
\]
|
||||
|
||||
进入 reasoning loop 后，\(E\) 的数值保持不变，但不执行 `detach`，梯度仍可回传到事件编码器。
|
||||
|
||||
Key 和 Value 同样每次 forward 只投影一次：
|
||||
|
||||
```python
|
||||
event_key_value = reasoning_core.project_event_memory(E)
|
||||
```
|
||||
|
||||
12 轮共享并复用该结果。
|
||||
|
||||
## 3. Trajectory State
|
||||
|
||||
每个查询维护 `n_trajectory` 个显式 trajectory slots；nano 默认使用6个：
|
||||
|
||||
\[
|
||||
S\in\mathbb{R}^{B\times Q\times n_{\mathrm{trajectory}}\times
|
||||
d_{\mathrm{trajectory}}}.
|
||||
\]
|
||||
|
||||
其中：
|
||||
|
||||
- all-future：\(Q=1\)；
|
||||
- next-token：\(Q=L\)，所有查询位置并行计算。
|
||||
|
||||
定义可学习原型：
|
||||
|
||||
\[
|
||||
P\in\mathbb{R}^{n_{\mathrm{trajectory}}\times d_{\mathrm{trajectory}}}.
|
||||
\]
|
||||
|
||||
查询上下文经过投影并 reshape：
|
||||
|
||||
\[
|
||||
C_Q
|
||||
=\operatorname{QueryProjection}(\text{query features})
|
||||
\in\mathbb{R}^{B\times Q\times n_{\mathrm{trajectory}}\times
|
||||
d_{\mathrm{trajectory}}},
|
||||
\]
|
||||
|
||||
\[
|
||||
S^{(0)}=P+C_Q.
|
||||
\]
|
||||
|
||||
all-future 的 query features 包含可学习 query token、查询年龄和性别；next-token 的 query features 使用当前位置的事件、时间和性别表示，以保留 token-level 预测语义。
|
||||
|
||||
## 4. 共享 Trajectory-to-Event Attention
|
||||
|
||||
Trajectory State 作为 Query，固定 Event Memory 作为 Key 和 Value：
|
||||
|
||||
\[
|
||||
Q^{(r)}
|
||||
=W_Q\operatorname{LN}_{A}(S^{(r)}),
|
||||
\]
|
||||
|
||||
\[
|
||||
K=W_KE,\qquad V=W_VE.
|
||||
\]
|
||||
|
||||
形状为：
|
||||
|
||||
```text
|
||||
Q: [B, query, trajectory, trajectory_dim]
|
||||
K: [B, trajectory, event, trajectory_dim]
|
||||
V: [B, trajectory, event, trajectory_dim]
|
||||
```
|
||||
|
||||
Attention：
|
||||
|
||||
\[
|
||||
\operatorname{score}_{b,q,h,l}
|
||||
=
|
||||
\frac{
|
||||
\left\langle Q_{b,q,h,:},K_{b,h,l,:}\right\rangle
|
||||
}{
|
||||
\sqrt{d_{\mathrm{trajectory}}}
|
||||
}.
|
||||
\]
|
||||
|
||||
每个 trajectory slot 独立读取整段 Event Memory。Attention 不包含跨 trajectory 的完整输出投影；trajectory 之间的交互只由后续 TrajMixer 完成。
|
||||
|
||||
外部 `padding_mask` 的语义固定为 `True = valid`。
|
||||
|
||||
all-future 的内部 mask 必须满足：
|
||||
|
||||
\[
|
||||
\operatorname{valid}_{b,q,l}
|
||||
=
|
||||
\operatorname{eventValid}_{b,l}
|
||||
\land
|
||||
(t_l\le t_q).
|
||||
\]
|
||||
|
||||
next-token 还必须对相同时间戳加入位置因果约束：
|
||||
|
||||
\[
|
||||
\operatorname{valid}_{b,q,l}
|
||||
=
|
||||
\operatorname{eventValid}_{b,l}
|
||||
\land
|
||||
\left[
|
||||
(t_l<t_q)
|
||||
\lor
|
||||
\left((t_l=t_q)\land(l\le q)\right)
|
||||
\right].
|
||||
\]
|
||||
|
||||
这可以阻止前一个 token 直接读到同时间的后续目标 token；它只改变并行
|
||||
Attention 的可见性矩阵，不沿疾病时间轴递归。
|
||||
|
||||
两种 mask 均不能读取未来事件或协变量。
|
||||
|
||||
全 masked query 的 Attention readout 必须显式返回零，不能产生 NaN。
|
||||
|
||||
## 5. 时间信息
|
||||
|
||||
所有模式都在 Event Memory 和 query context 中加入 age/time encoding。
|
||||
|
||||
当 `time_mode=relative` 时，shared cross-attention 额外使用：
|
||||
|
||||
- query-time 对 event-time 的 Cross-TimeRoPE；
|
||||
- query–event 时间差的 Gaussian RBF bias。
|
||||
|
||||
对应缓存形状为：
|
||||
|
||||
\[
|
||||
\text{RBF cache}\in\mathbb{R}^{B\times Q\times L\times n_{\mathrm{rbf}}}.
|
||||
\]
|
||||
|
||||
## 6. 共享 TrajMixer
|
||||
|
||||
TrajMixer 输入：
|
||||
|
||||
\[
|
||||
U\in\mathbb{R}^{B\times Q\times H\times D_h}.
|
||||
\]
|
||||
|
||||
它只沿 trajectory 轴交互，不使用普通全维度 FFN，也不包含旧版 Group Alignment。
|
||||
|
||||
对每个内部坐标 \(r\)：
|
||||
|
||||
\[
|
||||
G_r=U_rW_g^{(r)},\qquad
|
||||
V_r=U_rW_v^{(r)},
|
||||
\]
|
||||
|
||||
\[
|
||||
M_r
|
||||
=
|
||||
\operatorname{SiLU}(G_r)\odot V_r,
|
||||
\]
|
||||
|
||||
\[
|
||||
Y_r=M_rW_o^{(r)}.
|
||||
\]
|
||||
|
||||
参数形状：
|
||||
|
||||
```text
|
||||
W_g: [trajectory_dim, n_trajectory, traj_hidden]
|
||||
W_v: [trajectory_dim, n_trajectory, traj_hidden]
|
||||
W_o: [trajectory_dim, traj_hidden, n_trajectory]
|
||||
```
|
||||
|
||||
默认：
|
||||
|
||||
```text
|
||||
n_trajectory -> 4 * n_trajectory -> n_trajectory
|
||||
```
|
||||
|
||||
其中：
|
||||
|
||||
\[
|
||||
d_{\mathrm{trajHidden}}=4n_{\mathrm{trajectory}}.
|
||||
\]
|
||||
|
||||
## 7. 单轮共享核心
|
||||
|
||||
单轮计算：
|
||||
|
||||
\[
|
||||
R^{(r)}
|
||||
=
|
||||
A_\theta\left(
|
||||
\operatorname{LN}_A(S^{(r)}),E
|
||||
\right),
|
||||
\]
|
||||
|
||||
\[
|
||||
U^{(r)}
|
||||
=
|
||||
S^{(r)}
|
||||
+\alpha_A\operatorname{Dropout}(R^{(r)}),
|
||||
\]
|
||||
|
||||
\[
|
||||
S^{(r+1)}
|
||||
=
|
||||
U^{(r)}
|
||||
+\alpha_M\operatorname{Dropout}
|
||||
\left(
|
||||
M_\phi(\operatorname{LN}_M(U^{(r)}))
|
||||
\right).
|
||||
\]
|
||||
|
||||
残差统一使用加法。
|
||||
|
||||
## 8. 多轮参数共享
|
||||
|
||||
同一个核心重复运行：
|
||||
|
||||
```python
|
||||
for _ in range(n_reasoning_rounds):
|
||||
S = self.reasoning_core(...)
|
||||
```
|
||||
|
||||
所有轮次共享：
|
||||
|
||||
```text
|
||||
q_proj / k_proj / v_proj
|
||||
relative-time projection
|
||||
TrajMixer parameters
|
||||
LayerNorm parameters
|
||||
attn_scale / mixer_scale
|
||||
```
|
||||
|
||||
因此推理轮数不改变模型参数量：
|
||||
|
||||
\[
|
||||
A^{(1)}=\cdots=A^{(12)}=A_\theta,
|
||||
\]
|
||||
|
||||
\[
|
||||
M^{(1)}=\cdots=M^{(12)}=M_\phi.
|
||||
\]
|
||||
|
||||
但每轮 state 不同，因此 Query 与 Attention weights 也不同。
|
||||
|
||||
## 9. 稳定性设计
|
||||
|
||||
共享残差缩放初始化为：
|
||||
|
||||
\[
|
||||
\alpha_A=\alpha_M
|
||||
=
|
||||
\frac{1}{\sqrt{n_{\mathrm{reasoningRounds}}}}.
|
||||
\]
|
||||
|
||||
两个标量可学习，并由全部轮次共享。
|
||||
|
||||
第一版不加入：
|
||||
|
||||
```text
|
||||
round-specific parameters
|
||||
round embedding
|
||||
每轮独立 LayerNorm
|
||||
每轮独立 residual scale
|
||||
GRU 或其他时间递归
|
||||
```
|
||||
|
||||
## 10. 输出接口
|
||||
|
||||
推理结束后按固定顺序 flatten trajectory slots：
|
||||
|
||||
\[
|
||||
H
|
||||
=
|
||||
\operatorname{FinalNorm}
|
||||
\left(
|
||||
\operatorname{Flatten}(S^{(R)})
|
||||
\right).
|
||||
\]
|
||||
|
||||
- all-future 输出：`[B, d_model]`；
|
||||
- next-token 输出：`[B, L, d_model]`；
|
||||
- next-token 的 risk-head weight tying 保持不变；
|
||||
- Weibull 与 mixed heads 继续使用同一最终 hidden。
|
||||
|
||||
next-token 的 query 位置全部并行，只有 reasoning rounds 顺序执行，因此不存在沿疾病时间轴的状态递归。
|
||||
|
||||
## 11. 配置与 checkpoint 约束
|
||||
|
||||
训练配置必须写入：
|
||||
|
||||
```yaml
|
||||
model_architecture: event_trajectory_shared_v2
|
||||
model_size: nano
|
||||
d_model: 120
|
||||
n_trajectory: 6
|
||||
trajectory_dim: 20
|
||||
traj_hidden: 24
|
||||
n_reasoning_rounds: 12
|
||||
model_parameter_count: <runtime count>
|
||||
trainable_parameter_count: <runtime count>
|
||||
```
|
||||
|
||||
评估和导出入口必须同时验证：
|
||||
|
||||
1. `model_architecture` 完全匹配；
|
||||
2. `model_size` 属于 `nano / tiny / small / medium / huge`；
|
||||
3. `d_model`、`n_trajectory`、`trajectory_dim` 和 `traj_hidden`
|
||||
与对应规模预设完全匹配；
|
||||
4. checkpoint 包含一套且仅一套 `reasoning_core` 关键参数；
|
||||
5. checkpoint 内持久化的 `d_model`、`n_trajectory` 和
|
||||
`n_reasoning_rounds` 架构指纹与训练配置完全一致；
|
||||
6. 不接受旧 `traj_mixer_v2` checkpoint。
|
||||
|
||||
其中 `n_reasoning_rounds` 必须进入 checkpoint 架构指纹，因为改变轮数
|
||||
不会改变参数 shape，不能仅依赖 `load_state_dict(strict=True)` 检出错配。
|
||||
|
||||
## 12. 信息流
|
||||
|
||||
```text
|
||||
E ─────────────┬──────────────┬──────────────┬──────────────┐
|
||||
│ │ │ │
|
||||
▼ ▼ ▼ ▼
|
||||
S0 -> Shared Core -> S1 -> Shared Core -> S2 -> ... -> Shared Core -> S12
|
||||
同一套参数 同一套参数 同一套参数
|
||||
```
|
||||
|
||||
整体定义：
|
||||
|
||||
\[
|
||||
\boxed{
|
||||
\text{一个共享 Event–Trajectory 推理核心}
|
||||
\times
|
||||
\text{多轮状态依赖推理}
|
||||
}
|
||||
\]
|
||||
64
MODEL_ARCHITECTURES.md
Normal file
64
MODEL_ARCHITECTURES.md
Normal file
@@ -0,0 +1,64 @@
|
||||
# Model architectures
|
||||
|
||||
DeepHealth uses one codebase for both supported history-block architectures.
|
||||
Select the architecture explicitly when starting a training run:
|
||||
|
||||
| `model_architecture` | History block | Checkpoint fingerprint |
|
||||
| --- | --- | --- |
|
||||
| `transformer_ffn_v1` | Temporal attention + SwiGLU FFN | `blocks.*.mlp.w1/w2/w3` and `blocks.*.ln2` |
|
||||
| `traj_mixer_v5` | Temporal attention + TrajMixer | `blocks.*.mlp.intra_*`, `gate_proj`, and `output_proj` |
|
||||
|
||||
`transformer_ffn_v1` is the CLI default; pass `traj_mixer_v5` explicitly for
|
||||
TrajMixer runs.
|
||||
|
||||
## Training
|
||||
|
||||
Next-step example:
|
||||
|
||||
```powershell
|
||||
python train_next_step.py --model_architecture traj_mixer_v5 --n_layer 12
|
||||
```
|
||||
|
||||
All-future example:
|
||||
|
||||
```powershell
|
||||
python train_all_future.py --model_architecture transformer_ffn_v1 --n_layer 12
|
||||
```
|
||||
|
||||
New runs are separated by architecture:
|
||||
|
||||
```text
|
||||
runs/
|
||||
transformer_ffn_v1/
|
||||
<run_name>/
|
||||
traj_mixer_v5/
|
||||
<run_name>/
|
||||
```
|
||||
|
||||
Use `--runs_root` to place this structure under a different root. Existing run
|
||||
directories are not moved or renamed.
|
||||
|
||||
Each generated `train_config.json` records `model_architecture`, total parameter
|
||||
count, and trainable parameter count.
|
||||
|
||||
Both training entry points use the single `--n_layer` option to set the number
|
||||
of history backbone blocks. The same value is passed to `DeepHealth.n_layer`
|
||||
and saved as `n_layer` in `train_config.json`; it must be at least 1.
|
||||
|
||||
## Architecture validation
|
||||
|
||||
Evaluation resolves the architecture before constructing the model and always
|
||||
loads weights with `strict=True`.
|
||||
|
||||
- Every config must include an explicit `model_architecture` marker.
|
||||
- Checkpoint fingerprints are used to validate that the selected architecture
|
||||
matches the stored weights.
|
||||
- A config marker that conflicts with the checkpoint fingerprint raises an
|
||||
error instead of silently choosing one architecture.
|
||||
- Unsupported historical TrajMixer markers such as `traj_mixer_v2`,
|
||||
`traj_mixer_v3`, and `traj_mixer_v4` are rejected.
|
||||
- Checkpoints and configs created before architecture markers were introduced
|
||||
are intentionally unsupported.
|
||||
|
||||
Project code should use the architecture factory rather than instantiate a
|
||||
history block directly.
|
||||
73
README.md
73
README.md
@@ -6,7 +6,9 @@
|
||||
疾病序列 stream + 统一的额外信息 token stream
|
||||
```
|
||||
|
||||
疾病、死亡、checkup 事件仍然保存在事件序列里；性别单独保存在 `basic_info`；其他体检、暴露、生活方式等信息统一整理成 `(type, value, value_kind, time)` token。
|
||||
疾病和死亡事件保存在预处理事件文件中；性别单独保存在 `basic_info`；其他体检、暴露、生活方式等信息统一整理成 `(type, value, value_kind, time)` token。每个额外信息 token 自带测量时间并直接与疾病 token 拼接，不再生成或使用 assessment/checkup 事件 token。dataset 会无条件清除旧预处理文件中遗留的 `label=1` 事件。
|
||||
|
||||
所有连续额外信息强制使用训练子集拟合的 RobustScale（median/IQR）。center 和 scale 保存为模型 buffer，并用于验证集、测试集和推理；代码不提供未标准化模式，缺少 scaler buffer 的连续变量 checkpoint 不受支持。该规则同时适用于 `next_token` 和 `all_future`。
|
||||
|
||||
## 数据准备
|
||||
|
||||
@@ -28,7 +30,7 @@ python prepare_data.py
|
||||
- `ukb_event_data.npy`
|
||||
- 形状为 `(N, 3)`
|
||||
- 每行是 `(eid, days, label)`
|
||||
- 包含疾病、死亡、checkup 事件
|
||||
- 包含疾病和死亡事件
|
||||
|
||||
- `ukb_basic_info.csv`
|
||||
- index 为 `eid`
|
||||
@@ -59,12 +61,12 @@ python prepare_data.py
|
||||
`dataset.py` 提供两个 dataset：
|
||||
|
||||
- `NextStepHealthDataset`
|
||||
- 用于 next-token / next-time-point 监督
|
||||
- 对应 `Delphi2MLoss` 和 `UniqueTimeSetExponentialLoss`
|
||||
- 仅用于 absolute-time Delphi2M next-token 复现
|
||||
- 对应 `Delphi2MLoss`
|
||||
|
||||
- `AllFutureHealthDataset`
|
||||
- 用于 query-conditioned all-future 监督
|
||||
- 对应 `ExponentialLoss`、`WeibullLoss`、`MixedLoss`
|
||||
- 对应 `ExponentialLoss`、`WeibullLoss`
|
||||
|
||||
为了兼容旧训练入口：
|
||||
|
||||
@@ -173,7 +175,7 @@ extra-info 不再通过独立的 `BaselineEncoder` 或 `CrossAttention` 注入
|
||||
如果需要拿到完整 next-token 输出，可使用结构化返回：
|
||||
|
||||
```python
|
||||
out = model(..., target_mode="next_token", return_output=True)
|
||||
out = model(..., return_output=True)
|
||||
out.hidden # disease tokens + pooled extra-info readout tokens
|
||||
out.time_seq # 与 hidden 对齐的时间
|
||||
out.padding_mask # 与 hidden 对齐的有效位置
|
||||
@@ -187,8 +189,7 @@ model = DeepHealth(
|
||||
vocab_size=dataset.vocab_size,
|
||||
n_embd=120,
|
||||
n_head=10,
|
||||
n_hist_layer=12,
|
||||
n_tab_layer=4, # 兼容旧配置；当前不再创建独立 tabular transformer
|
||||
n_layer=12,
|
||||
n_types=dataset.n_types,
|
||||
n_cont_types=dataset.n_cont_types,
|
||||
n_categories=dataset.n_categories,
|
||||
@@ -204,50 +205,44 @@ model = DeepHealth(
|
||||
next-token 监督：
|
||||
|
||||
- `Delphi2MLoss`
|
||||
- `UniqueTimeSetExponentialLoss`
|
||||
|
||||
next-token 训练中，模型会请求 `return_output=True`，因此 loss 的预测位置包括：
|
||||
|
||||
- 原 disease token readout 位置
|
||||
- 同一时间点 extra hidden 池化后的 pooled extra-info readout token
|
||||
|
||||
pooled extra-info readout token 的监督目标在训练时动态构造：对 pooled extra-info readout token 的时间 `t`，寻找该患者 `t` 之后的下一个 disease 事件时间；`delphi2m` 使用第一个未来事件作为 next-token target，`uts` 使用下一唯一时间点上的事件集合做 multi-hot target。若该 extra-info 时间点之后没有未来 disease target，则该位置不参与 loss。
|
||||
pooled extra-info readout token 的监督目标在训练时动态构造：对 pooled extra-info readout token 的时间 `t`，寻找该患者 `t` 之后的第一个 disease 事件作为 next-token target。若该时间点之后没有未来 disease target，则该位置不参与 loss。
|
||||
|
||||
all-future / query-conditioned 监督：
|
||||
|
||||
- `ExponentialLoss`
|
||||
- `WeibullLoss`
|
||||
- `MixedLoss`
|
||||
|
||||
all-future 训练只读出 `t_query` 对应的 query hidden。展开的 extra-info tokens 作为主序列上下文输入，但不会被单独读出，也不会被纳入 loss 监督。
|
||||
|
||||
`UniqueTimeSetExponentialLoss` 的 observed term 固定使用 sum reduction，不再暴露旧的 `observed_reduction` 参数。
|
||||
|
||||
## 训练
|
||||
|
||||
当前提供两类训练入口：
|
||||
|
||||
- `train_next_step.py`
|
||||
- 使用 `NextStepHealthDataset`
|
||||
- `--target_mode delphi2m` 默认搭配 `Delphi2MLoss` + `token` readout
|
||||
- `--target_mode uts` 默认搭配 `UniqueTimeSetExponentialLoss` + `same_time_group_end` readout
|
||||
- 当前 next-token 训练只支持 exponential time loss
|
||||
- 仅用于 Delphi2M 复现，固定 `time_mode=absolute`、`target_mode=delphi2m`
|
||||
- 固定使用 `Delphi2MLoss` + `token` readout
|
||||
- next-token 训练只支持 exponential time loss
|
||||
- 展开的 extra-info tokens 进入主序列；读出端 pooled extra-info tokens 会加入 prediction/loss 监督
|
||||
- `train_all_future.py`
|
||||
- 使用 `AllFutureHealthDataset`
|
||||
- 不使用 readout，直接对 query hidden 计算风险
|
||||
- `--dist_mode exponential/weibull/mixed` 分别搭配 `ExponentialLoss`、`WeibullLoss`、`MixedLoss`
|
||||
- `--dist_mode exponential/weibull` 分别搭配 `ExponentialLoss`、`WeibullLoss`
|
||||
- 展开的 extra-info tokens 只作为 query 上下文，不单独监督
|
||||
|
||||
当前 `train_next_step.py` / `train_all_future.py` 支持所有已有训练目标定义的组合：
|
||||
|
||||
| 训练模式 | 时间模式 | 分布/监督 | 默认 loss/readout |
|
||||
| --- | --- | --- | --- |
|
||||
| `next_token` | `relative`, `absolute` | `target_mode=delphi2m`, `dist_mode=exponential` | `Delphi2MLoss` + `token` |
|
||||
| `next_token` | `relative`, `absolute` | `target_mode=uts`, `dist_mode=exponential` | `UniqueTimeSetExponentialLoss` + `same_time_group_end` |
|
||||
| `next_token` | `absolute` | `target_mode=delphi2m`, `dist_mode=exponential` | `Delphi2MLoss` + `token` |
|
||||
| `all_future` | `relative`, `absolute` | `dist_mode=exponential` | `ExponentialLoss`，无 readout |
|
||||
| `all_future` | `relative`, `absolute` | `dist_mode=weibull` | `WeibullLoss`，无 readout |
|
||||
| `all_future` | `relative`, `absolute` | `dist_mode=mixed` | `MixedLoss`，无 readout |
|
||||
|
||||
示例：
|
||||
|
||||
@@ -255,11 +250,9 @@ all-future 训练只读出 `t_query` 对应的 query hidden。展开的 extra-in
|
||||
python train_next_step.py \
|
||||
--data_prefix ukb \
|
||||
--labels_file labels.csv \
|
||||
--target_mode uts \
|
||||
--n_embd 120 \
|
||||
--n_head 10 \
|
||||
--n_hist_layer 12 \
|
||||
--n_tab_layer 4 \
|
||||
--n_layer 12 \
|
||||
--extra_pool_reduce mean
|
||||
```
|
||||
|
||||
@@ -274,6 +267,28 @@ python train_all_future.py \
|
||||
--extra_pool_reduce mean
|
||||
```
|
||||
|
||||
纯疾病历史的 T/O/S 消融仅用于
|
||||
`TrajMixer + all_future + relative + Weibull + extra_info_types=[]`：
|
||||
|
||||
```bash
|
||||
python train_all_future.py \
|
||||
--model_architecture traj_mixer_v5 \
|
||||
--time_mode relative \
|
||||
--dist_mode weibull \
|
||||
--extra_info_types_file extra_info_types_none.txt \
|
||||
--disease_history_mode ordered
|
||||
```
|
||||
|
||||
`--disease_history_mode` 的含义：
|
||||
|
||||
- `timed`：保留疾病事件顺序和真实患病时间（T）。
|
||||
- `ordered`：保留疾病事件顺序，以 `0, 1, ..., G-1` 的首发日期顺序组替代真实时间；同日首发疾病共享一个位置，query 位置为 `G`（O）。
|
||||
- `set`：疾病代码去重并排序，疾病和 query 的模型时间全部为 `0`，不保留顺序或患病时间（S）。
|
||||
|
||||
查询点、未来疾病、`future_dt`、exposure、Landmark 年龄与 AUC
|
||||
病例/对照始终使用真实时间。训练配置会记录 `disease_history_mode`；
|
||||
旧配置缺少该字段时按 `timed` 处理。
|
||||
|
||||
选择额外信息变量：
|
||||
|
||||
```bash
|
||||
@@ -286,6 +301,7 @@ python train_next_step.py --extra_info_types_file extra_info_types_smoking_alcoh
|
||||
|
||||
- `extra_info_types_file`：训练时使用的列表文件名
|
||||
- `extra_info_types`：解析后的实际 type id 列表，用于评估脚本复现变量选择
|
||||
- `disease_history_mode`：all-future 模型使用的 T/O/S 疾病历史表示
|
||||
- `extra_pool_reduce`：同一 `other_time` 的 extra-info tokens 池化方式，默认为 `mean`
|
||||
- `model_target_mode`、`time_mode`、`dist_mode`、`dataset_class`、`collate_fn`、`resolved_loss_name`：用于评估脚本重建模型和输入方式
|
||||
|
||||
@@ -350,7 +366,7 @@ python evaluate_auc.py \
|
||||
- 为每个患者和 landmark age 构造 landmark query 样本。
|
||||
- 根据模型模式插入 `<NO_EVENT>` token 或直接传 `t_query`，取 landmark/query hidden。
|
||||
- 对疾病 token 分块投影到 `risk_head`；`score_mode="risk"` 时会根据 `dist_mode` 把线性输出转换为固定 horizon 风险概率。
|
||||
- 分布转换规则与 all-future 训练损失一致：`exponential` 使用 `1 - exp(-rate * horizon)`；`weibull` 使用 `1 - exp(-rate * horizon ** rho)`；`mixed` 中普通疾病使用 exponential，死亡 endpoint 使用 Weibull death rho。
|
||||
- 分布转换规则与 all-future 训练损失一致：`exponential` 使用 `1 - exp(-rate * horizon)`；`weibull` 使用 `1 - exp(-rate * horizon ** rho)`。
|
||||
- `score_mode="eta"` 是诊断用排序分数，不使用 `rho`，因此不区分不同分布的风险曲线。
|
||||
- 按疾病、性别、landmark age、horizon 计算 incident disease AUC。
|
||||
- 可选择排除 horizon 内先于目标疾病发生的死亡竞争风险。
|
||||
@@ -386,7 +402,7 @@ python evaluate_auc_v2.py \
|
||||
| 评估口径 | next-step/token-level 预测点 | landmark fixed-horizon incident risk |
|
||||
| 查询位置 | next-token 用满足 offset 条件的最新 readout token；all-future 直接用该预测点年龄作为 `t_query` | next-token 用人工插入的 `<NO_EVENT>` landmark token；all-future 直接用 `t_query` |
|
||||
| 时间参数 | `offsets`：预测点至少早于目标事件多少年 | `landmark_*` 和 `horizons`：固定年龄点与未来窗口 |
|
||||
| score 与分布 | 使用 eta/logit 排序分数；不按 `dist_mode` 转换风险概率 | `score_mode="risk"` 按 `dist_mode` 区分 exponential / Weibull / mixed；`score_mode="eta"` 不区分分布 |
|
||||
| score 与分布 | 使用 eta/logit 排序分数；不按 `dist_mode` 转换风险概率 | `score_mode="risk"` 按 `dist_mode` 区分 exponential / Weibull；`score_mode="eta"` 不区分分布 |
|
||||
| 病例定义 | target table 中出现目标疾病的患者/事件 | landmark 后 horizon 内首次发生目标疾病 |
|
||||
| 对照定义 | 从未出现该疾病的患者的 eligible target occurrence | landmark 时未患病，且 horizon 内未发病并有足够随访 |
|
||||
| 分层 | sex + age bracket + offset | sex + landmark age + horizon |
|
||||
@@ -424,11 +440,6 @@ python evaluate_auc_v2.py \
|
||||
- `losses.py`
|
||||
- next-token 和 all-future losses
|
||||
|
||||
- `readouts.py`
|
||||
- token readout
|
||||
- same-time group readout
|
||||
- last-valid readout
|
||||
|
||||
- `evaluate_auc.py`
|
||||
- next-step/token-level 疾病 AUC 评估
|
||||
- 使用 prediction offset、sex、age bracket 分层
|
||||
@@ -439,4 +450,4 @@ python evaluate_auc_v2.py \
|
||||
- landmark fixed-horizon incident disease AUC 评估
|
||||
- next-token 模型通过插入 `<NO_EVENT>` landmark token 查询固定年龄点风险
|
||||
- all-future 模型直接通过 `t_query` 查询固定年龄点风险
|
||||
- `score_mode="risk"` 按 exponential / Weibull / mixed 分布计算固定 horizon 风险
|
||||
- `score_mode="risk"` 按 exponential / Weibull 分布计算固定 horizon 风险
|
||||
|
||||
542
backbones.py
542
backbones.py
@@ -4,6 +4,12 @@ import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from model_architectures import (
|
||||
TRAJ_MIXER_ARCHITECTURE,
|
||||
TRANSFORMER_FFN_ARCHITECTURE,
|
||||
resolve_model_architecture,
|
||||
)
|
||||
|
||||
|
||||
class TimeRoPE(nn.Module):
|
||||
def __init__(self, dim: int, base: float = 10000.0):
|
||||
@@ -29,14 +35,6 @@ class TimeRoPE(nn.Module):
|
||||
x2 = x[..., 1::2]
|
||||
return torch.stack((-x2, x1), dim=-1).flatten(-2)
|
||||
|
||||
@staticmethod
|
||||
def apply_single_from_cache(
|
||||
x: torch.Tensor,
|
||||
rope_cache: tuple[torch.Tensor, torch.Tensor],
|
||||
) -> torch.Tensor:
|
||||
cos, sin = rope_cache
|
||||
return x * cos + TimeRoPE._rotate_half(x) * sin
|
||||
|
||||
@staticmethod
|
||||
def apply_from_cache(
|
||||
q: torch.Tensor,
|
||||
@@ -93,280 +91,364 @@ class GaussianRBFTimeBasis(nn.Module):
|
||||
)
|
||||
return rbf_acts
|
||||
|
||||
def precompute_cross_cache(
|
||||
self,
|
||||
query_tau: torch.Tensor,
|
||||
key_tau: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""Return RBF activations for query-time minus event-time."""
|
||||
diff = query_tau.float().unsqueeze(2) - key_tau.float().unsqueeze(1)
|
||||
widths = self.log_widths.exp()
|
||||
return torch.exp(
|
||||
-0.5
|
||||
* (
|
||||
(diff.unsqueeze(-1) - self.centers)
|
||||
/ widths
|
||||
).square()
|
||||
)
|
||||
|
||||
|
||||
class TrajectoryCrossAttention(nn.Module):
|
||||
"""Shared trajectory-slot queries reading a fixed event memory."""
|
||||
|
||||
class TemporalAttention(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
d_model: int,
|
||||
n_trajectory: int,
|
||||
n_embd: int,
|
||||
n_head: int,
|
||||
n_rbf_bases: int = 16,
|
||||
use_time_rope: bool = False,
|
||||
use_rbf_bias: bool = False,
|
||||
dropout: float = 0.0,
|
||||
use_time_rope: bool = True,
|
||||
use_rbf_bias: bool = True,
|
||||
):
|
||||
super().__init__()
|
||||
if d_model <= 0 or n_trajectory <= 0:
|
||||
raise ValueError("d_model and n_trajectory must be positive")
|
||||
if d_model % n_trajectory != 0:
|
||||
raise ValueError(
|
||||
"d_model must be divisible by n_trajectory, got "
|
||||
f"{d_model} and {n_trajectory}"
|
||||
)
|
||||
self.d_model = d_model
|
||||
self.n_trajectory = n_trajectory
|
||||
self.trajectory_dim = d_model // n_trajectory
|
||||
self.scale = self.trajectory_dim ** -0.5
|
||||
assert n_embd % n_head == 0, "n_embd must be divisible by n_head"
|
||||
self.n_head = n_head
|
||||
self.d_head = n_embd // n_head
|
||||
self.scale = 1.0 / math.sqrt(self.d_head)
|
||||
self.use_time_rope = use_time_rope
|
||||
self.use_rbf_bias = use_rbf_bias
|
||||
|
||||
# q_proj acts on each slot independently and is shared across slots.
|
||||
self.q_proj = nn.Linear(
|
||||
self.trajectory_dim,
|
||||
self.trajectory_dim,
|
||||
bias=False,
|
||||
)
|
||||
self.k_proj = nn.Linear(d_model, d_model, bias=False)
|
||||
self.v_proj = nn.Linear(d_model, d_model, bias=False)
|
||||
if use_rbf_bias:
|
||||
self.rbf_proj = nn.Linear(
|
||||
n_rbf_bases,
|
||||
n_trajectory,
|
||||
bias=False,
|
||||
)
|
||||
self.time_bias_scale = nn.Parameter(torch.tensor(0.0))
|
||||
else:
|
||||
self.rbf_proj = None
|
||||
self.register_parameter("time_bias_scale", None)
|
||||
# QKV projection (fused for efficiency)
|
||||
self.qkv = nn.Linear(n_embd, 3 * n_embd, bias=False)
|
||||
# Output projection
|
||||
self.out_proj = nn.Linear(n_embd, n_embd, bias=False)
|
||||
|
||||
# Layer-specific projection from shared RBF basis activations to per-head attention bias.
|
||||
self.rbf_proj = nn.Linear(n_rbf_bases, n_head, bias=False)
|
||||
# Keep the initial RBF attention bias exactly zero through the
|
||||
# zero-initialized projection, while leaving that projection with a
|
||||
# live gradient from the first optimization step.
|
||||
self.time_bias_scale = nn.Parameter(torch.tensor(1.0))
|
||||
|
||||
self.resid_drop = nn.Dropout(dropout)
|
||||
self.reset_parameters()
|
||||
|
||||
def reset_parameters(self) -> None:
|
||||
nn.init.normal_(self.q_proj.weight, mean=0.0, std=0.02)
|
||||
nn.init.normal_(self.k_proj.weight, mean=0.0, std=0.02)
|
||||
nn.init.normal_(self.v_proj.weight, mean=0.0, std=0.02)
|
||||
if self.rbf_proj is not None:
|
||||
# The scalar gate starts at zero, so the relative-time bias still
|
||||
# starts disabled. A nonzero projection is necessary for the gate
|
||||
# itself to receive a gradient on the first optimization step.
|
||||
nn.init.xavier_uniform_(self.rbf_proj.weight)
|
||||
|
||||
def project_event_memory(
|
||||
self,
|
||||
event_memory: torch.Tensor,
|
||||
event_rope_cache: tuple[torch.Tensor, torch.Tensor] | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Project K/V once for reuse by every reasoning round."""
|
||||
batch_size, memory_len, _ = event_memory.shape
|
||||
key = self.k_proj(event_memory).reshape(
|
||||
batch_size,
|
||||
memory_len,
|
||||
self.n_trajectory,
|
||||
self.trajectory_dim,
|
||||
).transpose(1, 2)
|
||||
value = self.v_proj(event_memory).reshape(
|
||||
batch_size,
|
||||
memory_len,
|
||||
self.n_trajectory,
|
||||
self.trajectory_dim,
|
||||
).transpose(1, 2)
|
||||
if self.use_time_rope:
|
||||
if event_rope_cache is None:
|
||||
raise ValueError(
|
||||
"event_rope_cache is required when TimeRoPE is enabled"
|
||||
)
|
||||
key = TimeRoPE.apply_single_from_cache(key, event_rope_cache)
|
||||
return key, value
|
||||
"""Match the previous version's GPT-style weight initialization."""
|
||||
nn.init.normal_(self.qkv.weight, mean=0.0, std=0.02)
|
||||
nn.init.normal_(self.out_proj.weight, mean=0.0, std=0.02)
|
||||
nn.init.zeros_(self.rbf_proj.weight)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
trajectory_state: torch.Tensor,
|
||||
event_key_value: tuple[torch.Tensor, torch.Tensor],
|
||||
event_invalid_mask: torch.Tensor,
|
||||
query_rope_cache: tuple[torch.Tensor, torch.Tensor] | None = None,
|
||||
x: torch.Tensor,
|
||||
rope_cache: tuple[torch.Tensor, torch.Tensor] | None = None,
|
||||
rbf_cache: torch.Tensor | None = None,
|
||||
attn_mask: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Read memory for states shaped ``(B, Q, H, Dh)``."""
|
||||
if trajectory_state.ndim != 4:
|
||||
raise ValueError(
|
||||
"trajectory_state must have shape (B, Q, H, Dh), got "
|
||||
f"{tuple(trajectory_state.shape)}"
|
||||
)
|
||||
batch_size, n_query, n_trajectory, trajectory_dim = (
|
||||
trajectory_state.shape
|
||||
)
|
||||
if (n_trajectory, trajectory_dim) != (
|
||||
self.n_trajectory,
|
||||
self.trajectory_dim,
|
||||
):
|
||||
raise ValueError(
|
||||
"Unexpected trajectory shape: "
|
||||
f"{(n_trajectory, trajectory_dim)}"
|
||||
)
|
||||
key, value = event_key_value
|
||||
memory_len = key.size(2)
|
||||
if event_invalid_mask.shape != (batch_size, n_query, memory_len):
|
||||
raise ValueError(
|
||||
"event_invalid_mask must have shape "
|
||||
f"{(batch_size, n_query, memory_len)}, got "
|
||||
f"{tuple(event_invalid_mask.shape)}"
|
||||
)
|
||||
|
||||
query = self.q_proj(trajectory_state).transpose(1, 2)
|
||||
if self.use_time_rope:
|
||||
if query_rope_cache is None:
|
||||
raise ValueError(
|
||||
"query_rope_cache is required when TimeRoPE is enabled"
|
||||
)
|
||||
query = TimeRoPE.apply_single_from_cache(query, query_rope_cache)
|
||||
query = query.transpose(1, 2)
|
||||
|
||||
scores = torch.einsum("bqhd,bhld->bqhl", query, key) * self.scale
|
||||
assert rope_cache is not None, "rope_cache must be provided when use_time_rope is True"
|
||||
if self.use_rbf_bias:
|
||||
if rbf_cache is None or self.rbf_proj is None:
|
||||
raise ValueError(
|
||||
"rbf_cache is required when relative time bias is enabled"
|
||||
assert rbf_cache is not None, "rbf_cache must be provided when use_rbf_bias is True"
|
||||
|
||||
B, L, _ = x.shape
|
||||
H, D = self.n_head, self.d_head
|
||||
|
||||
# --- QKV ----------------------------------------------------------
|
||||
qkv = self.qkv(x).reshape(B, L, 3, H, D).permute(2, 0, 3, 1, 4)
|
||||
q, k, v = qkv.unbind(0) # each (B, H, L, D)
|
||||
|
||||
# --- Apply RoPE (from shared cache) --------------------------------
|
||||
if self.use_time_rope:
|
||||
q, k = TimeRoPE.apply_from_cache(q, k, rope_cache)
|
||||
|
||||
# Build additive attention bias mask: time bias + causal/padding mask.
|
||||
time_bias = None
|
||||
if self.use_rbf_bias:
|
||||
time_bias = self.rbf_proj(rbf_cache).permute(
|
||||
0, 3, 1, 2) # (B, H, L, L)
|
||||
time_bias = self.time_bias_scale.tanh() * time_bias
|
||||
|
||||
if time_bias is not None and attn_mask is not None:
|
||||
attn_bias = time_bias + attn_mask.to(time_bias.dtype)
|
||||
elif time_bias is not None:
|
||||
attn_bias = time_bias
|
||||
elif attn_mask is not None:
|
||||
attn_bias = attn_mask
|
||||
else:
|
||||
attn_bias = None
|
||||
|
||||
out = F.scaled_dot_product_attention(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
attn_mask=attn_bias,
|
||||
dropout_p=0.0,
|
||||
is_causal=False,
|
||||
scale=self.scale,
|
||||
)
|
||||
time_bias = self.rbf_proj(rbf_cache).permute(0, 1, 3, 2)
|
||||
scores = scores + self.time_bias_scale.tanh() * time_bias
|
||||
|
||||
mask = event_invalid_mask.unsqueeze(2)
|
||||
min_value = torch.finfo(scores.dtype).min
|
||||
masked_scores = scores.masked_fill(mask, min_value)
|
||||
weights = torch.softmax(masked_scores.float(), dim=-1).to(scores.dtype)
|
||||
weights = weights.masked_fill(mask, 0.0)
|
||||
denominator = weights.sum(dim=-1, keepdim=True)
|
||||
weights = weights / denominator.clamp_min(
|
||||
torch.finfo(weights.dtype).eps
|
||||
)
|
||||
return torch.einsum("bqhl,bhld->bqhd", weights, value)
|
||||
# --- Aggregate & project out --------------------------------------
|
||||
out = out.transpose(1, 2).reshape(B, L, H * D)
|
||||
return self.resid_drop(self.out_proj(out))
|
||||
|
||||
|
||||
class SharedTrajectoryMixer(nn.Module):
|
||||
"""SwiGLU interaction along the trajectory axis only."""
|
||||
|
||||
def __init__(self, n_trajectory: int, trajectory_dim: int):
|
||||
class SwiGLU(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
n_embd: int,
|
||||
hidden_dim: int | None = None,
|
||||
dropout: float = 0.0,
|
||||
bias: bool = True,
|
||||
):
|
||||
super().__init__()
|
||||
if n_trajectory <= 0 or trajectory_dim <= 0:
|
||||
raise ValueError("trajectory dimensions must be positive")
|
||||
self.n_trajectory = n_trajectory
|
||||
self.trajectory_dim = trajectory_dim
|
||||
self.traj_hidden = 4 * n_trajectory
|
||||
self.gate_proj = nn.Parameter(
|
||||
torch.empty(trajectory_dim, n_trajectory, self.traj_hidden)
|
||||
)
|
||||
self.value_proj = nn.Parameter(
|
||||
torch.empty(trajectory_dim, n_trajectory, self.traj_hidden)
|
||||
)
|
||||
self.output_proj = nn.Parameter(
|
||||
torch.empty(trajectory_dim, self.traj_hidden, n_trajectory)
|
||||
)
|
||||
hidden_dim = hidden_dim if hidden_dim is not None else int(
|
||||
n_embd * 2.5)
|
||||
|
||||
self.w1 = nn.Linear(n_embd, hidden_dim, bias=bias) # gate path
|
||||
self.w2 = nn.Linear(n_embd, hidden_dim, bias=bias) # value path
|
||||
# output projection
|
||||
self.w3 = nn.Linear(hidden_dim, n_embd, bias=bias)
|
||||
self.drop = nn.Dropout(dropout)
|
||||
self.reset_parameters()
|
||||
|
||||
def reset_parameters(self) -> None:
|
||||
for feature_idx in range(self.trajectory_dim):
|
||||
"""GPT-style parameter initialization for MLP paths."""
|
||||
nn.init.normal_(self.w1.weight, mean=0.0, std=0.02)
|
||||
nn.init.normal_(self.w2.weight, mean=0.0, std=0.02)
|
||||
nn.init.normal_(self.w3.weight, mean=0.0, std=0.02)
|
||||
if self.w1.bias is not None:
|
||||
nn.init.zeros_(self.w1.bias)
|
||||
nn.init.zeros_(self.w2.bias)
|
||||
nn.init.zeros_(self.w3.bias)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""``(B, L, n_embd) -> (B, L, n_embd)``."""
|
||||
return self.drop(self.w3(F.silu(self.w1(x)) * self.w2(x)))
|
||||
|
||||
|
||||
class TrajMixer(nn.Module):
|
||||
"""PreNorm gated mixing within and across latent trajectory groups.
|
||||
|
||||
The groups are contiguous partitions of the post-``W_O`` residual
|
||||
representation. They are deliberately not treated as attention heads.
|
||||
All operations are position-wise, so the sequence dimension remains fully
|
||||
parallel and no temporal information can leak between positions here.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
n_embd: int,
|
||||
n_head: int = 10,
|
||||
dropout: float = 0.0,
|
||||
):
|
||||
super().__init__()
|
||||
if n_embd <= 0:
|
||||
raise ValueError(f"n_embd must be > 0, got {n_embd}")
|
||||
if n_head <= 0:
|
||||
raise ValueError(f"n_head must be > 0, got {n_head}")
|
||||
if n_embd % n_head != 0:
|
||||
raise ValueError(
|
||||
f"n_embd must be divisible by n_head, got {n_embd} and {n_head}"
|
||||
)
|
||||
self.n_embd = n_embd
|
||||
self.n_group = n_head
|
||||
self.d_group = n_embd // n_head
|
||||
self.intra_hidden = 4 * self.d_group
|
||||
self.hidden_group = 4 * n_head
|
||||
|
||||
self.norm = nn.LayerNorm(self.n_embd)
|
||||
|
||||
self.intra_gate_proj = nn.Parameter(
|
||||
torch.empty(self.n_group, self.d_group, self.intra_hidden)
|
||||
)
|
||||
self.intra_value_proj = nn.Parameter(
|
||||
torch.empty(self.n_group, self.d_group, self.intra_hidden)
|
||||
)
|
||||
self.intra_output_proj = nn.Parameter(
|
||||
torch.empty(self.n_group, self.intra_hidden, self.d_group)
|
||||
)
|
||||
self.intra_gate_logits = nn.Parameter(
|
||||
torch.empty(self.n_group, self.d_group)
|
||||
)
|
||||
|
||||
self.gate_proj = nn.Parameter(
|
||||
torch.empty(self.d_group, self.n_group, self.hidden_group)
|
||||
)
|
||||
self.value_proj = nn.Parameter(
|
||||
torch.empty(self.d_group, self.n_group, self.hidden_group)
|
||||
)
|
||||
self.output_proj = nn.Parameter(
|
||||
torch.empty(self.d_group, self.hidden_group, self.n_group)
|
||||
)
|
||||
self.drop = nn.Dropout(dropout)
|
||||
self.reset_parameters()
|
||||
|
||||
def reset_parameters(self) -> None:
|
||||
for group_idx in range(self.n_group):
|
||||
nn.init.xavier_uniform_(self.intra_gate_proj[group_idx])
|
||||
nn.init.xavier_uniform_(self.intra_value_proj[group_idx])
|
||||
nn.init.xavier_uniform_(self.intra_output_proj[group_idx])
|
||||
nn.init.constant_(
|
||||
self.intra_gate_logits,
|
||||
math.log(0.1 / 0.9),
|
||||
)
|
||||
|
||||
for feature_idx in range(self.d_group):
|
||||
nn.init.xavier_uniform_(self.gate_proj[feature_idx])
|
||||
nn.init.xavier_uniform_(self.value_proj[feature_idx])
|
||||
nn.init.normal_(self.output_proj, mean=0.0, std=1e-3)
|
||||
|
||||
def forward(self, state: torch.Tensor) -> torch.Tensor:
|
||||
if state.shape[-2:] != (self.n_trajectory, self.trajectory_dim):
|
||||
raise ValueError(
|
||||
"Expected trailing trajectory shape "
|
||||
f"{(self.n_trajectory, self.trajectory_dim)}, got "
|
||||
f"{tuple(state.shape[-2:])}"
|
||||
def _intra_mix(self, grouped: torch.Tensor) -> torch.Tensor:
|
||||
"""Mix features independently inside each residual-space group."""
|
||||
intra_gate = torch.einsum(
|
||||
"blgd,gdh->blgh", grouped, self.intra_gate_proj
|
||||
)
|
||||
intra_value = torch.einsum(
|
||||
"blgd,gdh->blgh", grouped, self.intra_value_proj
|
||||
)
|
||||
intra_hidden = F.silu(intra_gate) * intra_value
|
||||
return torch.einsum(
|
||||
"blgh,ghd->blgd", intra_hidden, self.intra_output_proj
|
||||
)
|
||||
|
||||
def _cross_mix(self, grouped: torch.Tensor) -> torch.Tensor:
|
||||
"""Mix groups independently for each within-group coordinate."""
|
||||
gate = torch.einsum(
|
||||
"blgr,rgh->blhr", grouped, self.gate_proj
|
||||
)
|
||||
value = torch.einsum(
|
||||
"blgr,rgh->blhr", grouped, self.value_proj
|
||||
)
|
||||
gate = torch.einsum("...hr,rhk->...kr", state, self.gate_proj)
|
||||
value = torch.einsum("...hr,rhk->...kr", state, self.value_proj)
|
||||
hidden = F.silu(gate) * value
|
||||
return torch.einsum("...kr,rkh->...hr", hidden, self.output_proj)
|
||||
return torch.einsum(
|
||||
"blhr,rhg->blgr", hidden, self.output_proj
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""Apply one full-width PreNorm and one outer residual update."""
|
||||
if x.ndim != 3:
|
||||
raise ValueError(
|
||||
f"TrajMixer expects a 3D tensor, got shape {tuple(x.shape)}"
|
||||
)
|
||||
if x.size(-1) != self.n_embd:
|
||||
raise ValueError(
|
||||
f"Expected hidden size {self.n_embd}, got {x.size(-1)}"
|
||||
)
|
||||
|
||||
batch_size, seq_len, _ = x.shape
|
||||
grouped = self.norm(x).reshape(
|
||||
batch_size, seq_len, self.n_group, self.d_group
|
||||
)
|
||||
|
||||
intra_output = self._intra_mix(grouped)
|
||||
intra_gate = torch.sigmoid(self.intra_gate_logits).view(
|
||||
1, 1, self.n_group, self.d_group
|
||||
)
|
||||
mixed_input = grouped + intra_gate * intra_output
|
||||
|
||||
update = self._cross_mix(mixed_input).reshape(
|
||||
batch_size, seq_len, self.n_embd
|
||||
)
|
||||
return x + self.drop(update)
|
||||
|
||||
|
||||
class SharedEventTrajectoryCore(nn.Module):
|
||||
"""One parameter-shared reasoning core reused across all rounds."""
|
||||
|
||||
class TransformerFFNBlock(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
d_model: int,
|
||||
n_trajectory: int,
|
||||
n_reasoning_rounds: int,
|
||||
dropout: float = 0.0,
|
||||
n_rbf_bases: int = 16,
|
||||
n_embd: int,
|
||||
n_head: int,
|
||||
|
||||
attn_dropout: float = 0.0,
|
||||
mlp_dropout: float = 0.0,
|
||||
use_time_rope: bool = False,
|
||||
use_rbf_bias: bool = False,
|
||||
n_rbf_bases: int = 16,
|
||||
):
|
||||
super().__init__()
|
||||
if n_reasoning_rounds <= 0:
|
||||
raise ValueError("n_reasoning_rounds must be positive")
|
||||
if d_model <= 0 or n_trajectory <= 0:
|
||||
raise ValueError("d_model and n_trajectory must be positive")
|
||||
if d_model % n_trajectory != 0:
|
||||
raise ValueError("d_model must equal n_trajectory * trajectory_dim")
|
||||
trajectory_dim = d_model // n_trajectory
|
||||
self.norm_attn = nn.LayerNorm(trajectory_dim)
|
||||
self.cross_attention = TrajectoryCrossAttention(
|
||||
d_model=d_model,
|
||||
n_trajectory=n_trajectory,
|
||||
self.attn = TemporalAttention(
|
||||
n_embd=n_embd,
|
||||
n_head=n_head,
|
||||
n_rbf_bases=n_rbf_bases,
|
||||
dropout=attn_dropout,
|
||||
use_time_rope=use_time_rope,
|
||||
use_rbf_bias=use_rbf_bias,
|
||||
)
|
||||
self.norm_mixer = nn.LayerNorm(trajectory_dim)
|
||||
self.traj_mixer = SharedTrajectoryMixer(
|
||||
n_trajectory=n_trajectory,
|
||||
trajectory_dim=trajectory_dim,
|
||||
)
|
||||
initial_scale = 1.0 / math.sqrt(n_reasoning_rounds)
|
||||
self.attn_scale = nn.Parameter(torch.tensor(initial_scale))
|
||||
self.mixer_scale = nn.Parameter(torch.tensor(initial_scale))
|
||||
self.dropout = nn.Dropout(dropout)
|
||||
|
||||
def project_event_memory(
|
||||
self,
|
||||
event_memory: torch.Tensor,
|
||||
event_rope_cache: tuple[torch.Tensor, torch.Tensor] | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
return self.cross_attention.project_event_memory(
|
||||
event_memory,
|
||||
event_rope_cache=event_rope_cache,
|
||||
)
|
||||
self.mlp = SwiGLU(n_embd=n_embd, dropout=mlp_dropout)
|
||||
self.ln1 = nn.LayerNorm(n_embd)
|
||||
self.ln2 = nn.LayerNorm(n_embd)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
trajectory_state: torch.Tensor,
|
||||
event_key_value: tuple[torch.Tensor, torch.Tensor],
|
||||
event_invalid_mask: torch.Tensor,
|
||||
query_rope_cache: tuple[torch.Tensor, torch.Tensor] | None = None,
|
||||
x: torch.Tensor,
|
||||
rope_cache: tuple[torch.Tensor, torch.Tensor] | None = None,
|
||||
rbf_cache: torch.Tensor | None = None,
|
||||
attn_mask: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
readout = self.cross_attention(
|
||||
trajectory_state=self.norm_attn(trajectory_state),
|
||||
event_key_value=event_key_value,
|
||||
event_invalid_mask=event_invalid_mask,
|
||||
query_rope_cache=query_rope_cache,
|
||||
rbf_cache=rbf_cache,
|
||||
x = x + self.attn(self.ln1(x), rope_cache, rbf_cache, attn_mask)
|
||||
x = x + self.mlp(self.ln2(x))
|
||||
return x
|
||||
|
||||
|
||||
class TrajMixerBlock(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
n_embd: int,
|
||||
n_head: int,
|
||||
attn_dropout: float = 0.0,
|
||||
mlp_dropout: float = 0.0,
|
||||
use_time_rope: bool = False,
|
||||
use_rbf_bias: bool = False,
|
||||
n_rbf_bases: int = 16,
|
||||
):
|
||||
super().__init__()
|
||||
self.attn = TemporalAttention(
|
||||
n_embd=n_embd,
|
||||
n_head=n_head,
|
||||
n_rbf_bases=n_rbf_bases,
|
||||
dropout=attn_dropout,
|
||||
use_time_rope=use_time_rope,
|
||||
use_rbf_bias=use_rbf_bias,
|
||||
)
|
||||
updated = (
|
||||
trajectory_state
|
||||
+ self.attn_scale * self.dropout(readout)
|
||||
self.mlp = TrajMixer(
|
||||
n_embd=n_embd,
|
||||
n_head=n_head,
|
||||
dropout=mlp_dropout,
|
||||
)
|
||||
self.ln1 = nn.LayerNorm(n_embd)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
rope_cache: tuple[torch.Tensor, torch.Tensor] | None = None,
|
||||
rbf_cache: torch.Tensor | None = None,
|
||||
attn_mask: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
x = x + self.attn(self.ln1(x), rope_cache, rbf_cache, attn_mask)
|
||||
return self.mlp(x)
|
||||
|
||||
|
||||
def build_backbone_block(
|
||||
model_architecture: str,
|
||||
*,
|
||||
n_embd: int,
|
||||
n_head: int,
|
||||
attn_dropout: float = 0.0,
|
||||
mlp_dropout: float = 0.0,
|
||||
use_time_rope: bool = False,
|
||||
use_rbf_bias: bool = False,
|
||||
n_rbf_bases: int = 16,
|
||||
) -> nn.Module:
|
||||
"""Build one history block for a supported model architecture."""
|
||||
architecture = resolve_model_architecture(model_architecture)
|
||||
block_class: type[nn.Module]
|
||||
if architecture == TRANSFORMER_FFN_ARCHITECTURE:
|
||||
block_class = TransformerFFNBlock
|
||||
elif architecture == TRAJ_MIXER_ARCHITECTURE:
|
||||
block_class = TrajMixerBlock
|
||||
else: # pragma: no cover - guarded by resolve_model_architecture.
|
||||
raise ValueError(f"Unsupported model architecture: {architecture!r}")
|
||||
return block_class(
|
||||
n_embd=n_embd,
|
||||
n_head=n_head,
|
||||
attn_dropout=attn_dropout,
|
||||
mlp_dropout=mlp_dropout,
|
||||
use_time_rope=use_time_rope,
|
||||
use_rbf_bias=use_rbf_bias,
|
||||
n_rbf_bases=n_rbf_bases,
|
||||
)
|
||||
mixed = self.traj_mixer(self.norm_mixer(updated))
|
||||
return updated + self.mixer_scale * self.dropout(mixed)
|
||||
|
||||
|
||||
class TokenAutoDiscretization(nn.Module):
|
||||
|
||||
367
dataset.py
367
dataset.py
@@ -10,15 +10,173 @@ from torch.nn.utils.rnn import pad_sequence
|
||||
from torch.utils.data import Dataset
|
||||
|
||||
from targets import (
|
||||
CHECKUP_IDX,
|
||||
DAYS_PER_YEAR,
|
||||
NO_EVENT_IDX,
|
||||
PAD_IDX,
|
||||
build_all_targets,
|
||||
RESERVED_IDX,
|
||||
build_next_token_targets,
|
||||
)
|
||||
|
||||
DISEASE_HISTORY_MODE_TIMED = "timed"
|
||||
DISEASE_HISTORY_MODE_ORDERED = "ordered"
|
||||
DISEASE_HISTORY_MODE_SET = "set"
|
||||
DISEASE_HISTORY_MODES = (
|
||||
DISEASE_HISTORY_MODE_TIMED,
|
||||
DISEASE_HISTORY_MODE_ORDERED,
|
||||
DISEASE_HISTORY_MODE_SET,
|
||||
)
|
||||
|
||||
|
||||
ONE_DAY_YEARS = 1.0 / DAYS_PER_YEAR
|
||||
def normalize_disease_history_mode(mode: str | None) -> str:
|
||||
value = DISEASE_HISTORY_MODE_TIMED if mode is None else str(mode).lower()
|
||||
if value not in DISEASE_HISTORY_MODES:
|
||||
raise ValueError(
|
||||
"disease_history_mode must be one of "
|
||||
f"{list(DISEASE_HISTORY_MODES)}, got {mode!r}"
|
||||
)
|
||||
return value
|
||||
|
||||
|
||||
def transform_disease_history(
|
||||
event_seq: np.ndarray,
|
||||
actual_time_seq: np.ndarray,
|
||||
actual_t_query: float,
|
||||
disease_history_mode: str,
|
||||
) -> Tuple[np.ndarray, np.ndarray, np.float32]:
|
||||
"""
|
||||
Convert an already-truncated disease history into its model representation.
|
||||
|
||||
``timed`` keeps the real event/query times. ``ordered`` preserves the
|
||||
chronological event order but replaces calendar time with ordinal event-time
|
||||
groups. Diseases first recorded on the same day share one ordinal position.
|
||||
``set`` removes both time and order by sorting the unique disease codes and
|
||||
assigning every disease and the query the same model time.
|
||||
"""
|
||||
mode = normalize_disease_history_mode(disease_history_mode)
|
||||
events = np.asarray(event_seq, dtype=np.int64)
|
||||
times = np.asarray(actual_time_seq, dtype=np.float32)
|
||||
if events.ndim != 1 or times.ndim != 1 or events.shape != times.shape:
|
||||
raise ValueError(
|
||||
"event_seq and actual_time_seq must be aligned 1D arrays, got "
|
||||
f"{events.shape} and {times.shape}"
|
||||
)
|
||||
|
||||
if mode == DISEASE_HISTORY_MODE_TIMED:
|
||||
return events, times, np.float32(actual_t_query)
|
||||
|
||||
special = events <= NO_EVENT_IDX
|
||||
if np.any(special):
|
||||
raise ValueError(
|
||||
f"{mode} disease history must contain only disease events; "
|
||||
f"found special token ids {np.unique(events[special]).tolist()}"
|
||||
)
|
||||
|
||||
if mode == DISEASE_HISTORY_MODE_ORDERED:
|
||||
_, ordinal_groups = np.unique(times, return_inverse=True)
|
||||
model_times = ordinal_groups.astype(np.float32, copy=False)
|
||||
n_groups = int(model_times.max()) + 1 if model_times.size else 0
|
||||
return events, model_times, np.float32(n_groups)
|
||||
|
||||
set_events = np.unique(events)
|
||||
model_times = np.zeros(set_events.size, dtype=np.float32)
|
||||
return set_events, model_times, np.float32(0.0)
|
||||
|
||||
|
||||
def transform_disease_history_batch_at_position(
|
||||
event_seq: torch.Tensor,
|
||||
actual_time_seq: torch.Tensor,
|
||||
padding_mask: torch.Tensor,
|
||||
query_position: int,
|
||||
disease_history_mode: str,
|
||||
vocab_size: int,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Build a model-visible prefix for token-position all-future evaluation.
|
||||
|
||||
Actual event times remain outside this return value for AUC bookkeeping.
|
||||
For ordered/set modes, events after ``query_position`` are explicitly
|
||||
masked so collapsing time cannot expose future diseases.
|
||||
"""
|
||||
mode = normalize_disease_history_mode(disease_history_mode)
|
||||
if event_seq.ndim != 2 or actual_time_seq.shape != event_seq.shape:
|
||||
raise ValueError(
|
||||
"event_seq and actual_time_seq must be aligned 2D tensors, got "
|
||||
f"{tuple(event_seq.shape)} and {tuple(actual_time_seq.shape)}"
|
||||
)
|
||||
if padding_mask.shape != event_seq.shape:
|
||||
raise ValueError(
|
||||
"padding_mask must match event_seq, got "
|
||||
f"{tuple(padding_mask.shape)} and {tuple(event_seq.shape)}"
|
||||
)
|
||||
if query_position < 0 or query_position >= event_seq.size(1):
|
||||
raise ValueError(
|
||||
f"query_position={query_position} is outside sequence length "
|
||||
f"{event_seq.size(1)}"
|
||||
)
|
||||
|
||||
padding_mask = padding_mask.to(device=event_seq.device, dtype=torch.bool)
|
||||
if not torch.all(padding_mask[:, query_position]):
|
||||
raise ValueError("query_position must be valid for every batch row")
|
||||
|
||||
if mode == DISEASE_HISTORY_MODE_TIMED:
|
||||
return (
|
||||
event_seq,
|
||||
actual_time_seq,
|
||||
padding_mask,
|
||||
actual_time_seq[:, query_position],
|
||||
)
|
||||
|
||||
positions = torch.arange(
|
||||
event_seq.size(1),
|
||||
device=event_seq.device,
|
||||
)[None, :]
|
||||
history_mask = padding_mask & (positions <= query_position)
|
||||
visible_events = event_seq.masked_select(history_mask)
|
||||
if torch.any(visible_events <= NO_EVENT_IDX):
|
||||
special_ids = torch.unique(
|
||||
visible_events[visible_events <= NO_EVENT_IDX]
|
||||
).detach().cpu().tolist()
|
||||
raise ValueError(
|
||||
f"{mode} disease history must contain only disease events; "
|
||||
f"found special token ids {special_ids}"
|
||||
)
|
||||
|
||||
if mode == DISEASE_HISTORY_MODE_ORDERED:
|
||||
model_times = torch.zeros_like(actual_time_seq)
|
||||
model_t_query = torch.zeros(
|
||||
event_seq.size(0),
|
||||
device=actual_time_seq.device,
|
||||
dtype=actual_time_seq.dtype,
|
||||
)
|
||||
for row_idx in range(event_seq.size(0)):
|
||||
row_mask = history_mask[row_idx]
|
||||
_, ordinal_groups = torch.unique(
|
||||
actual_time_seq[row_idx, row_mask],
|
||||
sorted=True,
|
||||
return_inverse=True,
|
||||
)
|
||||
model_times[row_idx, row_mask] = ordinal_groups.to(
|
||||
dtype=actual_time_seq.dtype
|
||||
)
|
||||
model_t_query[row_idx] = float(
|
||||
int(ordinal_groups.max().item()) + 1
|
||||
if ordinal_groups.numel()
|
||||
else 0
|
||||
)
|
||||
return event_seq, model_times, history_mask, model_t_query
|
||||
|
||||
sentinel = torch.full_like(event_seq, int(vocab_size))
|
||||
sortable = torch.where(history_mask, event_seq, sentinel)
|
||||
set_events = torch.sort(sortable, dim=1).values
|
||||
set_mask = set_events != int(vocab_size)
|
||||
set_events = set_events.masked_fill(~set_mask, PAD_IDX)
|
||||
model_times = torch.zeros_like(actual_time_seq)
|
||||
model_t_query = torch.zeros(
|
||||
event_seq.size(0),
|
||||
device=actual_time_seq.device,
|
||||
dtype=actual_time_seq.dtype,
|
||||
)
|
||||
return set_events, model_times, set_mask, model_t_query
|
||||
|
||||
|
||||
def load_label_vocab(
|
||||
@@ -27,12 +185,12 @@ def load_label_vocab(
|
||||
) -> Tuple[Dict[str, int], Dict[int, str]]:
|
||||
label_id_to_code: Dict[int, str] = {
|
||||
PAD_IDX: "<PAD>",
|
||||
CHECKUP_IDX: "<CHECKUP>",
|
||||
RESERVED_IDX: "<RESERVED>",
|
||||
}
|
||||
if include_no_event:
|
||||
label_id_to_code[NO_EVENT_IDX] = "<NO_EVENT>"
|
||||
|
||||
offset = NO_EVENT_IDX + 1 if include_no_event else CHECKUP_IDX + 1
|
||||
offset = NO_EVENT_IDX + 1 if include_no_event else RESERVED_IDX + 1
|
||||
label_code_to_id: Dict[str, int] = {}
|
||||
with open(labels_file, encoding="utf-8") as f:
|
||||
for i, line in enumerate(f):
|
||||
@@ -86,13 +244,11 @@ class _ExpoBaseDataset(Dataset):
|
||||
data_prefix: str = "ukb",
|
||||
labels_file: str = "labels.csv",
|
||||
no_event_interval_years: float = 5.0,
|
||||
include_no_event_in_uts_target: bool = False,
|
||||
extra_info_types: Iterable[int] | None = None,
|
||||
) -> None:
|
||||
self.data_prefix = data_prefix
|
||||
self.labels_file = labels_file
|
||||
self.no_event_interval_years = float(no_event_interval_years)
|
||||
self.include_no_event_in_uts_target = bool(include_no_event_in_uts_target)
|
||||
self.requested_extra_info_types = (
|
||||
None
|
||||
if extra_info_types is None
|
||||
@@ -138,11 +294,6 @@ class _ExpoBaseDataset(Dataset):
|
||||
max_id_in_data += 1
|
||||
self.vocab_size = max(max_id_in_vocab, max_id_in_data) + 1
|
||||
|
||||
if not self.include_no_event_in_uts_target:
|
||||
self.ignored_uts_target_ids = {PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX}
|
||||
else:
|
||||
self.ignored_uts_target_ids = {PAD_IDX, CHECKUP_IDX}
|
||||
|
||||
def _prepare_sex(self, basic_table: pd.DataFrame, unique_eids: np.ndarray) -> None:
|
||||
sex_values = pd.to_numeric(basic_table["sex"], errors="coerce").to_numpy()
|
||||
if np.isnan(sex_values).any():
|
||||
@@ -262,6 +413,13 @@ class _ExpoBaseDataset(Dataset):
|
||||
times_days_raw = rows[:, 1].astype(np.float32)
|
||||
labels_raw = rows[:, 2].astype(np.int64)
|
||||
|
||||
# Label 1 was emitted as a CHECKUP event by older prepared files.
|
||||
# It is now an unused reserved slot and must never enter either the
|
||||
# next-token or all-future disease sequence.
|
||||
keep = labels_raw != RESERVED_IDX
|
||||
times_days_raw = times_days_raw[keep]
|
||||
labels_raw = labels_raw[keep]
|
||||
|
||||
if len(labels_raw) == 0:
|
||||
yield eid, times_days_raw, labels_raw
|
||||
continue
|
||||
@@ -289,12 +447,7 @@ class _ExpoBaseDataset(Dataset):
|
||||
|
||||
class NextStepHealthDataset(_ExpoBaseDataset):
|
||||
"""
|
||||
Dataset for next-token and next-time-point losses with unified other-info
|
||||
tokens.
|
||||
|
||||
Returned targets cover both:
|
||||
- Delphi2MLoss: target_event_seq, target_time_seq
|
||||
- UniqueTimeSetExponentialLoss: readout_mask, target_dt_unique, target_multi_hot
|
||||
Delphi2M next-token dataset with unified other-info tokens.
|
||||
"""
|
||||
|
||||
CACHE_VERSION = 3
|
||||
@@ -304,14 +457,12 @@ class NextStepHealthDataset(_ExpoBaseDataset):
|
||||
data_prefix: str = "ukb",
|
||||
labels_file: str = "labels.csv",
|
||||
no_event_interval_years: float = 5.0,
|
||||
include_no_event_in_uts_target: bool = False,
|
||||
extra_info_types: Iterable[int] | None = None,
|
||||
) -> None:
|
||||
super().__init__(
|
||||
data_prefix=data_prefix,
|
||||
labels_file=labels_file,
|
||||
no_event_interval_years=no_event_interval_years,
|
||||
include_no_event_in_uts_target=include_no_event_in_uts_target,
|
||||
extra_info_types=extra_info_types,
|
||||
)
|
||||
|
||||
@@ -326,23 +477,18 @@ class NextStepHealthDataset(_ExpoBaseDataset):
|
||||
if features is None:
|
||||
continue
|
||||
|
||||
target_pack = build_all_targets(
|
||||
targets = build_next_token_targets(
|
||||
labels=labels,
|
||||
times_days=times_days,
|
||||
vocab_size=self.vocab_size,
|
||||
ignored_uts_target_ids=self.ignored_uts_target_ids,
|
||||
require_sorted=True,
|
||||
)
|
||||
|
||||
self.samples.append({
|
||||
"eid": eid,
|
||||
"event_seq": target_pack.next_token.input_events,
|
||||
"time_seq": target_pack.next_token.input_times_years,
|
||||
"target_event_seq": target_pack.next_token.target_events,
|
||||
"target_time_seq": target_pack.next_token.target_times_years,
|
||||
"readout_mask": target_pack.unique_time_set.readout_mask,
|
||||
"target_dt_unique": target_pack.unique_time_set.target_dt_unique,
|
||||
"target_multi_hot": target_pack.unique_time_set.target_multi_hot,
|
||||
"event_seq": targets.input_events,
|
||||
"time_seq": targets.input_times_years,
|
||||
"target_event_seq": targets.target_events,
|
||||
"target_time_seq": targets.target_times_years,
|
||||
**features,
|
||||
})
|
||||
|
||||
@@ -361,9 +507,6 @@ class NextStepHealthDataset(_ExpoBaseDataset):
|
||||
"other_time": torch.from_numpy(s["other_time"]).float(),
|
||||
"target_event_seq": torch.from_numpy(s["target_event_seq"]).long(),
|
||||
"target_time_seq": torch.from_numpy(s["target_time_seq"]).float(),
|
||||
"readout_mask": torch.from_numpy(s["readout_mask"]).bool(),
|
||||
"target_dt_unique": torch.from_numpy(s["target_dt_unique"]).float(),
|
||||
"target_multi_hot": torch.from_numpy(s["target_multi_hot"]).bool(),
|
||||
}
|
||||
|
||||
|
||||
@@ -372,10 +515,10 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
|
||||
Dataset with unified other-info tokens and DeepHealthV2-style all-future
|
||||
targets.
|
||||
|
||||
Train samples one query time per patient at each __getitem__ call.
|
||||
Valid/test use random-but-fixed query points. For each patient with N real
|
||||
disease events, N - 2 query points are sampled from the eligible observed
|
||||
time range, with at least one future event after every query.
|
||||
Every split uses the same patient-equal query distribution: choose one
|
||||
eligible inter-event interval uniformly, then choose a time uniformly in
|
||||
that interval. Train resamples on every ``__getitem__`` call; valid/test
|
||||
keep one deterministic draw per patient.
|
||||
"""
|
||||
|
||||
CACHE_VERSION = 5
|
||||
@@ -386,11 +529,11 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
|
||||
labels_file: str = "labels.csv",
|
||||
split: Literal["train", "valid", "test"] = "train",
|
||||
no_event_interval_years: float = 5.0,
|
||||
include_no_event_in_uts_target: bool = False,
|
||||
min_history_events: int = 1,
|
||||
min_future_events: int = 1,
|
||||
validation_query_seed: int = 42,
|
||||
extra_info_types: Iterable[int] | None = None,
|
||||
disease_history_mode: str = DISEASE_HISTORY_MODE_TIMED,
|
||||
) -> None:
|
||||
if split not in {"train", "valid", "test"}:
|
||||
raise ValueError(f"split must be train/valid/test, got {split!r}")
|
||||
@@ -399,10 +542,21 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
|
||||
data_prefix=data_prefix,
|
||||
labels_file=labels_file,
|
||||
no_event_interval_years=no_event_interval_years,
|
||||
include_no_event_in_uts_target=include_no_event_in_uts_target,
|
||||
extra_info_types=extra_info_types,
|
||||
)
|
||||
|
||||
self.disease_history_mode = normalize_disease_history_mode(
|
||||
disease_history_mode
|
||||
)
|
||||
if (
|
||||
self.disease_history_mode != DISEASE_HISTORY_MODE_TIMED
|
||||
and self.extra_info_types
|
||||
):
|
||||
raise ValueError(
|
||||
f"disease_history_mode={self.disease_history_mode!r} is only "
|
||||
"supported with an explicitly empty extra-info selection"
|
||||
)
|
||||
|
||||
self.split = split
|
||||
self.min_history_events = int(min_history_events)
|
||||
self.min_future_events = int(min_future_events)
|
||||
@@ -434,6 +588,11 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
|
||||
**features,
|
||||
}
|
||||
|
||||
query_intervals = self._eligible_query_intervals(patient)
|
||||
if not query_intervals:
|
||||
continue
|
||||
patient["query_intervals"] = query_intervals
|
||||
|
||||
pidx = len(self.patients)
|
||||
self.patients.append(patient)
|
||||
|
||||
@@ -456,9 +615,9 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
|
||||
labels = patient["labels"]
|
||||
real_event_mask = ~np.isin(
|
||||
labels,
|
||||
np.array([PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX], dtype=np.int64),
|
||||
np.array([PAD_IDX, RESERVED_IDX, NO_EVENT_IDX], dtype=np.int64),
|
||||
)
|
||||
n_hist = int((times <= t_query).sum())
|
||||
n_hist = int(((times <= t_query) & real_event_mask).sum())
|
||||
n_future = int(((times > t_query) & real_event_mask).sum())
|
||||
return (
|
||||
n_hist >= self.min_history_events
|
||||
@@ -466,73 +625,76 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
|
||||
and patient["t_obs"] > t_query
|
||||
)
|
||||
|
||||
def _eligible_query_intervals(self, patient: Dict) -> List[Tuple[float, float]]:
|
||||
times = np.asarray(patient["times"], dtype=np.float32)
|
||||
labels = np.asarray(patient["labels"], dtype=np.int64)
|
||||
real_event_mask = ~np.isin(
|
||||
labels,
|
||||
np.array([PAD_IDX, RESERVED_IDX, NO_EVENT_IDX], dtype=np.int64),
|
||||
)
|
||||
unique_times = np.unique(times[real_event_mask])
|
||||
intervals: List[Tuple[float, float]] = []
|
||||
for j in range(1, len(unique_times)):
|
||||
left = float(unique_times[j - 1])
|
||||
right = float(unique_times[j])
|
||||
probe = float(
|
||||
np.nextafter(np.float32(right), np.float32(-np.inf))
|
||||
)
|
||||
if np.isfinite(left) and np.isfinite(probe) and probe >= left:
|
||||
if self._is_valid_query(patient, probe):
|
||||
intervals.append((left, right))
|
||||
return intervals
|
||||
|
||||
def sample_query(
|
||||
self,
|
||||
patient: Dict,
|
||||
rng,
|
||||
) -> float:
|
||||
intervals = patient.get("query_intervals")
|
||||
if intervals is None:
|
||||
intervals = self._eligible_query_intervals(patient)
|
||||
if not intervals:
|
||||
raise RuntimeError("Patient has no eligible all-future query interval.")
|
||||
|
||||
interval_idx = int(rng.randint(0, len(intervals)))
|
||||
left, right_event_time = intervals[interval_idx]
|
||||
right = float(
|
||||
np.nextafter(np.float32(right_event_time), np.float32(-np.inf))
|
||||
)
|
||||
if right <= left:
|
||||
t_query = float(left)
|
||||
else:
|
||||
t_query = float(rng.uniform(left, right))
|
||||
if not self._is_valid_query(patient, t_query):
|
||||
raise RuntimeError("Sampled an invalid all-future query time.")
|
||||
return t_query
|
||||
|
||||
def _sample_fixed_validation_queries(
|
||||
self,
|
||||
patient: Dict,
|
||||
rng: np.random.RandomState,
|
||||
) -> List[float]:
|
||||
times = np.asarray(patient["times"], dtype=np.float32)
|
||||
labels = np.asarray(patient["labels"], dtype=np.int64)
|
||||
real_event_mask = ~np.isin(
|
||||
labels,
|
||||
np.array([PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX], dtype=np.int64),
|
||||
)
|
||||
real_times = np.sort(times[real_event_mask].astype(np.float32, copy=False))
|
||||
n_real_events = int(real_times.size)
|
||||
n_queries = max(0, n_real_events - 2)
|
||||
if n_queries == 0:
|
||||
return []
|
||||
|
||||
min_hist = int(self.min_history_events)
|
||||
min_future = int(self.min_future_events)
|
||||
if n_real_events < min_hist + min_future:
|
||||
return []
|
||||
|
||||
left = float(real_times[min_hist - 1])
|
||||
right_event_time = float(real_times[n_real_events - min_future])
|
||||
right = np.nextafter(np.float32(right_event_time), np.float32(-np.inf))
|
||||
if not np.isfinite(left) or not np.isfinite(right) or float(right) <= left:
|
||||
return []
|
||||
|
||||
queries: List[float] = []
|
||||
max_attempts = max(100, n_queries * 50)
|
||||
for _ in range(max_attempts):
|
||||
if len(queries) >= n_queries:
|
||||
break
|
||||
t_query = float(rng.uniform(left, float(right)))
|
||||
if self._is_valid_query(patient, t_query):
|
||||
queries.append(t_query)
|
||||
|
||||
return queries
|
||||
return [self.sample_query(patient, rng)]
|
||||
|
||||
def _sample_train_query(self, patient: Dict) -> float:
|
||||
unique_times = np.unique(patient["times"])
|
||||
if len(unique_times) < 2:
|
||||
raise RuntimeError("Training patient has fewer than two unique times.")
|
||||
|
||||
j = np.random.randint(1, len(unique_times))
|
||||
left = float(unique_times[j - 1])
|
||||
right = float(unique_times[j])
|
||||
|
||||
if right - left <= ONE_DAY_YEARS:
|
||||
t_query = right - ONE_DAY_YEARS
|
||||
else:
|
||||
t_query = np.random.uniform(left, right - ONE_DAY_YEARS)
|
||||
|
||||
if not self._is_valid_query(patient, t_query):
|
||||
t_query = right - 1e-6
|
||||
return float(t_query)
|
||||
return self.sample_query(patient, np.random)
|
||||
|
||||
def _build_item(self, patient: Dict, t_query: float) -> Dict:
|
||||
times = patient["times"]
|
||||
labels = patient["labels"]
|
||||
hist = times <= t_query
|
||||
fut = times > t_query
|
||||
event_seq, model_time_seq, model_t_query = transform_disease_history(
|
||||
event_seq=labels[hist],
|
||||
actual_time_seq=times[hist],
|
||||
actual_t_query=t_query,
|
||||
disease_history_mode=self.disease_history_mode,
|
||||
)
|
||||
|
||||
return {
|
||||
"event_seq": torch.from_numpy(labels[hist]).long(),
|
||||
"time_seq": torch.from_numpy(times[hist]).float(),
|
||||
"t_query": torch.tensor(t_query, dtype=torch.float32),
|
||||
"event_seq": torch.from_numpy(event_seq).long(),
|
||||
"time_seq": torch.from_numpy(model_time_seq).float(),
|
||||
"t_query": torch.tensor(model_t_query, dtype=torch.float32),
|
||||
"future_targets": torch.from_numpy(labels[fut]).long(),
|
||||
"future_dt": torch.from_numpy(times[fut] - np.float32(t_query)).float(),
|
||||
"exposure": torch.tensor(np.float32(patient["t_obs"] - t_query), dtype=torch.float32),
|
||||
@@ -605,31 +767,12 @@ def next_step_collate_fn(batch: List[Dict]) -> Dict:
|
||||
batch_first=True,
|
||||
padding_value=0.0,
|
||||
)
|
||||
readout_mask = pad_sequence(
|
||||
[s["readout_mask"] for s in batch],
|
||||
batch_first=True,
|
||||
padding_value=False,
|
||||
)
|
||||
target_dt_unique = pad_sequence(
|
||||
[s["target_dt_unique"] for s in batch],
|
||||
batch_first=True,
|
||||
padding_value=0.0,
|
||||
)
|
||||
target_multi_hot = pad_sequence(
|
||||
[s["target_multi_hot"] for s in batch],
|
||||
batch_first=True,
|
||||
padding_value=False,
|
||||
)
|
||||
|
||||
out = {
|
||||
"event_seq": event_seq,
|
||||
"time_seq": time_seq,
|
||||
"padding_mask": event_seq > PAD_IDX,
|
||||
"target_event_seq": target_event_seq,
|
||||
"target_time_seq": target_time_seq,
|
||||
"readout_mask": readout_mask,
|
||||
"target_dt_unique": target_dt_unique,
|
||||
"target_multi_hot": target_multi_hot,
|
||||
}
|
||||
out.update(_collate_common_static(batch))
|
||||
return out
|
||||
|
||||
183
delphi2m_auc_report.py
Normal file
183
delphi2m_auc_report.py
Normal file
@@ -0,0 +1,183 @@
|
||||
"""Build Delphi2M-style sex-specific AUC reports.
|
||||
|
||||
The Delphi2M evaluation code uses 0.1 years for the no-gap evaluation. The
|
||||
published report displays that point as 0 months, while retaining the actual
|
||||
0.1-year evaluation period in this project's report output.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from typing import Dict, Optional
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
|
||||
|
||||
DEFAULT_DELPHI2M_PERIODS_YEARS = (0.1, 1.0, 5.0, 10.0)
|
||||
|
||||
_CHAPTER_SHORT_NAMES = {
|
||||
"I": "I. Infectious Diseases",
|
||||
"II": "II. Neoplasms",
|
||||
"III": "III. Blood & Immune Disorders",
|
||||
"IV": "IV. Metabolic Diseases",
|
||||
"V": "V. Mental Disorders",
|
||||
"VI": "VI. Nervous System Diseases",
|
||||
"VII": "VII. Eye Diseases",
|
||||
"VIII": "VIII. Ear Diseases",
|
||||
"IX": "IX. Circulatory Diseases",
|
||||
"X": "X. Respiratory Diseases",
|
||||
"XI": "XI. Digestive Diseases",
|
||||
"XII": "XII. Skin Diseases",
|
||||
"XIII": "XIII. Musculoskeletal Diseases",
|
||||
"XIV": "XIV. Genitourinary Diseases",
|
||||
"XV": "XV. Pregnancy & Childbirth",
|
||||
"XVI": "XVI. Perinatal Conditions",
|
||||
"XVII": "XVII. Congenital Abnormalities",
|
||||
"XVIII": "XVIII. Symptoms & Signs",
|
||||
"XIX": "XIX. Injury & Poisoning",
|
||||
"XX": "XX. External Causes",
|
||||
"XXI": "XXI. Health Services",
|
||||
"XXII": "XXII. Special Purposes",
|
||||
"Death": "Death",
|
||||
"Unmapped": "Unmapped",
|
||||
}
|
||||
|
||||
|
||||
def _is_no_gap(period_years: float) -> bool:
|
||||
return bool(np.isclose(float(period_years), 0.1, rtol=0.0, atol=1e-8))
|
||||
|
||||
|
||||
def _canonical_period_years(period_years: float) -> float:
|
||||
value = float(period_years)
|
||||
for canonical in DEFAULT_DELPHI2M_PERIODS_YEARS:
|
||||
if np.isclose(value, canonical, rtol=0.0, atol=1e-6):
|
||||
return float(canonical)
|
||||
return value
|
||||
|
||||
|
||||
def _gap_months(period_years: float) -> int:
|
||||
if _is_no_gap(period_years):
|
||||
return 0
|
||||
return int(round(float(period_years) * 12.0))
|
||||
|
||||
|
||||
def _gap_label(period_years: float) -> str:
|
||||
if _is_no_gap(period_years):
|
||||
return "No gap"
|
||||
value = float(period_years)
|
||||
value_text = f"{value:g}"
|
||||
unit = "year" if np.isclose(value, 1.0) else "years"
|
||||
return f"{value_text} {unit}"
|
||||
|
||||
|
||||
def _load_chapter_by_code(
|
||||
chapter_mapping_path: Optional[str | Path] = None,
|
||||
) -> Dict[str, str]:
|
||||
if chapter_mapping_path is None:
|
||||
chapter_mapping_path = Path(__file__).with_name(
|
||||
"icd10_chapter_organ_mapping.csv"
|
||||
)
|
||||
path = Path(chapter_mapping_path)
|
||||
if not path.exists():
|
||||
return {}
|
||||
|
||||
mapping = pd.read_csv(
|
||||
path,
|
||||
usecols=["code", "icd10_chapter"],
|
||||
dtype={"code": str, "icd10_chapter": str},
|
||||
)
|
||||
mapping["code"] = mapping["code"].str.strip()
|
||||
mapping["chapter"] = (
|
||||
mapping["icd10_chapter"]
|
||||
.str.strip()
|
||||
.map(_CHAPTER_SHORT_NAMES)
|
||||
.fillna("Unmapped")
|
||||
)
|
||||
return dict(zip(mapping["code"], mapping["chapter"]))
|
||||
|
||||
|
||||
def build_delphi2m_auc_report(
|
||||
df_unpooled: pd.DataFrame,
|
||||
*,
|
||||
period_col: str,
|
||||
chapter_mapping_path: Optional[str | Path] = None,
|
||||
) -> pd.DataFrame:
|
||||
"""Aggregate age strata by sex and return a Delphi2M-style AUC report.
|
||||
|
||||
Required input columns are ``token``, ``label_code``, ``sex``,
|
||||
``auc_delong``, and the supplied ``period_col`` (``offset`` or
|
||||
``horizon``). The output begins with the five columns used by Delphi2M
|
||||
Fig. 2e and then records the actual evaluation period and ICD-10 code.
|
||||
"""
|
||||
required = {"token", "label_code", "sex", "auc_delong", period_col}
|
||||
missing = sorted(required - set(df_unpooled.columns))
|
||||
if missing:
|
||||
raise ValueError(
|
||||
"Cannot build Delphi2M AUC report; missing columns: "
|
||||
+ ", ".join(missing)
|
||||
)
|
||||
|
||||
source = df_unpooled.loc[
|
||||
:,
|
||||
["token", "label_code", "sex", "auc_delong", period_col],
|
||||
].copy()
|
||||
source["sex"] = source["sex"].astype(str).str.strip().str.lower()
|
||||
source = source[source["sex"].isin(["female", "male"])]
|
||||
source["auc_delong"] = pd.to_numeric(
|
||||
source["auc_delong"], errors="coerce"
|
||||
)
|
||||
source[period_col] = pd.to_numeric(source[period_col], errors="coerce")
|
||||
source = source.dropna(subset=[period_col, "auc_delong"])
|
||||
source[period_col] = source[period_col].map(_canonical_period_years)
|
||||
|
||||
if source.empty:
|
||||
raise ValueError("Cannot build Delphi2M AUC report from empty AUC data.")
|
||||
|
||||
grouped = (
|
||||
source.groupby(
|
||||
["token", "label_code", period_col, "sex"],
|
||||
dropna=False,
|
||||
as_index=False,
|
||||
)
|
||||
.agg(auc=("auc_delong", "mean"))
|
||||
)
|
||||
report = (
|
||||
grouped.pivot(
|
||||
index=["token", "label_code", period_col],
|
||||
columns="sex",
|
||||
values="auc",
|
||||
)
|
||||
.reset_index()
|
||||
.rename_axis(columns=None)
|
||||
.rename(columns={"female": "Female", "male": "Male"})
|
||||
)
|
||||
for col in ["Female", "Male"]:
|
||||
if col not in report.columns:
|
||||
report[col] = np.nan
|
||||
|
||||
chapter_by_code = _load_chapter_by_code(chapter_mapping_path)
|
||||
report["chapter"] = (
|
||||
report["label_code"].astype(str).map(chapter_by_code).fillna("Unmapped")
|
||||
)
|
||||
report["Gap, months"] = report[period_col].map(_gap_months).astype("Int64")
|
||||
report["Gap label"] = report[period_col].map(_gap_label)
|
||||
report["icd10"] = pd.to_numeric(report["token"], errors="coerce").astype(
|
||||
"Int64"
|
||||
)
|
||||
|
||||
report = report.sort_values(
|
||||
["icd10", period_col], kind="stable", ignore_index=True
|
||||
)
|
||||
return report.loc[
|
||||
:,
|
||||
[
|
||||
"Gap, months",
|
||||
"chapter",
|
||||
"icd10",
|
||||
"Female",
|
||||
"Male",
|
||||
period_col,
|
||||
"Gap label",
|
||||
"label_code",
|
||||
],
|
||||
]
|
||||
326
eval_data.py
326
eval_data.py
@@ -1,22 +1,323 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Dict, Iterable, List
|
||||
import argparse
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Iterable, List, Optional, Tuple
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch.nn.utils.rnn import pad_sequence
|
||||
|
||||
from dataset import AllFutureHealthDataset, HealthDataset
|
||||
from dataset import (
|
||||
DISEASE_HISTORY_MODE_TIMED,
|
||||
AllFutureHealthDataset,
|
||||
HealthDataset,
|
||||
normalize_disease_history_mode,
|
||||
)
|
||||
from model_architectures import resolve_model_architecture
|
||||
from models import DeepHealth
|
||||
from targets import PAD_IDX
|
||||
|
||||
|
||||
def load_json_config(path: str | Path | None) -> Dict[str, Any]:
|
||||
if path is None:
|
||||
return {}
|
||||
config_path = Path(path)
|
||||
if not config_path.exists():
|
||||
return {}
|
||||
with config_path.open("r", encoding="utf-8") as file:
|
||||
return json.load(file)
|
||||
|
||||
|
||||
def cfg_get(
|
||||
args: argparse.Namespace | Dict[str, Any] | None,
|
||||
cfg: Dict[str, Any],
|
||||
name: str,
|
||||
default: Any,
|
||||
) -> Any:
|
||||
if args is not None:
|
||||
value = (
|
||||
args.get(name)
|
||||
if isinstance(args, dict)
|
||||
else getattr(args, name, None)
|
||||
)
|
||||
if value is not None:
|
||||
return value
|
||||
return cfg.get(name, default)
|
||||
|
||||
|
||||
def resolve_eval_device(device_arg: Optional[str]) -> torch.device:
|
||||
device_name = device_arg or ("cuda" if torch.cuda.is_available() else "cpu")
|
||||
device = torch.device(device_name)
|
||||
if device.type == "cuda" and not torch.cuda.is_available():
|
||||
raise RuntimeError(
|
||||
f"Requested device {device_name!r}, but CUDA is not available."
|
||||
)
|
||||
return device
|
||||
|
||||
|
||||
def validate_training_mode_config(cfg: Dict[str, Any]) -> None:
|
||||
model_target_mode = str(
|
||||
cfg.get("model_target_mode", "next_token")
|
||||
).lower()
|
||||
if model_target_mode not in {"next_token", "all_future"}:
|
||||
raise ValueError(
|
||||
"model_target_mode must be next_token or all_future, got "
|
||||
f"{model_target_mode!r}"
|
||||
)
|
||||
disease_history_mode = normalize_disease_history_mode(
|
||||
cfg.get("disease_history_mode", DISEASE_HISTORY_MODE_TIMED)
|
||||
)
|
||||
if disease_history_mode != DISEASE_HISTORY_MODE_TIMED:
|
||||
expected = {
|
||||
"model_target_mode": "all_future",
|
||||
"time_mode": "relative",
|
||||
"dist_mode": "weibull",
|
||||
"model_architecture": "traj_mixer_v5",
|
||||
}
|
||||
actual = {
|
||||
"model_target_mode": model_target_mode,
|
||||
"time_mode": str(cfg.get("time_mode", "")).lower(),
|
||||
"dist_mode": str(cfg.get("dist_mode", "")).lower(),
|
||||
"model_architecture": str(
|
||||
cfg.get("model_architecture", "")
|
||||
).lower(),
|
||||
}
|
||||
mismatches = [
|
||||
f"{name}={actual[name]!r} (expected {value!r})"
|
||||
for name, value in expected.items()
|
||||
if actual[name] != value
|
||||
]
|
||||
extra_info_types = cfg.get("extra_info_types", None)
|
||||
if extra_info_types != []:
|
||||
mismatches.append("extra_info_types must be []")
|
||||
if mismatches:
|
||||
raise ValueError(
|
||||
f"disease_history_mode={disease_history_mode!r} is only valid "
|
||||
"for the no-extra TrajMixer + all_future + relative + Weibull "
|
||||
"ablation; " + "; ".join(mismatches)
|
||||
)
|
||||
|
||||
if model_target_mode != "next_token":
|
||||
return
|
||||
|
||||
time_mode = str(cfg.get("time_mode", "")).lower()
|
||||
target_mode = str(cfg.get("target_mode", "")).lower()
|
||||
if time_mode != "absolute" or target_mode != "delphi2m":
|
||||
raise ValueError(
|
||||
"next_token is reserved for Delphi2M reproduction and requires "
|
||||
"time_mode='absolute' and target_mode='delphi2m'; got "
|
||||
f"time_mode={time_mode!r}, target_mode={target_mode!r}"
|
||||
)
|
||||
|
||||
|
||||
def split_indices(
|
||||
n: int,
|
||||
train_ratio: float,
|
||||
val_ratio: float,
|
||||
test_ratio: float,
|
||||
seed: int,
|
||||
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
|
||||
total = float(train_ratio) + float(val_ratio) + float(test_ratio)
|
||||
if not np.isclose(total, 1.0, atol=1e-6):
|
||||
raise ValueError(
|
||||
f"train/val/test ratios must sum to 1.0, got {total}"
|
||||
)
|
||||
indices = np.random.RandomState(seed).permutation(n)
|
||||
n_train = int(n * train_ratio)
|
||||
n_val = int(n * val_ratio)
|
||||
return (
|
||||
indices[:n_train],
|
||||
indices[n_train:n_train + n_val],
|
||||
indices[n_train + n_val:],
|
||||
)
|
||||
|
||||
|
||||
def select_indices_by_eid_file(
|
||||
dataset: Any,
|
||||
eid_file: str | Path,
|
||||
) -> Tuple[np.ndarray, Path]:
|
||||
"""Return dataset indices whose patient EIDs occur in ``eid_file``."""
|
||||
from train_util import load_eid_file
|
||||
|
||||
path = Path(eid_file)
|
||||
if not path.is_absolute():
|
||||
direct = Path.cwd() / path
|
||||
path = direct if direct.is_file() else Path(__file__).resolve().parent / path
|
||||
if not path.is_file():
|
||||
raise FileNotFoundError(f"EID split file not found: {path}")
|
||||
|
||||
samples = getattr(dataset, "samples", None)
|
||||
if samples is None:
|
||||
raise TypeError("EID-based evaluation requires dataset.samples")
|
||||
selected_eids = load_eid_file(path)
|
||||
indices = np.asarray(
|
||||
[
|
||||
index
|
||||
for index, sample in enumerate(samples)
|
||||
if int(sample["eid"]) in selected_eids
|
||||
],
|
||||
dtype=np.int64,
|
||||
)
|
||||
if indices.size == 0:
|
||||
raise ValueError(
|
||||
f"No dataset patients matched the EID split file: {path}"
|
||||
)
|
||||
return indices, path.resolve()
|
||||
|
||||
|
||||
def build_model_from_dataset(
|
||||
args: argparse.Namespace,
|
||||
cfg: Dict[str, Any],
|
||||
dataset: HealthDataset,
|
||||
state_dict: Optional[Dict[str, Any]] = None,
|
||||
) -> DeepHealth:
|
||||
validate_training_mode_config(cfg)
|
||||
model_target_mode = str(
|
||||
cfg_get(args, cfg, "model_target_mode", "next_token")
|
||||
).lower()
|
||||
if model_target_mode not in {"next_token", "all_future"}:
|
||||
raise ValueError(
|
||||
"model_target_mode must be next_token or all_future, got "
|
||||
f"{model_target_mode!r}"
|
||||
)
|
||||
risk_head_bias = bool(cfg_get(args, cfg, "risk_head_bias", False))
|
||||
if state_dict is not None:
|
||||
# The checkpoint schema is authoritative. This keeps all older
|
||||
# bias-free checkpoints loadable while restoring the new baseline bias.
|
||||
risk_head_bias = "risk_head.bias" in state_dict
|
||||
model_architecture = resolve_model_architecture(cfg, state_dict)
|
||||
continuous_value_center = None
|
||||
continuous_value_scale = None
|
||||
if dataset.n_cont_types > 0:
|
||||
scaling = str(cfg.get("continuous_value_scaling", "")).lower()
|
||||
if scaling != "robust":
|
||||
raise RuntimeError(
|
||||
"Continuous-variable checkpoints must declare "
|
||||
"continuous_value_scaling='robust'; unscaled checkpoints are "
|
||||
"not supported"
|
||||
)
|
||||
if state_dict is None:
|
||||
raise RuntimeError(
|
||||
"A checkpoint state_dict is required to restore RobustScale buffers"
|
||||
)
|
||||
center_key = "tokenizer.continuous_value_center"
|
||||
scale_key = "tokenizer.continuous_value_scale"
|
||||
missing = [
|
||||
key for key in (center_key, scale_key)
|
||||
if key not in state_dict
|
||||
]
|
||||
if missing:
|
||||
raise RuntimeError(
|
||||
"Checkpoint is missing required RobustScale buffers: "
|
||||
+ ", ".join(missing)
|
||||
)
|
||||
continuous_value_center = state_dict[center_key]
|
||||
continuous_value_scale = state_dict[scale_key]
|
||||
return DeepHealth(
|
||||
vocab_size=dataset.vocab_size,
|
||||
n_embd=int(cfg_get(args, cfg, "n_embd", 120)),
|
||||
n_head=int(cfg_get(args, cfg, "n_head", 10)),
|
||||
n_layer=int(cfg["n_layer"]),
|
||||
n_types=dataset.n_types,
|
||||
n_cont_types=dataset.n_cont_types,
|
||||
n_categories=dataset.n_categories,
|
||||
cont_type_ids=dataset.cont_type_ids,
|
||||
n_bins=int(cfg_get(args, cfg, "n_bins", 16)),
|
||||
continuous_value_center=continuous_value_center,
|
||||
continuous_value_scale=continuous_value_scale,
|
||||
extra_pool_reduce=str(
|
||||
cfg_get(args, cfg, "extra_pool_reduce", "mean")
|
||||
),
|
||||
target_mode=model_target_mode,
|
||||
time_mode=str(cfg_get(args, cfg, "time_mode", "absolute")),
|
||||
dist_mode=str(cfg_get(args, cfg, "dist_mode", "exponential")),
|
||||
dropout=float(cfg_get(args, cfg, "dropout", 0.0)),
|
||||
model_architecture=model_architecture,
|
||||
risk_head_bias=risk_head_bias,
|
||||
)
|
||||
|
||||
|
||||
def validate_dataset_metadata(
|
||||
dataset: HealthDataset,
|
||||
cfg: Dict[str, Any],
|
||||
) -> None:
|
||||
metadata = cfg.get("dataset_metadata")
|
||||
if not isinstance(metadata, dict):
|
||||
return
|
||||
|
||||
actual: Dict[str, Any] = {
|
||||
"vocab_size": int(dataset.vocab_size),
|
||||
"n_types": int(dataset.n_types),
|
||||
"n_cont_types": int(dataset.n_cont_types),
|
||||
"n_categories": int(dataset.n_categories),
|
||||
"cont_type_ids": [int(x) for x in dataset.cont_type_ids],
|
||||
"extra_info_types": [int(x) for x in dataset.extra_info_types],
|
||||
}
|
||||
mismatches = [
|
||||
f"{key}: train_config={metadata.get(key)!r}, current_dataset={value!r}"
|
||||
for key, value in actual.items()
|
||||
if key in metadata and metadata.get(key) != value
|
||||
]
|
||||
if mismatches:
|
||||
raise RuntimeError(
|
||||
"Current dataset metadata does not match train_config.json. "
|
||||
"Use the same prepared data and extra_info_types as training. "
|
||||
+ "; ".join(mismatches)
|
||||
)
|
||||
|
||||
|
||||
def build_first_occurrence_map(
|
||||
dataset: HealthDataset,
|
||||
subset_indices: np.ndarray,
|
||||
) -> Dict[int, Tuple[np.ndarray, np.ndarray]]:
|
||||
first_lists: Dict[int, List[Tuple[int, float]]] = {}
|
||||
for patient_id, dataset_index in enumerate(subset_indices.tolist()):
|
||||
sample = dataset.samples[int(dataset_index)]
|
||||
sequence_events = np.asarray(sample["event_seq"], dtype=np.int64)
|
||||
sequence_times = np.asarray(sample["time_seq"], dtype=np.float32)
|
||||
target_events = np.asarray(
|
||||
sample["target_event_seq"], dtype=np.int64
|
||||
)
|
||||
target_times = np.asarray(
|
||||
sample["target_time_seq"], dtype=np.float32
|
||||
)
|
||||
if sequence_events.size == 0 or target_events.size == 0:
|
||||
continue
|
||||
|
||||
full_events = np.concatenate(
|
||||
[sequence_events, target_events[-1:]]
|
||||
)
|
||||
full_times = np.concatenate([sequence_times, target_times[-1:]])
|
||||
unique_tokens, first_indices = np.unique(
|
||||
full_events, return_index=True
|
||||
)
|
||||
for token, first_index in zip(
|
||||
unique_tokens.tolist(), first_indices.tolist()
|
||||
):
|
||||
first_lists.setdefault(int(token), []).append(
|
||||
(patient_id, float(full_times[int(first_index)]))
|
||||
)
|
||||
|
||||
return {
|
||||
int(token): (
|
||||
np.asarray([patient for patient, _ in pairs], dtype=np.int32),
|
||||
np.asarray([time for _, time in pairs], dtype=np.float32),
|
||||
)
|
||||
for token, pairs in first_lists.items()
|
||||
if pairs
|
||||
}
|
||||
|
||||
|
||||
class AllFutureSequenceEvalDataset:
|
||||
"""
|
||||
Eval-only sequence view for all-future checkpoints.
|
||||
|
||||
All-future training uses the observed history, including CHECKUP state
|
||||
tokens, without reusing the next-step view that contains imputed
|
||||
<NO_EVENT> gap tokens.
|
||||
All-future training uses the observed history without reusing the
|
||||
next-step view that contains imputed <NO_EVENT> gap tokens. Legacy label-1
|
||||
assessment events are removed by the shared base dataset for every
|
||||
extra-info selection.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
@@ -26,6 +327,7 @@ class AllFutureSequenceEvalDataset:
|
||||
min_history_events: int = 1,
|
||||
min_future_events: int = 1,
|
||||
extra_info_types: Iterable[int] | None = None,
|
||||
disease_history_mode: str = DISEASE_HISTORY_MODE_TIMED,
|
||||
) -> None:
|
||||
base = AllFutureHealthDataset(
|
||||
data_prefix=data_prefix,
|
||||
@@ -34,6 +336,7 @@ class AllFutureSequenceEvalDataset:
|
||||
min_history_events=min_history_events,
|
||||
min_future_events=min_future_events,
|
||||
extra_info_types=extra_info_types,
|
||||
disease_history_mode=disease_history_mode,
|
||||
)
|
||||
|
||||
self.base = base
|
||||
@@ -45,6 +348,7 @@ class AllFutureSequenceEvalDataset:
|
||||
self.n_categories = base.n_categories
|
||||
self.cont_type_ids = base.cont_type_ids
|
||||
self.extra_info_types = base.extra_info_types
|
||||
self.disease_history_mode = base.disease_history_mode
|
||||
|
||||
self.samples: List[Dict[str, Any]] = []
|
||||
for patient in base.patients:
|
||||
@@ -52,7 +356,6 @@ class AllFutureSequenceEvalDataset:
|
||||
times = np.asarray(patient["times"], dtype=np.float32)
|
||||
if labels.size < 2:
|
||||
continue
|
||||
input_len = int(labels.size - 1)
|
||||
self.samples.append(
|
||||
{
|
||||
"eid": int(patient["eid"]),
|
||||
@@ -60,7 +363,6 @@ class AllFutureSequenceEvalDataset:
|
||||
"time_seq": times[:-1],
|
||||
"target_event_seq": labels[1:],
|
||||
"target_time_seq": times[1:],
|
||||
"readout_mask": np.ones(input_len, dtype=bool),
|
||||
"sex": int(patient["sex"]),
|
||||
"other_type": np.asarray(patient["other_type"], dtype=np.int64),
|
||||
"other_value": np.asarray(patient["other_value"], dtype=np.float32),
|
||||
@@ -79,7 +381,6 @@ class AllFutureSequenceEvalDataset:
|
||||
"time_seq": torch.from_numpy(s["time_seq"]).float(),
|
||||
"target_event_seq": torch.from_numpy(s["target_event_seq"]).long(),
|
||||
"target_time_seq": torch.from_numpy(s["target_time_seq"]).float(),
|
||||
"readout_mask": torch.from_numpy(s["readout_mask"]).bool(),
|
||||
"sex": torch.tensor(s["sex"], dtype=torch.long),
|
||||
"other_type": torch.from_numpy(s["other_type"]).long(),
|
||||
"other_value": torch.from_numpy(s["other_value"]).float(),
|
||||
@@ -94,10 +395,10 @@ def load_sequence_eval_dataset(
|
||||
data_prefix: str,
|
||||
labels_file: str,
|
||||
no_event_interval_years: float,
|
||||
include_no_event_in_uts_target: bool,
|
||||
min_history_events: int,
|
||||
min_future_events: int,
|
||||
extra_info_types: Iterable[int] | None,
|
||||
disease_history_mode: str = DISEASE_HISTORY_MODE_TIMED,
|
||||
):
|
||||
mode = str(model_target_mode).lower()
|
||||
if mode == "next_token":
|
||||
@@ -105,7 +406,6 @@ def load_sequence_eval_dataset(
|
||||
data_prefix=data_prefix,
|
||||
labels_file=labels_file,
|
||||
no_event_interval_years=no_event_interval_years,
|
||||
include_no_event_in_uts_target=include_no_event_in_uts_target,
|
||||
extra_info_types=extra_info_types,
|
||||
)
|
||||
if mode == "all_future":
|
||||
@@ -115,6 +415,7 @@ def load_sequence_eval_dataset(
|
||||
min_history_events=min_history_events,
|
||||
min_future_events=min_future_events,
|
||||
extra_info_types=extra_info_types,
|
||||
disease_history_mode=disease_history_mode,
|
||||
)
|
||||
raise ValueError(f"Unknown model_target_mode: {model_target_mode!r}")
|
||||
|
||||
@@ -132,9 +433,6 @@ def sequence_eval_collate_fn(batch: List[Dict[str, torch.Tensor]]) -> Dict[str,
|
||||
target_time_seq = pad_sequence(
|
||||
[s["target_time_seq"] for s in batch], batch_first=True, padding_value=0.0
|
||||
)
|
||||
readout_mask = pad_sequence(
|
||||
[s["readout_mask"] for s in batch], batch_first=True, padding_value=False
|
||||
)
|
||||
other_type = pad_sequence(
|
||||
[s["other_type"] for s in batch], batch_first=True, padding_value=0
|
||||
)
|
||||
@@ -154,7 +452,7 @@ def sequence_eval_collate_fn(batch: List[Dict[str, torch.Tensor]]) -> Dict[str,
|
||||
"padding_mask": event_seq > PAD_IDX,
|
||||
"target_event_seq": target_event_seq,
|
||||
"target_time_seq": target_time_seq,
|
||||
"readout_mask": readout_mask,
|
||||
"readout_mask": event_seq > PAD_IDX,
|
||||
"sex": torch.stack([s["sex"] for s in batch]),
|
||||
"other_type": other_type,
|
||||
"other_value": other_value,
|
||||
|
||||
342
evaluate_all_runs_linux.sh
Normal file
342
evaluate_all_runs_linux.sh
Normal file
@@ -0,0 +1,342 @@
|
||||
#!/usr/bin/env bash
|
||||
#
|
||||
# Recursively evaluate every completed run under runs/.
|
||||
#
|
||||
# A directory is considered a runnable run when it contains both:
|
||||
# - train_config.json
|
||||
# - best_model.pt
|
||||
#
|
||||
# Each missing report is scheduled independently:
|
||||
# - evaluate_auc.py
|
||||
# -> df_auc_delphi2m_report.csv
|
||||
# - evaluate_auc_v2.py
|
||||
# -> df_auc_landmark_delphi2m_report.csv
|
||||
#
|
||||
# Existing non-empty reports are skipped unless --force is supplied. Each
|
||||
# evaluation task uses one GPU.
|
||||
# Tasks on the same GPU run sequentially; different GPUs run in parallel.
|
||||
#
|
||||
# Examples:
|
||||
# bash evaluate_all_runs_linux.sh --gpus 0
|
||||
# bash evaluate_all_runs_linux.sh --gpus 0,1,2,3
|
||||
# bash evaluate_all_runs_linux.sh --gpus 0,1 --force
|
||||
# bash evaluate_all_runs_linux.sh --gpus 0,1 --dry-run
|
||||
#
|
||||
|
||||
set -uo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
|
||||
cd "$SCRIPT_DIR"
|
||||
|
||||
RUNS_ROOT="$SCRIPT_DIR/runs"
|
||||
LOG_ROOT="$SCRIPT_DIR/batch_logs/evaluate_all_runs"
|
||||
GPU_CSV="0"
|
||||
PYTHON_BIN="${PYTHON_BIN:-python}"
|
||||
NUM_WORKERS=4
|
||||
NUM_WORKERS_AUC=4
|
||||
FORCE=0
|
||||
DRY_RUN=0
|
||||
|
||||
TOKEN_REPORT="df_auc_delphi2m_report.csv"
|
||||
LANDMARK_REPORT="df_auc_landmark_delphi2m_report.csv"
|
||||
|
||||
usage() {
|
||||
cat <<'EOF'
|
||||
Usage:
|
||||
bash evaluate_all_runs_linux.sh [options]
|
||||
|
||||
Options:
|
||||
--gpus LIST Comma-separated GPU ids (default: 0).
|
||||
--runs-root PATH Root directory scanned recursively (default: ./runs).
|
||||
--log-root PATH Evaluation log root.
|
||||
--python PATH Python executable (default: $PYTHON_BIN or python).
|
||||
--num-workers N DataLoader workers per evaluation (default: 4).
|
||||
--num-workers-auc N CPU AUC workers per evaluation (default: 4).
|
||||
--force Recompute both AUC reports even when they exist.
|
||||
--dry-run Discover runs and print scheduled evaluations only.
|
||||
-h, --help Show this help message.
|
||||
|
||||
Completion files:
|
||||
evaluate_auc.py df_auc_delphi2m_report.csv
|
||||
evaluate_auc_v2.py df_auc_landmark_delphi2m_report.csv
|
||||
|
||||
Existing non-empty completion files are skipped independently unless --force
|
||||
is supplied.
|
||||
EOF
|
||||
}
|
||||
|
||||
while (($# > 0)); do
|
||||
case "$1" in
|
||||
--gpus)
|
||||
[[ $# -ge 2 ]] || {
|
||||
echo "ERROR: --gpus requires a value." >&2
|
||||
exit 2
|
||||
}
|
||||
GPU_CSV="$2"
|
||||
shift 2
|
||||
;;
|
||||
--runs-root)
|
||||
[[ $# -ge 2 ]] || {
|
||||
echo "ERROR: --runs-root requires a value." >&2
|
||||
exit 2
|
||||
}
|
||||
RUNS_ROOT="$2"
|
||||
shift 2
|
||||
;;
|
||||
--log-root)
|
||||
[[ $# -ge 2 ]] || {
|
||||
echo "ERROR: --log-root requires a value." >&2
|
||||
exit 2
|
||||
}
|
||||
LOG_ROOT="$2"
|
||||
shift 2
|
||||
;;
|
||||
--python)
|
||||
[[ $# -ge 2 ]] || {
|
||||
echo "ERROR: --python requires a value." >&2
|
||||
exit 2
|
||||
}
|
||||
PYTHON_BIN="$2"
|
||||
shift 2
|
||||
;;
|
||||
--num-workers)
|
||||
[[ $# -ge 2 ]] || {
|
||||
echo "ERROR: --num-workers requires a value." >&2
|
||||
exit 2
|
||||
}
|
||||
NUM_WORKERS="$2"
|
||||
shift 2
|
||||
;;
|
||||
--num-workers-auc)
|
||||
[[ $# -ge 2 ]] || {
|
||||
echo "ERROR: --num-workers-auc requires a value." >&2
|
||||
exit 2
|
||||
}
|
||||
NUM_WORKERS_AUC="$2"
|
||||
shift 2
|
||||
;;
|
||||
--force)
|
||||
FORCE=1
|
||||
shift
|
||||
;;
|
||||
--dry-run)
|
||||
DRY_RUN=1
|
||||
shift
|
||||
;;
|
||||
-h|--help)
|
||||
usage
|
||||
exit 0
|
||||
;;
|
||||
*)
|
||||
echo "ERROR: unknown argument: $1" >&2
|
||||
usage >&2
|
||||
exit 2
|
||||
;;
|
||||
esac
|
||||
done
|
||||
|
||||
[[ -n "$GPU_CSV" ]] || {
|
||||
echo "ERROR: --gpus must not be empty." >&2
|
||||
exit 2
|
||||
}
|
||||
[[ "$NUM_WORKERS" =~ ^[0-9]+$ ]] || {
|
||||
echo "ERROR: --num-workers must be a non-negative integer." >&2
|
||||
exit 2
|
||||
}
|
||||
[[ "$NUM_WORKERS_AUC" =~ ^[1-9][0-9]*$ ]] || {
|
||||
echo "ERROR: --num-workers-auc must be a positive integer." >&2
|
||||
exit 2
|
||||
}
|
||||
[[ -d "$RUNS_ROOT" ]] || {
|
||||
echo "ERROR: runs root does not exist: $RUNS_ROOT" >&2
|
||||
exit 2
|
||||
}
|
||||
[[ -f "$SCRIPT_DIR/evaluate_auc.py" ]] || {
|
||||
echo "ERROR: missing evaluator: $SCRIPT_DIR/evaluate_auc.py" >&2
|
||||
exit 2
|
||||
}
|
||||
[[ -f "$SCRIPT_DIR/evaluate_auc_v2.py" ]] || {
|
||||
echo "ERROR: missing evaluator: $SCRIPT_DIR/evaluate_auc_v2.py" >&2
|
||||
exit 2
|
||||
}
|
||||
command -v "$PYTHON_BIN" >/dev/null 2>&1 || {
|
||||
echo "ERROR: Python executable not found: $PYTHON_BIN" >&2
|
||||
exit 2
|
||||
}
|
||||
|
||||
RUNS_ROOT="$(cd -- "$RUNS_ROOT" && pwd)"
|
||||
if [[ "$LOG_ROOT" != /* ]]; then
|
||||
LOG_ROOT="$SCRIPT_DIR/$LOG_ROOT"
|
||||
fi
|
||||
|
||||
IFS=',' read -r -a GPU_IDS <<< "$GPU_CSV"
|
||||
declare -A SEEN_GPUS=()
|
||||
for gpu in "${GPU_IDS[@]}"; do
|
||||
[[ -n "$gpu" && "$gpu" =~ ^[A-Za-z0-9._:-]+$ ]] || {
|
||||
echo "ERROR: invalid GPU id: $gpu" >&2
|
||||
exit 2
|
||||
}
|
||||
[[ -z "${SEEN_GPUS[$gpu]+x}" ]] || {
|
||||
echo "ERROR: duplicate GPU id: $gpu" >&2
|
||||
exit 2
|
||||
}
|
||||
SEEN_GPUS["$gpu"]=1
|
||||
done
|
||||
|
||||
declare -a JOB_RUN_DIRS=()
|
||||
declare -a JOB_EVALUATORS=()
|
||||
declare -a JOB_REPORTS=()
|
||||
declare -a JOB_LOG_FILES=()
|
||||
|
||||
add_job() {
|
||||
local run_dir="$1"
|
||||
local evaluator="$2"
|
||||
local report_name="$3"
|
||||
local relative_run="${run_dir#"$RUNS_ROOT"/}"
|
||||
local evaluator_name="${evaluator%.py}"
|
||||
|
||||
JOB_RUN_DIRS+=("$run_dir")
|
||||
JOB_EVALUATORS+=("$evaluator")
|
||||
JOB_REPORTS+=("$report_name")
|
||||
JOB_LOG_FILES+=("$LOG_ROOT/$relative_run/$evaluator_name.log")
|
||||
}
|
||||
|
||||
run_count=0
|
||||
incomplete_run_count=0
|
||||
skipped_token_count=0
|
||||
skipped_landmark_count=0
|
||||
|
||||
while IFS= read -r -d '' config_path; do
|
||||
run_dir="${config_path%/train_config.json}"
|
||||
((run_count += 1))
|
||||
|
||||
if [[ ! -f "$run_dir/best_model.pt" ]]; then
|
||||
echo "[SKIP] Incomplete run without best_model.pt: $run_dir"
|
||||
((incomplete_run_count += 1))
|
||||
continue
|
||||
fi
|
||||
|
||||
if ((!FORCE)) && [[ -s "$run_dir/$TOKEN_REPORT" ]]; then
|
||||
((skipped_token_count += 1))
|
||||
else
|
||||
add_job "$run_dir" "evaluate_auc.py" "$TOKEN_REPORT"
|
||||
fi
|
||||
|
||||
if ((!FORCE)) && [[ -s "$run_dir/$LANDMARK_REPORT" ]]; then
|
||||
((skipped_landmark_count += 1))
|
||||
else
|
||||
add_job "$run_dir" "evaluate_auc_v2.py" "$LANDMARK_REPORT"
|
||||
fi
|
||||
done < <(find "$RUNS_ROOT" -type f -name "train_config.json" -print0)
|
||||
|
||||
if ((!DRY_RUN)); then
|
||||
mkdir -p "$LOG_ROOT"
|
||||
fi
|
||||
|
||||
print_command() {
|
||||
printf '%q ' "$@"
|
||||
printf '\n'
|
||||
}
|
||||
|
||||
run_job() {
|
||||
local job_index="$1"
|
||||
local gpu="$2"
|
||||
local run_dir="${JOB_RUN_DIRS[$job_index]}"
|
||||
local evaluator="${JOB_EVALUATORS[$job_index]}"
|
||||
local report_name="${JOB_REPORTS[$job_index]}"
|
||||
local log_file="${JOB_LOG_FILES[$job_index]}"
|
||||
local -a command=(
|
||||
"$PYTHON_BIN"
|
||||
-u
|
||||
"$SCRIPT_DIR/$evaluator"
|
||||
--run_path "$run_dir"
|
||||
--output_path "$run_dir"
|
||||
--device cuda
|
||||
--num_workers "$NUM_WORKERS"
|
||||
--num_workers_auc "$NUM_WORKERS_AUC"
|
||||
)
|
||||
|
||||
echo "[$(date '+%F %T')] START evaluator=$evaluator gpu=$gpu"
|
||||
echo " run=$run_dir"
|
||||
echo " report=$report_name"
|
||||
|
||||
if ((DRY_RUN)); then
|
||||
printf ' CUDA_VISIBLE_DEVICES=%q ' "$gpu"
|
||||
print_command "${command[@]}"
|
||||
return 0
|
||||
fi
|
||||
|
||||
mkdir -p "$(dirname -- "$log_file")"
|
||||
if CUDA_VISIBLE_DEVICES="$gpu" PYTHONUNBUFFERED=1 \
|
||||
"${command[@]}" >"$log_file" 2>&1; then
|
||||
if [[ -s "$run_dir/$report_name" ]]; then
|
||||
echo "[$(date '+%F %T')] DONE evaluator=$evaluator gpu=$gpu"
|
||||
return 0
|
||||
fi
|
||||
|
||||
echo "[$(date '+%F %T')] FAIL evaluator=$evaluator gpu=$gpu" >&2
|
||||
echo " Evaluator exited successfully but did not create: $run_dir/$report_name" >&2
|
||||
echo " See: $log_file" >&2
|
||||
return 1
|
||||
else
|
||||
local exit_code=$?
|
||||
echo "[$(date '+%F %T')] FAIL evaluator=$evaluator gpu=$gpu exit=$exit_code" >&2
|
||||
echo " See: $log_file" >&2
|
||||
return "$exit_code"
|
||||
fi
|
||||
}
|
||||
|
||||
worker() {
|
||||
local slot="$1"
|
||||
local gpu="${GPU_IDS[$slot]}"
|
||||
local job_index
|
||||
local failed=0
|
||||
|
||||
for ((job_index = slot; job_index < ${#JOB_RUN_DIRS[@]}; job_index += ${#GPU_IDS[@]})); do
|
||||
run_job "$job_index" "$gpu" || failed=1
|
||||
done
|
||||
return "$failed"
|
||||
}
|
||||
|
||||
force_label="no"
|
||||
if ((FORCE)); then
|
||||
force_label="yes"
|
||||
fi
|
||||
|
||||
echo "Runs root: $RUNS_ROOT"
|
||||
echo "GPUs: ${GPU_IDS[*]}"
|
||||
echo "Force recompute: $force_label"
|
||||
echo "Runs discovered: $run_count"
|
||||
echo "Incomplete runs skipped: $incomplete_run_count"
|
||||
echo "Existing token reports skipped: $skipped_token_count"
|
||||
echo "Existing landmark reports skipped: $skipped_landmark_count"
|
||||
echo "Scheduled evaluation tasks: ${#JOB_RUN_DIRS[@]}"
|
||||
echo "Log root: $LOG_ROOT"
|
||||
echo
|
||||
|
||||
if ((${#JOB_RUN_DIRS[@]} == 0)); then
|
||||
echo "No AUC evaluation tasks are scheduled."
|
||||
exit 0
|
||||
fi
|
||||
|
||||
declare -a WORKER_PIDS=()
|
||||
for ((slot = 0; slot < ${#GPU_IDS[@]}; slot++)); do
|
||||
worker "$slot" &
|
||||
WORKER_PIDS+=("$!")
|
||||
done
|
||||
|
||||
overall_status=0
|
||||
for pid in "${WORKER_PIDS[@]}"; do
|
||||
wait "$pid" || overall_status=1
|
||||
done
|
||||
|
||||
if ((overall_status != 0)); then
|
||||
echo "One or more evaluation tasks failed. Inspect logs under: $LOG_ROOT" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
if ((DRY_RUN)); then
|
||||
echo "Dry run completed successfully."
|
||||
else
|
||||
echo "All scheduled AUC evaluations completed successfully."
|
||||
fi
|
||||
415
evaluate_auc.py
415
evaluate_auc.py
@@ -7,7 +7,8 @@ This script follows the logic of the Delphi evaluation script supplied by the us
|
||||
at least `offset` years before the target time;
|
||||
3. run model inference by disease chunks to avoid materializing all logits;
|
||||
4. compute AUC separately by sex and age bracket;
|
||||
5. aggregate age brackets with DeLong variance.
|
||||
5. average age-bracket AUCs within each sex and write a Delphi2M-style
|
||||
Female/Male report.
|
||||
|
||||
Efficiency notes:
|
||||
- transformer/readout inference is executed once and cached;
|
||||
@@ -18,7 +19,7 @@ Efficiency notes:
|
||||
avoiding repeated pickling of arrays for every disease.
|
||||
|
||||
Run from the DeepHealth code directory containing dataset.py, models.py,
|
||||
readouts.py, and train_config.json-compatible checkpoints/configs.
|
||||
and train_config.json-compatible checkpoints/configs.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -38,15 +39,32 @@ import torch
|
||||
from torch.utils.data import DataLoader, Subset
|
||||
from tqdm.auto import tqdm
|
||||
|
||||
from dataset import HealthDataset
|
||||
from eval_data import load_sequence_eval_dataset, sequence_eval_collate_fn
|
||||
from models import (
|
||||
DeepHealth,
|
||||
validate_event_trajectory_config,
|
||||
validate_event_trajectory_state_dict,
|
||||
from dataset import (
|
||||
DISEASE_HISTORY_MODE_TIMED,
|
||||
HealthDataset,
|
||||
normalize_disease_history_mode,
|
||||
transform_disease_history_batch_at_position,
|
||||
)
|
||||
from readouts import build_readout
|
||||
from targets import PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX
|
||||
from delphi2m_auc_report import (
|
||||
DEFAULT_DELPHI2M_PERIODS_YEARS,
|
||||
build_delphi2m_auc_report,
|
||||
)
|
||||
from eval_data import (
|
||||
build_first_occurrence_map,
|
||||
build_model_from_dataset,
|
||||
cfg_get,
|
||||
load_json_config,
|
||||
load_sequence_eval_dataset,
|
||||
resolve_eval_device,
|
||||
select_indices_by_eid_file,
|
||||
sequence_eval_collate_fn,
|
||||
split_indices,
|
||||
validate_training_mode_config,
|
||||
validate_dataset_metadata,
|
||||
)
|
||||
from model_architectures import resolve_model_architecture
|
||||
from models import DeepHealth
|
||||
from targets import NO_EVENT_IDX, PAD_IDX, RESERVED_IDX
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -151,56 +169,7 @@ def get_auc_delong_var(control_scores: np.ndarray, case_scores: np.ndarray) -> T
|
||||
# Disease selection
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
SPECIAL_TOKENS = {PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX}
|
||||
|
||||
|
||||
def _build_first_occurrence_maps(
|
||||
dataset: HealthDataset,
|
||||
subset_indices: np.ndarray,
|
||||
) -> Tuple[Dict[int, Tuple[np.ndarray, np.ndarray]], np.ndarray, np.ndarray, np.ndarray]:
|
||||
patient_count = len(subset_indices)
|
||||
followup_end = np.full(patient_count, -np.inf, dtype=np.float32)
|
||||
death_time = np.full(patient_count, np.inf, dtype=np.float32)
|
||||
sex = np.full(patient_count, -1, dtype=np.int8)
|
||||
|
||||
first_lists: Dict[int, List[Tuple[int, float]]] = {}
|
||||
|
||||
for patient_id, dataset_index in enumerate(subset_indices.tolist()):
|
||||
s = dataset.samples[int(dataset_index)]
|
||||
|
||||
seq_event = np.asarray(s["event_seq"], dtype=np.int64)
|
||||
seq_time = np.asarray(s["time_seq"], dtype=np.float32)
|
||||
tgt_event = np.asarray(s["target_event_seq"], dtype=np.int64)
|
||||
tgt_time = np.asarray(s["target_time_seq"], dtype=np.float32)
|
||||
|
||||
if seq_event.size == 0 or tgt_event.size == 0:
|
||||
continue
|
||||
|
||||
full_event = np.concatenate([seq_event, tgt_event[-1:]])
|
||||
full_time = np.concatenate([seq_time, tgt_time[-1:]])
|
||||
|
||||
sex[patient_id] = int(s["sex"])
|
||||
followup_end[patient_id] = np.max(full_time).astype(np.float32)
|
||||
|
||||
uniq_tokens, first_idx = np.unique(full_event, return_index=True)
|
||||
|
||||
for token, idx in zip(uniq_tokens.tolist(), first_idx.tolist()):
|
||||
token = int(token)
|
||||
event_time = float(full_time[int(idx)])
|
||||
if token not in first_lists:
|
||||
first_lists[token] = []
|
||||
first_lists[token].append((patient_id, event_time))
|
||||
|
||||
packed: Dict[int, Tuple[np.ndarray, np.ndarray]] = {}
|
||||
for token, pairs in first_lists.items():
|
||||
if not pairs:
|
||||
continue
|
||||
packed[int(token)] = (
|
||||
np.asarray([p for p, _ in pairs], dtype=np.int32),
|
||||
np.asarray([t for _, t in pairs], dtype=np.float32),
|
||||
)
|
||||
|
||||
return packed, followup_end, death_time, sex
|
||||
SPECIAL_TOKENS = {PAD_IDX, RESERVED_IDX, NO_EVENT_IDX}
|
||||
|
||||
|
||||
def _get_death_token_ids(dataset: HealthDataset) -> List[int]:
|
||||
@@ -262,83 +231,6 @@ def select_disease_tokens(
|
||||
# Dataset/split/model helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def load_json_config(path: Optional[str]) -> Dict[str, Any]:
|
||||
if path is None:
|
||||
return {}
|
||||
p = Path(path)
|
||||
if not p.exists():
|
||||
return {}
|
||||
with p.open("r", encoding="utf-8") as f:
|
||||
return json.load(f)
|
||||
|
||||
|
||||
def cfg_get(args: argparse.Namespace | Dict[str, Any] | None, cfg: Dict[str, Any], name: str, default: Any) -> Any:
|
||||
"""Get a value from CLI args first, then train_config.json, then default.
|
||||
|
||||
This helper intentionally accepts either an argparse.Namespace or a dict.
|
||||
The earlier version passed cfg as both args and cfg, then tried to access
|
||||
args.eval_split, which fails because dict has no attributes.
|
||||
"""
|
||||
val = None
|
||||
if args is not None:
|
||||
if isinstance(args, dict):
|
||||
val = args.get(name, None)
|
||||
else:
|
||||
val = getattr(args, name, None)
|
||||
if val is not None:
|
||||
return val
|
||||
return cfg.get(name, default)
|
||||
|
||||
|
||||
def resolve_eval_device(device_arg: Optional[str]) -> torch.device:
|
||||
"""Resolve evaluation device without inheriting train_config.json device."""
|
||||
device_name = device_arg or ("cuda" if torch.cuda.is_available() else "cpu")
|
||||
device = torch.device(device_name)
|
||||
if device.type == "cuda" and not torch.cuda.is_available():
|
||||
raise RuntimeError(
|
||||
f"Requested device {device_name!r}, but CUDA is not available."
|
||||
)
|
||||
return device
|
||||
|
||||
|
||||
def split_indices(n: int, train_ratio: float, val_ratio: float, test_ratio: float, seed: int) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
|
||||
total = train_ratio + val_ratio + test_ratio
|
||||
if not np.isclose(total, 1.0, atol=1e-6):
|
||||
raise ValueError(f"train/val/test ratios must sum to 1.0, got {total}")
|
||||
rng = np.random.RandomState(seed)
|
||||
idx = rng.permutation(n)
|
||||
n_train = int(n * train_ratio)
|
||||
n_val = int(n * val_ratio)
|
||||
return idx[:n_train], idx[n_train:n_train + n_val], idx[n_train + n_val:]
|
||||
|
||||
|
||||
def build_model_from_dataset(args: argparse.Namespace, cfg: Dict[str, Any], dataset: HealthDataset) -> DeepHealth:
|
||||
validate_event_trajectory_config(cfg)
|
||||
model_target_mode = str(cfg_get(
|
||||
args, cfg, "model_target_mode", "next_token")).lower()
|
||||
if model_target_mode not in {"next_token", "all_future"}:
|
||||
raise ValueError(
|
||||
f"model_target_mode must be next_token or all_future, got {model_target_mode!r}"
|
||||
)
|
||||
return DeepHealth(
|
||||
vocab_size=dataset.vocab_size,
|
||||
model_size=str(cfg_get(args, cfg, "model_size", "nano")),
|
||||
n_reasoning_rounds=int(
|
||||
cfg_get(args, cfg, "n_reasoning_rounds", 12)
|
||||
),
|
||||
n_types=dataset.n_types,
|
||||
n_cont_types=dataset.n_cont_types,
|
||||
n_categories=dataset.n_categories,
|
||||
cont_type_ids=dataset.cont_type_ids,
|
||||
n_bins=int(cfg_get(args, cfg, "n_bins", 16)),
|
||||
extra_pool_reduce=str(cfg_get(args, cfg, "extra_pool_reduce", "mean")),
|
||||
target_mode=model_target_mode,
|
||||
time_mode=str(cfg_get(args, cfg, "time_mode", "relative")),
|
||||
dist_mode=str(cfg_get(args, cfg, "dist_mode", "exponential")),
|
||||
dropout=float(cfg_get(args, cfg, "dropout", 0.0)),
|
||||
)
|
||||
|
||||
|
||||
def _extract_state_dict(ckpt: Any) -> Dict[str, Any]:
|
||||
if isinstance(ckpt, dict) and "model" in ckpt:
|
||||
return ckpt["model"]
|
||||
@@ -358,32 +250,25 @@ def load_checkpoint_state_dict(checkpoint_path: str, map_location: str | torch.d
|
||||
|
||||
def resolve_dist_mode_for_checkpoint(cfg_dist_mode: str, state_dict: Dict[str, Any]) -> str:
|
||||
mode = str(cfg_dist_mode).lower()
|
||||
if mode not in {"exponential", "weibull"}:
|
||||
raise ValueError(
|
||||
f"Unsupported dist_mode={mode!r}; expected exponential or weibull."
|
||||
)
|
||||
has_rho_head = any(str(k).startswith("rho_head.")
|
||||
for k in state_dict.keys())
|
||||
has_rho_death_head = any(str(k).startswith("rho_death_head.")
|
||||
for k in state_dict.keys())
|
||||
|
||||
if has_rho_head and mode != "weibull":
|
||||
print(
|
||||
"[WARN] Checkpoint contains rho_head weights; overriding dist_mode to 'weibull' for evaluation.")
|
||||
return "weibull"
|
||||
if has_rho_death_head and mode != "mixed":
|
||||
print(
|
||||
"[WARN] Checkpoint contains rho_death_head weights; overriding dist_mode to 'mixed' for evaluation.")
|
||||
return "mixed"
|
||||
if (not has_rho_head) and mode == "weibull":
|
||||
print(
|
||||
"[WARN] dist_mode is 'weibull' but checkpoint has no rho_head weights; overriding dist_mode to 'exponential'.")
|
||||
return "exponential"
|
||||
if (not has_rho_death_head) and mode == "mixed":
|
||||
print(
|
||||
"[WARN] dist_mode is 'mixed' but checkpoint has no rho_death_head weights; overriding dist_mode to 'exponential'.")
|
||||
return "exponential"
|
||||
if mode == "weibull" and not has_rho_head:
|
||||
raise RuntimeError(
|
||||
"Weibull checkpoint is missing rho_head parameters."
|
||||
)
|
||||
if mode == "exponential" and has_rho_head:
|
||||
raise RuntimeError(
|
||||
"Exponential checkpoint unexpectedly contains rho_head parameters."
|
||||
)
|
||||
return mode
|
||||
|
||||
|
||||
def load_model_state(
|
||||
model: torch.nn.Module,
|
||||
model: DeepHealth,
|
||||
checkpoint_path: str,
|
||||
device: torch.device,
|
||||
state_dict: Optional[Dict[str, Any]] = None,
|
||||
@@ -391,16 +276,15 @@ def load_model_state(
|
||||
state = state_dict if state_dict is not None else load_checkpoint_state_dict(
|
||||
checkpoint_path, map_location=device)
|
||||
|
||||
validate_event_trajectory_state_dict(
|
||||
state,
|
||||
expected_d_model=model.d_model,
|
||||
expected_n_trajectory=model.n_trajectory,
|
||||
expected_n_reasoning_rounds=model.n_reasoning_rounds,
|
||||
)
|
||||
resolve_model_architecture(model.model_architecture, state)
|
||||
model.load_state_dict(state, strict=True)
|
||||
|
||||
|
||||
def make_eval_subset(dataset: HealthDataset, args: argparse.Namespace | Dict[str, Any] | None, cfg: Dict[str, Any]) -> Tuple[Subset, np.ndarray]:
|
||||
def make_eval_subset(
|
||||
dataset: HealthDataset,
|
||||
args: argparse.Namespace | Dict[str, Any] | None,
|
||||
cfg: Dict[str, Any],
|
||||
) -> Tuple[Subset, np.ndarray]:
|
||||
train_ratio = float(cfg_get(args, cfg, "train_ratio", 0.7))
|
||||
val_ratio = float(cfg_get(args, cfg, "val_ratio", 0.15))
|
||||
test_ratio = float(cfg_get(args, cfg, "test_ratio", 0.15))
|
||||
@@ -408,52 +292,43 @@ def make_eval_subset(dataset: HealthDataset, args: argparse.Namespace | Dict[str
|
||||
eval_split = str(cfg_get(args, cfg, "eval_split", "test")).lower()
|
||||
dataset_subset_size = cfg_get(args, cfg, "dataset_subset_size", None)
|
||||
|
||||
if eval_split in {"valid", "validation"}:
|
||||
eval_split = "val"
|
||||
if eval_split not in {"train", "val", "test", "all"}:
|
||||
raise ValueError(
|
||||
"eval_split must be one of train/val/test/all, got "
|
||||
f"{eval_split!r}"
|
||||
)
|
||||
|
||||
test_eid_file = cfg_get(
|
||||
args,
|
||||
cfg,
|
||||
"test_eid_file",
|
||||
"ukb_test_eid.csv",
|
||||
)
|
||||
if eval_split == "test" and test_eid_file not in {None, ""}:
|
||||
indices, eid_path = select_indices_by_eid_file(
|
||||
dataset,
|
||||
str(test_eid_file),
|
||||
)
|
||||
print(f"Test split source: EID file {eid_path}")
|
||||
else:
|
||||
train_idx, val_idx, test_idx = split_indices(
|
||||
len(dataset), train_ratio, val_ratio, test_ratio, seed)
|
||||
len(dataset), train_ratio, val_ratio, test_ratio, seed
|
||||
)
|
||||
split_map = {
|
||||
"train": train_idx,
|
||||
"val": val_idx,
|
||||
"valid": val_idx,
|
||||
"validation": val_idx,
|
||||
"test": test_idx,
|
||||
"all": np.arange(len(dataset)),
|
||||
}
|
||||
if eval_split not in split_map:
|
||||
raise ValueError(
|
||||
f"eval_split must be one of {sorted(split_map)}, got {eval_split!r}")
|
||||
|
||||
indices = split_map[eval_split]
|
||||
|
||||
if dataset_subset_size is not None and int(dataset_subset_size) > 0:
|
||||
indices = indices[: int(dataset_subset_size)]
|
||||
return Subset(dataset, indices.tolist()), np.asarray(indices, dtype=np.int64)
|
||||
|
||||
|
||||
def validate_dataset_metadata(dataset: HealthDataset, cfg: Dict[str, Any]) -> None:
|
||||
meta = cfg.get("dataset_metadata")
|
||||
if not isinstance(meta, dict):
|
||||
return
|
||||
|
||||
actual: Dict[str, Any] = {
|
||||
"vocab_size": int(dataset.vocab_size),
|
||||
"n_types": int(dataset.n_types),
|
||||
"n_cont_types": int(dataset.n_cont_types),
|
||||
"n_categories": int(dataset.n_categories),
|
||||
"cont_type_ids": [int(x) for x in dataset.cont_type_ids],
|
||||
"extra_info_types": [int(x) for x in dataset.extra_info_types],
|
||||
}
|
||||
mismatches = [
|
||||
f"{key}: train_config={meta.get(key)!r}, current_dataset={value!r}"
|
||||
for key, value in actual.items()
|
||||
if key in meta and meta.get(key) != value
|
||||
]
|
||||
if mismatches:
|
||||
raise RuntimeError(
|
||||
"Current dataset metadata does not match train_config.json. "
|
||||
"Use the same prepared data and extra_info_types as training. "
|
||||
+ "; ".join(mismatches)
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Batched inference + cached hidden states
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -478,8 +353,7 @@ def infer_readout_hidden(
|
||||
loader: DataLoader,
|
||||
device: torch.device,
|
||||
model_target_mode: str,
|
||||
readout_name: str,
|
||||
readout_reduce: str,
|
||||
disease_history_mode: str,
|
||||
use_amp: bool,
|
||||
hidden_cache_dtype: str = "float16",
|
||||
) -> Tuple[np.ndarray, Dict[str, np.ndarray]]:
|
||||
@@ -489,15 +363,16 @@ def infer_readout_hidden(
|
||||
raise ValueError(
|
||||
f"model_target_mode must be next_token or all_future, got {model_target_mode!r}"
|
||||
)
|
||||
|
||||
readout = None
|
||||
if model_target_mode == "next_token" and readout_name == "same_time_group_end":
|
||||
readout = build_readout("same_time_group_end",
|
||||
reduce=readout_reduce).to(device)
|
||||
elif model_target_mode == "next_token":
|
||||
readout = build_readout(readout_name).to(device)
|
||||
if readout is not None:
|
||||
readout.eval()
|
||||
disease_history_mode = normalize_disease_history_mode(
|
||||
disease_history_mode
|
||||
)
|
||||
if (
|
||||
model_target_mode != "all_future"
|
||||
and disease_history_mode != DISEASE_HISTORY_MODE_TIMED
|
||||
):
|
||||
raise ValueError(
|
||||
"ordered/set disease history is only supported for all_future models"
|
||||
)
|
||||
|
||||
hidden_parts: List[np.ndarray] = []
|
||||
arrays: Dict[str, List[np.ndarray]] = {
|
||||
@@ -533,7 +408,7 @@ def infer_readout_hidden(
|
||||
hidden = torch.zeros(
|
||||
batch_size,
|
||||
seq_len,
|
||||
model.d_model,
|
||||
model.n_embd,
|
||||
device=event_seq.device,
|
||||
dtype=torch.float32,
|
||||
)
|
||||
@@ -541,17 +416,29 @@ def infer_readout_hidden(
|
||||
active = padding_mask[:, pos].bool()
|
||||
if not active.any():
|
||||
continue
|
||||
hidden_pos = model(
|
||||
(
|
||||
model_event_seq,
|
||||
model_time_seq,
|
||||
model_padding_mask,
|
||||
model_t_query,
|
||||
) = transform_disease_history_batch_at_position(
|
||||
event_seq=event_seq[active],
|
||||
time_seq=time_seq[active],
|
||||
sex=batch_dev["sex"][active],
|
||||
actual_time_seq=time_seq[active],
|
||||
padding_mask=padding_mask[active],
|
||||
t_query=time_seq[active, pos],
|
||||
query_position=pos,
|
||||
disease_history_mode=disease_history_mode,
|
||||
vocab_size=model.vocab_size,
|
||||
)
|
||||
hidden_pos = model(
|
||||
event_seq=model_event_seq,
|
||||
time_seq=model_time_seq,
|
||||
sex=batch_dev["sex"][active],
|
||||
padding_mask=model_padding_mask,
|
||||
t_query=model_t_query,
|
||||
other_type=batch_dev["other_type"][active],
|
||||
other_value=batch_dev["other_value"][active],
|
||||
other_value_kind=batch_dev["other_value_kind"][active],
|
||||
other_time=batch_dev["other_time"][active],
|
||||
target_mode="all_future",
|
||||
)
|
||||
hidden[active, pos, :] = hidden_pos.float()
|
||||
readout_mask_np = batch["padding_mask"].cpu().numpy()
|
||||
@@ -565,16 +452,9 @@ def infer_readout_hidden(
|
||||
other_value=batch_dev["other_value"],
|
||||
other_value_kind=batch_dev["other_value_kind"],
|
||||
other_time=batch_dev["other_time"],
|
||||
target_mode="next_token",
|
||||
)
|
||||
ro = readout(
|
||||
hidden=hidden_raw,
|
||||
time_seq=time_seq,
|
||||
padding_mask=padding_mask,
|
||||
readout_mask=batch_dev["readout_mask"],
|
||||
)
|
||||
hidden = ro.hidden
|
||||
readout_mask_np = ro.readout_mask.detach().cpu().numpy()
|
||||
hidden = hidden_raw
|
||||
readout_mask_np = padding_mask.detach().cpu().numpy()
|
||||
|
||||
h = hidden.detach().cpu().numpy().astype(out_dtype, copy=False)
|
||||
hidden_parts.append(h)
|
||||
@@ -1063,8 +943,7 @@ def evaluate_auc_pipeline(
|
||||
offsets: Sequence[float],
|
||||
device: torch.device,
|
||||
model_target_mode: str,
|
||||
readout_name: str,
|
||||
readout_reduce: str,
|
||||
disease_history_mode: str,
|
||||
num_workers_auc: int,
|
||||
use_amp: bool,
|
||||
auc_task_chunk_size: int = 0,
|
||||
@@ -1113,7 +992,7 @@ def evaluate_auc_pipeline(
|
||||
sex_items = [("female", 0), ("male", 1)]
|
||||
all_rows: List[Dict[str, Any]] = []
|
||||
|
||||
valid_target_min_id = CHECKUP_IDX if NO_EVENT_IDX >= dataset.vocab_size else CHECKUP_IDX
|
||||
valid_target_min_id = RESERVED_IDX
|
||||
# If NO_EVENT exists and should not be a disease/control target, require target > NO_EVENT_IDX.
|
||||
if NO_EVENT_IDX in dataset.label_id_to_code and dataset.label_id_to_code.get(NO_EVENT_IDX) == "<NO_EVENT>":
|
||||
valid_target_min_id = NO_EVENT_IDX
|
||||
@@ -1123,8 +1002,7 @@ def evaluate_auc_pipeline(
|
||||
loader=loader,
|
||||
device=device,
|
||||
model_target_mode=model_target_mode,
|
||||
readout_name=readout_name,
|
||||
readout_reduce=readout_reduce,
|
||||
disease_history_mode=disease_history_mode,
|
||||
use_amp=use_amp,
|
||||
hidden_cache_dtype=hidden_cache_dtype,
|
||||
)
|
||||
@@ -1169,30 +1047,23 @@ def evaluate_auc_pipeline(
|
||||
df_auc_unpooled["label_code"] = df_auc_unpooled["token"].map(
|
||||
dataset.label_id_to_code)
|
||||
|
||||
print("Using DeLong method to calculate AUC confidence intervals.")
|
||||
grouped = df_auc_unpooled.groupby(
|
||||
["token", "label_code", "offset"], dropna=False, as_index=False)
|
||||
df_auc = grouped.agg(
|
||||
auc=("auc_delong", "mean"),
|
||||
n_strata=("auc_delong", "size"),
|
||||
n_diseased=("n_diseased", "sum"),
|
||||
n_healthy=("n_healthy", "sum"),
|
||||
auc_variance_sum=("auc_variance_delong", "sum"),
|
||||
print(
|
||||
"Building Delphi2M-style report: mean AUC across age strata, "
|
||||
"reported separately for Female and Male."
|
||||
)
|
||||
df_auc["auc_variance_delong"] = (
|
||||
df_auc["auc_variance_sum"]
|
||||
/ (df_auc["n_strata"].clip(lower=1).astype(np.float64) ** 2)
|
||||
df_report = build_delphi2m_auc_report(
|
||||
df_auc_unpooled,
|
||||
period_col="offset",
|
||||
)
|
||||
df_auc = df_auc.drop(columns=["auc_variance_sum"])
|
||||
|
||||
if output_path is not None:
|
||||
out_dir = Path(output_path)
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
df_auc.to_csv(out_dir / "df_both.csv", index=False)
|
||||
df_auc_unpooled.to_csv(
|
||||
out_dir / "df_auc_unpooled.csv", index=False)
|
||||
report_path = out_dir / "df_auc_delphi2m_report.csv"
|
||||
df_report.to_csv(report_path, index=False)
|
||||
print(f"Saved Delphi2M-style AUC report: {report_path}")
|
||||
|
||||
return df_auc_unpooled, df_auc
|
||||
return df_auc_unpooled, df_report
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -1242,8 +1113,18 @@ def make_auc_offsets(args: argparse.Namespace, cfg: Dict[str, Any]) -> List[floa
|
||||
if explicit_offsets is not None:
|
||||
base_offsets = explicit_offsets
|
||||
else:
|
||||
next_token_offset = float(cfg_get(args, cfg, "offset", 0.1))
|
||||
base_offsets = [next_token_offset, 1.0, 5.0, 10.0]
|
||||
next_token_offset = float(
|
||||
cfg_get(
|
||||
args,
|
||||
cfg,
|
||||
"offset",
|
||||
DEFAULT_DELPHI2M_PERIODS_YEARS[0],
|
||||
)
|
||||
)
|
||||
base_offsets = [
|
||||
next_token_offset,
|
||||
*DEFAULT_DELPHI2M_PERIODS_YEARS[1:],
|
||||
]
|
||||
|
||||
offsets: List[float] = []
|
||||
seen = set()
|
||||
@@ -1270,6 +1151,15 @@ def main() -> None:
|
||||
choices=["train", "val", "valid",
|
||||
"validation", "test", "all"],
|
||||
help="Evaluation split. Defaults to 'test' unless cfg contains eval_split.")
|
||||
parser.add_argument(
|
||||
"--test_eid_file",
|
||||
type=str,
|
||||
default=None,
|
||||
help=(
|
||||
"Patient EID file for the test split. Defaults to train_config.json "
|
||||
"or ukb_test_eid.csv. Set to an empty value to use ratio splitting."
|
||||
),
|
||||
)
|
||||
parser.add_argument("--dataset_subset_size", type=int, default=None,
|
||||
help="Optional number of patients from the selected split.")
|
||||
parser.add_argument("--batch_size", type=int, default=None,
|
||||
@@ -1291,9 +1181,9 @@ def main() -> None:
|
||||
parser.add_argument("--filter_min_total", type=int, default=None,
|
||||
help="Minimum metadata count for disease selection; default 0.")
|
||||
parser.add_argument("--offset", type=float, default=None,
|
||||
help="Next-token prediction offset in years; preserved and evaluated alongside 1, 5, and 10 years by default.")
|
||||
help="Next-token prediction offset in years; 0.1 is Delphi2M no gap and is evaluated alongside 1, 5, and 10 years by default.")
|
||||
parser.add_argument("--offsets", type=str, default=None,
|
||||
help="Comma-separated prediction offsets in years. Overrides the default set of offset,1,5,10.")
|
||||
help="Comma-separated prediction offsets in years. Overrides the default set of 0.1,1,5,10.")
|
||||
parser.add_argument("--age_start", type=float, default=None)
|
||||
parser.add_argument("--age_stop", type=float, default=None)
|
||||
parser.add_argument("--age_step", type=float, default=None)
|
||||
@@ -1313,6 +1203,7 @@ def main() -> None:
|
||||
raise FileNotFoundError(f"best_model.pt not found in {run_path}")
|
||||
|
||||
cfg = load_json_config(str(config_path))
|
||||
validate_training_mode_config(cfg)
|
||||
|
||||
if args.output_path is None:
|
||||
args.output_path = str(run_path)
|
||||
@@ -1321,9 +1212,7 @@ def main() -> None:
|
||||
data_prefix = cfg.get("data_prefix", "ukb")
|
||||
labels_file = cfg.get("labels_file", "labels.csv")
|
||||
no_event_interval_years = cfg.get("no_event_interval_years", 5.0)
|
||||
include_no_event = cfg.get("include_no_event_in_uts_target", False)
|
||||
|
||||
target_mode = cfg.get("target_mode", "uts")
|
||||
model_target_mode = str(cfg.get("model_target_mode", "next_token")).lower()
|
||||
if model_target_mode not in {"next_token", "all_future"}:
|
||||
raise ValueError(
|
||||
@@ -1331,9 +1220,9 @@ def main() -> None:
|
||||
f"got {model_target_mode!r}"
|
||||
)
|
||||
dist_mode_cfg = cfg.get("dist_mode", "exponential")
|
||||
readout_name = cfg.get(
|
||||
"readout_name", "same_time_group_end" if target_mode == "uts" else "token")
|
||||
readout_reduce = cfg.get("readout_reduce", "mean")
|
||||
disease_history_mode = normalize_disease_history_mode(
|
||||
cfg.get("disease_history_mode", DISEASE_HISTORY_MODE_TIMED)
|
||||
)
|
||||
|
||||
device = resolve_eval_device(args.device)
|
||||
if device.type == "cuda":
|
||||
@@ -1345,10 +1234,10 @@ def main() -> None:
|
||||
data_prefix=data_prefix,
|
||||
labels_file=labels_file,
|
||||
no_event_interval_years=no_event_interval_years,
|
||||
include_no_event_in_uts_target=include_no_event,
|
||||
min_history_events=int(cfg.get("all_future_min_history_events", 1)),
|
||||
min_future_events=int(cfg.get("all_future_min_future_events", 1)),
|
||||
extra_info_types=parse_int_list(cfg.get("extra_info_types", None)),
|
||||
disease_history_mode=disease_history_mode,
|
||||
)
|
||||
validate_dataset_metadata(dataset, cfg)
|
||||
|
||||
@@ -1374,14 +1263,20 @@ def main() -> None:
|
||||
cfg = dict(cfg)
|
||||
cfg["dist_mode"] = dist_mode
|
||||
cfg["model_target_mode"] = model_target_mode
|
||||
model_architecture = resolve_model_architecture(cfg, state_dict)
|
||||
cfg["model_architecture"] = model_architecture
|
||||
print(f"Resolved dist_mode for evaluation: {dist_mode}")
|
||||
print(f"Resolved model architecture: {model_architecture}")
|
||||
print(f"Model target mode for AUC: {model_target_mode}")
|
||||
print(f"Disease history mode for AUC: {disease_history_mode}")
|
||||
print(
|
||||
"AUC score semantics: evaluate_auc.py uses disease-specific eta/logit scores; "
|
||||
"dist_mode affects model loading but is not converted to horizon-specific risk probability."
|
||||
)
|
||||
|
||||
model = build_model_from_dataset(args, cfg, dataset).to(device)
|
||||
model = build_model_from_dataset(
|
||||
args, cfg, dataset, state_dict=state_dict
|
||||
).to(device)
|
||||
load_model_state(model, str(model_ckpt_path),
|
||||
device, state_dict=state_dict)
|
||||
model.eval()
|
||||
@@ -1396,8 +1291,9 @@ def main() -> None:
|
||||
if disease_spec is None:
|
||||
disease_spec = cfg.get("disease_tokens", None)
|
||||
diseases = parse_int_list(disease_spec)
|
||||
first_occurrence_by_token, _, _, _ = _build_first_occurrence_maps(
|
||||
dataset, subset_indices)
|
||||
first_occurrence_by_token = build_first_occurrence_map(
|
||||
dataset, subset_indices
|
||||
)
|
||||
include_death = bool(cfg_get(args, cfg, "include_death", True))
|
||||
exclude_death = bool(cfg_get(args, cfg, "exclude_death", False))
|
||||
auc_offsets = make_auc_offsets(args, cfg)
|
||||
@@ -1417,8 +1313,7 @@ def main() -> None:
|
||||
offsets=auc_offsets,
|
||||
device=device,
|
||||
model_target_mode=model_target_mode,
|
||||
readout_name=readout_name,
|
||||
readout_reduce=readout_reduce,
|
||||
disease_history_mode=disease_history_mode,
|
||||
num_workers_auc=int(cfg_get(args, cfg, "num_workers_auc", max(
|
||||
1, (os.cpu_count() or 2) - 1))),
|
||||
use_amp=bool(cfg_get(args, cfg, "use_amp", False)),
|
||||
|
||||
@@ -1,7 +1,10 @@
|
||||
"""Evaluate landmark fixed-horizon incident disease AUC for DeepHealth.
|
||||
|
||||
This script supports DeepHealth fixed-horizon risk scores for exponential,
|
||||
Weibull, and mixed all-future distributions.
|
||||
This script supports DeepHealth fixed-horizon risk scores for exponential and
|
||||
Weibull all-future distributions.
|
||||
|
||||
The default horizons are 0.1, 1, 5, and 10 years. As in Delphi2M, 0.1 years
|
||||
is reported as the no-gap evaluation.
|
||||
|
||||
Landmark querying depends on the model target mode saved in train_config.json:
|
||||
- next_token: insert a <NO_EVENT> token at landmark age and read it out;
|
||||
@@ -27,43 +30,33 @@ from torch.nn.utils.rnn import pad_sequence
|
||||
from torch.utils.data import DataLoader, Dataset
|
||||
from tqdm.auto import tqdm
|
||||
|
||||
from dataset import HealthDataset
|
||||
from eval_data import load_sequence_eval_dataset
|
||||
from models import (
|
||||
DeepHealth,
|
||||
validate_event_trajectory_config,
|
||||
validate_event_trajectory_state_dict,
|
||||
from dataset import (
|
||||
DISEASE_HISTORY_MODE_TIMED,
|
||||
HealthDataset,
|
||||
normalize_disease_history_mode,
|
||||
transform_disease_history,
|
||||
)
|
||||
from readouts import build_readout
|
||||
from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX
|
||||
from delphi2m_auc_report import (
|
||||
DEFAULT_DELPHI2M_PERIODS_YEARS,
|
||||
build_delphi2m_auc_report,
|
||||
)
|
||||
from eval_data import (
|
||||
build_first_occurrence_map,
|
||||
build_model_from_dataset,
|
||||
cfg_get,
|
||||
load_json_config,
|
||||
load_sequence_eval_dataset,
|
||||
resolve_eval_device,
|
||||
select_indices_by_eid_file,
|
||||
split_indices,
|
||||
validate_training_mode_config,
|
||||
validate_dataset_metadata,
|
||||
)
|
||||
from model_architectures import resolve_model_architecture
|
||||
from models import DeepHealth
|
||||
from targets import NO_EVENT_IDX, PAD_IDX, RESERVED_IDX
|
||||
|
||||
SPECIAL_TOKENS = {PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX}
|
||||
_TARGET_AWARE_MODES = {"target_aware", "delphi2m", "d2m"}
|
||||
|
||||
|
||||
def load_json_config(path: Path) -> Dict[str, Any]:
|
||||
if not path.exists():
|
||||
return {}
|
||||
with path.open("r", encoding="utf-8") as f:
|
||||
return json.load(f)
|
||||
|
||||
|
||||
def cfg_get(args: argparse.Namespace, cfg: Dict[str, Any], name: str, default: Any) -> Any:
|
||||
value = getattr(args, name, None)
|
||||
if value is not None:
|
||||
return value
|
||||
return cfg.get(name, default)
|
||||
|
||||
|
||||
def resolve_eval_device(device_arg: Optional[str]) -> torch.device:
|
||||
"""Resolve evaluation device without inheriting train_config.json device."""
|
||||
device_name = device_arg or ("cuda" if torch.cuda.is_available() else "cpu")
|
||||
device = torch.device(device_name)
|
||||
if device.type == "cuda" and not torch.cuda.is_available():
|
||||
raise RuntimeError(
|
||||
f"Requested device {device_name!r}, but CUDA is not available."
|
||||
)
|
||||
return device
|
||||
SPECIAL_TOKENS = {PAD_IDX, RESERVED_IDX, NO_EVENT_IDX}
|
||||
|
||||
|
||||
def parse_int_list(value: Any) -> Optional[List[int]]:
|
||||
@@ -104,18 +97,11 @@ def parse_float_list(value: Any) -> Optional[List[float]]:
|
||||
return [float(x.strip()) for x in text.split(",") if x.strip()]
|
||||
|
||||
|
||||
def split_indices(n: int, train_ratio: float, val_ratio: float, test_ratio: float, seed: int) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
|
||||
total = float(train_ratio) + float(val_ratio) + float(test_ratio)
|
||||
if not np.isclose(total, 1.0, atol=1e-6):
|
||||
raise ValueError(f"train/val/test ratios must sum to 1.0, got {total}")
|
||||
rng = np.random.RandomState(int(seed))
|
||||
idx = rng.permutation(int(n))
|
||||
n_train = int(n * train_ratio)
|
||||
n_val = int(n * val_ratio)
|
||||
return idx[:n_train], idx[n_train:n_train + n_val], idx[n_train + n_val:]
|
||||
|
||||
|
||||
def make_eval_indices(dataset: HealthDataset, args: argparse.Namespace, cfg: Dict[str, Any]) -> np.ndarray:
|
||||
def make_eval_indices(
|
||||
dataset: HealthDataset,
|
||||
args: argparse.Namespace,
|
||||
cfg: Dict[str, Any],
|
||||
) -> np.ndarray:
|
||||
train_ratio = float(cfg_get(args, cfg, "train_ratio", 0.7))
|
||||
val_ratio = float(cfg_get(args, cfg, "val_ratio", 0.15))
|
||||
test_ratio = float(cfg_get(args, cfg, "test_ratio", 0.15))
|
||||
@@ -124,6 +110,22 @@ def make_eval_indices(dataset: HealthDataset, args: argparse.Namespace, cfg: Dic
|
||||
if eval_split in {"valid", "validation"}:
|
||||
eval_split = "val"
|
||||
|
||||
if eval_split not in {"train", "val", "test", "all"}:
|
||||
raise ValueError(f"Unsupported eval_split={eval_split!r}")
|
||||
|
||||
test_eid_file = cfg_get(
|
||||
args,
|
||||
cfg,
|
||||
"test_eid_file",
|
||||
"ukb_test_eid.csv",
|
||||
)
|
||||
if eval_split == "test" and test_eid_file not in {None, ""}:
|
||||
indices, eid_path = select_indices_by_eid_file(
|
||||
dataset,
|
||||
str(test_eid_file),
|
||||
)
|
||||
print(f"Test split source: EID file {eid_path}")
|
||||
else:
|
||||
train_idx, val_idx, test_idx = split_indices(
|
||||
len(dataset), train_ratio, val_ratio, test_ratio, seed
|
||||
)
|
||||
@@ -133,8 +135,6 @@ def make_eval_indices(dataset: HealthDataset, args: argparse.Namespace, cfg: Dic
|
||||
"test": test_idx,
|
||||
"all": np.arange(len(dataset), dtype=np.int64),
|
||||
}
|
||||
if eval_split not in split_map:
|
||||
raise ValueError(f"Unsupported eval_split={eval_split!r}")
|
||||
indices = split_map[eval_split]
|
||||
|
||||
subset_size = cfg_get(args, cfg, "dataset_subset_size", None)
|
||||
@@ -156,92 +156,24 @@ def load_checkpoint_state_dict(checkpoint_path: Path, map_location: str | torch.
|
||||
|
||||
def resolve_dist_mode_for_checkpoint(cfg_dist_mode: str, state_dict: Dict[str, Any]) -> str:
|
||||
mode = str(cfg_dist_mode).lower()
|
||||
if mode not in {"exponential", "weibull"}:
|
||||
raise ValueError(
|
||||
f"Unsupported dist_mode={mode!r}; expected exponential or weibull."
|
||||
)
|
||||
has_rho_head = any(str(k).startswith("rho_head.")
|
||||
for k in state_dict.keys())
|
||||
has_rho_death_head = any(str(k).startswith("rho_death_head.")
|
||||
for k in state_dict.keys())
|
||||
if has_rho_head:
|
||||
if mode != "weibull":
|
||||
print(
|
||||
"[WARN] Checkpoint contains rho_head weights; overriding dist_mode to 'weibull' for evaluation.")
|
||||
return "weibull"
|
||||
if has_rho_death_head:
|
||||
if mode != "mixed":
|
||||
print(
|
||||
"[WARN] Checkpoint contains rho_death_head weights; overriding dist_mode to 'mixed' for evaluation.")
|
||||
return "mixed"
|
||||
if mode == "weibull":
|
||||
print(
|
||||
"[WARN] dist_mode is 'weibull' but checkpoint has no rho_head weights; overriding dist_mode to 'exponential'.")
|
||||
return "exponential"
|
||||
if mode == "mixed":
|
||||
print(
|
||||
"[WARN] dist_mode is 'mixed' but checkpoint has no rho_death_head weights; overriding dist_mode to 'exponential'.")
|
||||
return "exponential"
|
||||
return mode if mode in {"exponential", "weibull", "mixed"} else "exponential"
|
||||
|
||||
|
||||
def build_model_from_dataset(args: argparse.Namespace, cfg: Dict[str, Any], dataset: HealthDataset) -> DeepHealth:
|
||||
validate_event_trajectory_config(cfg)
|
||||
model_target_mode = str(cfg_get(
|
||||
args, cfg, "model_target_mode", "next_token")).lower()
|
||||
if model_target_mode not in {"next_token", "all_future"}:
|
||||
raise ValueError(
|
||||
f"model_target_mode must be next_token or all_future, got {model_target_mode!r}"
|
||||
)
|
||||
return DeepHealth(
|
||||
vocab_size=dataset.vocab_size,
|
||||
model_size=str(cfg_get(args, cfg, "model_size", "nano")),
|
||||
n_reasoning_rounds=int(
|
||||
cfg_get(args, cfg, "n_reasoning_rounds", 12)
|
||||
),
|
||||
n_types=dataset.n_types,
|
||||
n_cont_types=dataset.n_cont_types,
|
||||
n_categories=dataset.n_categories,
|
||||
cont_type_ids=dataset.cont_type_ids,
|
||||
n_bins=int(cfg_get(args, cfg, "n_bins", 16)),
|
||||
extra_pool_reduce=str(cfg_get(args, cfg, "extra_pool_reduce", "mean")),
|
||||
target_mode=model_target_mode,
|
||||
time_mode=str(cfg_get(args, cfg, "time_mode", "relative")),
|
||||
dist_mode=str(cfg_get(args, cfg, "dist_mode", "exponential")),
|
||||
dropout=float(cfg_get(args, cfg, "dropout", 0.0)),
|
||||
)
|
||||
|
||||
|
||||
def load_model_state(model: torch.nn.Module, state_dict: Dict[str, Any]) -> None:
|
||||
validate_event_trajectory_state_dict(
|
||||
state_dict,
|
||||
expected_d_model=model.d_model,
|
||||
expected_n_trajectory=model.n_trajectory,
|
||||
expected_n_reasoning_rounds=model.n_reasoning_rounds,
|
||||
)
|
||||
model.load_state_dict(state_dict, strict=True)
|
||||
|
||||
|
||||
def validate_dataset_metadata(dataset: HealthDataset, cfg: Dict[str, Any]) -> None:
|
||||
meta = cfg.get("dataset_metadata")
|
||||
if not isinstance(meta, dict):
|
||||
return
|
||||
|
||||
actual: Dict[str, Any] = {
|
||||
"vocab_size": int(dataset.vocab_size),
|
||||
"n_types": int(dataset.n_types),
|
||||
"n_cont_types": int(dataset.n_cont_types),
|
||||
"n_categories": int(dataset.n_categories),
|
||||
"cont_type_ids": [int(x) for x in dataset.cont_type_ids],
|
||||
"extra_info_types": [int(x) for x in dataset.extra_info_types],
|
||||
}
|
||||
mismatches = [
|
||||
f"{key}: train_config={meta.get(key)!r}, current_dataset={value!r}"
|
||||
for key, value in actual.items()
|
||||
if key in meta and meta.get(key) != value
|
||||
]
|
||||
if mismatches:
|
||||
if mode == "weibull" and not has_rho_head:
|
||||
raise RuntimeError("Weibull checkpoint is missing rho_head parameters.")
|
||||
if mode == "exponential" and has_rho_head:
|
||||
raise RuntimeError(
|
||||
"Current dataset metadata does not match train_config.json. "
|
||||
"Use the same prepared data and extra_info_types as training. "
|
||||
+ "; ".join(mismatches)
|
||||
"Exponential checkpoint unexpectedly contains rho_head parameters."
|
||||
)
|
||||
return mode
|
||||
|
||||
|
||||
def load_model_state(model: DeepHealth, state_dict: Dict[str, Any]) -> None:
|
||||
resolve_model_architecture(model.model_architecture, state_dict)
|
||||
model.load_state_dict(state_dict, strict=True)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -335,44 +267,6 @@ def _first_existing_column(df: pd.DataFrame, candidates: Sequence[str]) -> Optio
|
||||
return None
|
||||
|
||||
|
||||
def build_metadata_for_merge(dataset: HealthDataset, labels_meta: Optional[pd.DataFrame]) -> pd.DataFrame:
|
||||
base_rows = []
|
||||
for token, code in dataset.label_id_to_code.items():
|
||||
token = int(token)
|
||||
code_text = str(code)
|
||||
if token in SPECIAL_TOKENS or code_text.startswith("<"):
|
||||
continue
|
||||
base_rows.append({"token": token, "label_code": code_text})
|
||||
base = pd.DataFrame(base_rows)
|
||||
if labels_meta is None or labels_meta.empty:
|
||||
return base
|
||||
|
||||
meta = labels_meta.copy()
|
||||
code_col = _first_existing_column(
|
||||
meta, ["Name", "code", "ICD10", "icd10", "label", "token", "disease_code"])
|
||||
if code_col is not None:
|
||||
meta["_label_code"] = meta[code_col].astype(
|
||||
str).map(lambda s: s.split()[0].strip())
|
||||
merged = base.merge(meta, left_on="label_code",
|
||||
right_on="_label_code", how="left")
|
||||
return merged.drop(columns=["_label_code"], errors="ignore")
|
||||
|
||||
if "index" in meta.columns:
|
||||
idx = pd.to_numeric(meta["index"], errors="coerce")
|
||||
has_no_event = (
|
||||
NO_EVENT_IDX in dataset.label_id_to_code
|
||||
and dataset.label_id_to_code.get(NO_EVENT_IDX) == "<NO_EVENT>"
|
||||
)
|
||||
if has_no_event:
|
||||
idx = idx.where(idx < NO_EVENT_IDX, idx + 1)
|
||||
meta["_index_int"] = idx.astype("Int64")
|
||||
merged = base.merge(meta, left_on="token",
|
||||
right_on="_index_int", how="left")
|
||||
return merged.drop(columns=["_index_int"], errors="ignore")
|
||||
|
||||
return base
|
||||
|
||||
|
||||
def _metadata_count_map(dataset: HealthDataset, labels_meta: Optional[pd.DataFrame]) -> Dict[int, float]:
|
||||
if labels_meta is None or labels_meta.empty or "count" not in labels_meta.columns:
|
||||
return {}
|
||||
@@ -405,58 +299,10 @@ def _metadata_count_map(dataset: HealthDataset, labels_meta: Optional[pd.DataFra
|
||||
return out
|
||||
|
||||
|
||||
def _get_death_token_ids(dataset: HealthDataset, labels_meta: Optional[pd.DataFrame]) -> List[int]:
|
||||
def _get_death_token_ids(dataset: HealthDataset) -> List[int]:
|
||||
return [int(dataset.vocab_size) - 1]
|
||||
|
||||
|
||||
def _build_first_occurrence_maps(
|
||||
dataset: HealthDataset,
|
||||
subset_indices: np.ndarray,
|
||||
) -> Tuple[Dict[int, Tuple[np.ndarray, np.ndarray]], np.ndarray, np.ndarray, np.ndarray]:
|
||||
patient_count = len(subset_indices)
|
||||
followup_end = np.full(patient_count, -np.inf, dtype=np.float32)
|
||||
death_time = np.full(patient_count, np.inf, dtype=np.float32)
|
||||
sex = np.full(patient_count, -1, dtype=np.int8)
|
||||
|
||||
first_lists: Dict[int, List[Tuple[int, float]]] = {}
|
||||
|
||||
for patient_id, dataset_index in enumerate(subset_indices.tolist()):
|
||||
s = dataset.samples[int(dataset_index)]
|
||||
|
||||
seq_event = np.asarray(s["event_seq"], dtype=np.int64)
|
||||
seq_time = np.asarray(s["time_seq"], dtype=np.float32)
|
||||
tgt_event = np.asarray(s["target_event_seq"], dtype=np.int64)
|
||||
tgt_time = np.asarray(s["target_time_seq"], dtype=np.float32)
|
||||
|
||||
if seq_event.size == 0 or tgt_event.size == 0:
|
||||
continue
|
||||
|
||||
full_event = np.concatenate([seq_event, tgt_event[-1:]])
|
||||
full_time = np.concatenate([seq_time, tgt_time[-1:]])
|
||||
|
||||
sex[patient_id] = int(s["sex"])
|
||||
followup_end[patient_id] = np.max(full_time).astype(np.float32)
|
||||
|
||||
uniq_tokens, first_idx = np.unique(full_event, return_index=True)
|
||||
for token, idx in zip(uniq_tokens.tolist(), first_idx.tolist()):
|
||||
token = int(token)
|
||||
event_time = float(full_time[int(idx)])
|
||||
if token not in first_lists:
|
||||
first_lists[token] = []
|
||||
first_lists[token].append((patient_id, event_time))
|
||||
|
||||
packed: Dict[int, Tuple[np.ndarray, np.ndarray]] = {}
|
||||
for token, pairs in first_lists.items():
|
||||
if not pairs:
|
||||
continue
|
||||
packed[int(token)] = (
|
||||
np.asarray([p for p, _ in pairs], dtype=np.int32),
|
||||
np.asarray([t for _, t in pairs], dtype=np.float32),
|
||||
)
|
||||
|
||||
return packed, followup_end, death_time, sex
|
||||
|
||||
|
||||
def select_disease_tokens(
|
||||
dataset: HealthDataset,
|
||||
labels_meta: Optional[pd.DataFrame],
|
||||
@@ -502,22 +348,31 @@ class LandmarkDataset(Dataset):
|
||||
dataset: HealthDataset,
|
||||
subset_indices: np.ndarray,
|
||||
landmark_ages: np.ndarray,
|
||||
attn_mask_mode: str,
|
||||
model_target_mode: str,
|
||||
min_history_events: int,
|
||||
first_occurrence_by_token: Dict[int, Tuple[np.ndarray, np.ndarray]],
|
||||
death_token_ids: Sequence[int],
|
||||
disease_history_mode: str = DISEASE_HISTORY_MODE_TIMED,
|
||||
) -> None:
|
||||
self.dataset = dataset
|
||||
self.subset_indices = np.asarray(subset_indices, dtype=np.int64)
|
||||
self.landmark_ages = np.asarray(landmark_ages, dtype=np.float32)
|
||||
self.attn_mask_mode = str(attn_mask_mode).lower()
|
||||
self.model_target_mode = str(model_target_mode).lower()
|
||||
if self.model_target_mode not in {"next_token", "all_future"}:
|
||||
raise ValueError(
|
||||
"model_target_mode must be next_token or all_future, got "
|
||||
f"{self.model_target_mode!r}"
|
||||
)
|
||||
self.disease_history_mode = normalize_disease_history_mode(
|
||||
disease_history_mode
|
||||
)
|
||||
if (
|
||||
self.model_target_mode != "all_future"
|
||||
and self.disease_history_mode != DISEASE_HISTORY_MODE_TIMED
|
||||
):
|
||||
raise ValueError(
|
||||
"ordered/set disease history is only supported for all_future models"
|
||||
)
|
||||
self.min_history_events = int(min_history_events)
|
||||
|
||||
self.first_occurrence_by_token = first_occurrence_by_token
|
||||
@@ -572,7 +427,7 @@ class LandmarkDataset(Dataset):
|
||||
prefix_times = full_time[prefix_mask]
|
||||
|
||||
valid_history_mask = ~np.isin(prefix_events, np.array(
|
||||
[PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX], dtype=np.int64))
|
||||
[PAD_IDX, RESERVED_IDX, NO_EVENT_IDX], dtype=np.int64))
|
||||
if valid_history_mask.sum() < self.min_history_events:
|
||||
continue
|
||||
|
||||
@@ -589,18 +444,25 @@ class LandmarkDataset(Dataset):
|
||||
np.array([np.float32(landmark_age)], dtype=np.float32),
|
||||
]
|
||||
)
|
||||
if self.attn_mask_mode in _TARGET_AWARE_MODES:
|
||||
time_seq_landmark[-1] = np.nextafter(
|
||||
np.float32(landmark_age), np.float32(np.inf), dtype=np.float32
|
||||
np.float32(landmark_age),
|
||||
np.float32(np.inf),
|
||||
dtype=np.float32,
|
||||
)
|
||||
landmark_pos = int(len(event_seq_landmark) - 1)
|
||||
readout_mask = np.zeros(len(event_seq_landmark), dtype=bool)
|
||||
readout_mask[-1] = True
|
||||
else:
|
||||
event_seq_landmark = prefix_events.astype(
|
||||
np.int64, copy=False)
|
||||
time_seq_landmark = prefix_times.astype(
|
||||
np.float32, copy=False)
|
||||
(
|
||||
event_seq_landmark,
|
||||
time_seq_landmark,
|
||||
model_t_query,
|
||||
) = transform_disease_history(
|
||||
event_seq=prefix_events,
|
||||
actual_time_seq=prefix_times,
|
||||
actual_t_query=landmark_age,
|
||||
disease_history_mode=self.disease_history_mode,
|
||||
)
|
||||
landmark_pos = int(len(event_seq_landmark) - 1)
|
||||
readout_mask = np.zeros(len(event_seq_landmark), dtype=bool)
|
||||
|
||||
@@ -613,7 +475,11 @@ class LandmarkDataset(Dataset):
|
||||
"followup_end_time": np.float32(followup_end),
|
||||
"death_time": np.float32(self.patient_death_time[patient_id]),
|
||||
"landmark_pos": landmark_pos,
|
||||
"t_query": np.float32(landmark_age),
|
||||
"t_query": (
|
||||
np.float32(landmark_age)
|
||||
if self.model_target_mode == "next_token"
|
||||
else model_t_query
|
||||
),
|
||||
"event_seq": event_seq_landmark,
|
||||
"time_seq": time_seq_landmark,
|
||||
"readout_mask": readout_mask,
|
||||
@@ -705,8 +571,6 @@ def infer_landmark_hidden(
|
||||
loader: DataLoader,
|
||||
device: torch.device,
|
||||
model_target_mode: str,
|
||||
readout_name: str,
|
||||
readout_reduce: str,
|
||||
use_amp: bool,
|
||||
hidden_cache_dtype: str,
|
||||
) -> Tuple[np.ndarray, Dict[str, np.ndarray]]:
|
||||
@@ -716,15 +580,6 @@ def infer_landmark_hidden(
|
||||
f"model_target_mode must be next_token or all_future, got {model_target_mode!r}"
|
||||
)
|
||||
|
||||
readout = None
|
||||
if model_target_mode == "next_token" and readout_name == "same_time_group_end":
|
||||
readout = build_readout("same_time_group_end",
|
||||
reduce=readout_reduce).to(device)
|
||||
elif model_target_mode == "next_token":
|
||||
readout = build_readout(readout_name).to(device)
|
||||
if readout is not None:
|
||||
readout.eval()
|
||||
|
||||
hidden_parts: List[np.ndarray] = []
|
||||
arrays = {
|
||||
"patient_id": [],
|
||||
@@ -761,7 +616,6 @@ def infer_landmark_hidden(
|
||||
other_value=batch_dev["other_value"],
|
||||
other_value_kind=batch_dev["other_value_kind"],
|
||||
other_time=batch_dev["other_time"],
|
||||
target_mode="all_future",
|
||||
)
|
||||
else:
|
||||
hidden = model(
|
||||
@@ -773,18 +627,11 @@ def infer_landmark_hidden(
|
||||
other_value=batch_dev["other_value"],
|
||||
other_value_kind=batch_dev["other_value_kind"],
|
||||
other_time=batch_dev["other_time"],
|
||||
target_mode="next_token",
|
||||
)
|
||||
readout_out = readout(
|
||||
hidden=hidden,
|
||||
time_seq=batch_dev["time_seq"],
|
||||
padding_mask=batch_dev["padding_mask"],
|
||||
readout_mask=batch_dev["readout_mask"],
|
||||
)
|
||||
landmark_hidden = readout_out.hidden.gather(
|
||||
landmark_hidden = hidden.gather(
|
||||
1,
|
||||
batch_dev["landmark_pos"].long()[:, None, None].expand(
|
||||
-1, 1, readout_out.hidden.shape[-1]
|
||||
-1, 1, hidden.shape[-1]
|
||||
),
|
||||
).squeeze(1)
|
||||
|
||||
@@ -833,24 +680,11 @@ def project_distribution_chunk(
|
||||
device=device, dtype=compute_dtype)
|
||||
rho_weight = None
|
||||
rho_bias = None
|
||||
death_rho_weight = None
|
||||
death_rho_bias = None
|
||||
mixed_death_cols: List[int] = []
|
||||
death_idx = int(getattr(model, "death_idx", getattr(model, "vocab_size", 0) - 1))
|
||||
|
||||
if dist_mode == "weibull":
|
||||
rho_weight = model.rho_head.weight[disease_ids].detach().to(
|
||||
device=device, dtype=compute_dtype)
|
||||
rho_bias = model.rho_head.bias[disease_ids].detach().to(
|
||||
device=device, dtype=compute_dtype)
|
||||
elif dist_mode == "mixed":
|
||||
mixed_death_cols = [j for j, token in enumerate(disease_ids)
|
||||
if int(token) == death_idx]
|
||||
if mixed_death_cols:
|
||||
death_rho_weight = model.rho_death_head.weight.detach().to(
|
||||
device=device, dtype=compute_dtype)
|
||||
death_rho_bias = model.rho_death_head.bias.detach().to(
|
||||
device=device, dtype=compute_dtype)
|
||||
|
||||
out_parts: List[np.ndarray] = []
|
||||
rho_parts: List[np.ndarray] = []
|
||||
@@ -865,14 +699,6 @@ def project_distribution_chunk(
|
||||
if dist_mode == "weibull":
|
||||
assert rho_weight is not None and rho_bias is not None
|
||||
rho = F.softplus(torch.matmul(h, rho_weight.t()) + rho_bias) + 1e-6
|
||||
elif dist_mode == "mixed" and mixed_death_cols:
|
||||
assert death_rho_weight is not None and death_rho_bias is not None
|
||||
rho = torch.ones_like(logits)
|
||||
death_rho = F.softplus(
|
||||
torch.matmul(h, death_rho_weight.t()).squeeze(-1) + death_rho_bias.squeeze(0)
|
||||
) + 1e-6
|
||||
for col in mixed_death_cols:
|
||||
rho[:, int(col)] = death_rho
|
||||
|
||||
out_parts.append(logits.float().cpu(
|
||||
).numpy().astype(np.float32, copy=False))
|
||||
@@ -909,7 +735,6 @@ def _init_worker(
|
||||
exclude_death_competing: bool,
|
||||
death_token_ids: np.ndarray,
|
||||
dist_mode: str,
|
||||
model_death_idx: int,
|
||||
) -> None:
|
||||
os.environ.setdefault("OMP_NUM_THREADS", "1")
|
||||
os.environ.setdefault("MKL_NUM_THREADS", "1")
|
||||
@@ -934,7 +759,6 @@ def _init_worker(
|
||||
"exclude_death_competing": bool(exclude_death_competing),
|
||||
"death_token_ids": set(int(x) for x in np.asarray(death_token_ids, dtype=np.int64).tolist()),
|
||||
"dist_mode": str(dist_mode).lower(),
|
||||
"model_death_idx": int(model_death_idx),
|
||||
"first_time_cache": {},
|
||||
}
|
||||
)
|
||||
@@ -960,8 +784,6 @@ def _score_to_probability(
|
||||
score_mode: str,
|
||||
horizon: float,
|
||||
dist_mode: str,
|
||||
token: int,
|
||||
death_idx: int,
|
||||
) -> np.ndarray:
|
||||
if score_mode == "eta":
|
||||
return logits.astype(np.float64, copy=False)
|
||||
@@ -973,11 +795,6 @@ def _score_to_probability(
|
||||
raise RuntimeError("Weibull risk scoring requires rho parameters.")
|
||||
exposure = np.power(np.float32(horizon), rho.astype(np.float32, copy=False))
|
||||
return (-np.expm1(-rate * exposure)).astype(np.float64, copy=False)
|
||||
if dist_mode == "mixed" and int(token) == int(death_idx):
|
||||
if rho is None:
|
||||
raise RuntimeError("Mixed death risk scoring requires death rho parameters.")
|
||||
exposure = np.power(np.float32(horizon), rho.astype(np.float32, copy=False))
|
||||
return (-np.expm1(-rate * exposure)).astype(np.float64, copy=False)
|
||||
return (-np.expm1(-rate * np.float32(horizon))).astype(np.float64, copy=False)
|
||||
|
||||
|
||||
@@ -994,7 +811,6 @@ def _eval_token(task: Tuple[int, int, str]) -> List[Dict[str, Any]]:
|
||||
rho_chunk = _WORKER["rho_chunk"]
|
||||
rho_token = None if rho_chunk is None else rho_chunk[:, int(j)]
|
||||
dist_mode = _WORKER["dist_mode"]
|
||||
model_death_idx = int(_WORKER["model_death_idx"])
|
||||
|
||||
first_time_patient = _first_time_by_patient(token)
|
||||
is_death_target = token in _WORKER["death_token_ids"]
|
||||
@@ -1053,8 +869,6 @@ def _eval_token(task: Tuple[int, int, str]) -> List[Dict[str, Any]]:
|
||||
score_mode=score_mode,
|
||||
horizon=horizon,
|
||||
dist_mode=dist_mode,
|
||||
token=token,
|
||||
death_idx=model_death_idx,
|
||||
)
|
||||
control_scores = _score_to_probability(
|
||||
logits_token[idx[control_idx]],
|
||||
@@ -1062,8 +876,6 @@ def _eval_token(task: Tuple[int, int, str]) -> List[Dict[str, Any]]:
|
||||
score_mode=score_mode,
|
||||
horizon=horizon,
|
||||
dist_mode=dist_mode,
|
||||
token=token,
|
||||
death_idx=model_death_idx,
|
||||
)
|
||||
|
||||
auc, auc_var = get_auc_delong_var(case_scores, control_scores)
|
||||
@@ -1112,7 +924,6 @@ def evaluate_landmark_auc(
|
||||
loader: DataLoader,
|
||||
landmark_dataset: LandmarkDataset,
|
||||
output_path: Path,
|
||||
labels_meta: Optional[pd.DataFrame],
|
||||
disease_ids: Sequence[int],
|
||||
disease_chunk_size: int,
|
||||
score_mode: str,
|
||||
@@ -1120,8 +931,6 @@ def evaluate_landmark_auc(
|
||||
horizons: np.ndarray,
|
||||
device: torch.device,
|
||||
model_target_mode: str,
|
||||
readout_name: str,
|
||||
readout_reduce: str,
|
||||
num_workers_auc: int,
|
||||
auc_task_chunk_size: int,
|
||||
min_cases: int,
|
||||
@@ -1129,7 +938,6 @@ def evaluate_landmark_auc(
|
||||
use_amp: bool,
|
||||
hidden_cache_dtype: str,
|
||||
logit_batch_size: int,
|
||||
meta_info: Dict[str, Any],
|
||||
) -> Tuple[pd.DataFrame, pd.DataFrame]:
|
||||
model.eval().to(device)
|
||||
|
||||
@@ -1138,8 +946,6 @@ def evaluate_landmark_auc(
|
||||
loader=loader,
|
||||
device=device,
|
||||
model_target_mode=model_target_mode,
|
||||
readout_name=readout_name,
|
||||
readout_reduce=readout_reduce,
|
||||
use_amp=use_amp,
|
||||
hidden_cache_dtype=hidden_cache_dtype,
|
||||
)
|
||||
@@ -1187,8 +993,6 @@ def evaluate_landmark_auc(
|
||||
death_token_ids=np.asarray(
|
||||
landmark_dataset.death_token_ids, dtype=np.int64),
|
||||
dist_mode=dist_mode,
|
||||
model_death_idx=int(getattr(
|
||||
model, "death_idx", getattr(model, "vocab_size", 1) - 1)),
|
||||
)
|
||||
nested = [_eval_token(t) for t in tqdm(
|
||||
tasks, desc=f"AUC chunk {chunk_idx}", leave=False, dynamic_ncols=True)]
|
||||
@@ -1216,8 +1020,6 @@ def evaluate_landmark_auc(
|
||||
np.asarray(landmark_dataset.death_token_ids,
|
||||
dtype=np.int64),
|
||||
dist_mode,
|
||||
int(getattr(
|
||||
model, "death_idx", getattr(model, "vocab_size", 1) - 1)),
|
||||
),
|
||||
) as ex:
|
||||
nested = list(
|
||||
@@ -1246,54 +1048,21 @@ def evaluate_landmark_auc(
|
||||
df_unpooled["label_code"] = df_unpooled["token"].map(
|
||||
landmark_dataset.dataset.label_id_to_code)
|
||||
|
||||
for k, v in meta_info.items():
|
||||
df_unpooled[k] = v
|
||||
|
||||
meta_table = build_metadata_for_merge(landmark_dataset.dataset, labels_meta)
|
||||
df_unpooled = df_unpooled.merge(
|
||||
meta_table, on=["token", "label_code"], how="left")
|
||||
|
||||
grouped = df_unpooled.groupby(
|
||||
["token", "label_code", "horizon"], dropna=False, as_index=False)
|
||||
df_merged = grouped.agg(
|
||||
auc=("auc_delong", "mean"),
|
||||
n_strata=("auc_delong", "size"),
|
||||
n_diseased=("n_diseased", "sum"),
|
||||
n_healthy=("n_healthy", "sum"),
|
||||
auc_variance_sum=("auc_variance_delong", "sum"),
|
||||
print(
|
||||
"Building Delphi2M-style report: mean AUC across landmark-age "
|
||||
"strata, reported separately for Female and Male."
|
||||
)
|
||||
df_merged["auc_variance_delong"] = (
|
||||
df_merged["auc_variance_sum"]
|
||||
/ (df_merged["n_strata"].clip(lower=1).astype(np.float64) ** 2)
|
||||
df_report = build_delphi2m_auc_report(
|
||||
df_unpooled,
|
||||
period_col="horizon",
|
||||
)
|
||||
df_merged = df_merged.drop(columns=["auc_variance_sum"])
|
||||
|
||||
keep_meta = [
|
||||
c for c in [
|
||||
"model_ckpt_path",
|
||||
"config_path",
|
||||
"target_mode",
|
||||
"model_target_mode",
|
||||
"dist_mode",
|
||||
"time_mode",
|
||||
"attn_mask_mode",
|
||||
"readout_name",
|
||||
"landmark_query_mode",
|
||||
"landmark_token_mode",
|
||||
"score_mode",
|
||||
"eval_split",
|
||||
]
|
||||
if c in df_unpooled.columns
|
||||
]
|
||||
for col in keep_meta:
|
||||
df_merged[col] = meta_info[col]
|
||||
|
||||
output_path.mkdir(parents=True, exist_ok=True)
|
||||
df_unpooled.to_csv(
|
||||
output_path / "df_auc_landmark_unpooled.csv", index=False)
|
||||
df_merged.to_csv(output_path / "df_auc_landmark.csv", index=False)
|
||||
report_path = output_path / "df_auc_landmark_delphi2m_report.csv"
|
||||
df_report.to_csv(report_path, index=False)
|
||||
print(f"Saved Delphi2M-style landmark AUC report: {report_path}")
|
||||
|
||||
return df_unpooled, df_merged
|
||||
return df_unpooled, df_report
|
||||
|
||||
|
||||
def main() -> None:
|
||||
@@ -1304,6 +1073,15 @@ def main() -> None:
|
||||
parser.add_argument("--output_path", type=str, default=None)
|
||||
parser.add_argument("--eval_split", type=str, default="test",
|
||||
choices=["train", "val", "valid", "validation", "test", "all"])
|
||||
parser.add_argument(
|
||||
"--test_eid_file",
|
||||
type=str,
|
||||
default=None,
|
||||
help=(
|
||||
"Patient EID file for the test split. Defaults to train_config.json "
|
||||
"or ukb_test_eid.csv. Set to an empty value to use ratio splitting."
|
||||
),
|
||||
)
|
||||
parser.add_argument("--dataset_subset_size", type=int, default=None)
|
||||
|
||||
parser.add_argument("--batch_size", type=int, default=None)
|
||||
@@ -1319,7 +1097,12 @@ def main() -> None:
|
||||
parser.add_argument("--landmark_start", type=float, default=None)
|
||||
parser.add_argument("--landmark_stop", type=float, default=None)
|
||||
parser.add_argument("--landmark_step", type=float, default=None)
|
||||
parser.add_argument("--horizons", type=str, default=None)
|
||||
parser.add_argument(
|
||||
"--horizons",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Comma-separated horizons in years; defaults to 0.1,1,5,10, where 0.1 is Delphi2M no gap.",
|
||||
)
|
||||
|
||||
parser.add_argument("--min_cases", type=int, default=None)
|
||||
parser.add_argument("--min_history_events", type=int, default=None)
|
||||
@@ -1346,14 +1129,12 @@ def main() -> None:
|
||||
raise FileNotFoundError(f"best_model.pt not found in {run_path}")
|
||||
|
||||
cfg = load_json_config(config_path)
|
||||
validate_training_mode_config(cfg)
|
||||
|
||||
data_prefix = cfg.get("data_prefix", "ukb")
|
||||
labels_file = cfg.get("labels_file", "labels.csv")
|
||||
no_event_interval_years = cfg.get("no_event_interval_years", 5.0)
|
||||
include_no_event_in_uts_target = cfg.get(
|
||||
"include_no_event_in_uts_target", False)
|
||||
|
||||
target_mode = cfg.get("target_mode", "uts")
|
||||
model_target_mode = str(cfg.get("model_target_mode", "next_token")).lower()
|
||||
if model_target_mode not in {"next_token", "all_future"}:
|
||||
raise ValueError(
|
||||
@@ -1361,12 +1142,9 @@ def main() -> None:
|
||||
f"got {model_target_mode!r}"
|
||||
)
|
||||
dist_mode_cfg = str(cfg.get("dist_mode", "exponential"))
|
||||
attn_mask_mode = str(cfg.get(
|
||||
"attn_mask_mode", "non_strict_time" if target_mode == "uts" else "target_aware"))
|
||||
readout_name = str(cfg.get(
|
||||
"readout_name", "same_time_group_end" if target_mode == "uts" else "token"))
|
||||
readout_reduce = str(cfg.get("readout_reduce", "mean"))
|
||||
time_mode = str(cfg.get("time_mode", "relative"))
|
||||
disease_history_mode = normalize_disease_history_mode(
|
||||
cfg.get("disease_history_mode", DISEASE_HISTORY_MODE_TIMED)
|
||||
)
|
||||
|
||||
output_path = Path(
|
||||
cfg_get(args, cfg, "output_path", None)
|
||||
@@ -1389,10 +1167,10 @@ def main() -> None:
|
||||
data_prefix=data_prefix,
|
||||
labels_file=labels_file,
|
||||
no_event_interval_years=float(no_event_interval_years),
|
||||
include_no_event_in_uts_target=bool(include_no_event_in_uts_target),
|
||||
min_history_events=int(cfg.get("all_future_min_history_events", 1)),
|
||||
min_future_events=int(cfg.get("all_future_min_future_events", 1)),
|
||||
extra_info_types=parse_int_list(cfg.get("extra_info_types", None)),
|
||||
disease_history_mode=disease_history_mode,
|
||||
)
|
||||
validate_dataset_metadata(dataset, cfg)
|
||||
|
||||
@@ -1411,8 +1189,9 @@ def main() -> None:
|
||||
|
||||
subset_indices = make_eval_indices(dataset, args, cfg)
|
||||
|
||||
first_occurrence_by_token, _, _, _ = _build_first_occurrence_maps(
|
||||
dataset, subset_indices)
|
||||
first_occurrence_by_token = build_first_occurrence_map(
|
||||
dataset, subset_indices
|
||||
)
|
||||
|
||||
disease_requested = parse_int_list(
|
||||
cfg_get(args, cfg, "diseases_of_interest", None))
|
||||
@@ -1439,8 +1218,9 @@ def main() -> None:
|
||||
"Landmark ages are empty. Check landmark_start/landmark_stop/landmark_step.")
|
||||
|
||||
horizons = np.asarray(
|
||||
parse_float_list(cfg_get(args, cfg, "horizons", "1,5,10")) or [
|
||||
1.0, 5.0, 10.0],
|
||||
parse_float_list(
|
||||
cfg_get(args, cfg, "horizons", "0.1,1,5,10")
|
||||
) or list(DEFAULT_DELPHI2M_PERIODS_YEARS),
|
||||
dtype=np.float32,
|
||||
)
|
||||
if horizons.size == 0:
|
||||
@@ -1452,10 +1232,6 @@ def main() -> None:
|
||||
|
||||
state_dict = load_checkpoint_state_dict(model_ckpt_path, map_location="cpu")
|
||||
dist_mode = resolve_dist_mode_for_checkpoint(dist_mode_cfg, state_dict)
|
||||
if dist_mode not in {"exponential", "weibull", "mixed"}:
|
||||
raise ValueError(
|
||||
f"Unsupported dist_mode={dist_mode!r}; expected exponential, weibull, or mixed."
|
||||
)
|
||||
|
||||
if score_mode == "eta":
|
||||
print(
|
||||
@@ -1463,12 +1239,17 @@ def main() -> None:
|
||||
|
||||
cfg_model = dict(cfg)
|
||||
cfg_model["dist_mode"] = dist_mode
|
||||
model_architecture = resolve_model_architecture(cfg_model, state_dict)
|
||||
cfg_model["model_architecture"] = model_architecture
|
||||
print(f"Resolved model architecture: {model_architecture}")
|
||||
|
||||
device = resolve_eval_device(args.device)
|
||||
if device.type == "cuda":
|
||||
torch.backends.cudnn.benchmark = True
|
||||
|
||||
model = build_model_from_dataset(args, cfg_model, dataset).to(device)
|
||||
model = build_model_from_dataset(
|
||||
args, cfg_model, dataset, state_dict=state_dict
|
||||
).to(device)
|
||||
|
||||
if (
|
||||
model_target_mode == "next_token"
|
||||
@@ -1496,17 +1277,17 @@ def main() -> None:
|
||||
"Please use a checkpoint trained with the same no-event vocabulary configuration."
|
||||
)
|
||||
|
||||
death_token_ids = _get_death_token_ids(dataset, labels_meta)
|
||||
death_token_ids = _get_death_token_ids(dataset)
|
||||
min_history_events = int(cfg_get(args, cfg, "min_history_events", 1))
|
||||
landmark_dataset = LandmarkDataset(
|
||||
dataset=dataset,
|
||||
subset_indices=subset_indices,
|
||||
landmark_ages=landmark_ages,
|
||||
attn_mask_mode=attn_mask_mode,
|
||||
model_target_mode=model_target_mode,
|
||||
min_history_events=min_history_events,
|
||||
first_occurrence_by_token=first_occurrence_by_token,
|
||||
death_token_ids=death_token_ids,
|
||||
disease_history_mode=disease_history_mode,
|
||||
)
|
||||
|
||||
batch_size = int(cfg_get(args, cfg, "batch_size", 128))
|
||||
@@ -1531,8 +1312,6 @@ def main() -> None:
|
||||
if model_target_mode == "next_token"
|
||||
else "direct_t_query"
|
||||
)
|
||||
score_mode_out = f"{landmark_query_mode}_{score_mode}"
|
||||
|
||||
num_workers_auc = int(
|
||||
cfg_get(args, cfg, "num_workers_auc", max(1, (os.cpu_count() or 2) - 1)))
|
||||
auc_task_chunk_size = int(cfg_get(args, cfg, "auc_task_chunk_size", 0))
|
||||
@@ -1549,6 +1328,7 @@ def main() -> None:
|
||||
print(f"Number of selected patients: {len(subset_indices)}")
|
||||
print(f"No-event support: {bool(has_no_event)}")
|
||||
print(f"Model target mode: {model_target_mode}")
|
||||
print(f"Disease history mode: {disease_history_mode}")
|
||||
print(f"Landmark query mode: {landmark_query_mode}")
|
||||
print(
|
||||
"Landmark token mode: no_event"
|
||||
@@ -1564,27 +1344,11 @@ def main() -> None:
|
||||
print(f"AUC workers: {num_workers_auc}")
|
||||
print(f"Output path: {output_path}")
|
||||
|
||||
meta_info = {
|
||||
"score_mode": score_mode_out,
|
||||
"eval_split": eval_split,
|
||||
"model_ckpt_path": str(model_ckpt_path),
|
||||
"config_path": str(config_path),
|
||||
"target_mode": str(target_mode),
|
||||
"model_target_mode": str(model_target_mode),
|
||||
"dist_mode": str(dist_mode),
|
||||
"time_mode": str(time_mode),
|
||||
"attn_mask_mode": str(attn_mask_mode),
|
||||
"readout_name": str(readout_name),
|
||||
"landmark_query_mode": landmark_query_mode,
|
||||
"landmark_token_mode": "no_event" if model_target_mode == "next_token" else "none",
|
||||
}
|
||||
|
||||
evaluate_landmark_auc(
|
||||
model=model,
|
||||
loader=loader,
|
||||
landmark_dataset=landmark_dataset,
|
||||
output_path=output_path,
|
||||
labels_meta=labels_meta,
|
||||
disease_ids=disease_ids,
|
||||
disease_chunk_size=disease_chunk_size,
|
||||
score_mode=score_mode,
|
||||
@@ -1592,8 +1356,6 @@ def main() -> None:
|
||||
horizons=horizons,
|
||||
device=device,
|
||||
model_target_mode=model_target_mode,
|
||||
readout_name=readout_name,
|
||||
readout_reduce=readout_reduce,
|
||||
num_workers_auc=num_workers_auc,
|
||||
auc_task_chunk_size=auc_task_chunk_size,
|
||||
min_cases=min_cases,
|
||||
@@ -1601,7 +1363,6 @@ def main() -> None:
|
||||
use_amp=use_amp,
|
||||
hidden_cache_dtype=hidden_cache_dtype,
|
||||
logit_batch_size=logit_batch_size,
|
||||
meta_info=meta_info,
|
||||
)
|
||||
|
||||
|
||||
|
||||
1980
evaluate_calibration.py
Normal file
1980
evaluate_calibration.py
Normal file
File diff suppressed because it is too large
Load Diff
404
evaluate_calibration_all_runs_linux.sh
Normal file
404
evaluate_calibration_all_runs_linux.sh
Normal file
@@ -0,0 +1,404 @@
|
||||
#!/usr/bin/env bash
|
||||
#
|
||||
# Recursively evaluate calibration for every completed all_future run.
|
||||
#
|
||||
# A runnable run contains:
|
||||
# - train_config.json with model_target_mode="all_future"
|
||||
# - best_model.pt
|
||||
#
|
||||
# calibration_evaluation_summary.json is written last by the evaluator and is
|
||||
# used as the completion marker. Existing non-empty markers are skipped unless
|
||||
# --force is supplied. next_token/Delphi2M runs are intentionally skipped.
|
||||
#
|
||||
# Jobs assigned to the same GPU run sequentially; different GPUs run in
|
||||
# parallel.
|
||||
#
|
||||
# Examples:
|
||||
# bash evaluate_calibration_all_runs_linux.sh --gpus 0
|
||||
# bash evaluate_calibration_all_runs_linux.sh --gpus 0,1,2,3
|
||||
# bash evaluate_calibration_all_runs_linux.sh --gpus 0,1 --dry-run
|
||||
#
|
||||
|
||||
set -uo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
|
||||
cd "$SCRIPT_DIR"
|
||||
|
||||
RUNS_ROOT="$SCRIPT_DIR/runs"
|
||||
LOG_ROOT="$SCRIPT_DIR/batch_logs/evaluate_calibration_all_runs"
|
||||
GPU_CSV="0"
|
||||
PYTHON_BIN="${PYTHON_BIN:-python}"
|
||||
NUM_WORKERS=4
|
||||
NUM_WORKERS_CALIBRATION=0
|
||||
BATCH_SIZE=128
|
||||
DISEASE_CHUNK_SIZE=64
|
||||
HORIZONS=""
|
||||
USE_AMP=0
|
||||
FORCE=0
|
||||
DRY_RUN=0
|
||||
|
||||
COMPLETION_FILE="calibration_evaluation_summary.json"
|
||||
EVALUATOR="$SCRIPT_DIR/evaluate_calibration.py"
|
||||
|
||||
usage() {
|
||||
cat <<'EOF'
|
||||
Usage:
|
||||
bash evaluate_calibration_all_runs_linux.sh [options]
|
||||
|
||||
Options:
|
||||
--gpus LIST Comma-separated GPU ids (default: 0).
|
||||
--runs-root PATH Root directory scanned recursively
|
||||
(default: ./runs).
|
||||
--log-root PATH Evaluation log root.
|
||||
--python PATH Python executable
|
||||
(default: $PYTHON_BIN or python).
|
||||
--num-workers N DataLoader workers per job (default: 4).
|
||||
--num-workers-calibration N CPU calibration workers per job. Default: 0,
|
||||
which divides all logical CPUs across GPUs.
|
||||
--batch-size N Evaluation batch size (default: 128).
|
||||
--disease-chunk-size N Disease projection chunk size (default: 64).
|
||||
--horizons LIST Optional comma-separated horizons in years.
|
||||
--use-amp Force CUDA automatic mixed precision.
|
||||
--force Recompute runs with an existing completion file.
|
||||
--dry-run Discover and print pending jobs only.
|
||||
-h, --help Show this help message.
|
||||
|
||||
Only all_future runs are evaluated. next_token runs are skipped.
|
||||
Completion marker: calibration_evaluation_summary.json
|
||||
EOF
|
||||
}
|
||||
|
||||
while (($# > 0)); do
|
||||
case "$1" in
|
||||
--gpus)
|
||||
[[ $# -ge 2 ]] || {
|
||||
echo "ERROR: --gpus requires a value." >&2
|
||||
exit 2
|
||||
}
|
||||
GPU_CSV="$2"
|
||||
shift 2
|
||||
;;
|
||||
--runs-root)
|
||||
[[ $# -ge 2 ]] || {
|
||||
echo "ERROR: --runs-root requires a value." >&2
|
||||
exit 2
|
||||
}
|
||||
RUNS_ROOT="$2"
|
||||
shift 2
|
||||
;;
|
||||
--log-root)
|
||||
[[ $# -ge 2 ]] || {
|
||||
echo "ERROR: --log-root requires a value." >&2
|
||||
exit 2
|
||||
}
|
||||
LOG_ROOT="$2"
|
||||
shift 2
|
||||
;;
|
||||
--python)
|
||||
[[ $# -ge 2 ]] || {
|
||||
echo "ERROR: --python requires a value." >&2
|
||||
exit 2
|
||||
}
|
||||
PYTHON_BIN="$2"
|
||||
shift 2
|
||||
;;
|
||||
--num-workers)
|
||||
[[ $# -ge 2 ]] || {
|
||||
echo "ERROR: --num-workers requires a value." >&2
|
||||
exit 2
|
||||
}
|
||||
NUM_WORKERS="$2"
|
||||
shift 2
|
||||
;;
|
||||
--num-workers-calibration)
|
||||
[[ $# -ge 2 ]] || {
|
||||
echo "ERROR: --num-workers-calibration requires a value." >&2
|
||||
exit 2
|
||||
}
|
||||
NUM_WORKERS_CALIBRATION="$2"
|
||||
shift 2
|
||||
;;
|
||||
--batch-size)
|
||||
[[ $# -ge 2 ]] || {
|
||||
echo "ERROR: --batch-size requires a value." >&2
|
||||
exit 2
|
||||
}
|
||||
BATCH_SIZE="$2"
|
||||
shift 2
|
||||
;;
|
||||
--disease-chunk-size)
|
||||
[[ $# -ge 2 ]] || {
|
||||
echo "ERROR: --disease-chunk-size requires a value." >&2
|
||||
exit 2
|
||||
}
|
||||
DISEASE_CHUNK_SIZE="$2"
|
||||
shift 2
|
||||
;;
|
||||
--horizons)
|
||||
[[ $# -ge 2 ]] || {
|
||||
echo "ERROR: --horizons requires a value." >&2
|
||||
exit 2
|
||||
}
|
||||
HORIZONS="$2"
|
||||
shift 2
|
||||
;;
|
||||
--use-amp)
|
||||
USE_AMP=1
|
||||
shift
|
||||
;;
|
||||
--force)
|
||||
FORCE=1
|
||||
shift
|
||||
;;
|
||||
--dry-run)
|
||||
DRY_RUN=1
|
||||
shift
|
||||
;;
|
||||
-h|--help)
|
||||
usage
|
||||
exit 0
|
||||
;;
|
||||
*)
|
||||
echo "ERROR: unknown argument: $1" >&2
|
||||
usage >&2
|
||||
exit 2
|
||||
;;
|
||||
esac
|
||||
done
|
||||
|
||||
[[ -n "$GPU_CSV" ]] || {
|
||||
echo "ERROR: --gpus must not be empty." >&2
|
||||
exit 2
|
||||
}
|
||||
[[ "$NUM_WORKERS" =~ ^[0-9]+$ ]] || {
|
||||
echo "ERROR: --num-workers must be a non-negative integer." >&2
|
||||
exit 2
|
||||
}
|
||||
[[ "$NUM_WORKERS_CALIBRATION" =~ ^[0-9]+$ ]] || {
|
||||
echo "ERROR: --num-workers-calibration must be a non-negative integer." >&2
|
||||
exit 2
|
||||
}
|
||||
[[ "$BATCH_SIZE" =~ ^[1-9][0-9]*$ ]] || {
|
||||
echo "ERROR: --batch-size must be a positive integer." >&2
|
||||
exit 2
|
||||
}
|
||||
[[ "$DISEASE_CHUNK_SIZE" =~ ^[1-9][0-9]*$ ]] || {
|
||||
echo "ERROR: --disease-chunk-size must be a positive integer." >&2
|
||||
exit 2
|
||||
}
|
||||
[[ -d "$RUNS_ROOT" ]] || {
|
||||
echo "ERROR: runs root does not exist: $RUNS_ROOT" >&2
|
||||
exit 2
|
||||
}
|
||||
[[ -f "$EVALUATOR" ]] || {
|
||||
echo "ERROR: missing evaluator: $EVALUATOR" >&2
|
||||
exit 2
|
||||
}
|
||||
command -v "$PYTHON_BIN" >/dev/null 2>&1 || {
|
||||
echo "ERROR: Python executable not found: $PYTHON_BIN" >&2
|
||||
exit 2
|
||||
}
|
||||
|
||||
RUNS_ROOT="$(cd -- "$RUNS_ROOT" && pwd)"
|
||||
if [[ "$LOG_ROOT" != /* ]]; then
|
||||
LOG_ROOT="$SCRIPT_DIR/$LOG_ROOT"
|
||||
fi
|
||||
|
||||
IFS=',' read -r -a GPU_IDS <<< "$GPU_CSV"
|
||||
declare -A SEEN_GPUS=()
|
||||
for gpu in "${GPU_IDS[@]}"; do
|
||||
[[ -n "$gpu" && "$gpu" =~ ^[A-Za-z0-9._:-]+$ ]] || {
|
||||
echo "ERROR: invalid GPU id: $gpu" >&2
|
||||
exit 2
|
||||
}
|
||||
[[ -z "${SEEN_GPUS[$gpu]+x}" ]] || {
|
||||
echo "ERROR: duplicate GPU id: $gpu" >&2
|
||||
exit 2
|
||||
}
|
||||
SEEN_GPUS["$gpu"]=1
|
||||
done
|
||||
|
||||
if ((NUM_WORKERS_CALIBRATION == 0)); then
|
||||
TOTAL_CPUS="$(nproc)"
|
||||
NUM_WORKERS_CALIBRATION=$(( (TOTAL_CPUS + ${#GPU_IDS[@]} - 1) / ${#GPU_IDS[@]} ))
|
||||
if ((NUM_WORKERS_CALIBRATION < 1)); then
|
||||
NUM_WORKERS_CALIBRATION=1
|
||||
fi
|
||||
fi
|
||||
|
||||
declare -a JOB_RUN_DIRS=()
|
||||
declare -a JOB_LOG_FILES=()
|
||||
|
||||
add_job() {
|
||||
local run_dir="$1"
|
||||
local relative_run
|
||||
if [[ "$run_dir" == "$RUNS_ROOT" ]]; then
|
||||
relative_run="_root"
|
||||
else
|
||||
relative_run="${run_dir#"$RUNS_ROOT"/}"
|
||||
fi
|
||||
JOB_RUN_DIRS+=("$run_dir")
|
||||
JOB_LOG_FILES+=("$LOG_ROOT/$relative_run/evaluate_calibration.log")
|
||||
}
|
||||
|
||||
run_count=0
|
||||
incomplete_count=0
|
||||
next_token_count=0
|
||||
invalid_config_count=0
|
||||
existing_count=0
|
||||
|
||||
while IFS= read -r -d '' config_path; do
|
||||
run_dir="${config_path%/train_config.json}"
|
||||
((run_count += 1))
|
||||
|
||||
if [[ ! -f "$run_dir/best_model.pt" ]]; then
|
||||
echo "[SKIP] Incomplete run without best_model.pt: $run_dir"
|
||||
((incomplete_count += 1))
|
||||
continue
|
||||
fi
|
||||
|
||||
if ! target_mode="$(
|
||||
"$PYTHON_BIN" -c \
|
||||
'import json,sys; print(str(json.load(open(sys.argv[1], encoding="utf-8")).get("model_target_mode", "next_token")).lower())' \
|
||||
"$config_path"
|
||||
)"; then
|
||||
echo "[SKIP] Invalid train_config.json: $config_path" >&2
|
||||
((invalid_config_count += 1))
|
||||
continue
|
||||
fi
|
||||
|
||||
if [[ "$target_mode" != "all_future" ]]; then
|
||||
echo "[SKIP] model_target_mode=$target_mode: $run_dir"
|
||||
((next_token_count += 1))
|
||||
continue
|
||||
fi
|
||||
|
||||
if ((FORCE == 0)) && [[ -s "$run_dir/$COMPLETION_FILE" ]]; then
|
||||
((existing_count += 1))
|
||||
continue
|
||||
fi
|
||||
|
||||
add_job "$run_dir"
|
||||
done < <(find "$RUNS_ROOT" -type f -name "train_config.json" -print0)
|
||||
|
||||
if ((DRY_RUN == 0)); then
|
||||
mkdir -p "$LOG_ROOT"
|
||||
fi
|
||||
|
||||
print_command() {
|
||||
printf '%q ' "$@"
|
||||
printf '\n'
|
||||
}
|
||||
|
||||
run_job() {
|
||||
local job_index="$1"
|
||||
local gpu="$2"
|
||||
local run_dir="${JOB_RUN_DIRS[$job_index]}"
|
||||
local log_file="${JOB_LOG_FILES[$job_index]}"
|
||||
local -a command=(
|
||||
"$PYTHON_BIN"
|
||||
-u
|
||||
"$EVALUATOR"
|
||||
--run_path "$run_dir"
|
||||
--output_path "$run_dir"
|
||||
--eval_split test
|
||||
--device cuda
|
||||
--num_workers "$NUM_WORKERS"
|
||||
--num_workers_calibration "$NUM_WORKERS_CALIBRATION"
|
||||
--batch_size "$BATCH_SIZE"
|
||||
--disease_chunk_size "$DISEASE_CHUNK_SIZE"
|
||||
)
|
||||
|
||||
if [[ -n "$HORIZONS" ]]; then
|
||||
command+=(--horizons "$HORIZONS")
|
||||
fi
|
||||
if ((USE_AMP)); then
|
||||
command+=(--use_amp)
|
||||
fi
|
||||
if ((FORCE)); then
|
||||
command+=(--force)
|
||||
fi
|
||||
|
||||
echo "[$(date '+%F %T')] START gpu=$gpu"
|
||||
echo " run=$run_dir"
|
||||
|
||||
if ((DRY_RUN)); then
|
||||
printf ' CUDA_VISIBLE_DEVICES=%q ' "$gpu"
|
||||
print_command "${command[@]}"
|
||||
return 0
|
||||
fi
|
||||
|
||||
mkdir -p "$(dirname -- "$log_file")"
|
||||
if CUDA_VISIBLE_DEVICES="$gpu" PYTHONUNBUFFERED=1 \
|
||||
"${command[@]}" >"$log_file" 2>&1; then
|
||||
if [[ -s "$run_dir/$COMPLETION_FILE" ]]; then
|
||||
echo "[$(date '+%F %T')] DONE gpu=$gpu"
|
||||
return 0
|
||||
fi
|
||||
echo "[$(date '+%F %T')] FAIL gpu=$gpu" >&2
|
||||
echo " Missing completion marker: $run_dir/$COMPLETION_FILE" >&2
|
||||
echo " See: $log_file" >&2
|
||||
return 1
|
||||
else
|
||||
local exit_code=$?
|
||||
echo "[$(date '+%F %T')] FAIL gpu=$gpu exit=$exit_code" >&2
|
||||
echo " See: $log_file" >&2
|
||||
return "$exit_code"
|
||||
fi
|
||||
}
|
||||
|
||||
worker() {
|
||||
local slot="$1"
|
||||
local gpu="${GPU_IDS[$slot]}"
|
||||
local job_index
|
||||
local failed=0
|
||||
|
||||
for ((
|
||||
job_index = slot;
|
||||
job_index < ${#JOB_RUN_DIRS[@]};
|
||||
job_index += ${#GPU_IDS[@]}
|
||||
)); do
|
||||
run_job "$job_index" "$gpu" || failed=1
|
||||
done
|
||||
return "$failed"
|
||||
}
|
||||
|
||||
echo "Runs root: $RUNS_ROOT"
|
||||
echo "GPUs: ${GPU_IDS[*]}"
|
||||
echo "Calibration CPU workers per GPU job: $NUM_WORKERS_CALIBRATION"
|
||||
echo "Runs discovered: $run_count"
|
||||
echo "Incomplete runs skipped: $incomplete_count"
|
||||
echo "next_token runs skipped: $next_token_count"
|
||||
echo "Invalid configs skipped: $invalid_config_count"
|
||||
echo "Existing calibration results skipped: $existing_count"
|
||||
echo "Pending all_future evaluations: ${#JOB_RUN_DIRS[@]}"
|
||||
echo "Log root: $LOG_ROOT"
|
||||
echo
|
||||
|
||||
if ((${#JOB_RUN_DIRS[@]} == 0)); then
|
||||
echo "No pending all_future calibration evaluations."
|
||||
exit 0
|
||||
fi
|
||||
|
||||
declare -a WORKER_PIDS=()
|
||||
for ((slot = 0; slot < ${#GPU_IDS[@]}; slot++)); do
|
||||
worker "$slot" &
|
||||
WORKER_PIDS+=("$!")
|
||||
done
|
||||
|
||||
overall_status=0
|
||||
for pid in "${WORKER_PIDS[@]}"; do
|
||||
wait "$pid" || overall_status=1
|
||||
done
|
||||
|
||||
if ((overall_status != 0)); then
|
||||
echo "One or more calibration evaluations failed." >&2
|
||||
echo "Inspect logs under: $LOG_ROOT" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
if ((DRY_RUN)); then
|
||||
echo "Dry run completed successfully."
|
||||
else
|
||||
echo "All pending all_future calibration evaluations completed."
|
||||
fi
|
||||
@@ -1,814 +0,0 @@
|
||||
"""Compute landmark future death and incident system-disease risks.
|
||||
|
||||
For each selected patient and landmark age, this script computes:
|
||||
|
||||
* 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.
|
||||
|
||||
Death is always token vocab_size - 1. Disease groups are read from
|
||||
icd10_chapter_organ_mapping.csv.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional, Sequence
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import torch
|
||||
from torch.nn.utils.rnn import pad_sequence
|
||||
from torch.utils.data import DataLoader, Dataset
|
||||
from tqdm.auto import tqdm
|
||||
|
||||
from dataset import HealthDataset
|
||||
from eval_data import load_sequence_eval_dataset
|
||||
from evaluate_auc_v2 import (
|
||||
LandmarkDataset,
|
||||
build_model_from_dataset,
|
||||
cfg_get,
|
||||
load_checkpoint_state_dict,
|
||||
load_json_config,
|
||||
load_model_state,
|
||||
make_eval_indices,
|
||||
resolve_dist_mode_for_checkpoint,
|
||||
resolve_eval_device,
|
||||
validate_dataset_metadata,
|
||||
)
|
||||
from future_risk import (
|
||||
death_risk_from_probabilities,
|
||||
new_disease_risk_from_probabilities,
|
||||
probabilities_from_logits,
|
||||
)
|
||||
from models import DeepHealth
|
||||
from readouts import build_readout
|
||||
from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX
|
||||
from train_util import load_eid_file, load_extra_info_types_file
|
||||
|
||||
|
||||
SPECIAL_TOKENS = {PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX}
|
||||
|
||||
|
||||
def parse_int_list(value: Any) -> Optional[List[int]]:
|
||||
if value is None:
|
||||
return None
|
||||
if isinstance(value, (list, tuple, np.ndarray)):
|
||||
return [int(x) for x in value]
|
||||
text = str(value).strip()
|
||||
if text == "":
|
||||
return None
|
||||
if text.startswith("["):
|
||||
values = json.loads(text)
|
||||
if not isinstance(values, list):
|
||||
raise ValueError(
|
||||
f"Expected a JSON list, got {type(values).__name__}")
|
||||
return [int(x) for x in values]
|
||||
return [int(x.strip()) for x in text.split(",") if x.strip()]
|
||||
|
||||
|
||||
def load_extra_info_types(value: Any) -> Optional[List[int]]:
|
||||
if value is None:
|
||||
return None
|
||||
text = str(value)
|
||||
path = Path(text)
|
||||
if path.exists():
|
||||
return load_extra_info_types_file(text)
|
||||
return parse_int_list(value)
|
||||
|
||||
|
||||
def make_landmark_ages(start: float, stop: float, step: float) -> np.ndarray:
|
||||
if step <= 0:
|
||||
raise ValueError("landmark_step must be positive")
|
||||
if stop < start:
|
||||
raise ValueError("landmark_stop must be >= landmark_start")
|
||||
# Include stop when it lands on the grid, e.g. 40,45,...,80.
|
||||
return np.arange(start, stop + step * 0.5, step, dtype=np.float32)
|
||||
|
||||
|
||||
def build_first_occurrence_maps_for_landmarks(
|
||||
dataset: HealthDataset,
|
||||
subset_indices: np.ndarray,
|
||||
) -> Dict[int, tuple[np.ndarray, np.ndarray]]:
|
||||
first_lists: Dict[int, list[tuple[int, float]]] = {}
|
||||
for patient_id, dataset_index in enumerate(np.asarray(subset_indices, dtype=np.int64).tolist()):
|
||||
s = dataset.samples[int(dataset_index)]
|
||||
seq_event = np.asarray(s["event_seq"], dtype=np.int64)
|
||||
seq_time = np.asarray(s["time_seq"], dtype=np.float32)
|
||||
tgt_event = np.asarray(s["target_event_seq"], dtype=np.int64)
|
||||
tgt_time = np.asarray(s["target_time_seq"], dtype=np.float32)
|
||||
if seq_event.size == 0 or tgt_event.size == 0:
|
||||
continue
|
||||
|
||||
full_event = np.concatenate([seq_event, tgt_event[-1:]])
|
||||
full_time = np.concatenate([seq_time, tgt_time[-1:]])
|
||||
uniq_tokens, first_idx = np.unique(full_event, return_index=True)
|
||||
for token, idx in zip(uniq_tokens.tolist(), first_idx.tolist()):
|
||||
token = int(token)
|
||||
if token in SPECIAL_TOKENS:
|
||||
continue
|
||||
first_lists.setdefault(token, []).append(
|
||||
(patient_id, float(full_time[int(idx)])))
|
||||
|
||||
return {
|
||||
int(token): (
|
||||
np.asarray([p for p, _ in pairs], dtype=np.int32),
|
||||
np.asarray([t for _, t in pairs], dtype=np.float32),
|
||||
)
|
||||
for token, pairs in first_lists.items()
|
||||
if pairs
|
||||
}
|
||||
|
||||
|
||||
def normalize_eval_split(args: argparse.Namespace, cfg: Dict[str, Any]) -> str:
|
||||
eval_split = str(cfg_get(args, cfg, "eval_split", "test")).lower()
|
||||
if eval_split in {"valid", "validation"}:
|
||||
return "val"
|
||||
if eval_split not in {"train", "val", "test", "all"}:
|
||||
raise ValueError(f"Unsupported eval_split={eval_split!r}")
|
||||
return eval_split
|
||||
|
||||
|
||||
def load_eval_sequence_dataset(
|
||||
args: argparse.Namespace,
|
||||
cfg: Dict[str, Any],
|
||||
) -> tuple[Any, np.ndarray, str, str]:
|
||||
eval_split = normalize_eval_split(args, cfg)
|
||||
model_target_mode = str(cfg.get("model_target_mode", "next_token")).lower()
|
||||
data_prefix = str(cfg.get("data_prefix", "ukb"))
|
||||
labels_file = str(cfg.get("labels_file", "labels.csv"))
|
||||
no_event_interval_years = float(cfg.get("no_event_interval_years", 5.0))
|
||||
include_no_event_in_uts_target = bool(
|
||||
cfg.get("include_no_event_in_uts_target", False))
|
||||
extra_info_types = load_extra_info_types(args.extra_info_types)
|
||||
if extra_info_types is None:
|
||||
extra_info_types = parse_int_list(cfg.get("extra_info_types", None))
|
||||
|
||||
print("Loading one sequence eval dataset...")
|
||||
dataset = load_sequence_eval_dataset(
|
||||
model_target_mode=model_target_mode,
|
||||
data_prefix=data_prefix,
|
||||
labels_file=labels_file,
|
||||
no_event_interval_years=no_event_interval_years,
|
||||
include_no_event_in_uts_target=include_no_event_in_uts_target,
|
||||
min_history_events=int(cfg.get("all_future_min_history_events", 1)),
|
||||
min_future_events=int(cfg.get("all_future_min_future_events", 1)),
|
||||
extra_info_types=extra_info_types,
|
||||
)
|
||||
|
||||
train_eid_file = cfg_get(args, cfg, "train_eid_file", "ukb_train_eid.csv")
|
||||
val_eid_file = cfg_get(args, cfg, "val_eid_file", "ukb_val_eid.csv")
|
||||
test_eid_file = cfg_get(args, cfg, "test_eid_file", "ukb_test_eid.csv")
|
||||
split_files_exist = all(
|
||||
Path(str(path)).exists()
|
||||
for path in (train_eid_file, val_eid_file, test_eid_file)
|
||||
)
|
||||
|
||||
if eval_split != "all" and split_files_exist:
|
||||
split_files = {
|
||||
"train": train_eid_file,
|
||||
"val": val_eid_file,
|
||||
"test": test_eid_file,
|
||||
}
|
||||
selected_eids = load_eid_file(split_files[eval_split])
|
||||
out = np.asarray(
|
||||
[
|
||||
idx
|
||||
for idx, sample in enumerate(dataset.samples)
|
||||
if int(sample["eid"]) in selected_eids
|
||||
],
|
||||
dtype=np.int64,
|
||||
)
|
||||
if out.size == 0:
|
||||
raise ValueError(
|
||||
f"No samples found for eval_split={eval_split!r} using {split_files[eval_split]}"
|
||||
)
|
||||
split_source = "eid_files"
|
||||
else:
|
||||
if eval_split == "all":
|
||||
out = np.arange(len(dataset.samples), dtype=np.int64)
|
||||
split_source = "all"
|
||||
else:
|
||||
out = make_eval_indices(dataset, args, cfg)
|
||||
split_source = "ratio_split"
|
||||
|
||||
subset_size = cfg_get(args, cfg, "dataset_subset_size", None)
|
||||
if subset_size is not None and int(subset_size) > 0:
|
||||
out = out[: int(subset_size)]
|
||||
return dataset, np.asarray(out, dtype=np.int64), eval_split, split_source
|
||||
|
||||
|
||||
def load_organ_groups(
|
||||
path: Path,
|
||||
*,
|
||||
vocab_size: int,
|
||||
) -> tuple[dict[str, list[int]], dict[str, str], dict[int, str]]:
|
||||
table = pd.read_csv(path)
|
||||
required = {"token_id", "organ_system", "organ_system_label", "is_death"}
|
||||
missing = required - set(table.columns)
|
||||
if missing:
|
||||
raise ValueError(f"{path} is missing columns: {sorted(missing)}")
|
||||
|
||||
death_idx = int(vocab_size) - 1
|
||||
groups: dict[str, list[int]] = {}
|
||||
labels: dict[str, str] = {}
|
||||
token_to_group: dict[int, str] = {}
|
||||
for row in table.itertuples(index=False):
|
||||
token = int(getattr(row, "token_id"))
|
||||
if token in SPECIAL_TOKENS or token == death_idx:
|
||||
continue
|
||||
if token < 0 or token >= int(vocab_size):
|
||||
continue
|
||||
if int(getattr(row, "is_death")) == 1:
|
||||
continue
|
||||
group = str(getattr(row, "organ_system"))
|
||||
label = str(getattr(row, "organ_system_label"))
|
||||
groups.setdefault(group, []).append(token)
|
||||
labels[group] = label
|
||||
token_to_group[token] = group
|
||||
|
||||
groups = {k: sorted(set(v)) for k, v in groups.items() if v}
|
||||
return groups, labels, token_to_group
|
||||
|
||||
|
||||
class IndexedLandmarkDataset(Dataset):
|
||||
def __init__(self, base: LandmarkDataset) -> None:
|
||||
self.base = base
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.base)
|
||||
|
||||
def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
|
||||
item = dict(self.base[idx])
|
||||
item["row_idx"] = torch.tensor(int(idx), dtype=torch.long)
|
||||
return item
|
||||
|
||||
|
||||
def collate_indexed_landmark_fn(batch: List[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]:
|
||||
event_seq = pad_sequence(
|
||||
[x["event_seq"] for x in batch], batch_first=True, padding_value=PAD_IDX
|
||||
)
|
||||
time_seq = pad_sequence(
|
||||
[x["time_seq"] for x in batch], batch_first=True, padding_value=0.0
|
||||
)
|
||||
readout_mask = pad_sequence(
|
||||
[x["readout_mask"] for x in batch], batch_first=True, padding_value=False
|
||||
)
|
||||
other_type = pad_sequence(
|
||||
[x["other_type"] for x in batch], batch_first=True, padding_value=0
|
||||
)
|
||||
other_value = pad_sequence(
|
||||
[x["other_value"] for x in batch], batch_first=True, padding_value=0.0
|
||||
)
|
||||
other_value_kind = pad_sequence(
|
||||
[x["other_value_kind"] for x in batch], batch_first=True, padding_value=0
|
||||
)
|
||||
other_time = pad_sequence(
|
||||
[x["other_time"] for x in batch], batch_first=True, padding_value=0.0
|
||||
)
|
||||
return {
|
||||
"event_seq": event_seq,
|
||||
"time_seq": time_seq,
|
||||
"padding_mask": event_seq > PAD_IDX,
|
||||
"readout_mask": readout_mask,
|
||||
"sex": torch.stack([x["sex"] for x in batch]),
|
||||
"other_type": other_type,
|
||||
"other_value": other_value,
|
||||
"other_value_kind": other_value_kind,
|
||||
"other_time": other_time,
|
||||
"landmark_pos": torch.stack([x["landmark_pos"] for x in batch]),
|
||||
"t_query": torch.stack([x["t_query"] for x in batch]),
|
||||
"patient_id": torch.stack([x["patient_id"] for x in batch]),
|
||||
"landmark_age": torch.stack([x["landmark_age"] for x in batch]),
|
||||
"followup_end_time": torch.stack([x["followup_end_time"] for x in batch]),
|
||||
"death_time": torch.stack([x["death_time"] for x in batch]),
|
||||
"row_idx": torch.stack([x["row_idx"] for x in batch]),
|
||||
}
|
||||
|
||||
|
||||
def build_group_ablated_slice(
|
||||
batch: Dict[str, torch.Tensor],
|
||||
token_ids: Sequence[int],
|
||||
row_indices: torch.Tensor,
|
||||
) -> Dict[str, torch.Tensor]:
|
||||
"""Build one fixed-width ablated slice without rebuilding variable-length rows."""
|
||||
event_seq = batch["event_seq"]
|
||||
|
||||
out: Dict[str, torch.Tensor] = {}
|
||||
out["event_seq"] = event_seq[row_indices].clone()
|
||||
out["time_seq"] = batch["time_seq"][row_indices]
|
||||
out["readout_mask"] = batch["readout_mask"][row_indices].clone()
|
||||
out["padding_mask"] = batch["padding_mask"][row_indices].bool().clone()
|
||||
out["landmark_pos"] = batch["landmark_pos"][row_indices].clone()
|
||||
|
||||
seq_len = int(event_seq.shape[1])
|
||||
positions = torch.arange(seq_len, device=event_seq.device)[None, :]
|
||||
ids = torch.as_tensor(token_ids, dtype=event_seq.dtype,
|
||||
device=event_seq.device)
|
||||
remove = torch.isin(out["event_seq"], ids) & out["padding_mask"]
|
||||
out["event_seq"] = torch.where(
|
||||
remove,
|
||||
torch.full_like(out["event_seq"], PAD_IDX),
|
||||
out["event_seq"],
|
||||
)
|
||||
out["padding_mask"] &= ~remove
|
||||
out["readout_mask"] &= ~remove
|
||||
|
||||
has_valid = out["padding_mask"].any(dim=1)
|
||||
if not bool(has_valid.all().item()):
|
||||
empty_rows = torch.nonzero(~has_valid, as_tuple=False).flatten()
|
||||
out["event_seq"][empty_rows, 0] = CHECKUP_IDX
|
||||
out["time_seq"][empty_rows, 0] = batch["t_query"][row_indices[empty_rows]].to(
|
||||
dtype=out["time_seq"].dtype
|
||||
)
|
||||
out["padding_mask"][empty_rows, 0] = True
|
||||
out["readout_mask"][empty_rows, 0] = True
|
||||
out["landmark_pos"][empty_rows] = 0
|
||||
|
||||
has_readout = out["readout_mask"].any(dim=1)
|
||||
if not bool(has_readout.all().item()):
|
||||
rows = torch.nonzero(~has_readout, as_tuple=False).flatten()
|
||||
local_valid = out["padding_mask"][rows]
|
||||
last_pos = torch.where(
|
||||
local_valid,
|
||||
positions.expand(local_valid.shape[0], -1),
|
||||
torch.zeros_like(positions.expand(local_valid.shape[0], -1)),
|
||||
).amax(dim=1)
|
||||
out["readout_mask"][rows] = False
|
||||
out["readout_mask"][rows, last_pos] = True
|
||||
out["landmark_pos"][rows] = last_pos.to(dtype=out["landmark_pos"].dtype)
|
||||
|
||||
repeated_keys = (
|
||||
"sex",
|
||||
"other_type",
|
||||
"other_value",
|
||||
"other_value_kind",
|
||||
"other_time",
|
||||
"t_query",
|
||||
"patient_id",
|
||||
"landmark_age",
|
||||
"followup_end_time",
|
||||
"death_time",
|
||||
"row_idx",
|
||||
)
|
||||
for key in repeated_keys:
|
||||
out[key] = batch[key][row_indices]
|
||||
return out
|
||||
|
||||
|
||||
def concat_tensor_batches(chunks: Sequence[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]:
|
||||
return {
|
||||
key: torch.cat([chunk[key] for chunk in chunks], dim=0)
|
||||
for key in chunks[0]
|
||||
}
|
||||
|
||||
|
||||
def iter_group_ablated_batches(
|
||||
batch: Dict[str, torch.Tensor],
|
||||
group_names: Sequence[str],
|
||||
organ_groups: dict[str, list[int]],
|
||||
occurred: torch.Tensor,
|
||||
max_batch_size: int,
|
||||
):
|
||||
"""Yield ablated chunks as soon as enough rows are available for a forward pass."""
|
||||
pending_batches: list[Dict[str, torch.Tensor]] = []
|
||||
pending_groups: list[str] = []
|
||||
pending_rows: list[int] = []
|
||||
pending_n = 0
|
||||
|
||||
for group in group_names:
|
||||
ids = torch.as_tensor(
|
||||
organ_groups[group], dtype=torch.long, device=occurred.device)
|
||||
if ids.numel() == 0:
|
||||
continue
|
||||
active_rows = torch.nonzero(
|
||||
occurred[:, ids].any(dim=1), as_tuple=False).flatten()
|
||||
if active_rows.numel() == 0:
|
||||
continue
|
||||
|
||||
row_offset = 0
|
||||
while row_offset < int(active_rows.numel()):
|
||||
capacity = int(max_batch_size) - pending_n
|
||||
row_stop = min(int(active_rows.numel()), row_offset + capacity)
|
||||
row_indices = active_rows[row_offset:row_stop].to(
|
||||
device=batch["event_seq"].device)
|
||||
chunk = build_group_ablated_slice(
|
||||
batch=batch,
|
||||
token_ids=organ_groups[group],
|
||||
row_indices=row_indices,
|
||||
)
|
||||
chunk_n = int(row_indices.numel())
|
||||
pending_batches.append(chunk)
|
||||
pending_groups.extend([group] * chunk_n)
|
||||
pending_rows.extend(int(x)
|
||||
for x in row_indices.detach().cpu().tolist())
|
||||
pending_n += chunk_n
|
||||
row_offset = row_stop
|
||||
|
||||
if pending_n >= int(max_batch_size):
|
||||
yield concat_tensor_batches(pending_batches), pending_groups, pending_rows
|
||||
pending_batches = []
|
||||
pending_groups = []
|
||||
pending_rows = []
|
||||
pending_n = 0
|
||||
|
||||
if pending_batches:
|
||||
yield concat_tensor_batches(pending_batches), pending_groups, pending_rows
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def infer_landmark_hidden(
|
||||
*,
|
||||
model: DeepHealth,
|
||||
batch: Dict[str, torch.Tensor],
|
||||
device: torch.device,
|
||||
model_target_mode: str,
|
||||
readout_name: str,
|
||||
readout_reduce: str,
|
||||
) -> torch.Tensor:
|
||||
batch_dev = {
|
||||
k: (v.to(device, non_blocking=True) if isinstance(v, torch.Tensor) else v)
|
||||
for k, v in batch.items()
|
||||
}
|
||||
if model_target_mode == "all_future":
|
||||
return model(
|
||||
event_seq=batch_dev["event_seq"].long(),
|
||||
time_seq=batch_dev["time_seq"].float(),
|
||||
sex=batch_dev["sex"].long(),
|
||||
padding_mask=batch_dev["padding_mask"].bool(),
|
||||
t_query=batch_dev["t_query"].float(),
|
||||
other_type=batch_dev["other_type"].long(),
|
||||
other_value=batch_dev["other_value"].float(),
|
||||
other_value_kind=batch_dev["other_value_kind"].long(),
|
||||
other_time=batch_dev["other_time"].float(),
|
||||
target_mode="all_future",
|
||||
)
|
||||
|
||||
hidden = model(
|
||||
event_seq=batch_dev["event_seq"].long(),
|
||||
time_seq=batch_dev["time_seq"].float(),
|
||||
sex=batch_dev["sex"].long(),
|
||||
padding_mask=batch_dev["padding_mask"].bool(),
|
||||
other_type=batch_dev["other_type"].long(),
|
||||
other_value=batch_dev["other_value"].float(),
|
||||
other_value_kind=batch_dev["other_value_kind"].long(),
|
||||
other_time=batch_dev["other_time"].float(),
|
||||
target_mode="next_token",
|
||||
)
|
||||
readout = build_readout(readout_name, reduce=readout_reduce)
|
||||
readout_out = readout(
|
||||
hidden=hidden,
|
||||
time_seq=batch_dev["time_seq"].float(),
|
||||
padding_mask=batch_dev["padding_mask"].bool(),
|
||||
readout_mask=batch_dev["readout_mask"].bool(),
|
||||
)
|
||||
return readout_out.hidden.gather(
|
||||
1,
|
||||
batch_dev["landmark_pos"].long()[:, None, None].expand(
|
||||
-1, 1, readout_out.hidden.shape[-1]
|
||||
),
|
||||
).squeeze(1)
|
||||
|
||||
|
||||
def make_occurred_mask(
|
||||
event_seq: torch.Tensor,
|
||||
*,
|
||||
vocab_size: int,
|
||||
device: torch.device,
|
||||
) -> torch.Tensor:
|
||||
occurred = torch.zeros(event_seq.shape[0], int(
|
||||
vocab_size), dtype=torch.bool, device=device)
|
||||
valid = (event_seq >= 0) & (event_seq < int(vocab_size))
|
||||
safe = event_seq.clamp(min=0, max=int(vocab_size) - 1).to(device)
|
||||
occurred.scatter_(1, safe, valid.to(device))
|
||||
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,
|
||||
*,
|
||||
death_idx: int,
|
||||
token_to_group: dict[int, str],
|
||||
group_names: Sequence[str],
|
||||
) -> tuple[int, dict[str, int]]:
|
||||
unique_tokens = {
|
||||
int(token)
|
||||
for token in np.asarray(tokens, dtype=np.int64).tolist()
|
||||
if int(token) not in SPECIAL_TOKENS and int(token) != int(death_idx)
|
||||
}
|
||||
total = len(unique_tokens)
|
||||
out = {group: 0 for group in group_names}
|
||||
for token in unique_tokens:
|
||||
group = token_to_group.get(token)
|
||||
if group in out:
|
||||
out[group] += 1
|
||||
return total, out
|
||||
|
||||
|
||||
def output_name_for_run(run_path: Path, eval_split: str, tau: float) -> Path:
|
||||
return run_path / f"future_risk_{eval_split}_tau{tau:g}y.csv"
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
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)
|
||||
parser.add_argument("--organ_mapping_path", type=str,
|
||||
default="icd10_chapter_organ_mapping.csv")
|
||||
parser.add_argument("--eval_split", type=str, default=None)
|
||||
parser.add_argument("--dataset_subset_size", type=int, default=None)
|
||||
parser.add_argument("--train_eid_file", type=str, default=None)
|
||||
parser.add_argument("--val_eid_file", type=str, default=None)
|
||||
parser.add_argument("--test_eid_file", type=str, default=None)
|
||||
parser.add_argument("--landmark_start", type=float, default=40.0)
|
||||
parser.add_argument("--landmark_stop", type=float, default=80.0)
|
||||
parser.add_argument("--landmark_step", type=float, default=5.0)
|
||||
parser.add_argument("--tau", type=float, default=5.0)
|
||||
parser.add_argument("--min_history_events", type=int, default=None)
|
||||
parser.add_argument("--batch_size", type=int, default=None)
|
||||
parser.add_argument(
|
||||
"--attribution_batch_size",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Forward batch size for expanded organ/system ablation queries.",
|
||||
)
|
||||
parser.add_argument("--num_workers", type=int, default=None)
|
||||
parser.add_argument("--device", type=str, default=None)
|
||||
parser.add_argument("--extra_info_types", type=str, default=None)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
run_path = Path(args.run_path)
|
||||
config_path = run_path / "train_config.json"
|
||||
checkpoint_path = run_path / "best_model.pt"
|
||||
if not config_path.exists():
|
||||
raise FileNotFoundError(f"train_config.json not found: {config_path}")
|
||||
if not checkpoint_path.exists():
|
||||
raise FileNotFoundError(f"best_model.pt not found: {checkpoint_path}")
|
||||
|
||||
cfg = load_json_config(config_path)
|
||||
model_target_mode = str(cfg.get("model_target_mode", "next_token")).lower()
|
||||
if model_target_mode not in {"next_token", "all_future"}:
|
||||
raise ValueError(
|
||||
f"Unsupported model_target_mode: {model_target_mode!r}")
|
||||
|
||||
target_mode = str(cfg.get("target_mode", "uts"))
|
||||
attn_mask_mode = str(
|
||||
cfg.get("attn_mask_mode", "non_strict_time" if target_mode ==
|
||||
"uts" else "target_aware")
|
||||
)
|
||||
readout_name = str(cfg.get(
|
||||
"readout_name", "same_time_group_end" if target_mode == "uts" else "token"))
|
||||
readout_reduce = str(cfg.get("readout_reduce", "mean"))
|
||||
|
||||
dataset, subset_indices, eval_split, split_source = load_eval_sequence_dataset(
|
||||
args,
|
||||
cfg,
|
||||
)
|
||||
validate_dataset_metadata(dataset, cfg)
|
||||
|
||||
landmark_ages = make_landmark_ages(
|
||||
float(args.landmark_start),
|
||||
float(args.landmark_stop),
|
||||
float(args.landmark_step),
|
||||
)
|
||||
tau = float(args.tau)
|
||||
if tau < 0:
|
||||
raise ValueError("tau must be non-negative")
|
||||
|
||||
first_occurrence_by_token = build_first_occurrence_maps_for_landmarks(
|
||||
dataset,
|
||||
subset_indices,
|
||||
)
|
||||
death_idx = int(dataset.vocab_size) - 1
|
||||
landmark_dataset = LandmarkDataset(
|
||||
dataset=dataset,
|
||||
subset_indices=subset_indices,
|
||||
landmark_ages=landmark_ages,
|
||||
attn_mask_mode=attn_mask_mode,
|
||||
model_target_mode=model_target_mode,
|
||||
min_history_events=int(cfg_get(args, cfg, "min_history_events", 1)),
|
||||
first_occurrence_by_token=first_occurrence_by_token,
|
||||
death_token_ids=[death_idx],
|
||||
)
|
||||
|
||||
organ_groups, organ_labels, token_to_group = load_organ_groups(
|
||||
Path(args.organ_mapping_path),
|
||||
vocab_size=int(dataset.vocab_size),
|
||||
)
|
||||
group_names = sorted(organ_groups)
|
||||
|
||||
state_dict = load_checkpoint_state_dict(checkpoint_path, map_location="cpu")
|
||||
dist_mode = resolve_dist_mode_for_checkpoint(
|
||||
str(cfg.get("dist_mode", "exponential")), state_dict)
|
||||
cfg_model = dict(cfg)
|
||||
cfg_model["dist_mode"] = dist_mode
|
||||
device = resolve_eval_device(args.device)
|
||||
model = build_model_from_dataset(args, cfg_model, dataset).to(device)
|
||||
load_model_state(model, state_dict)
|
||||
model.eval()
|
||||
|
||||
batch_size = int(cfg_get(args, cfg, "batch_size", 128))
|
||||
attribution_batch_size = int(
|
||||
cfg_get(args, cfg, "attribution_batch_size",
|
||||
max(batch_size * 4, batch_size))
|
||||
)
|
||||
if attribution_batch_size <= 0:
|
||||
raise ValueError("attribution_batch_size must be positive")
|
||||
num_workers = int(cfg_get(args, cfg, "num_workers", 4))
|
||||
loader = DataLoader(
|
||||
IndexedLandmarkDataset(landmark_dataset),
|
||||
batch_size=batch_size,
|
||||
shuffle=False,
|
||||
collate_fn=collate_indexed_landmark_fn,
|
||||
num_workers=num_workers,
|
||||
pin_memory=device.type == "cuda",
|
||||
persistent_workers=num_workers > 0,
|
||||
prefetch_factor=2 if num_workers > 0 else None,
|
||||
)
|
||||
|
||||
output_path = Path(args.output_path) if args.output_path else output_name_for_run(
|
||||
run_path, eval_split, tau)
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
print(f"Eval split: {eval_split}")
|
||||
print(f"Split source: {split_source}")
|
||||
print(f"Selected patients: {len(subset_indices)}")
|
||||
print(f"Landmark ages: {landmark_ages.tolist()}")
|
||||
print(f"Tau: {tau:g} years")
|
||||
print(f"Dist mode: {dist_mode}")
|
||||
print(f"Device: {device}")
|
||||
print(f"Death token: {death_idx}")
|
||||
print(f"Organ/system groups: {len(group_names)}")
|
||||
print(f"Landmark rows: {len(landmark_dataset)}")
|
||||
print(f"Attribution batch size: {attribution_batch_size}")
|
||||
print(f"Output: {output_path}")
|
||||
|
||||
rows: list[dict[str, Any]] = []
|
||||
for batch in tqdm(loader, desc="Future risks", dynamic_ncols=True):
|
||||
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,
|
||||
)
|
||||
occurred = make_occurred_mask(
|
||||
batch["event_seq"].to(device),
|
||||
vocab_size=int(dataset.vocab_size),
|
||||
device=device,
|
||||
)
|
||||
|
||||
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,
|
||||
organ_groups[group],
|
||||
).detach().cpu().numpy()
|
||||
|
||||
group_mortality_attr_prob: dict[str, np.ndarray] = {}
|
||||
group_mortality_attr_hazard: dict[str, np.ndarray] = {}
|
||||
batch_n = int(batch["event_seq"].shape[0])
|
||||
zeros = np.zeros(batch_n, dtype=np.float32)
|
||||
for group in group_names:
|
||||
group_mortality_attr_prob[group] = zeros.copy()
|
||||
group_mortality_attr_hazard[group] = zeros.copy()
|
||||
|
||||
for ablated_chunk, chunk_groups, chunk_rows in iter_group_ablated_batches(
|
||||
batch=batch,
|
||||
group_names=group_names,
|
||||
organ_groups=organ_groups,
|
||||
occurred=occurred,
|
||||
max_batch_size=attribution_batch_size,
|
||||
):
|
||||
ablated_death_risk = death_risk_for_batch(
|
||||
model=model,
|
||||
batch=ablated_chunk,
|
||||
device=device,
|
||||
model_target_mode=model_target_mode,
|
||||
readout_name=readout_name,
|
||||
readout_reduce=readout_reduce,
|
||||
dist_mode=dist_mode,
|
||||
tau=tau,
|
||||
)
|
||||
row_tensor = torch.as_tensor(
|
||||
chunk_rows, dtype=torch.long, device=device)
|
||||
ablated_death_hazard = mortality_hazard_from_risk(
|
||||
ablated_death_risk)
|
||||
attr_prob = (
|
||||
death_risk_tensor[row_tensor] - ablated_death_risk
|
||||
).detach().cpu().numpy()
|
||||
attr_hazard = (
|
||||
death_hazard_tensor[row_tensor] - ablated_death_hazard
|
||||
).detach().cpu().numpy()
|
||||
for local_idx, (group, row_idx) in enumerate(zip(chunk_groups, chunk_rows)):
|
||||
group_mortality_attr_prob[group][row_idx] = attr_prob[local_idx]
|
||||
group_mortality_attr_hazard[group][row_idx] = attr_hazard[local_idx]
|
||||
|
||||
row_indices = batch["row_idx"].cpu().numpy().astype(np.int64)
|
||||
for j, row_idx in enumerate(row_indices.tolist()):
|
||||
meta = landmark_dataset.rows[int(row_idx)]
|
||||
dataset_index = int(meta["dataset_index"])
|
||||
sample = dataset.samples[dataset_index]
|
||||
hist_tokens = np.asarray(meta["event_seq"], dtype=np.int64)
|
||||
total_count, group_counts = historical_counts_by_group(
|
||||
hist_tokens,
|
||||
death_idx=death_idx,
|
||||
token_to_group=token_to_group,
|
||||
group_names=group_names,
|
||||
)
|
||||
|
||||
out: dict[str, Any] = {
|
||||
"patient_id": int(meta["patient_id"]),
|
||||
"dataset_index": dataset_index,
|
||||
"eid": int(sample.get("eid", -1)),
|
||||
"sex": int(meta["sex"]),
|
||||
"landmark_age": float(meta["landmark_age"]),
|
||||
"tau": tau,
|
||||
"followup_end_time": float(meta["followup_end_time"]),
|
||||
"history_disease_count": int(total_count),
|
||||
"death_risk": float(death_risk[j]),
|
||||
}
|
||||
for group in group_names:
|
||||
out[f"history_count__{group}"] = int(group_counts[group])
|
||||
out[f"new_disease_risk__{group}"] = float(group_risk[group][j])
|
||||
if int(group_counts[group]) == 0:
|
||||
group_mortality_attr_prob[group][j] = 0.0
|
||||
group_mortality_attr_hazard[group][j] = 0.0
|
||||
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)
|
||||
df.to_csv(output_path, index=False)
|
||||
print(f"Wrote {len(df)} rows to {output_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,950 +0,0 @@
|
||||
"""Evaluate extra-info attribution to death and disease distribution parameters.
|
||||
|
||||
For each landmark query, this script scans selected extra-info types that are
|
||||
available at or before the query age. For each such type it re-runs the model
|
||||
with that extra-info type removed and summarizes:
|
||||
|
||||
* death distribution parameters before and after ablation;
|
||||
* disease distribution parameters before and after ablation, by ICD-10
|
||||
chapter-derived organ/system groups.
|
||||
|
||||
Death is always token vocab_size - 1.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import re
|
||||
from concurrent.futures import ProcessPoolExecutor, as_completed
|
||||
from pathlib import Path
|
||||
from typing import Any, Sequence
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch.utils.data import DataLoader
|
||||
from tqdm.auto import tqdm
|
||||
|
||||
from evaluate_auc_v2 import (
|
||||
build_model_from_dataset,
|
||||
cfg_get,
|
||||
load_checkpoint_state_dict,
|
||||
load_json_config,
|
||||
load_model_state,
|
||||
resolve_dist_mode_for_checkpoint,
|
||||
resolve_eval_device,
|
||||
validate_dataset_metadata,
|
||||
)
|
||||
from landmark_eval_utils import (
|
||||
IndexedLandmarkDataset,
|
||||
LandmarkDataset,
|
||||
build_first_occurrence_maps_for_landmarks,
|
||||
collate_indexed_landmark_fn,
|
||||
infer_landmark_hidden,
|
||||
load_eval_sequence_dataset,
|
||||
load_organ_groups,
|
||||
make_landmark_ages,
|
||||
)
|
||||
|
||||
EXTRA_KEY_COLUMNS = [
|
||||
"selected_extra_info_type_id",
|
||||
"selected_extra_info_var_name",
|
||||
"selected_extra_info_full_name",
|
||||
"landmark_age",
|
||||
"sex",
|
||||
]
|
||||
|
||||
DEATH_PARAMETER_COLUMNS = [
|
||||
"original_death_lambda",
|
||||
"ablated_death_lambda",
|
||||
"original_death_scale",
|
||||
"ablated_death_scale",
|
||||
"original_death_shape",
|
||||
"ablated_death_shape",
|
||||
]
|
||||
|
||||
DISEASE_PARAMETER_KEY_COLUMNS = [
|
||||
*EXTRA_KEY_COLUMNS,
|
||||
"target_group",
|
||||
"target_group_label",
|
||||
]
|
||||
|
||||
DISEASE_PARAMETER_COLUMNS = [
|
||||
"original_disease_lambda",
|
||||
"ablated_disease_lambda",
|
||||
"original_disease_scale",
|
||||
"ablated_disease_scale",
|
||||
"original_disease_shape",
|
||||
"ablated_disease_shape",
|
||||
]
|
||||
|
||||
|
||||
def parse_int_list(value: Any) -> list[int] | None:
|
||||
if value is None:
|
||||
return None
|
||||
if isinstance(value, (list, tuple, np.ndarray)):
|
||||
return [int(x) for x in value]
|
||||
text = str(value).strip()
|
||||
if text == "":
|
||||
return None
|
||||
if text.startswith("["):
|
||||
raw = json.loads(text)
|
||||
if not isinstance(raw, list):
|
||||
raise ValueError("Expected JSON list for integer list")
|
||||
return [int(x) for x in raw]
|
||||
return [int(x.strip()) for x in re.split(r"[,;\s]+", text) if x.strip()]
|
||||
|
||||
|
||||
def load_extra_info_metadata(
|
||||
*,
|
||||
dataset_extra_info_types: Sequence[int],
|
||||
search_root: Path = Path("."),
|
||||
) -> dict[int, dict[str, Any]]:
|
||||
metadata: dict[int, dict[str, Any]] = {
|
||||
int(type_id): {
|
||||
"type_id": int(type_id),
|
||||
"var_name": f"extra_info_{int(type_id)}",
|
||||
"full_name": f"extra-info type {int(type_id)}",
|
||||
}
|
||||
for type_id in dataset_extra_info_types
|
||||
}
|
||||
|
||||
line_re = re.compile(r"^\s*(\d+)\s*#\s*([^|#]+?)(?:\s*\|\s*(.*?))?\s*$")
|
||||
for path in sorted(search_root.glob("extra_info_types*.txt")):
|
||||
for line in path.read_text(encoding="utf-8").splitlines():
|
||||
match = line_re.match(line)
|
||||
if not match:
|
||||
continue
|
||||
type_id = int(match.group(1))
|
||||
if type_id not in metadata:
|
||||
continue
|
||||
var_name = match.group(2).strip()
|
||||
full_name = (match.group(3) or var_name).strip()
|
||||
metadata[type_id] = {
|
||||
"type_id": type_id,
|
||||
"var_name": var_name,
|
||||
"full_name": full_name,
|
||||
}
|
||||
|
||||
return metadata
|
||||
|
||||
|
||||
def resolve_extra_info_types(
|
||||
value: str | None,
|
||||
*,
|
||||
dataset_extra_info_types: Sequence[int],
|
||||
metadata: dict[int, dict[str, Any]],
|
||||
) -> list[int]:
|
||||
available = [int(x) for x in dataset_extra_info_types]
|
||||
if value is None or str(value).strip() == "":
|
||||
return available
|
||||
|
||||
out: list[int] = []
|
||||
seen: set[int] = set()
|
||||
by_var = {
|
||||
str(meta.get("var_name", "")).lower(): int(type_id)
|
||||
for type_id, meta in metadata.items()
|
||||
}
|
||||
by_full = {
|
||||
str(meta.get("full_name", "")).lower(): int(type_id)
|
||||
for type_id, meta in metadata.items()
|
||||
}
|
||||
for part in re.split(r"[,;]+", str(value)):
|
||||
text = part.strip()
|
||||
if not text:
|
||||
continue
|
||||
if text.isdigit() or (text.startswith("-") and text[1:].isdigit()):
|
||||
type_id = int(text)
|
||||
else:
|
||||
lower = text.lower()
|
||||
if lower in by_var:
|
||||
type_id = by_var[lower]
|
||||
elif lower in by_full:
|
||||
type_id = by_full[lower]
|
||||
else:
|
||||
matches = [
|
||||
int(t)
|
||||
for t, meta in metadata.items()
|
||||
if lower in str(meta.get("var_name", "")).lower()
|
||||
or lower in str(meta.get("full_name", "")).lower()
|
||||
]
|
||||
if len(matches) != 1:
|
||||
raise ValueError(
|
||||
f"--extra_info={text!r} matched {len(matches)} types; "
|
||||
"use a type id or exact variable name."
|
||||
)
|
||||
type_id = matches[0]
|
||||
if type_id not in available:
|
||||
raise ValueError(
|
||||
f"extra-info type {type_id} is not available in this dataset/run"
|
||||
)
|
||||
if type_id not in seen:
|
||||
out.append(type_id)
|
||||
seen.add(type_id)
|
||||
return out
|
||||
|
||||
|
||||
def death_distribution_parameters(
|
||||
model,
|
||||
hidden: torch.Tensor,
|
||||
*,
|
||||
dist_mode: str,
|
||||
eps: float = 1e-8,
|
||||
) -> tuple[str, torch.Tensor]:
|
||||
logits = model.calc_risk(hidden)
|
||||
death_idx = int(logits.shape[1]) - 1
|
||||
death_lambda = F.softplus(logits[:, death_idx]) + float(eps)
|
||||
|
||||
if dist_mode == "exponential":
|
||||
nan = torch.full_like(death_lambda, float("nan"))
|
||||
return "exponential", torch.stack([death_lambda, nan, nan], dim=1)
|
||||
|
||||
if dist_mode == "weibull":
|
||||
rho = model.calc_weibull_rho(hidden)[:, death_idx].to(dtype=death_lambda.dtype)
|
||||
elif dist_mode == "mixed":
|
||||
rho = model.calc_death_rho(hidden).to(dtype=death_lambda.dtype)
|
||||
else:
|
||||
raise ValueError(f"Unsupported dist_mode={dist_mode!r}")
|
||||
|
||||
shape = rho.clamp_min(float(eps))
|
||||
scale = torch.pow(death_lambda.clamp_min(float(eps)), -1.0 / shape)
|
||||
nan = torch.full_like(death_lambda, float("nan"))
|
||||
return "weibull", torch.stack([nan, scale, shape], dim=1)
|
||||
|
||||
|
||||
def parameter_pair_block(original: torch.Tensor, ablated: torch.Tensor) -> torch.Tensor:
|
||||
return torch.stack(
|
||||
[
|
||||
original[:, 0],
|
||||
ablated[:, 0],
|
||||
original[:, 1],
|
||||
ablated[:, 1],
|
||||
original[:, 2],
|
||||
ablated[:, 2],
|
||||
],
|
||||
dim=1,
|
||||
)
|
||||
|
||||
|
||||
def all_disease_parameter_pair_block(
|
||||
*,
|
||||
original_logits: torch.Tensor,
|
||||
ablated_logits: torch.Tensor,
|
||||
dist_mode: str,
|
||||
original_rho: torch.Tensor | None = None,
|
||||
ablated_rho: torch.Tensor | None = None,
|
||||
eps: float = 1e-8,
|
||||
) -> torch.Tensor:
|
||||
original_lambda = F.softplus(original_logits) + float(eps)
|
||||
ablated_lambda = F.softplus(ablated_logits) + float(eps)
|
||||
|
||||
if dist_mode in {"exponential", "mixed"}:
|
||||
nan = torch.full_like(original_lambda, float("nan"))
|
||||
return torch.stack(
|
||||
[
|
||||
original_lambda,
|
||||
ablated_lambda,
|
||||
nan,
|
||||
nan,
|
||||
nan,
|
||||
nan,
|
||||
],
|
||||
dim=2,
|
||||
)
|
||||
|
||||
if dist_mode == "weibull":
|
||||
if original_rho is None or ablated_rho is None:
|
||||
raise ValueError("rho tensors are required for weibull disease parameters")
|
||||
original_shape = original_rho.to(dtype=original_lambda.dtype).clamp_min(float(eps))
|
||||
ablated_shape = ablated_rho.to(dtype=ablated_lambda.dtype).clamp_min(float(eps))
|
||||
original_scale = torch.pow(original_lambda.clamp_min(float(eps)), -1.0 / original_shape)
|
||||
ablated_scale = torch.pow(ablated_lambda.clamp_min(float(eps)), -1.0 / ablated_shape)
|
||||
nan = torch.full_like(original_lambda, float("nan"))
|
||||
return torch.stack(
|
||||
[
|
||||
nan,
|
||||
nan,
|
||||
original_scale,
|
||||
ablated_scale,
|
||||
original_shape,
|
||||
ablated_shape,
|
||||
],
|
||||
dim=2,
|
||||
)
|
||||
|
||||
raise ValueError(f"Unsupported dist_mode={dist_mode!r}")
|
||||
|
||||
|
||||
def grouped_parameter_stats(
|
||||
values: torch.Tensor,
|
||||
group_token_mask: torch.Tensor,
|
||||
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
|
||||
finite = torch.isfinite(values)
|
||||
values64 = values.to(dtype=torch.float64)
|
||||
safe_values = torch.where(finite, values64, torch.zeros_like(values64))
|
||||
mask = group_token_mask.to(device=values.device, dtype=torch.float64)
|
||||
sums = torch.einsum("nvc,gv->ngc", safe_values, mask)
|
||||
sumsq = torch.einsum("nvc,gv->ngc", safe_values * safe_values, mask)
|
||||
counts = torch.einsum("nvc,gv->ngc", finite.to(dtype=torch.float64), mask)
|
||||
return (
|
||||
sums.detach().cpu().numpy().astype(np.float64, copy=False),
|
||||
sumsq.detach().cpu().numpy().astype(np.float64, copy=False),
|
||||
counts.detach().cpu().numpy().astype(np.float64, copy=False),
|
||||
)
|
||||
|
||||
|
||||
def build_extra_info_ablated_slice(
|
||||
batch: dict[str, torch.Tensor],
|
||||
*,
|
||||
row_indices: torch.Tensor,
|
||||
extra_info_type_id: int,
|
||||
) -> dict[str, torch.Tensor]:
|
||||
out: dict[str, torch.Tensor] = {}
|
||||
repeated_keys = (
|
||||
"event_seq",
|
||||
"time_seq",
|
||||
"padding_mask",
|
||||
"readout_mask",
|
||||
"sex",
|
||||
"landmark_pos",
|
||||
"t_query",
|
||||
"patient_id",
|
||||
"landmark_age",
|
||||
"followup_end_time",
|
||||
"death_time",
|
||||
"row_idx",
|
||||
)
|
||||
for key in repeated_keys:
|
||||
out[key] = batch[key][row_indices]
|
||||
|
||||
out["other_type"] = batch["other_type"][row_indices].clone()
|
||||
out["other_value"] = batch["other_value"][row_indices].clone()
|
||||
out["other_value_kind"] = batch["other_value_kind"][row_indices].clone()
|
||||
out["other_time"] = batch["other_time"][row_indices].clone()
|
||||
|
||||
remove = out["other_type"] == int(extra_info_type_id)
|
||||
out["other_type"][remove] = 0
|
||||
out["other_value"][remove] = 0
|
||||
out["other_value_kind"][remove] = 0
|
||||
out["other_time"][remove] = 0
|
||||
return out
|
||||
|
||||
|
||||
def concat_tensor_batches(chunks: Sequence[dict[str, torch.Tensor]]) -> dict[str, torch.Tensor]:
|
||||
return {key: torch.cat([chunk[key] for chunk in chunks], dim=0) for key in chunks[0]}
|
||||
|
||||
|
||||
def iter_extra_info_ablated_batches(
|
||||
batch: dict[str, torch.Tensor],
|
||||
*,
|
||||
selected_extra_info_types: Sequence[int],
|
||||
max_batch_size: int,
|
||||
):
|
||||
pending_batches: list[dict[str, torch.Tensor]] = []
|
||||
pending_types: list[int] = []
|
||||
pending_rows: list[int] = []
|
||||
pending_n = 0
|
||||
|
||||
other_type = batch["other_type"]
|
||||
visible = other_type > 0
|
||||
visible &= batch["other_time"] <= batch["t_query"][:, None].to(batch["other_time"].dtype)
|
||||
|
||||
for type_id in selected_extra_info_types:
|
||||
active_rows = torch.nonzero(
|
||||
((other_type == int(type_id)) & visible).any(dim=1),
|
||||
as_tuple=False,
|
||||
).flatten()
|
||||
if active_rows.numel() == 0:
|
||||
continue
|
||||
|
||||
row_offset = 0
|
||||
while row_offset < int(active_rows.numel()):
|
||||
capacity = int(max_batch_size) - pending_n
|
||||
row_stop = min(int(active_rows.numel()), row_offset + capacity)
|
||||
row_indices = active_rows[row_offset:row_stop].to(device=batch["event_seq"].device)
|
||||
chunk = build_extra_info_ablated_slice(
|
||||
batch,
|
||||
row_indices=row_indices,
|
||||
extra_info_type_id=int(type_id),
|
||||
)
|
||||
chunk_n = int(row_indices.numel())
|
||||
pending_batches.append(chunk)
|
||||
pending_types.extend([int(type_id)] * chunk_n)
|
||||
pending_rows.extend(int(x) for x in row_indices.detach().cpu().tolist())
|
||||
pending_n += chunk_n
|
||||
row_offset = row_stop
|
||||
|
||||
if pending_n >= int(max_batch_size):
|
||||
yield concat_tensor_batches(pending_batches), pending_types, pending_rows
|
||||
pending_batches = []
|
||||
pending_types = []
|
||||
pending_rows = []
|
||||
pending_n = 0
|
||||
|
||||
if pending_batches:
|
||||
yield concat_tensor_batches(pending_batches), pending_types, pending_rows
|
||||
|
||||
|
||||
def finite_float64(values: Any) -> np.ndarray:
|
||||
arr = np.asarray(values, dtype=np.float64)
|
||||
return arr[np.isfinite(arr)]
|
||||
|
||||
|
||||
def update_death_summary(
|
||||
summary: dict[tuple[Any, ...], dict[str, float]],
|
||||
*,
|
||||
key_rows: pd.DataFrame,
|
||||
values: np.ndarray,
|
||||
) -> None:
|
||||
if key_rows.empty:
|
||||
return
|
||||
table = key_rows.copy()
|
||||
for idx, column in enumerate(DEATH_PARAMETER_COLUMNS):
|
||||
table[column] = values[:, idx]
|
||||
|
||||
for key, group in table.groupby(EXTRA_KEY_COLUMNS, dropna=False, sort=False):
|
||||
if not isinstance(key, tuple):
|
||||
key = (key,)
|
||||
acc = summary.setdefault(
|
||||
key,
|
||||
{
|
||||
"n": 0.0,
|
||||
**{f"count__{col}": 0.0 for col in DEATH_PARAMETER_COLUMNS},
|
||||
**{f"sum__{col}": 0.0 for col in DEATH_PARAMETER_COLUMNS},
|
||||
**{f"sumsq__{col}": 0.0 for col in DEATH_PARAMETER_COLUMNS},
|
||||
},
|
||||
)
|
||||
acc["n"] += float(len(group))
|
||||
for column in DEATH_PARAMETER_COLUMNS:
|
||||
vals = finite_float64(pd.to_numeric(group[column], errors="coerce"))
|
||||
acc[f"count__{column}"] += float(vals.size)
|
||||
acc[f"sum__{column}"] += float(vals.sum())
|
||||
acc[f"sumsq__{column}"] += float(np.square(vals).sum())
|
||||
|
||||
|
||||
def update_disease_parameter_summary_from_group_stats(
|
||||
summary: dict[tuple[Any, ...], dict[str, float]],
|
||||
*,
|
||||
key_rows: pd.DataFrame,
|
||||
group_names: Sequence[str],
|
||||
group_labels: Sequence[str],
|
||||
sums: np.ndarray,
|
||||
sumsq: np.ndarray,
|
||||
counts: np.ndarray,
|
||||
) -> None:
|
||||
if key_rows.empty or sums.size == 0:
|
||||
return
|
||||
rows = key_rows.reset_index(drop=True)
|
||||
for row_idx, row in rows.iterrows():
|
||||
base_key = tuple(row[column] for column in EXTRA_KEY_COLUMNS)
|
||||
for group_idx, (group, label) in enumerate(zip(group_names, group_labels)):
|
||||
count_row = counts[int(row_idx), int(group_idx)]
|
||||
n_add = float(np.nanmax(count_row)) if count_row.size else 0.0
|
||||
if n_add <= 0:
|
||||
continue
|
||||
full_key = (*base_key, str(group), str(label))
|
||||
acc = summary.setdefault(
|
||||
full_key,
|
||||
{
|
||||
"n": 0.0,
|
||||
**{f"count__{col}": 0.0 for col in DISEASE_PARAMETER_COLUMNS},
|
||||
**{f"sum__{col}": 0.0 for col in DISEASE_PARAMETER_COLUMNS},
|
||||
**{f"sumsq__{col}": 0.0 for col in DISEASE_PARAMETER_COLUMNS},
|
||||
},
|
||||
)
|
||||
acc["n"] += n_add
|
||||
for col_idx, column in enumerate(DISEASE_PARAMETER_COLUMNS):
|
||||
count = float(counts[int(row_idx), int(group_idx), int(col_idx)])
|
||||
if count <= 0:
|
||||
continue
|
||||
acc[f"count__{column}"] += count
|
||||
acc[f"sum__{column}"] += float(sums[int(row_idx), int(group_idx), int(col_idx)])
|
||||
acc[f"sumsq__{column}"] += float(sumsq[int(row_idx), int(group_idx), int(col_idx)])
|
||||
|
||||
|
||||
def merge_summary_dict(
|
||||
dst: dict[tuple[Any, ...], dict[str, float]],
|
||||
src: dict[tuple[Any, ...], dict[str, float]],
|
||||
) -> None:
|
||||
for key, src_acc in src.items():
|
||||
dst_acc = dst.setdefault(key, {name: 0.0 for name in src_acc})
|
||||
for name, value in src_acc.items():
|
||||
dst_acc[name] = dst_acc.get(name, 0.0) + float(value)
|
||||
|
||||
|
||||
def reduce_attribution_chunk_bundle(
|
||||
payload: tuple[
|
||||
list[tuple[pd.DataFrame, np.ndarray]],
|
||||
list[tuple[pd.DataFrame, np.ndarray, np.ndarray, np.ndarray]],
|
||||
list[str],
|
||||
list[str],
|
||||
],
|
||||
) -> tuple[dict[tuple[Any, ...], dict[str, float]], dict[tuple[Any, ...], dict[str, float]]]:
|
||||
death_items, disease_items, group_names, group_labels = payload
|
||||
death_summary: dict[tuple[Any, ...], dict[str, float]] = {}
|
||||
disease_summary: dict[tuple[Any, ...], dict[str, float]] = {}
|
||||
|
||||
for key_rows, values in death_items:
|
||||
update_death_summary(
|
||||
death_summary,
|
||||
key_rows=key_rows,
|
||||
values=values,
|
||||
)
|
||||
|
||||
for key_rows, sums, sumsq, counts in disease_items:
|
||||
update_disease_parameter_summary_from_group_stats(
|
||||
disease_summary,
|
||||
key_rows=key_rows,
|
||||
group_names=group_names,
|
||||
group_labels=group_labels,
|
||||
sums=sums,
|
||||
sumsq=sumsq,
|
||||
counts=counts,
|
||||
)
|
||||
|
||||
return death_summary, disease_summary
|
||||
|
||||
|
||||
def reduce_attribution_chunks(
|
||||
*,
|
||||
death_key_chunks: list[pd.DataFrame],
|
||||
death_value_chunks: list[np.ndarray],
|
||||
disease_stat_chunks: list[tuple[pd.DataFrame, np.ndarray, np.ndarray, np.ndarray]],
|
||||
group_names: list[str],
|
||||
group_labels: list[str],
|
||||
cpu_reduce_workers: int,
|
||||
) -> tuple[dict[tuple[Any, ...], dict[str, float]], dict[tuple[Any, ...], dict[str, float]]]:
|
||||
n_chunks = max(len(death_key_chunks), len(disease_stat_chunks))
|
||||
if n_chunks == 0:
|
||||
return {}, {}
|
||||
|
||||
worker_count = max(1, min(int(cpu_reduce_workers), n_chunks))
|
||||
if worker_count == 1:
|
||||
return reduce_attribution_chunk_bundle(
|
||||
(
|
||||
list(zip(death_key_chunks, death_value_chunks)),
|
||||
disease_stat_chunks,
|
||||
group_names,
|
||||
group_labels,
|
||||
)
|
||||
)
|
||||
|
||||
bundles = []
|
||||
for worker_idx in range(worker_count):
|
||||
start = worker_idx * n_chunks // worker_count
|
||||
stop = (worker_idx + 1) * n_chunks // worker_count
|
||||
if start >= stop:
|
||||
continue
|
||||
death_items = [
|
||||
(death_key_chunks[i], death_value_chunks[i])
|
||||
for i in range(start, min(stop, len(death_key_chunks)))
|
||||
]
|
||||
disease_items = disease_stat_chunks[start:min(stop, len(disease_stat_chunks))]
|
||||
bundles.append((death_items, disease_items, group_names, group_labels))
|
||||
|
||||
merged_death: dict[tuple[Any, ...], dict[str, float]] = {}
|
||||
merged_disease: dict[tuple[Any, ...], dict[str, float]] = {}
|
||||
with ProcessPoolExecutor(max_workers=len(bundles)) as executor:
|
||||
futures = [executor.submit(reduce_attribution_chunk_bundle, bundle) for bundle in bundles]
|
||||
for future in tqdm(as_completed(futures), total=len(futures), desc="CPU summary reduction", dynamic_ncols=True):
|
||||
death_part, disease_part = future.result()
|
||||
merge_summary_dict(merged_death, death_part)
|
||||
merge_summary_dict(merged_disease, disease_part)
|
||||
|
||||
return merged_death, merged_disease
|
||||
|
||||
|
||||
def write_death_summary_csv(
|
||||
path: Path,
|
||||
summary: dict[tuple[Any, ...], dict[str, float]],
|
||||
*,
|
||||
death_distribution: str,
|
||||
) -> int:
|
||||
rows: list[dict[str, Any]] = []
|
||||
for key, acc in summary.items():
|
||||
n = int(acc["n"])
|
||||
row = {column: value for column, value in zip(EXTRA_KEY_COLUMNS, key)}
|
||||
row["n"] = n
|
||||
row["death_distribution"] = death_distribution
|
||||
for column in DEATH_PARAMETER_COLUMNS:
|
||||
count = int(acc[f"count__{column}"])
|
||||
mean = acc[f"sum__{column}"] / count if count > 0 else np.nan
|
||||
second = acc[f"sumsq__{column}"] / count if count > 0 else np.nan
|
||||
row[f"mean__{column}"] = mean
|
||||
row[f"var__{column}"] = second - mean * mean if count > 0 else np.nan
|
||||
rows.append(row)
|
||||
columns = [
|
||||
*EXTRA_KEY_COLUMNS,
|
||||
"n",
|
||||
"death_distribution",
|
||||
*[
|
||||
name
|
||||
for column in DEATH_PARAMETER_COLUMNS
|
||||
for name in (f"mean__{column}", f"var__{column}")
|
||||
],
|
||||
]
|
||||
pd.DataFrame(rows, columns=columns).sort_values(
|
||||
["selected_extra_info_type_id", "landmark_age", "sex"],
|
||||
kind="mergesort",
|
||||
).to_csv(path, index=False)
|
||||
return len(rows)
|
||||
|
||||
|
||||
def write_disease_parameter_summary_csv(
|
||||
path: Path,
|
||||
summary: dict[tuple[Any, ...], dict[str, float]],
|
||||
) -> int:
|
||||
rows: list[dict[str, Any]] = []
|
||||
for key, acc in summary.items():
|
||||
n = int(acc["n"])
|
||||
row = {column: value for column, value in zip(DISEASE_PARAMETER_KEY_COLUMNS, key)}
|
||||
row["n"] = n
|
||||
for column in DISEASE_PARAMETER_COLUMNS:
|
||||
count = int(acc[f"count__{column}"])
|
||||
mean = acc[f"sum__{column}"] / count if count > 0 else np.nan
|
||||
second = acc[f"sumsq__{column}"] / count if count > 0 else np.nan
|
||||
row[f"mean__{column}"] = mean
|
||||
row[f"var__{column}"] = second - mean * mean if count > 0 else np.nan
|
||||
rows.append(row)
|
||||
columns = [
|
||||
*DISEASE_PARAMETER_KEY_COLUMNS,
|
||||
"n",
|
||||
*[
|
||||
name
|
||||
for column in DISEASE_PARAMETER_COLUMNS
|
||||
for name in (f"mean__{column}", f"var__{column}")
|
||||
],
|
||||
]
|
||||
pd.DataFrame(rows, columns=columns).sort_values(
|
||||
["selected_extra_info_type_id", "target_group", "landmark_age", "sex"],
|
||||
kind="mergesort",
|
||||
).to_csv(path, index=False)
|
||||
return len(rows)
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Compute extra-info ablation attribution for death and disease distribution parameters."
|
||||
)
|
||||
parser.add_argument("--run_path", type=str, required=True)
|
||||
parser.add_argument(
|
||||
"--extra_info",
|
||||
type=str,
|
||||
default=None,
|
||||
help=(
|
||||
"Optional type id, variable name, exact full name, or comma-separated list. "
|
||||
"If omitted, scan all extra-info types available in the run."
|
||||
),
|
||||
)
|
||||
parser.add_argument("--output_dir", type=str, default=None)
|
||||
parser.add_argument("--organ_mapping_path", type=str, default="icd10_chapter_organ_mapping.csv")
|
||||
parser.add_argument("--eval_split", type=str, default=None)
|
||||
parser.add_argument("--dataset_subset_size", type=int, default=None)
|
||||
parser.add_argument("--train_eid_file", type=str, default=None)
|
||||
parser.add_argument("--val_eid_file", type=str, default=None)
|
||||
parser.add_argument("--test_eid_file", type=str, default=None)
|
||||
parser.add_argument("--landmark_start", type=float, default=40.0)
|
||||
parser.add_argument("--landmark_stop", type=float, default=80.0)
|
||||
parser.add_argument("--landmark_step", type=float, default=5.0)
|
||||
parser.add_argument("--min_history_events", type=int, default=None)
|
||||
parser.add_argument("--batch_size", type=int, default=None)
|
||||
parser.add_argument(
|
||||
"--attribution_batch_size",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Forward batch size for expanded extra-info ablation queries.",
|
||||
)
|
||||
parser.add_argument("--num_workers", type=int, default=None)
|
||||
parser.add_argument(
|
||||
"--cpu_reduce_workers",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Worker processes for CPU-side summary reduction. Defaults to --num_workers.",
|
||||
)
|
||||
parser.add_argument("--device", type=str, default=None)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
# Dataset extra-info types must reproduce the checkpoint training config.
|
||||
# --extra_info only filters which already-trained types are ablated.
|
||||
args.extra_info_types = None
|
||||
run_path = Path(args.run_path)
|
||||
config_path = run_path / "train_config.json"
|
||||
checkpoint_path = run_path / "best_model.pt"
|
||||
if not config_path.exists():
|
||||
raise FileNotFoundError(f"train_config.json not found: {config_path}")
|
||||
if not checkpoint_path.exists():
|
||||
raise FileNotFoundError(f"best_model.pt not found: {checkpoint_path}")
|
||||
|
||||
cfg = load_json_config(config_path)
|
||||
model_target_mode = str(cfg.get("model_target_mode", "next_token")).lower()
|
||||
if model_target_mode not in {"next_token", "all_future"}:
|
||||
raise ValueError(f"Unsupported model_target_mode: {model_target_mode!r}")
|
||||
|
||||
target_mode = str(cfg.get("target_mode", "uts"))
|
||||
attn_mask_mode = str(
|
||||
cfg.get("attn_mask_mode", "non_strict_time" if target_mode == "uts" else "target_aware")
|
||||
)
|
||||
readout_name = str(
|
||||
cfg.get("readout_name", "same_time_group_end" if target_mode == "uts" else "token")
|
||||
)
|
||||
readout_reduce = str(cfg.get("readout_reduce", "mean"))
|
||||
|
||||
dataset, subset_indices, eval_split, split_source = load_eval_sequence_dataset(args, cfg)
|
||||
validate_dataset_metadata(dataset, cfg)
|
||||
|
||||
extra_metadata = load_extra_info_metadata(
|
||||
dataset_extra_info_types=dataset.extra_info_types,
|
||||
search_root=Path("."),
|
||||
)
|
||||
selected_extra_info_types = resolve_extra_info_types(
|
||||
args.extra_info,
|
||||
dataset_extra_info_types=dataset.extra_info_types,
|
||||
metadata=extra_metadata,
|
||||
)
|
||||
if not selected_extra_info_types:
|
||||
raise ValueError("No extra-info types selected for attribution")
|
||||
|
||||
landmark_ages = make_landmark_ages(
|
||||
float(args.landmark_start),
|
||||
float(args.landmark_stop),
|
||||
float(args.landmark_step),
|
||||
)
|
||||
first_occurrence_by_token = build_first_occurrence_maps_for_landmarks(
|
||||
dataset,
|
||||
subset_indices,
|
||||
)
|
||||
death_idx = int(dataset.vocab_size) - 1
|
||||
landmark_dataset = LandmarkDataset(
|
||||
dataset=dataset,
|
||||
subset_indices=subset_indices,
|
||||
landmark_ages=landmark_ages,
|
||||
attn_mask_mode=attn_mask_mode,
|
||||
model_target_mode=model_target_mode,
|
||||
min_history_events=int(cfg_get(args, cfg, "min_history_events", 1)),
|
||||
first_occurrence_by_token=first_occurrence_by_token,
|
||||
death_token_ids=[death_idx],
|
||||
)
|
||||
|
||||
organ_groups, organ_labels, _token_to_group = load_organ_groups(
|
||||
Path(args.organ_mapping_path),
|
||||
vocab_size=int(dataset.vocab_size),
|
||||
)
|
||||
all_disease_tokens = sorted(
|
||||
{
|
||||
int(token)
|
||||
for tokens in organ_groups.values()
|
||||
for token in tokens
|
||||
if int(token) != death_idx
|
||||
}
|
||||
)
|
||||
risk_groups = {
|
||||
"all_modeled_diseases": all_disease_tokens,
|
||||
**{group: tokens for group, tokens in sorted(organ_groups.items())},
|
||||
}
|
||||
risk_group_labels = {
|
||||
"all_modeled_diseases": "All modeled diseases",
|
||||
**organ_labels,
|
||||
}
|
||||
group_names = list(risk_groups.keys())
|
||||
group_labels = [str(risk_group_labels[group]) for group in group_names]
|
||||
|
||||
state_dict = load_checkpoint_state_dict(checkpoint_path, map_location="cpu")
|
||||
dist_mode = resolve_dist_mode_for_checkpoint(str(cfg.get("dist_mode", "exponential")), state_dict)
|
||||
death_distribution_name = "exponential" if dist_mode == "exponential" else "weibull"
|
||||
cfg_model = dict(cfg)
|
||||
cfg_model["dist_mode"] = dist_mode
|
||||
device = resolve_eval_device(args.device)
|
||||
model = build_model_from_dataset(args, cfg_model, dataset).to(device)
|
||||
load_model_state(model, state_dict)
|
||||
model.eval()
|
||||
|
||||
group_token_mask = torch.zeros(
|
||||
(len(group_names), int(dataset.vocab_size)),
|
||||
dtype=torch.float32,
|
||||
device=device,
|
||||
)
|
||||
for group_idx, group in enumerate(group_names):
|
||||
valid_tokens = [
|
||||
int(token)
|
||||
for token in risk_groups[group]
|
||||
if 0 <= int(token) < int(dataset.vocab_size) and int(token) != death_idx
|
||||
]
|
||||
if valid_tokens:
|
||||
group_token_mask[group_idx, torch.as_tensor(valid_tokens, dtype=torch.long, device=device)] = 1.0
|
||||
|
||||
batch_size = int(cfg_get(args, cfg, "batch_size", 128))
|
||||
attribution_batch_size = int(
|
||||
cfg_get(args, cfg, "attribution_batch_size", max(batch_size * 32, 4096))
|
||||
)
|
||||
if attribution_batch_size <= 0:
|
||||
raise ValueError("attribution_batch_size must be positive")
|
||||
|
||||
num_workers = int(cfg_get(args, cfg, "num_workers", 4))
|
||||
cpu_reduce_workers = int(
|
||||
args.cpu_reduce_workers
|
||||
if args.cpu_reduce_workers is not None
|
||||
else max(1, num_workers)
|
||||
)
|
||||
if cpu_reduce_workers <= 0:
|
||||
raise ValueError("--cpu_reduce_workers must be positive")
|
||||
loader = DataLoader(
|
||||
IndexedLandmarkDataset(landmark_dataset),
|
||||
batch_size=batch_size,
|
||||
shuffle=False,
|
||||
collate_fn=collate_indexed_landmark_fn,
|
||||
num_workers=num_workers,
|
||||
pin_memory=device.type == "cuda",
|
||||
persistent_workers=num_workers > 0,
|
||||
prefetch_factor=2 if num_workers > 0 else None,
|
||||
)
|
||||
|
||||
output_dir = (
|
||||
Path(args.output_dir)
|
||||
if args.output_dir
|
||||
else run_path / f"extra_info_attribution_{eval_split}"
|
||||
)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
print(f"Eval split: {eval_split}")
|
||||
print(f"Split source: {split_source}")
|
||||
print(f"Selected patients: {len(subset_indices)}")
|
||||
print(f"Landmark ages: {landmark_ages.tolist()}")
|
||||
print(f"Dist mode: {dist_mode}")
|
||||
print(f"Device: {device}")
|
||||
print(f"Death token: {death_idx}")
|
||||
print(f"Extra-info types: {selected_extra_info_types}")
|
||||
print(f"Landmark rows: {len(landmark_dataset)}")
|
||||
print(f"Attribution batch size: {attribution_batch_size}")
|
||||
print(f"CPU reduce workers: {cpu_reduce_workers}")
|
||||
print(f"Output directory: {output_dir}")
|
||||
|
||||
death_key_chunks: list[pd.DataFrame] = []
|
||||
death_value_chunks: list[np.ndarray] = []
|
||||
disease_stat_chunks: list[tuple[pd.DataFrame, np.ndarray, np.ndarray, np.ndarray]] = []
|
||||
|
||||
for batch in tqdm(loader, desc="Extra-info attribution", dynamic_ncols=True):
|
||||
batch_dev = {
|
||||
k: (v.to(device, non_blocking=True) if isinstance(v, torch.Tensor) else v)
|
||||
for k, v in batch.items()
|
||||
}
|
||||
with torch.no_grad():
|
||||
hidden = infer_landmark_hidden(
|
||||
model=model,
|
||||
batch=batch_dev,
|
||||
device=device,
|
||||
model_target_mode=model_target_mode,
|
||||
readout_name=readout_name,
|
||||
readout_reduce=readout_reduce,
|
||||
)
|
||||
_death_distribution, original_death_params = death_distribution_parameters(
|
||||
model,
|
||||
hidden,
|
||||
dist_mode=dist_mode,
|
||||
)
|
||||
original_logits = model.calc_risk(hidden)
|
||||
original_rho = model.calc_weibull_rho(hidden) if dist_mode == "weibull" else None
|
||||
|
||||
for ablated_batch, type_ids, local_rows in iter_extra_info_ablated_batches(
|
||||
batch_dev,
|
||||
selected_extra_info_types=selected_extra_info_types,
|
||||
max_batch_size=attribution_batch_size,
|
||||
):
|
||||
row_tensor = torch.as_tensor(local_rows, dtype=torch.long, device=device)
|
||||
with torch.no_grad():
|
||||
ablated_hidden = infer_landmark_hidden(
|
||||
model=model,
|
||||
batch=ablated_batch,
|
||||
device=device,
|
||||
model_target_mode=model_target_mode,
|
||||
readout_name=readout_name,
|
||||
readout_reduce=readout_reduce,
|
||||
)
|
||||
_ablated_distribution, ablated_death_params = death_distribution_parameters(
|
||||
model,
|
||||
ablated_hidden,
|
||||
dist_mode=dist_mode,
|
||||
)
|
||||
ablated_logits = model.calc_risk(ablated_hidden)
|
||||
ablated_rho = model.calc_weibull_rho(ablated_hidden) if dist_mode == "weibull" else None
|
||||
|
||||
key_rows = []
|
||||
for type_id, local_row in zip(type_ids, local_rows):
|
||||
meta = extra_metadata[int(type_id)]
|
||||
key_rows.append(
|
||||
{
|
||||
"selected_extra_info_type_id": int(type_id),
|
||||
"selected_extra_info_var_name": str(meta.get("var_name", "")),
|
||||
"selected_extra_info_full_name": str(meta.get("full_name", "")),
|
||||
"landmark_age": float(batch["landmark_age"][int(local_row)].item()),
|
||||
"sex": int(batch["sex"][int(local_row)].item()),
|
||||
}
|
||||
)
|
||||
key_table = pd.DataFrame(key_rows, columns=EXTRA_KEY_COLUMNS)
|
||||
value_block = parameter_pair_block(
|
||||
original_death_params[row_tensor],
|
||||
ablated_death_params,
|
||||
).detach().cpu().numpy()
|
||||
death_key_chunks.append(key_table)
|
||||
death_value_chunks.append(value_block)
|
||||
|
||||
disease_values = all_disease_parameter_pair_block(
|
||||
original_logits=original_logits[row_tensor],
|
||||
ablated_logits=ablated_logits,
|
||||
dist_mode=dist_mode,
|
||||
original_rho=None if original_rho is None else original_rho[row_tensor],
|
||||
ablated_rho=ablated_rho,
|
||||
)
|
||||
sums, sumsq, counts = grouped_parameter_stats(
|
||||
disease_values,
|
||||
group_token_mask,
|
||||
)
|
||||
disease_stat_chunks.append((key_table, sums, sumsq, counts))
|
||||
|
||||
death_summary, disease_parameter_summary = reduce_attribution_chunks(
|
||||
death_key_chunks=death_key_chunks,
|
||||
death_value_chunks=death_value_chunks,
|
||||
disease_stat_chunks=disease_stat_chunks,
|
||||
group_names=group_names,
|
||||
group_labels=group_labels,
|
||||
cpu_reduce_workers=cpu_reduce_workers,
|
||||
)
|
||||
|
||||
death_summary_path = output_dir / "summary_extra_info_death_parameters.csv"
|
||||
disease_summary_path = output_dir / "summary_extra_info_disease_parameters.csv"
|
||||
death_rows = write_death_summary_csv(
|
||||
death_summary_path,
|
||||
death_summary,
|
||||
death_distribution=death_distribution_name,
|
||||
)
|
||||
disease_rows = write_disease_parameter_summary_csv(
|
||||
disease_summary_path,
|
||||
disease_parameter_summary,
|
||||
)
|
||||
manifest = {
|
||||
"death_summary_file": death_summary_path.name,
|
||||
"disease_parameter_summary_file": disease_summary_path.name,
|
||||
"death_summary_rows": int(death_rows),
|
||||
"disease_parameter_summary_rows": int(disease_rows),
|
||||
"eval_split": eval_split,
|
||||
"split_source": split_source,
|
||||
"dist_mode": dist_mode,
|
||||
"landmark_start": float(args.landmark_start),
|
||||
"landmark_stop": float(args.landmark_stop),
|
||||
"landmark_step": float(args.landmark_step),
|
||||
"selected_extra_info_types": [
|
||||
extra_metadata[int(type_id)] for type_id in selected_extra_info_types
|
||||
],
|
||||
}
|
||||
with (output_dir / "manifest.json").open("w", encoding="utf-8") as f:
|
||||
json.dump(manifest, f, ensure_ascii=False, indent=2)
|
||||
|
||||
print(f"Wrote {death_rows} death summary rows to {death_summary_path}")
|
||||
print(f"Wrote {disease_rows} disease-parameter summary rows to {disease_summary_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,7 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from evaluate_auc_v2 import main
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,817 +0,0 @@
|
||||
"""Compute per-disease attribution to predicted mortality distribution parameters.
|
||||
|
||||
For each selected patient and landmark age, this script keeps only rows where
|
||||
each scanned disease token has already occurred in the history. It then deletes
|
||||
that historical disease token, re-queries the model, and reports the original
|
||||
and ablated fitted death distribution parameters. If --disease is omitted, all
|
||||
disease tokens in the mapping are scanned.
|
||||
|
||||
Death is always token vocab_size - 1.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch.utils.data import DataLoader
|
||||
from tqdm.auto import tqdm
|
||||
|
||||
from evaluate_auc_v2 import (
|
||||
build_model_from_dataset,
|
||||
cfg_get,
|
||||
load_checkpoint_state_dict,
|
||||
load_json_config,
|
||||
load_model_state,
|
||||
resolve_dist_mode_for_checkpoint,
|
||||
resolve_eval_device,
|
||||
validate_dataset_metadata,
|
||||
)
|
||||
from landmark_eval_utils import (
|
||||
IndexedLandmarkDataset,
|
||||
LandmarkDataset,
|
||||
build_first_occurrence_maps_for_landmarks,
|
||||
collate_indexed_landmark_fn,
|
||||
historical_counts_by_group,
|
||||
infer_landmark_hidden,
|
||||
load_eval_sequence_dataset,
|
||||
load_organ_groups,
|
||||
make_landmark_ages,
|
||||
)
|
||||
from targets import CHECKUP_IDX, PAD_IDX
|
||||
|
||||
|
||||
OUTPUT_COLUMNS = [
|
||||
"patient_id",
|
||||
"dataset_index",
|
||||
"eid",
|
||||
"sex",
|
||||
"landmark_age",
|
||||
"followup_end_time",
|
||||
"history_disease_count",
|
||||
"selected_disease_history_count",
|
||||
"selected_disease_token_id",
|
||||
"selected_disease_code",
|
||||
"selected_disease_name",
|
||||
"selected_disease_organ_system",
|
||||
"selected_disease_organ_system_label",
|
||||
"history_count__selected_organ_system",
|
||||
"death_distribution",
|
||||
"original_death_lambda",
|
||||
"ablated_death_lambda",
|
||||
"original_death_scale",
|
||||
"ablated_death_scale",
|
||||
"original_death_shape",
|
||||
"ablated_death_shape",
|
||||
]
|
||||
|
||||
SUMMARY_KEY_COLUMNS = [
|
||||
"selected_disease_token_id",
|
||||
"selected_disease_code",
|
||||
"selected_disease_name",
|
||||
"selected_disease_organ_system",
|
||||
"selected_disease_organ_system_label",
|
||||
"landmark_age",
|
||||
"sex",
|
||||
]
|
||||
|
||||
SUMMARY_MEAN_COLUMNS = [
|
||||
"history_disease_count",
|
||||
"selected_disease_history_count",
|
||||
"history_count__selected_organ_system",
|
||||
]
|
||||
|
||||
SUMMARY_PARAMETER_COLUMNS = [
|
||||
"original_death_lambda",
|
||||
"ablated_death_lambda",
|
||||
"original_death_scale",
|
||||
"ablated_death_scale",
|
||||
"original_death_shape",
|
||||
"ablated_death_shape",
|
||||
]
|
||||
|
||||
|
||||
def write_compressed_npz_table(path: Path, table: pd.DataFrame) -> int:
|
||||
table = table.reindex(columns=OUTPUT_COLUMNS)
|
||||
arrays: dict[str, np.ndarray] = {
|
||||
"__columns__": np.asarray(OUTPUT_COLUMNS, dtype="U"),
|
||||
}
|
||||
for column in OUTPUT_COLUMNS:
|
||||
values = table[column] if column in table else pd.Series([], dtype=object)
|
||||
if values.dtype == object:
|
||||
arrays[column] = values.fillna("").astype(str).to_numpy(dtype="U")
|
||||
else:
|
||||
arrays[column] = values.to_numpy()
|
||||
np.savez_compressed(path, **arrays)
|
||||
return int(len(table))
|
||||
|
||||
|
||||
def normalize_output_dir(path: Path) -> Path:
|
||||
if path.suffix:
|
||||
return path.with_suffix(path.suffix + "_shards")
|
||||
return path
|
||||
|
||||
|
||||
def write_manifest(
|
||||
output_dir: Path,
|
||||
*,
|
||||
rows: int,
|
||||
shards: list[dict[str, Any]],
|
||||
summary_file: str,
|
||||
scanned_diseases: list[dict[str, Any]],
|
||||
eval_split: str,
|
||||
dist_mode: str,
|
||||
landmark_start: float,
|
||||
landmark_stop: float,
|
||||
landmark_step: float,
|
||||
) -> None:
|
||||
payload = {
|
||||
"format": "compressed_npz_shards",
|
||||
"columns": OUTPUT_COLUMNS,
|
||||
"rows": int(rows),
|
||||
"shards": shards,
|
||||
"summary_file": summary_file,
|
||||
"scanned_diseases": scanned_diseases,
|
||||
"eval_split": eval_split,
|
||||
"dist_mode": str(dist_mode),
|
||||
"landmark_start": float(landmark_start),
|
||||
"landmark_stop": float(landmark_stop),
|
||||
"landmark_step": float(landmark_step),
|
||||
}
|
||||
with (output_dir / "manifest.json").open("w", encoding="utf-8") as f:
|
||||
json.dump(payload, f, ensure_ascii=False, indent=2)
|
||||
|
||||
|
||||
def update_summary_accumulator(
|
||||
summary: dict[tuple[Any, ...], dict[str, float]],
|
||||
table: pd.DataFrame,
|
||||
) -> None:
|
||||
if table.empty:
|
||||
return
|
||||
grouped = table.groupby(SUMMARY_KEY_COLUMNS, dropna=False, sort=False)
|
||||
for key, group in grouped:
|
||||
if not isinstance(key, tuple):
|
||||
key = (key,)
|
||||
acc = summary.setdefault(
|
||||
key,
|
||||
{
|
||||
"n": 0.0,
|
||||
**{column: 0.0 for column in SUMMARY_MEAN_COLUMNS},
|
||||
**{f"count__{column}": 0.0 for column in SUMMARY_PARAMETER_COLUMNS},
|
||||
**{f"sum__{column}": 0.0 for column in SUMMARY_PARAMETER_COLUMNS},
|
||||
**{f"sumsq__{column}": 0.0 for column in SUMMARY_PARAMETER_COLUMNS},
|
||||
},
|
||||
)
|
||||
n = int(len(group))
|
||||
acc["n"] += float(n)
|
||||
for column in SUMMARY_MEAN_COLUMNS:
|
||||
acc[column] += float(pd.to_numeric(group[column], errors="coerce").sum())
|
||||
for column in SUMMARY_PARAMETER_COLUMNS:
|
||||
values = pd.to_numeric(group[column], errors="coerce").dropna()
|
||||
acc[f"count__{column}"] += float(len(values))
|
||||
acc[f"sum__{column}"] += float(values.sum())
|
||||
acc[f"sumsq__{column}"] += float((values * values).sum())
|
||||
|
||||
|
||||
def write_summary_csv(
|
||||
path: Path,
|
||||
summary: dict[tuple[Any, ...], dict[str, float]],
|
||||
) -> int:
|
||||
rows: list[dict[str, Any]] = []
|
||||
for key, acc in summary.items():
|
||||
n = int(acc["n"])
|
||||
out = {column: value for column, value in zip(SUMMARY_KEY_COLUMNS, key)}
|
||||
out["n"] = n
|
||||
for column in SUMMARY_MEAN_COLUMNS:
|
||||
out[f"mean__{column}"] = acc[column] / n if n > 0 else np.nan
|
||||
for column in SUMMARY_PARAMETER_COLUMNS:
|
||||
count = int(acc[f"count__{column}"])
|
||||
mean = acc[f"sum__{column}"] / count if count > 0 else np.nan
|
||||
second = acc[f"sumsq__{column}"] / count if count > 0 else np.nan
|
||||
out[f"mean__{column}"] = mean
|
||||
out[f"var__{column}"] = second - mean * mean if count > 0 else np.nan
|
||||
rows.append(out)
|
||||
|
||||
columns = [
|
||||
*SUMMARY_KEY_COLUMNS,
|
||||
"n",
|
||||
*[f"mean__{column}" for column in SUMMARY_MEAN_COLUMNS],
|
||||
*[
|
||||
name
|
||||
for column in SUMMARY_PARAMETER_COLUMNS
|
||||
for name in (f"mean__{column}", f"var__{column}")
|
||||
],
|
||||
]
|
||||
pd.DataFrame(rows, columns=columns).sort_values(
|
||||
["selected_disease_token_id", "landmark_age", "sex"],
|
||||
kind="mergesort",
|
||||
).to_csv(path, index=False)
|
||||
return len(rows)
|
||||
|
||||
|
||||
def build_disease_ablated_slice(
|
||||
batch: Dict[str, torch.Tensor],
|
||||
row_indices: torch.Tensor,
|
||||
token_ids: torch.Tensor,
|
||||
) -> Dict[str, torch.Tensor]:
|
||||
"""Build an ablated slice for aligned (row, disease_token) pairs."""
|
||||
event_seq = batch["event_seq"]
|
||||
row_indices = row_indices.to(device=event_seq.device, dtype=torch.long)
|
||||
token_ids = token_ids.to(device=event_seq.device, dtype=event_seq.dtype)
|
||||
|
||||
out: Dict[str, torch.Tensor] = {}
|
||||
out["event_seq"] = event_seq[row_indices].clone()
|
||||
out["time_seq"] = batch["time_seq"][row_indices]
|
||||
out["readout_mask"] = batch["readout_mask"][row_indices].clone()
|
||||
out["padding_mask"] = batch["padding_mask"][row_indices].bool().clone()
|
||||
out["landmark_pos"] = batch["landmark_pos"][row_indices].clone()
|
||||
|
||||
seq_len = int(event_seq.shape[1])
|
||||
positions = torch.arange(seq_len, device=event_seq.device)[None, :]
|
||||
remove = (out["event_seq"] == token_ids[:, None]) & out["padding_mask"]
|
||||
out["event_seq"] = torch.where(
|
||||
remove,
|
||||
torch.full_like(out["event_seq"], PAD_IDX),
|
||||
out["event_seq"],
|
||||
)
|
||||
out["padding_mask"] &= ~remove
|
||||
out["readout_mask"] &= ~remove
|
||||
|
||||
has_valid = out["padding_mask"].any(dim=1)
|
||||
empty_rows = ~has_valid
|
||||
out["event_seq"][empty_rows, 0] = CHECKUP_IDX
|
||||
out["time_seq"][empty_rows, 0] = batch["t_query"][row_indices[empty_rows]].to(
|
||||
dtype=out["time_seq"].dtype
|
||||
)
|
||||
out["padding_mask"][empty_rows, 0] = True
|
||||
out["readout_mask"][empty_rows, 0] = True
|
||||
out["landmark_pos"][empty_rows] = 0
|
||||
|
||||
has_readout = out["readout_mask"].any(dim=1)
|
||||
missing_readout = ~has_readout
|
||||
local_valid = out["padding_mask"]
|
||||
last_pos = torch.where(
|
||||
local_valid,
|
||||
positions.expand(local_valid.shape[0], -1),
|
||||
torch.zeros_like(positions.expand(local_valid.shape[0], -1)),
|
||||
).amax(dim=1)
|
||||
out["readout_mask"][missing_readout] = False
|
||||
out["readout_mask"][missing_readout, last_pos[missing_readout]] = True
|
||||
out["landmark_pos"][missing_readout] = last_pos[missing_readout].to(
|
||||
dtype=out["landmark_pos"].dtype
|
||||
)
|
||||
|
||||
repeated_keys = (
|
||||
"sex",
|
||||
"other_type",
|
||||
"other_value",
|
||||
"other_value_kind",
|
||||
"other_time",
|
||||
"t_query",
|
||||
"patient_id",
|
||||
"landmark_age",
|
||||
"followup_end_time",
|
||||
"death_time",
|
||||
"row_idx",
|
||||
)
|
||||
for key in repeated_keys:
|
||||
out[key] = batch[key][row_indices]
|
||||
return out
|
||||
|
||||
|
||||
def load_disease_metadata(
|
||||
mapping_path: Path,
|
||||
*,
|
||||
vocab_size: int,
|
||||
) -> dict[int, dict[str, Any]]:
|
||||
if not mapping_path.exists():
|
||||
raise FileNotFoundError(f"Disease mapping file not found: {mapping_path}")
|
||||
table = pd.read_csv(mapping_path)
|
||||
required = {"token_id", "code", "name", "is_death"}
|
||||
missing = required - set(table.columns)
|
||||
if missing:
|
||||
raise ValueError(f"{mapping_path} is missing columns: {sorted(missing)}")
|
||||
|
||||
death_idx = int(vocab_size) - 1
|
||||
out: dict[int, dict[str, Any]] = {}
|
||||
for row in table.itertuples(index=False):
|
||||
token = int(getattr(row, "token_id"))
|
||||
if token < 0 or token >= int(vocab_size) or token == death_idx:
|
||||
continue
|
||||
if int(getattr(row, "is_death")) == 1:
|
||||
continue
|
||||
meta = {
|
||||
"token_id": token,
|
||||
"code": str(getattr(row, "code")),
|
||||
"name": str(getattr(row, "name")),
|
||||
}
|
||||
for column in (
|
||||
"icd10_chapter",
|
||||
"icd10_chapter_title",
|
||||
"organ_system",
|
||||
"organ_system_label",
|
||||
):
|
||||
if hasattr(row, column):
|
||||
meta[column] = str(getattr(row, column))
|
||||
out[token] = meta
|
||||
return out
|
||||
|
||||
|
||||
def resolve_disease_token(
|
||||
value: str,
|
||||
metadata: dict[int, dict[str, Any]],
|
||||
) -> tuple[int, dict[str, Any]]:
|
||||
text = str(value).strip()
|
||||
if text == "":
|
||||
raise ValueError("--disease must not be empty")
|
||||
|
||||
if text.isdigit() or (text.startswith("-") and text[1:].isdigit()):
|
||||
token = int(text)
|
||||
if token not in metadata:
|
||||
raise ValueError(f"Disease token_id {token} was not found in the mapping")
|
||||
return token, metadata[token]
|
||||
|
||||
lower = text.lower()
|
||||
exact = [
|
||||
(token, meta)
|
||||
for token, meta in metadata.items()
|
||||
if str(meta.get("code", "")).lower() == lower
|
||||
or str(meta.get("name", "")).lower() == lower
|
||||
]
|
||||
if len(exact) == 1:
|
||||
return exact[0]
|
||||
if len(exact) > 1:
|
||||
raise ValueError(f"--disease={value!r} matched multiple diseases exactly")
|
||||
|
||||
contains = [
|
||||
(token, meta)
|
||||
for token, meta in metadata.items()
|
||||
if lower in str(meta.get("code", "")).lower()
|
||||
or lower in str(meta.get("name", "")).lower()
|
||||
]
|
||||
if len(contains) == 1:
|
||||
return contains[0]
|
||||
if not contains:
|
||||
raise ValueError(f"--disease={value!r} did not match any disease token")
|
||||
preview = ", ".join(
|
||||
f"{token}:{meta.get('code')} ({meta.get('name')})"
|
||||
for token, meta in contains[:10]
|
||||
)
|
||||
raise ValueError(
|
||||
f"--disease={value!r} matched {len(contains)} diseases; use token_id or code. "
|
||||
f"First matches: {preview}"
|
||||
)
|
||||
|
||||
|
||||
def resolve_disease_tokens(
|
||||
value: str | None,
|
||||
metadata: dict[int, dict[str, Any]],
|
||||
) -> list[tuple[int, dict[str, Any]]]:
|
||||
if value is None or str(value).strip() == "":
|
||||
return [(token, metadata[token]) for token in sorted(metadata)]
|
||||
out: list[tuple[int, dict[str, Any]]] = []
|
||||
seen: set[int] = set()
|
||||
for part in str(value).split(","):
|
||||
token, meta = resolve_disease_token(part, metadata)
|
||||
if token not in seen:
|
||||
out.append((token, meta))
|
||||
seen.add(token)
|
||||
return out
|
||||
|
||||
|
||||
def death_distribution_parameters(
|
||||
model,
|
||||
hidden: torch.Tensor,
|
||||
*,
|
||||
dist_mode: str,
|
||||
eps: float = 1e-8,
|
||||
) -> tuple[str, torch.Tensor]:
|
||||
"""Return death distribution parameters with columns matching PARAMETER_VALUE_COLUMNS."""
|
||||
logits = model.calc_risk(hidden)
|
||||
death_idx = int(logits.shape[1]) - 1
|
||||
death_lambda = F.softplus(logits[:, death_idx]) + float(eps)
|
||||
|
||||
if dist_mode == "exponential":
|
||||
nan = torch.full_like(death_lambda, float("nan"))
|
||||
return "exponential", torch.stack([death_lambda, nan, nan], dim=1)
|
||||
|
||||
if dist_mode == "weibull":
|
||||
rho = model.calc_weibull_rho(hidden)[:, death_idx].to(dtype=death_lambda.dtype)
|
||||
elif dist_mode == "mixed":
|
||||
rho = model.calc_death_rho(hidden).to(dtype=death_lambda.dtype)
|
||||
else:
|
||||
raise ValueError(f"Unsupported dist_mode={dist_mode!r}")
|
||||
|
||||
shape = rho.clamp_min(float(eps))
|
||||
scale = torch.pow(death_lambda.clamp_min(float(eps)), -1.0 / shape)
|
||||
nan = torch.full_like(death_lambda, float("nan"))
|
||||
return "weibull", torch.stack([nan, scale, shape], dim=1)
|
||||
|
||||
|
||||
def parameter_pair_block(original: torch.Tensor, ablated: torch.Tensor) -> torch.Tensor:
|
||||
return torch.stack(
|
||||
[
|
||||
original[:, 0],
|
||||
ablated[:, 0],
|
||||
original[:, 1],
|
||||
ablated[:, 1],
|
||||
original[:, 2],
|
||||
ablated[:, 2],
|
||||
],
|
||||
dim=1,
|
||||
)
|
||||
|
||||
|
||||
def output_name_for_run(run_path: Path, eval_split: str, *, all_diseases: bool) -> Path:
|
||||
scope = "all_diseases" if all_diseases else "selected_diseases"
|
||||
return run_path / f"single_disease_mortality_parameters_{eval_split}_{scope}"
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Compute per-disease model attribution to mortality distribution parameters."
|
||||
)
|
||||
parser.add_argument("--run_path", type=str, required=True)
|
||||
parser.add_argument(
|
||||
"--disease",
|
||||
type=str,
|
||||
default=None,
|
||||
help=(
|
||||
"Optional disease token_id, ICD-10 code, exact name, unambiguous name "
|
||||
"substring, or comma-separated list. If omitted, scan all disease tokens."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output_path",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Output directory for compressed .npz shards.",
|
||||
)
|
||||
parser.add_argument("--organ_mapping_path", type=str, default="icd10_chapter_organ_mapping.csv")
|
||||
parser.add_argument("--eval_split", type=str, default=None)
|
||||
parser.add_argument("--dataset_subset_size", type=int, default=None)
|
||||
parser.add_argument("--train_eid_file", type=str, default=None)
|
||||
parser.add_argument("--val_eid_file", type=str, default=None)
|
||||
parser.add_argument("--test_eid_file", type=str, default=None)
|
||||
parser.add_argument("--landmark_start", type=float, default=40.0)
|
||||
parser.add_argument("--landmark_stop", type=float, default=80.0)
|
||||
parser.add_argument("--landmark_step", type=float, default=5.0)
|
||||
parser.add_argument("--min_history_events", type=int, default=None)
|
||||
parser.add_argument("--batch_size", type=int, default=None)
|
||||
parser.add_argument(
|
||||
"--attribution_batch_size",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Forward batch size for disease-token ablation queries.",
|
||||
)
|
||||
parser.add_argument("--num_workers", type=int, default=None)
|
||||
parser.add_argument("--device", type=str, default=None)
|
||||
parser.add_argument("--extra_info_types", type=str, default=None)
|
||||
parser.add_argument(
|
||||
"--shard_rows",
|
||||
type=int,
|
||||
default=200_000,
|
||||
help="Approximate number of detailed rows to buffer before writing one .npz shard.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
run_path = Path(args.run_path)
|
||||
config_path = run_path / "train_config.json"
|
||||
checkpoint_path = run_path / "best_model.pt"
|
||||
if not config_path.exists():
|
||||
raise FileNotFoundError(f"train_config.json not found: {config_path}")
|
||||
if not checkpoint_path.exists():
|
||||
raise FileNotFoundError(f"best_model.pt not found: {checkpoint_path}")
|
||||
|
||||
cfg = load_json_config(config_path)
|
||||
model_target_mode = str(cfg.get("model_target_mode", "next_token")).lower()
|
||||
if model_target_mode not in {"next_token", "all_future"}:
|
||||
raise ValueError(f"Unsupported model_target_mode: {model_target_mode!r}")
|
||||
|
||||
target_mode = str(cfg.get("target_mode", "uts"))
|
||||
attn_mask_mode = str(
|
||||
cfg.get("attn_mask_mode", "non_strict_time" if target_mode == "uts" else "target_aware")
|
||||
)
|
||||
readout_name = str(
|
||||
cfg.get("readout_name", "same_time_group_end" if target_mode == "uts" else "token")
|
||||
)
|
||||
readout_reduce = str(cfg.get("readout_reduce", "mean"))
|
||||
|
||||
dataset, subset_indices, eval_split, split_source = load_eval_sequence_dataset(args, cfg)
|
||||
validate_dataset_metadata(dataset, cfg)
|
||||
|
||||
metadata = load_disease_metadata(
|
||||
Path(args.organ_mapping_path),
|
||||
vocab_size=int(dataset.vocab_size),
|
||||
)
|
||||
scanned_disease_items = resolve_disease_tokens(args.disease, metadata)
|
||||
if not scanned_disease_items:
|
||||
raise ValueError("No diseases selected for attribution")
|
||||
scanned_disease_tokens = [token for token, _meta in scanned_disease_items]
|
||||
|
||||
landmark_ages = make_landmark_ages(
|
||||
float(args.landmark_start),
|
||||
float(args.landmark_stop),
|
||||
float(args.landmark_step),
|
||||
)
|
||||
|
||||
first_occurrence_by_token = build_first_occurrence_maps_for_landmarks(
|
||||
dataset,
|
||||
subset_indices,
|
||||
)
|
||||
death_idx = int(dataset.vocab_size) - 1
|
||||
landmark_dataset = LandmarkDataset(
|
||||
dataset=dataset,
|
||||
subset_indices=subset_indices,
|
||||
landmark_ages=landmark_ages,
|
||||
attn_mask_mode=attn_mask_mode,
|
||||
model_target_mode=model_target_mode,
|
||||
min_history_events=int(cfg_get(args, cfg, "min_history_events", 1)),
|
||||
first_occurrence_by_token=first_occurrence_by_token,
|
||||
death_token_ids=[death_idx],
|
||||
)
|
||||
|
||||
organ_groups, _organ_labels, token_to_group = load_organ_groups(
|
||||
Path(args.organ_mapping_path),
|
||||
vocab_size=int(dataset.vocab_size),
|
||||
)
|
||||
group_names = sorted(organ_groups)
|
||||
|
||||
state_dict = load_checkpoint_state_dict(checkpoint_path, map_location="cpu")
|
||||
dist_mode = resolve_dist_mode_for_checkpoint(str(cfg.get("dist_mode", "exponential")), state_dict)
|
||||
death_distribution_name = "exponential" if dist_mode == "exponential" else "weibull"
|
||||
cfg_model = dict(cfg)
|
||||
cfg_model["dist_mode"] = dist_mode
|
||||
device = resolve_eval_device(args.device)
|
||||
selected_token_mask = np.zeros(int(dataset.vocab_size), dtype=bool)
|
||||
selected_token_mask[np.asarray(scanned_disease_tokens, dtype=np.int64)] = True
|
||||
model = build_model_from_dataset(args, cfg_model, dataset).to(device)
|
||||
load_model_state(model, state_dict)
|
||||
model.eval()
|
||||
|
||||
batch_size = int(cfg_get(args, cfg, "batch_size", 128))
|
||||
attribution_batch_size = int(
|
||||
cfg_get(args, cfg, "attribution_batch_size", max(batch_size * 32, 4096))
|
||||
)
|
||||
if attribution_batch_size <= 0:
|
||||
raise ValueError("attribution_batch_size must be positive")
|
||||
if int(args.shard_rows) <= 0:
|
||||
raise ValueError("--shard_rows must be positive")
|
||||
|
||||
num_workers = int(cfg_get(args, cfg, "num_workers", 4))
|
||||
loader = DataLoader(
|
||||
IndexedLandmarkDataset(landmark_dataset),
|
||||
batch_size=batch_size,
|
||||
shuffle=False,
|
||||
collate_fn=collate_indexed_landmark_fn,
|
||||
num_workers=num_workers,
|
||||
pin_memory=device.type == "cuda",
|
||||
persistent_workers=num_workers > 0,
|
||||
prefetch_factor=2 if num_workers > 0 else None,
|
||||
)
|
||||
|
||||
output_path = (
|
||||
Path(args.output_path)
|
||||
if args.output_path
|
||||
else output_name_for_run(
|
||||
run_path,
|
||||
eval_split,
|
||||
all_diseases=args.disease is None or str(args.disease).strip() == "",
|
||||
)
|
||||
)
|
||||
output_dir = normalize_output_dir(output_path)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
print(f"Eval split: {eval_split}")
|
||||
print(f"Split source: {split_source}")
|
||||
print(f"Selected patients: {len(subset_indices)}")
|
||||
print(f"Landmark ages: {landmark_ages.tolist()}")
|
||||
print(f"Dist mode: {dist_mode}")
|
||||
print(f"Device: {device}")
|
||||
print(f"Death token: {death_idx}")
|
||||
if len(scanned_disease_items) == len(metadata):
|
||||
print(f"Diseases: all mapped diseases ({len(scanned_disease_items)})")
|
||||
else:
|
||||
preview = ", ".join(
|
||||
f"{token}:{meta.get('code')}" for token, meta in scanned_disease_items[:10]
|
||||
)
|
||||
print(f"Diseases: {len(scanned_disease_items)} selected ({preview})")
|
||||
print(f"Landmark rows: {len(landmark_dataset)}")
|
||||
print(f"Attribution batch size: {attribution_batch_size}")
|
||||
print(f"Output directory: {output_dir}")
|
||||
|
||||
written_rows = 0
|
||||
shard_index = 0
|
||||
shards: list[dict[str, Any]] = []
|
||||
row_base_cache: dict[int, dict[str, Any]] = {}
|
||||
result_row_idx_chunks: list[np.ndarray] = []
|
||||
result_disease_token_chunks: list[np.ndarray] = []
|
||||
result_value_chunks: list[np.ndarray] = []
|
||||
|
||||
def get_row_base(row_idx: int) -> dict[str, Any]:
|
||||
cached = row_base_cache.get(row_idx)
|
||||
if cached is not None:
|
||||
return cached
|
||||
|
||||
meta = landmark_dataset.rows[int(row_idx)]
|
||||
dataset_index = int(meta["dataset_index"])
|
||||
sample = dataset.samples[dataset_index]
|
||||
hist_tokens = np.asarray(meta["event_seq"], dtype=np.int64)
|
||||
unique_tokens, token_counts = np.unique(hist_tokens, return_counts=True)
|
||||
total_count, group_counts = historical_counts_by_group(
|
||||
hist_tokens,
|
||||
death_idx=death_idx,
|
||||
token_to_group=token_to_group,
|
||||
group_names=group_names,
|
||||
)
|
||||
cached = {
|
||||
"patient_id": int(meta["patient_id"]),
|
||||
"dataset_index": dataset_index,
|
||||
"eid": int(sample.get("eid", -1)),
|
||||
"sex": int(meta["sex"]),
|
||||
"landmark_age": float(meta["landmark_age"]),
|
||||
"followup_end_time": float(meta["followup_end_time"]),
|
||||
"history_disease_count": int(total_count),
|
||||
"_hist_tokens": hist_tokens,
|
||||
"_token_counts": {
|
||||
int(token): int(count)
|
||||
for token, count in zip(unique_tokens.tolist(), token_counts.tolist())
|
||||
},
|
||||
"_group_counts": group_counts,
|
||||
}
|
||||
row_base_cache[row_idx] = cached
|
||||
return cached
|
||||
|
||||
for batch in tqdm(loader, desc="Per-disease mortality attribution", dynamic_ncols=True):
|
||||
batch_dev = {
|
||||
k: (v.to(device, non_blocking=True) if isinstance(v, torch.Tensor) else v)
|
||||
for k, v in batch.items()
|
||||
}
|
||||
hidden = infer_landmark_hidden(
|
||||
model=model,
|
||||
batch=batch_dev,
|
||||
device=device,
|
||||
model_target_mode=model_target_mode,
|
||||
readout_name=readout_name,
|
||||
readout_reduce=readout_reduce,
|
||||
)
|
||||
with torch.no_grad():
|
||||
_death_distribution, original_params = death_distribution_parameters(
|
||||
model,
|
||||
hidden,
|
||||
dist_mode=dist_mode,
|
||||
)
|
||||
event_np = batch["event_seq"].numpy()
|
||||
valid_event = (event_np >= 0) & (event_np < int(dataset.vocab_size))
|
||||
selected_event = np.zeros_like(valid_event, dtype=bool)
|
||||
selected_event[valid_event] = selected_token_mask[event_np[valid_event]]
|
||||
pair_row_np, pair_pos_np = np.nonzero(selected_event)
|
||||
if pair_row_np.size == 0:
|
||||
continue
|
||||
pair_disease_np = event_np[pair_row_np, pair_pos_np].astype(np.int64, copy=False)
|
||||
|
||||
pair_offset = 0
|
||||
while pair_offset < int(pair_row_np.shape[0]):
|
||||
pair_stop = min(int(pair_row_np.shape[0]), pair_offset + int(attribution_batch_size))
|
||||
local_rows_np = pair_row_np[pair_offset:pair_stop].astype(np.int64, copy=False)
|
||||
disease_tokens_np = pair_disease_np[pair_offset:pair_stop]
|
||||
local_rows = torch.as_tensor(local_rows_np, dtype=torch.long, device=device)
|
||||
disease_token_ids = torch.as_tensor(disease_tokens_np, dtype=torch.long, device=device)
|
||||
ablated_chunk = build_disease_ablated_slice(
|
||||
batch=batch_dev,
|
||||
row_indices=local_rows,
|
||||
token_ids=disease_token_ids,
|
||||
)
|
||||
with torch.no_grad():
|
||||
ablated_hidden = infer_landmark_hidden(
|
||||
model=model,
|
||||
batch=ablated_chunk,
|
||||
device=device,
|
||||
model_target_mode=model_target_mode,
|
||||
readout_name=readout_name,
|
||||
readout_reduce=readout_reduce,
|
||||
)
|
||||
_ablated_distribution, ablated_params = death_distribution_parameters(
|
||||
model,
|
||||
ablated_hidden,
|
||||
dist_mode=dist_mode,
|
||||
)
|
||||
value_block = parameter_pair_block(
|
||||
original_params[local_rows],
|
||||
ablated_params,
|
||||
).detach().cpu().numpy()
|
||||
row_ids = batch["row_idx"][local_rows_np].numpy().astype(np.int64, copy=False)
|
||||
disease_tokens_list = disease_tokens_np
|
||||
result_row_idx_chunks.append(row_ids)
|
||||
result_disease_token_chunks.append(disease_tokens_list)
|
||||
result_value_chunks.append(value_block)
|
||||
pair_offset = pair_stop
|
||||
|
||||
if result_value_chunks:
|
||||
all_row_ids = np.concatenate(result_row_idx_chunks).astype(np.int64, copy=False)
|
||||
all_disease_tokens = np.concatenate(result_disease_token_chunks).astype(
|
||||
np.int64,
|
||||
copy=False,
|
||||
)
|
||||
all_values = np.concatenate(result_value_chunks, axis=0)
|
||||
|
||||
rows: list[dict[str, Any]] = []
|
||||
for i, (row_idx, disease_token) in enumerate(
|
||||
zip(all_row_ids.tolist(), all_disease_tokens.tolist())
|
||||
):
|
||||
disease_token = int(disease_token)
|
||||
disease_meta = metadata[disease_token]
|
||||
row_base = get_row_base(int(row_idx))
|
||||
group_counts = row_base["_group_counts"]
|
||||
disease_history_count = int(row_base["_token_counts"].get(disease_token, 0))
|
||||
if disease_history_count <= 0:
|
||||
raise RuntimeError(
|
||||
"Internal mismatch: occurred mask selected disease "
|
||||
f"{disease_token} for row {row_idx}, but cached history has count 0"
|
||||
)
|
||||
|
||||
rows.append(
|
||||
{
|
||||
"patient_id": row_base["patient_id"],
|
||||
"dataset_index": row_base["dataset_index"],
|
||||
"eid": row_base["eid"],
|
||||
"sex": row_base["sex"],
|
||||
"landmark_age": row_base["landmark_age"],
|
||||
"followup_end_time": row_base["followup_end_time"],
|
||||
"history_disease_count": row_base["history_disease_count"],
|
||||
"selected_disease_history_count": disease_history_count,
|
||||
"selected_disease_token_id": int(disease_token),
|
||||
"selected_disease_code": str(disease_meta.get("code", "")),
|
||||
"selected_disease_name": str(disease_meta.get("name", "")),
|
||||
"selected_disease_organ_system": str(disease_meta.get("organ_system", "")),
|
||||
"selected_disease_organ_system_label": str(
|
||||
disease_meta.get("organ_system_label", "")
|
||||
),
|
||||
"history_count__selected_organ_system": int(
|
||||
group_counts.get(str(disease_meta.get("organ_system", "")), 0)
|
||||
),
|
||||
"death_distribution": death_distribution_name,
|
||||
"original_death_lambda": float(all_values[i, 0]),
|
||||
"ablated_death_lambda": float(all_values[i, 1]),
|
||||
"original_death_scale": float(all_values[i, 2]),
|
||||
"ablated_death_scale": float(all_values[i, 3]),
|
||||
"original_death_shape": float(all_values[i, 4]),
|
||||
"ablated_death_shape": float(all_values[i, 5]),
|
||||
}
|
||||
)
|
||||
|
||||
result_table = pd.DataFrame(rows).reindex(columns=OUTPUT_COLUMNS)
|
||||
written_rows = int(len(result_table))
|
||||
|
||||
summary_accumulator: dict[tuple[Any, ...], dict[str, float]] = {}
|
||||
update_summary_accumulator(summary_accumulator, result_table)
|
||||
|
||||
for start in range(0, written_rows, int(args.shard_rows)):
|
||||
stop = min(written_rows, start + int(args.shard_rows))
|
||||
shard_name = f"part-{shard_index:06d}.npz"
|
||||
shard_path = output_dir / shard_name
|
||||
shard_rows = write_compressed_npz_table(
|
||||
shard_path,
|
||||
result_table.iloc[start:stop],
|
||||
)
|
||||
shards.append({"file": shard_name, "rows": int(shard_rows)})
|
||||
shard_index += 1
|
||||
else:
|
||||
result_table = pd.DataFrame(columns=OUTPUT_COLUMNS)
|
||||
summary_accumulator = {}
|
||||
|
||||
if not shards:
|
||||
empty_path = output_dir / "part-000000.npz"
|
||||
write_compressed_npz_table(empty_path, pd.DataFrame(columns=OUTPUT_COLUMNS))
|
||||
shards.append({"file": empty_path.name, "rows": 0})
|
||||
summary_path = output_dir / "summary_by_disease_age_sex.csv"
|
||||
summary_rows = write_summary_csv(summary_path, summary_accumulator)
|
||||
write_manifest(
|
||||
output_dir,
|
||||
rows=written_rows,
|
||||
shards=shards,
|
||||
summary_file=summary_path.name,
|
||||
scanned_diseases=[
|
||||
{"token_id": int(token), **{k: v for k, v in meta.items() if k != "token_id"}}
|
||||
for token, meta in scanned_disease_items
|
||||
],
|
||||
eval_split=eval_split,
|
||||
dist_mode=dist_mode,
|
||||
landmark_start=float(args.landmark_start),
|
||||
landmark_stop=float(args.landmark_stop),
|
||||
landmark_step=float(args.landmark_step),
|
||||
)
|
||||
print(f"Wrote {written_rows} rows in {len(shards)} shard(s) to {output_dir}")
|
||||
print(f"Wrote {summary_rows} summary rows to {summary_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,7 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from evaluate_auc import main
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,104 +0,0 @@
|
||||
"""Read and query calendar-dated disease-event arrays from prepare_event_dates.py."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Iterable
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
|
||||
|
||||
REQUIRED_FIELDS = {"eid", "event_date", "token"}
|
||||
|
||||
|
||||
def load_event_dates(path: str | Path) -> np.ndarray:
|
||||
"""Load and validate the structured ``.npy`` event array."""
|
||||
events = np.load(path)
|
||||
if events.dtype.names is None or not REQUIRED_FIELDS.issubset(events.dtype.names):
|
||||
raise ValueError(
|
||||
"Expected a structured .npy with eid, event_date, token fields. "
|
||||
"Create it with prepare_event_dates.py."
|
||||
)
|
||||
return events
|
||||
|
||||
|
||||
def load_token_labels(labels_file: str | Path) -> dict[int, str]:
|
||||
"""Load token -> human-readable code using the project label convention."""
|
||||
labels = {1: "CHECKUP"}
|
||||
with Path(labels_file).open(encoding="utf-8") as handle:
|
||||
for index, line in enumerate(handle):
|
||||
code = line.strip().split(" ", maxsplit=1)[0]
|
||||
if code:
|
||||
labels[index + 2] = code
|
||||
return labels
|
||||
|
||||
|
||||
@dataclass
|
||||
class EventDateIndex:
|
||||
"""Small in-memory query wrapper for exposure-linkage and cohort scripts."""
|
||||
|
||||
events: np.ndarray
|
||||
token_labels: dict[int, str] | None = None
|
||||
|
||||
@classmethod
|
||||
def from_files(
|
||||
cls,
|
||||
event_file: str | Path,
|
||||
labels_file: str | Path | None = None,
|
||||
) -> "EventDateIndex":
|
||||
labels = load_token_labels(labels_file) if labels_file is not None else None
|
||||
return cls(load_event_dates(event_file), labels)
|
||||
|
||||
def to_frame(self, events: np.ndarray | None = None) -> pd.DataFrame:
|
||||
"""Convert records to a convenient, calendar-dated DataFrame."""
|
||||
data = self.events if events is None else events
|
||||
frame = pd.DataFrame(
|
||||
{
|
||||
"eid": data["eid"].astype("int64"),
|
||||
"event_date": pd.to_datetime(data["event_date"]),
|
||||
"token": data["token"].astype("int32"),
|
||||
}
|
||||
)
|
||||
if self.token_labels is not None:
|
||||
frame["label_code"] = frame["token"].map(self.token_labels).fillna("UNKNOWN")
|
||||
return frame.sort_values(["eid", "event_date", "token"], kind="stable").reset_index(drop=True)
|
||||
|
||||
def for_eid(self, eid: int) -> pd.DataFrame:
|
||||
"""Return every stored disease/death event for one participant."""
|
||||
return self.to_frame(self.events[self.events["eid"] == int(eid)])
|
||||
|
||||
def between(
|
||||
self,
|
||||
start: str | pd.Timestamp,
|
||||
end: str | pd.Timestamp,
|
||||
*,
|
||||
eids: Iterable[int] | None = None,
|
||||
tokens: Iterable[int] | None = None,
|
||||
) -> pd.DataFrame:
|
||||
"""Query events in an inclusive calendar-date interval."""
|
||||
start_day = np.datetime64(pd.Timestamp(start).date(), "D")
|
||||
end_day = np.datetime64(pd.Timestamp(end).date(), "D")
|
||||
mask = (self.events["event_date"] >= start_day) & (self.events["event_date"] <= end_day)
|
||||
if eids is not None:
|
||||
mask &= np.isin(self.events["eid"], list(eids))
|
||||
if tokens is not None:
|
||||
mask &= np.isin(self.events["token"], list(tokens))
|
||||
return self.to_frame(self.events[mask])
|
||||
|
||||
def anchors_before(self, eid: int, date: str | pd.Timestamp) -> pd.DataFrame:
|
||||
"""Return a participant's event history strictly before an exposure anchor."""
|
||||
day = np.datetime64(pd.Timestamp(date).date(), "D")
|
||||
mask = (self.events["eid"] == int(eid)) & (self.events["event_date"] < day)
|
||||
return self.to_frame(self.events[mask])
|
||||
|
||||
def first_event(self, token: int) -> pd.DataFrame:
|
||||
"""Return each participant's first date for a requested token."""
|
||||
selected = self.events[self.events["token"] == int(token)]
|
||||
# Arrays produced by prepare_event_dates.py are already deduplicated;
|
||||
# sorting makes this safe for externally produced compatible arrays too.
|
||||
order = np.lexsort((selected["event_date"], selected["eid"]))
|
||||
selected = selected[order]
|
||||
_, first = np.unique(selected["eid"], return_index=True)
|
||||
return self.to_frame(selected[first])
|
||||
@@ -1,325 +0,0 @@
|
||||
"""Export landmark risk logits and hidden states for t_query ages.
|
||||
|
||||
This script follows evaluate_event_free_survival.py's data loading,
|
||||
landmark construction, checkpoint loading, and readout logic, but only exports:
|
||||
|
||||
* all token/disease risk logits from ``model.calc_risk(hidden)``;
|
||||
* the corresponding landmark hidden state.
|
||||
|
||||
The two large arrays are saved separately as .npy files. Row metadata is saved
|
||||
as a CSV with matching row order.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Any, Optional
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import torch
|
||||
from torch.utils.data import DataLoader
|
||||
from tqdm.auto import tqdm
|
||||
|
||||
from evaluate_auc_v2 import (
|
||||
LandmarkDataset,
|
||||
build_model_from_dataset,
|
||||
cfg_get,
|
||||
load_checkpoint_state_dict,
|
||||
load_json_config,
|
||||
load_model_state,
|
||||
resolve_dist_mode_for_checkpoint,
|
||||
resolve_eval_device,
|
||||
validate_dataset_metadata,
|
||||
)
|
||||
from evaluate_event_free_survival import (
|
||||
IndexedLandmarkDataset,
|
||||
build_first_occurrence_maps_for_landmarks,
|
||||
collate_indexed_landmark_fn,
|
||||
infer_landmark_hidden,
|
||||
load_eval_sequence_dataset,
|
||||
make_landmark_ages,
|
||||
)
|
||||
|
||||
|
||||
def numpy_float_dtype(name: str) -> np.dtype:
|
||||
key = str(name).lower()
|
||||
if key in {"float16", "fp16", "half"}:
|
||||
return np.dtype(np.float16)
|
||||
if key in {"float32", "fp32", "single"}:
|
||||
return np.dtype(np.float32)
|
||||
raise ValueError(f"dtype must be float16 or float32, got {name!r}")
|
||||
|
||||
|
||||
def output_paths_for_run(
|
||||
run_path: Path,
|
||||
eval_split: str,
|
||||
landmark_start: float,
|
||||
landmark_stop: float,
|
||||
landmark_step: float,
|
||||
output_dir: Optional[str],
|
||||
) -> tuple[Path, Path, Path, Path]:
|
||||
suffix = f"{eval_split}_t{landmark_start:g}-{landmark_stop:g}_step{landmark_step:g}"
|
||||
base_dir = Path(output_dir) if output_dir else run_path
|
||||
return (
|
||||
base_dir / f"tquery_logits_{suffix}.npy",
|
||||
base_dir / f"tquery_hidden_{suffix}.npy",
|
||||
base_dir / f"tquery_metadata_{suffix}.csv",
|
||||
base_dir / f"tquery_export_config_{suffix}.json",
|
||||
)
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Export landmark risk logits and hidden states for t_query ages."
|
||||
)
|
||||
parser.add_argument("--run_path", type=str, required=True)
|
||||
parser.add_argument(
|
||||
"--output_dir",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Directory for output files. Defaults to run_path.",
|
||||
)
|
||||
parser.add_argument("--logits_path", type=str, default=None)
|
||||
parser.add_argument("--hidden_path", type=str, default=None)
|
||||
parser.add_argument("--metadata_path", type=str, default=None)
|
||||
parser.add_argument("--export_config_path", type=str, default=None)
|
||||
parser.add_argument("--eval_split", type=str, default=None)
|
||||
parser.add_argument("--dataset_subset_size", type=int, default=None)
|
||||
parser.add_argument("--train_eid_file", type=str, default=None)
|
||||
parser.add_argument("--val_eid_file", type=str, default=None)
|
||||
parser.add_argument("--test_eid_file", type=str, default=None)
|
||||
parser.add_argument("--landmark_start", type=float, default=40.0)
|
||||
parser.add_argument("--landmark_stop", type=float, default=80.0)
|
||||
parser.add_argument(
|
||||
"--landmark_step",
|
||||
type=float,
|
||||
default=1.0,
|
||||
help="t_query grid step in years. Default exports every integer age 40..80.",
|
||||
)
|
||||
parser.add_argument("--min_history_events", type=int, default=None)
|
||||
parser.add_argument("--batch_size", type=int, default=None)
|
||||
parser.add_argument("--num_workers", type=int, default=None)
|
||||
parser.add_argument("--device", type=str, default=None)
|
||||
parser.add_argument("--extra_info_types", type=str, default=None)
|
||||
parser.add_argument(
|
||||
"--logits_dtype",
|
||||
type=str,
|
||||
default="float32",
|
||||
choices=["float16", "float32"],
|
||||
)
|
||||
parser.add_argument(
|
||||
"--hidden_dtype",
|
||||
type=str,
|
||||
default="float32",
|
||||
choices=["float16", "float32"],
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
run_path = Path(args.run_path)
|
||||
config_path = run_path / "train_config.json"
|
||||
checkpoint_path = run_path / "best_model.pt"
|
||||
if not config_path.exists():
|
||||
raise FileNotFoundError(f"train_config.json not found: {config_path}")
|
||||
if not checkpoint_path.exists():
|
||||
raise FileNotFoundError(f"best_model.pt not found: {checkpoint_path}")
|
||||
|
||||
cfg = load_json_config(config_path)
|
||||
model_target_mode = str(cfg.get("model_target_mode", "next_token")).lower()
|
||||
if model_target_mode not in {"next_token", "all_future"}:
|
||||
raise ValueError(f"Unsupported model_target_mode: {model_target_mode!r}")
|
||||
|
||||
target_mode = str(cfg.get("target_mode", "uts"))
|
||||
attn_mask_mode = str(
|
||||
cfg.get(
|
||||
"attn_mask_mode",
|
||||
"non_strict_time" if target_mode == "uts" else "target_aware",
|
||||
)
|
||||
)
|
||||
readout_name = str(
|
||||
cfg.get("readout_name", "same_time_group_end" if target_mode == "uts" else "token")
|
||||
)
|
||||
readout_reduce = str(cfg.get("readout_reduce", "mean"))
|
||||
|
||||
dataset, subset_indices, eval_split, split_source = load_eval_sequence_dataset(
|
||||
args,
|
||||
cfg,
|
||||
)
|
||||
validate_dataset_metadata(dataset, cfg)
|
||||
|
||||
landmark_ages = make_landmark_ages(
|
||||
float(args.landmark_start),
|
||||
float(args.landmark_stop),
|
||||
float(args.landmark_step),
|
||||
)
|
||||
first_occurrence_by_token = build_first_occurrence_maps_for_landmarks(
|
||||
dataset,
|
||||
subset_indices,
|
||||
)
|
||||
death_idx = int(dataset.vocab_size) - 1
|
||||
landmark_dataset = LandmarkDataset(
|
||||
dataset=dataset,
|
||||
subset_indices=subset_indices,
|
||||
landmark_ages=landmark_ages,
|
||||
attn_mask_mode=attn_mask_mode,
|
||||
model_target_mode=model_target_mode,
|
||||
min_history_events=int(cfg_get(args, cfg, "min_history_events", 1)),
|
||||
first_occurrence_by_token=first_occurrence_by_token,
|
||||
death_token_ids=[death_idx],
|
||||
)
|
||||
|
||||
state_dict = load_checkpoint_state_dict(checkpoint_path, map_location="cpu")
|
||||
dist_mode = resolve_dist_mode_for_checkpoint(
|
||||
str(cfg.get("dist_mode", "exponential")),
|
||||
state_dict,
|
||||
)
|
||||
cfg_model = dict(cfg)
|
||||
cfg_model["dist_mode"] = dist_mode
|
||||
device = resolve_eval_device(args.device)
|
||||
model = build_model_from_dataset(args, cfg_model, dataset).to(device)
|
||||
load_model_state(model, state_dict)
|
||||
model.eval()
|
||||
|
||||
default_logits_path, default_hidden_path, default_metadata_path, default_config_path = (
|
||||
output_paths_for_run(
|
||||
run_path=run_path,
|
||||
eval_split=eval_split,
|
||||
landmark_start=float(args.landmark_start),
|
||||
landmark_stop=float(args.landmark_stop),
|
||||
landmark_step=float(args.landmark_step),
|
||||
output_dir=args.output_dir,
|
||||
)
|
||||
)
|
||||
logits_path = Path(args.logits_path) if args.logits_path else default_logits_path
|
||||
hidden_path = Path(args.hidden_path) if args.hidden_path else default_hidden_path
|
||||
metadata_path = Path(args.metadata_path) if args.metadata_path else default_metadata_path
|
||||
export_config_path = (
|
||||
Path(args.export_config_path) if args.export_config_path else default_config_path
|
||||
)
|
||||
for path in (logits_path, hidden_path, metadata_path, export_config_path):
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
n_rows = len(landmark_dataset)
|
||||
vocab_size = int(dataset.vocab_size)
|
||||
hidden_dim = int(model.d_model)
|
||||
logits_dtype = numpy_float_dtype(args.logits_dtype)
|
||||
hidden_dtype = numpy_float_dtype(args.hidden_dtype)
|
||||
|
||||
logits_memmap = np.lib.format.open_memmap(
|
||||
logits_path,
|
||||
mode="w+",
|
||||
dtype=logits_dtype,
|
||||
shape=(n_rows, vocab_size),
|
||||
)
|
||||
hidden_memmap = np.lib.format.open_memmap(
|
||||
hidden_path,
|
||||
mode="w+",
|
||||
dtype=hidden_dtype,
|
||||
shape=(n_rows, hidden_dim),
|
||||
)
|
||||
|
||||
batch_size = int(cfg_get(args, cfg, "batch_size", 128))
|
||||
num_workers = int(cfg_get(args, cfg, "num_workers", 4))
|
||||
loader = DataLoader(
|
||||
IndexedLandmarkDataset(landmark_dataset),
|
||||
batch_size=batch_size,
|
||||
shuffle=False,
|
||||
collate_fn=collate_indexed_landmark_fn,
|
||||
num_workers=num_workers,
|
||||
pin_memory=device.type == "cuda",
|
||||
persistent_workers=num_workers > 0,
|
||||
prefetch_factor=2 if num_workers > 0 else None,
|
||||
)
|
||||
|
||||
print(f"Eval split: {eval_split}")
|
||||
print(f"Split source: {split_source}")
|
||||
print(f"Selected patients: {len(subset_indices)}")
|
||||
print(f"t_query ages: {landmark_ages.tolist()}")
|
||||
print(f"Dist mode: {dist_mode}")
|
||||
print(f"Device: {device}")
|
||||
print(f"Landmark rows: {n_rows}")
|
||||
print(f"Logits: {logits_path} shape={(n_rows, vocab_size)} dtype={logits_dtype}")
|
||||
print(f"Hidden: {hidden_path} shape={(n_rows, hidden_dim)} dtype={hidden_dtype}")
|
||||
print(f"Metadata: {metadata_path}")
|
||||
|
||||
meta_rows: list[dict[str, Any]] = []
|
||||
written = 0
|
||||
with torch.no_grad():
|
||||
for batch in tqdm(loader, desc="Export logits/hidden", dynamic_ncols=True):
|
||||
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)
|
||||
row_indices = batch["row_idx"].detach().cpu().numpy().astype(np.int64)
|
||||
if not np.array_equal(row_indices, np.arange(written, written + len(row_indices))):
|
||||
raise RuntimeError("DataLoader row order changed; export requires shuffle=False.")
|
||||
|
||||
batch_n = int(logits.shape[0])
|
||||
logits_memmap[written : written + batch_n] = (
|
||||
logits.detach().cpu().numpy().astype(logits_dtype, copy=False)
|
||||
)
|
||||
hidden_memmap[written : written + batch_n] = (
|
||||
hidden.detach().cpu().numpy().astype(hidden_dtype, copy=False)
|
||||
)
|
||||
|
||||
for row_idx in row_indices.tolist():
|
||||
meta = landmark_dataset.rows[int(row_idx)]
|
||||
sample = dataset.samples[int(meta["dataset_index"])]
|
||||
meta_rows.append(
|
||||
{
|
||||
"row_index": int(row_idx),
|
||||
"patient_id": int(meta["patient_id"]),
|
||||
"dataset_index": int(meta["dataset_index"]),
|
||||
"eid": int(sample.get("eid", -1)),
|
||||
"sex": int(meta["sex"]),
|
||||
"t_query": float(meta["t_query"]),
|
||||
"landmark_age": float(meta["landmark_age"]),
|
||||
"followup_end_time": float(meta["followup_end_time"]),
|
||||
"death_time": float(meta["death_time"]),
|
||||
}
|
||||
)
|
||||
written += batch_n
|
||||
|
||||
logits_memmap.flush()
|
||||
hidden_memmap.flush()
|
||||
pd.DataFrame(meta_rows).to_csv(metadata_path, index=False)
|
||||
|
||||
export_config = {
|
||||
"run_path": str(run_path),
|
||||
"eval_split": eval_split,
|
||||
"split_source": split_source,
|
||||
"model_target_mode": model_target_mode,
|
||||
"target_mode": target_mode,
|
||||
"attn_mask_mode": attn_mask_mode,
|
||||
"readout_name": readout_name,
|
||||
"readout_reduce": readout_reduce,
|
||||
"dist_mode": dist_mode,
|
||||
"landmark_ages": [float(x) for x in landmark_ages.tolist()],
|
||||
"n_rows": int(n_rows),
|
||||
"vocab_size": int(vocab_size),
|
||||
"hidden_dim": int(hidden_dim),
|
||||
"death_token": int(death_idx),
|
||||
"logits_path": str(logits_path),
|
||||
"hidden_path": str(hidden_path),
|
||||
"metadata_path": str(metadata_path),
|
||||
"logits_dtype": str(logits_dtype),
|
||||
"hidden_dtype": str(hidden_dtype),
|
||||
}
|
||||
with export_config_path.open("w", encoding="utf-8") as f:
|
||||
json.dump(export_config, f, indent=2)
|
||||
|
||||
print(f"Wrote {written} rows.")
|
||||
print(f"Wrote export config: {export_config_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,517 +0,0 @@
|
||||
"""Export Weibull shape-parameter statistics on the test split.
|
||||
|
||||
The script is intended for all_future checkpoints with dist_mode="weibull" or
|
||||
dist_mode="mixed". For full Weibull models it reads rho_head[Death]; for mixed
|
||||
models it reads rho_death_head. For full Weibull models it also exports disease
|
||||
token rho summaries, which are the main evidence for whether risk/hazard changes
|
||||
with horizon instead of following an exponential shape.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import contextlib
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Iterable, List, Optional
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import torch.multiprocessing as torch_mp
|
||||
from torch.utils.data import DataLoader
|
||||
from tqdm.auto import tqdm
|
||||
|
||||
from eval_data import load_sequence_eval_dataset
|
||||
from evaluate_auc_v2 import (
|
||||
LandmarkDataset,
|
||||
_build_first_occurrence_maps,
|
||||
_get_death_token_ids,
|
||||
build_model_from_dataset,
|
||||
cfg_get,
|
||||
collate_landmark_fn,
|
||||
load_checkpoint_state_dict,
|
||||
load_json_config,
|
||||
load_model_state,
|
||||
make_eval_indices,
|
||||
parse_float_list,
|
||||
parse_int_list,
|
||||
resolve_dist_mode_for_checkpoint,
|
||||
resolve_eval_device,
|
||||
validate_dataset_metadata,
|
||||
)
|
||||
|
||||
try:
|
||||
torch_mp.set_sharing_strategy("file_system")
|
||||
except RuntimeError:
|
||||
pass
|
||||
|
||||
|
||||
def quantile_summary(df: pd.DataFrame, group_cols: List[str], value_cols: List[str]) -> pd.DataFrame:
|
||||
probs = [0.01, 0.05, 0.25, 0.50, 0.75, 0.95, 0.99]
|
||||
rows: List[Dict[str, Any]] = []
|
||||
grouped = [((), df)] if not group_cols else df.groupby(group_cols, dropna=False)
|
||||
|
||||
for key, g in grouped:
|
||||
if not isinstance(key, tuple):
|
||||
key = (key,)
|
||||
base = {col: val for col, val in zip(group_cols, key)}
|
||||
base["n"] = int(len(g))
|
||||
for col in value_cols:
|
||||
x = pd.to_numeric(g[col], errors="coerce").to_numpy(dtype=np.float64)
|
||||
x = x[np.isfinite(x)]
|
||||
if x.size == 0:
|
||||
continue
|
||||
row = dict(base)
|
||||
row["variable"] = col
|
||||
row["mean"] = float(np.mean(x))
|
||||
row["std"] = float(np.std(x, ddof=1)) if x.size > 1 else 0.0
|
||||
row["min"] = float(np.min(x))
|
||||
row["max"] = float(np.max(x))
|
||||
for p in probs:
|
||||
row[f"p{int(p * 100):02d}"] = float(np.quantile(x, p))
|
||||
rows.append(row)
|
||||
return pd.DataFrame(rows)
|
||||
|
||||
|
||||
def load_labels_meta(path: Optional[str]) -> Optional[pd.DataFrame]:
|
||||
if path is None:
|
||||
return None
|
||||
fp = Path(path)
|
||||
if not fp.exists():
|
||||
return None
|
||||
return pd.read_csv(fp)
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def infer_landmark_hidden_local(
|
||||
model,
|
||||
loader: DataLoader,
|
||||
device: torch.device,
|
||||
use_amp: bool,
|
||||
hidden_cache_dtype: str,
|
||||
) -> tuple[np.ndarray, Dict[str, np.ndarray]]:
|
||||
"""Minimal all_future landmark hidden inference for parameter export."""
|
||||
out_dtype = np.float32 if str(hidden_cache_dtype).lower() == "float32" else np.float16
|
||||
hidden_parts: List[np.ndarray] = []
|
||||
arrays: Dict[str, List[np.ndarray]] = {
|
||||
"patient_id": [],
|
||||
"sex": [],
|
||||
"landmark_age": [],
|
||||
"followup_end_time": [],
|
||||
"death_time": [],
|
||||
}
|
||||
amp_enabled = bool(use_amp and device.type == "cuda")
|
||||
|
||||
for batch in tqdm(loader, desc="Landmark hidden", dynamic_ncols=True):
|
||||
batch_dev = {
|
||||
k: (v.to(device, non_blocking=True) if isinstance(v, torch.Tensor) else v)
|
||||
for k, v in batch.items()
|
||||
}
|
||||
amp_ctx = (
|
||||
torch.autocast(device_type=device.type, dtype=torch.float16)
|
||||
if amp_enabled
|
||||
else contextlib.nullcontext()
|
||||
)
|
||||
with amp_ctx:
|
||||
hidden = model(
|
||||
event_seq=batch_dev["event_seq"],
|
||||
time_seq=batch_dev["time_seq"],
|
||||
sex=batch_dev["sex"],
|
||||
padding_mask=batch_dev["padding_mask"],
|
||||
t_query=batch_dev["t_query"],
|
||||
other_type=batch_dev["other_type"],
|
||||
other_value=batch_dev["other_value"],
|
||||
other_value_kind=batch_dev["other_value_kind"],
|
||||
other_time=batch_dev["other_time"],
|
||||
target_mode="all_future",
|
||||
)
|
||||
hidden_parts.append(hidden.detach().cpu().numpy().astype(out_dtype, copy=False))
|
||||
for key in arrays:
|
||||
arrays[key].append(batch[key].cpu().numpy())
|
||||
|
||||
hidden_all = np.concatenate(hidden_parts, axis=0)
|
||||
row_arrays = {key: np.concatenate(parts, axis=0) for key, parts in arrays.items()}
|
||||
return hidden_all, row_arrays
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def project_death_params(
|
||||
model,
|
||||
hidden_all: np.ndarray,
|
||||
dist_mode: str,
|
||||
device: torch.device,
|
||||
batch_size: int,
|
||||
use_amp: bool,
|
||||
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
|
||||
death_idx = int(getattr(model, "death_idx", getattr(model, "vocab_size", hidden_all.shape[0]) - 1))
|
||||
if not hasattr(model, "vocab_size"):
|
||||
death_idx = int(model.risk_head.out_features - 1)
|
||||
|
||||
compute_dtype = torch.float16 if (device.type == "cuda" and use_amp) else torch.float32
|
||||
risk_w = model.risk_head.weight[death_idx: death_idx + 1].detach().to(device=device, dtype=compute_dtype)
|
||||
risk_b = None
|
||||
if model.risk_head.bias is not None:
|
||||
risk_b = model.risk_head.bias[death_idx: death_idx + 1].detach().to(device=device, dtype=compute_dtype)
|
||||
|
||||
if dist_mode == "weibull":
|
||||
rho_w = model.rho_head.weight[death_idx: death_idx + 1].detach().to(device=device, dtype=compute_dtype)
|
||||
rho_b = model.rho_head.bias[death_idx: death_idx + 1].detach().to(device=device, dtype=compute_dtype)
|
||||
elif dist_mode == "mixed":
|
||||
rho_w = model.rho_death_head.weight.detach().to(device=device, dtype=compute_dtype)
|
||||
rho_b = model.rho_death_head.bias.detach().to(device=device, dtype=compute_dtype)
|
||||
else:
|
||||
raise ValueError("Death Weibull parameter export requires dist_mode='weibull' or 'mixed'.")
|
||||
|
||||
logits_out: List[np.ndarray] = []
|
||||
rate_out: List[np.ndarray] = []
|
||||
rho_out: List[np.ndarray] = []
|
||||
|
||||
for start in tqdm(range(0, hidden_all.shape[0], batch_size), desc="Death eta/rho", dynamic_ncols=True):
|
||||
end = min(start + batch_size, hidden_all.shape[0])
|
||||
h = torch.from_numpy(hidden_all[start:end]).to(device=device, dtype=compute_dtype, non_blocking=True)
|
||||
logits = F.linear(h, risk_w, risk_b).squeeze(-1)
|
||||
rate = F.softplus(logits) + 1e-8
|
||||
rho = F.softplus(F.linear(h, rho_w, rho_b).squeeze(-1)) + 1e-6
|
||||
logits_out.append(logits.float().cpu().numpy())
|
||||
rate_out.append(rate.float().cpu().numpy())
|
||||
rho_out.append(rho.float().cpu().numpy())
|
||||
del h, logits, rate, rho
|
||||
|
||||
return (
|
||||
np.concatenate(logits_out).astype(np.float32, copy=False),
|
||||
np.concatenate(rate_out).astype(np.float32, copy=False),
|
||||
np.concatenate(rho_out).astype(np.float32, copy=False),
|
||||
)
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def export_all_token_rho_summary(
|
||||
model,
|
||||
hidden_all: np.ndarray,
|
||||
dataset,
|
||||
device: torch.device,
|
||||
output_dir: Path,
|
||||
token_chunk_size: int,
|
||||
row_batch_size: int,
|
||||
use_amp: bool,
|
||||
horizons: np.ndarray,
|
||||
) -> None:
|
||||
if not hasattr(model, "rho_head"):
|
||||
print("[INFO] Skipping all-token rho summary because this is not a full Weibull model.")
|
||||
return
|
||||
|
||||
special = {0, 1, 2}
|
||||
token_ids = [
|
||||
int(t)
|
||||
for t, code in dataset.label_id_to_code.items()
|
||||
if int(t) not in special and not str(code).startswith("<")
|
||||
]
|
||||
token_ids = sorted(set(token_ids))
|
||||
death_idx = int(getattr(model, "death_idx", getattr(model, "vocab_size", len(token_ids)) - 1))
|
||||
if not hasattr(model, "vocab_size"):
|
||||
death_idx = int(model.risk_head.out_features - 1)
|
||||
compute_dtype = torch.float16 if (device.type == "cuda" and use_amp) else torch.float32
|
||||
rows: List[Dict[str, Any]] = []
|
||||
|
||||
for chunk_start in tqdm(range(0, len(token_ids), token_chunk_size), desc="All-token rho chunks", dynamic_ncols=True):
|
||||
chunk = token_ids[chunk_start: chunk_start + token_chunk_size]
|
||||
w = model.rho_head.weight[chunk].detach().to(device=device, dtype=compute_dtype)
|
||||
b = model.rho_head.bias[chunk].detach().to(device=device, dtype=compute_dtype)
|
||||
vals_parts: List[np.ndarray] = []
|
||||
for row_start in range(0, hidden_all.shape[0], row_batch_size):
|
||||
row_end = min(row_start + row_batch_size, hidden_all.shape[0])
|
||||
h = torch.from_numpy(hidden_all[row_start:row_end]).to(device=device, dtype=compute_dtype, non_blocking=True)
|
||||
rho = F.softplus(F.linear(h, w, b)) + 1e-6
|
||||
vals_parts.append(rho.float().cpu().numpy())
|
||||
del h, rho
|
||||
vals = np.concatenate(vals_parts, axis=0)
|
||||
for j, token in enumerate(chunk):
|
||||
x = vals[:, j].astype(np.float64, copy=False)
|
||||
row = {
|
||||
"token": int(token),
|
||||
"label_code": dataset.label_id_to_code.get(int(token), ""),
|
||||
"endpoint_type": "death" if int(token) == int(death_idx) else "disease",
|
||||
"n_landmark_rows": int(x.size),
|
||||
"rho_mean": float(np.mean(x)),
|
||||
"rho_std": float(np.std(x, ddof=1)) if x.size > 1 else 0.0,
|
||||
"rho_minus_one_mean": float(np.mean(x - 1.0)),
|
||||
"frac_rho_gt_1": float(np.mean(x > 1.0)),
|
||||
"frac_rho_lt_1": float(np.mean(x < 1.0)),
|
||||
"frac_rho_gt_1_1": float(np.mean(x > 1.1)),
|
||||
"frac_rho_lt_0_9": float(np.mean(x < 0.9)),
|
||||
"rho_p01": float(np.quantile(x, 0.01)),
|
||||
"rho_p05": float(np.quantile(x, 0.05)),
|
||||
"rho_p25": float(np.quantile(x, 0.25)),
|
||||
"rho_p50": float(np.quantile(x, 0.50)),
|
||||
"rho_p75": float(np.quantile(x, 0.75)),
|
||||
"rho_p95": float(np.quantile(x, 0.95)),
|
||||
"rho_p99": float(np.quantile(x, 0.99)),
|
||||
}
|
||||
for horizon in horizons.tolist():
|
||||
h = float(horizon)
|
||||
if h <= 0:
|
||||
continue
|
||||
# Shape-only time scaling. For rho=1 this equals 1, i.e. an
|
||||
# exponential model with constant instantaneous hazard.
|
||||
inst_scale = np.power(h, x - 1.0)
|
||||
cumhaz_scale = np.power(h, x)
|
||||
row[f"instant_hazard_scale_h{h:g}y_vs_1y_mean"] = float(np.mean(inst_scale))
|
||||
row[f"instant_hazard_scale_h{h:g}y_vs_1y_p50"] = float(np.quantile(inst_scale, 0.50))
|
||||
row[f"cumhaz_scale_h{h:g}y_mean"] = float(np.mean(cumhaz_scale))
|
||||
row[f"cumhaz_scale_h{h:g}y_p50"] = float(np.quantile(cumhaz_scale, 0.50))
|
||||
rows.append(row)
|
||||
del vals, vals_parts
|
||||
|
||||
out = pd.DataFrame(rows)
|
||||
out.to_csv(output_dir / "all_token_weibull_shape_summary.csv", index=False)
|
||||
out[out["endpoint_type"] == "disease"].to_csv(
|
||||
output_dir / "disease_token_weibull_shape_summary.csv", index=False
|
||||
)
|
||||
out[out["endpoint_type"] == "death"].to_csv(
|
||||
output_dir / "death_token_weibull_shape_summary.csv", index=False
|
||||
)
|
||||
|
||||
disease = out[out["endpoint_type"] == "disease"].copy()
|
||||
if not disease.empty:
|
||||
pd.DataFrame([
|
||||
{
|
||||
"n_tokens": int(len(disease)),
|
||||
"rho_mean_across_tokens": float(disease["rho_mean"].mean()),
|
||||
"rho_median_across_tokens": float(disease["rho_p50"].median()),
|
||||
"tokens_with_mean_rho_gt_1": int((disease["rho_mean"] > 1.0).sum()),
|
||||
"tokens_with_mean_rho_lt_1": int((disease["rho_mean"] < 1.0).sum()),
|
||||
"frac_tokens_with_mean_rho_gt_1": float((disease["rho_mean"] > 1.0).mean()),
|
||||
"frac_tokens_with_mean_rho_lt_1": float((disease["rho_mean"] < 1.0).mean()),
|
||||
"tokens_with_mean_rho_gt_1_1": int((disease["rho_mean"] > 1.1).sum()),
|
||||
"tokens_with_mean_rho_lt_0_9": int((disease["rho_mean"] < 0.9).sum()),
|
||||
}
|
||||
]).to_csv(output_dir / "disease_weibull_shape_overall_summary.csv", index=False)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="Export test-split Weibull shape parameter statistics.")
|
||||
parser.add_argument("--run_path", type=str, required=True)
|
||||
parser.add_argument("--output_path", type=str, default=None)
|
||||
parser.add_argument("--eval_split", type=str, default="test", choices=["test", "val", "valid", "validation", "train", "all"])
|
||||
parser.add_argument("--landmark_start", type=float, default=None)
|
||||
parser.add_argument("--landmark_stop", type=float, default=None)
|
||||
parser.add_argument("--landmark_step", type=float, default=None)
|
||||
parser.add_argument("--horizons", type=str, default=None)
|
||||
parser.add_argument("--batch_size", type=int, default=None)
|
||||
parser.add_argument(
|
||||
"--num_workers",
|
||||
type=int,
|
||||
default=0,
|
||||
help=(
|
||||
"DataLoader workers. Default 0 avoids Linux multiprocessing "
|
||||
"'received 0 items of ancdata' failures on shared filesystems."
|
||||
),
|
||||
)
|
||||
parser.add_argument("--device", type=str, default=None)
|
||||
parser.add_argument("--use_amp", action=argparse.BooleanOptionalAction, default=None)
|
||||
parser.add_argument("--hidden_cache_dtype", type=str, default="float32", choices=["float16", "float32"])
|
||||
parser.add_argument(
|
||||
"--include_all_token_rho_summary",
|
||||
action=argparse.BooleanOptionalAction,
|
||||
default=True,
|
||||
help=(
|
||||
"For full Weibull models, export disease/death token rho summaries. "
|
||||
"Use --no-include_all_token_rho_summary to skip the heavier token projection."
|
||||
),
|
||||
)
|
||||
parser.add_argument("--token_chunk_size", type=int, default=32)
|
||||
parser.add_argument("--row_batch_size", type=int, default=512)
|
||||
args = parser.parse_args()
|
||||
|
||||
run_path = Path(args.run_path)
|
||||
config_path = run_path / "train_config.json"
|
||||
ckpt_path = run_path / "best_model.pt"
|
||||
if not config_path.exists():
|
||||
raise FileNotFoundError(config_path)
|
||||
if not ckpt_path.exists():
|
||||
raise FileNotFoundError(ckpt_path)
|
||||
|
||||
cfg = load_json_config(config_path)
|
||||
model_target_mode = str(cfg.get("model_target_mode", "next_token")).lower()
|
||||
if model_target_mode != "all_future":
|
||||
raise ValueError("This export is intended for all_future checkpoints.")
|
||||
|
||||
data_prefix = cfg.get("data_prefix", "ukb")
|
||||
labels_file = cfg.get("labels_file", "labels.csv")
|
||||
no_event_interval_years = cfg.get("no_event_interval_years", 5.0)
|
||||
include_no_event_in_uts_target = cfg.get("include_no_event_in_uts_target", False)
|
||||
|
||||
dataset = load_sequence_eval_dataset(
|
||||
model_target_mode=model_target_mode,
|
||||
data_prefix=data_prefix,
|
||||
labels_file=labels_file,
|
||||
no_event_interval_years=float(no_event_interval_years),
|
||||
include_no_event_in_uts_target=bool(include_no_event_in_uts_target),
|
||||
min_history_events=int(cfg.get("all_future_min_history_events", 1)),
|
||||
min_future_events=int(cfg.get("all_future_min_future_events", 1)),
|
||||
extra_info_types=parse_int_list(cfg.get("extra_info_types", None)),
|
||||
)
|
||||
validate_dataset_metadata(dataset, cfg)
|
||||
|
||||
subset_indices = make_eval_indices(dataset, args, cfg)
|
||||
first_occurrence_by_token, _, _, _ = _build_first_occurrence_maps(dataset, subset_indices)
|
||||
|
||||
landmark_start = float(cfg_get(args, cfg, "landmark_start", 40.0))
|
||||
landmark_stop = float(cfg_get(args, cfg, "landmark_stop", 80.0))
|
||||
landmark_step = float(cfg_get(args, cfg, "landmark_step", 5.0))
|
||||
landmark_ages = np.arange(landmark_start, landmark_stop, landmark_step, dtype=np.float32)
|
||||
if landmark_ages.size == 0:
|
||||
raise ValueError("No landmark ages produced.")
|
||||
|
||||
horizons = np.asarray(
|
||||
parse_float_list(cfg_get(args, cfg, "horizons", "1,5,10")) or [1.0, 5.0, 10.0],
|
||||
dtype=np.float32,
|
||||
)
|
||||
if horizons.size == 0:
|
||||
raise ValueError("No horizons provided.")
|
||||
|
||||
state_dict = load_checkpoint_state_dict(ckpt_path, map_location="cpu")
|
||||
dist_mode = resolve_dist_mode_for_checkpoint(str(cfg.get("dist_mode", "exponential")), state_dict)
|
||||
if dist_mode not in {"weibull", "mixed"}:
|
||||
raise ValueError(
|
||||
f"Resolved dist_mode={dist_mode!r}; expected 'weibull' or 'mixed' for Weibull shape export."
|
||||
)
|
||||
|
||||
cfg_model = dict(cfg)
|
||||
cfg_model["dist_mode"] = dist_mode
|
||||
device = resolve_eval_device(args.device)
|
||||
model = build_model_from_dataset(args, cfg_model, dataset).to(device)
|
||||
load_model_state(model, state_dict)
|
||||
model.eval()
|
||||
|
||||
death_token_ids = _get_death_token_ids(dataset, None)
|
||||
death_idx = int(death_token_ids[0])
|
||||
attn_mask_mode = str(cfg.get("attn_mask_mode", "target_aware"))
|
||||
landmark_dataset = LandmarkDataset(
|
||||
dataset=dataset,
|
||||
subset_indices=subset_indices,
|
||||
landmark_ages=landmark_ages,
|
||||
attn_mask_mode=attn_mask_mode,
|
||||
model_target_mode=model_target_mode,
|
||||
min_history_events=int(cfg_get(args, cfg, "min_history_events", 1)),
|
||||
first_occurrence_by_token=first_occurrence_by_token,
|
||||
death_token_ids=death_token_ids,
|
||||
)
|
||||
|
||||
batch_size = int(cfg_get(args, cfg, "batch_size", 128))
|
||||
num_workers = int(cfg_get(args, cfg, "num_workers", 0))
|
||||
loader_kwargs = {
|
||||
"batch_size": batch_size,
|
||||
"shuffle": False,
|
||||
"collate_fn": collate_landmark_fn,
|
||||
"num_workers": num_workers,
|
||||
"pin_memory": device.type == "cuda",
|
||||
}
|
||||
if num_workers > 0:
|
||||
loader_kwargs["persistent_workers"] = True
|
||||
loader_kwargs["prefetch_factor"] = 2
|
||||
loader = DataLoader(landmark_dataset, **loader_kwargs)
|
||||
|
||||
use_amp = bool(cfg_get(args, cfg, "use_amp", False))
|
||||
hidden_all, row_arrays = infer_landmark_hidden_local(
|
||||
model=model,
|
||||
loader=loader,
|
||||
device=device,
|
||||
use_amp=use_amp,
|
||||
hidden_cache_dtype=str(args.hidden_cache_dtype),
|
||||
)
|
||||
eta, rate, rho = project_death_params(
|
||||
model=model,
|
||||
hidden_all=hidden_all,
|
||||
dist_mode=dist_mode,
|
||||
device=device,
|
||||
batch_size=int(args.row_batch_size),
|
||||
use_amp=use_amp,
|
||||
)
|
||||
|
||||
rows = pd.DataFrame({
|
||||
"patient_id": row_arrays["patient_id"].astype(np.int64),
|
||||
"sex": row_arrays["sex"].astype(np.int64),
|
||||
"sex_label": np.where(row_arrays["sex"].astype(np.int64) == 0, "female", "male"),
|
||||
"landmark_age": row_arrays["landmark_age"].astype(np.float32),
|
||||
"followup_end_time": row_arrays["followup_end_time"].astype(np.float32),
|
||||
"death_time": row_arrays["death_time"].astype(np.float32),
|
||||
"death_eta": eta,
|
||||
"death_rate": rate,
|
||||
"death_rho": rho,
|
||||
})
|
||||
for horizon in horizons.tolist():
|
||||
h = float(horizon)
|
||||
cumulative_hazard = rows["death_rate"].to_numpy(dtype=np.float64) * np.power(h, rows["death_rho"].to_numpy(dtype=np.float64))
|
||||
rows[f"death_cumhaz_h{h:g}y"] = cumulative_hazard
|
||||
rows[f"death_risk_h{h:g}y"] = -np.expm1(-cumulative_hazard)
|
||||
rows[f"death_observed_h{h:g}y"] = (
|
||||
(rows["death_time"].to_numpy(dtype=np.float64) > rows["landmark_age"].to_numpy(dtype=np.float64))
|
||||
& (rows["death_time"].to_numpy(dtype=np.float64) <= rows["landmark_age"].to_numpy(dtype=np.float64) + h)
|
||||
).astype(np.int8)
|
||||
|
||||
output_dir = Path(args.output_path) if args.output_path else run_path / "weibull_death_parameter_stats_test"
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
rows.to_csv(output_dir / "death_weibull_parameters_by_landmark.csv", index=False)
|
||||
value_cols = ["death_eta", "death_rate", "death_rho"]
|
||||
for horizon in horizons.tolist():
|
||||
h = float(horizon)
|
||||
value_cols.extend([f"death_cumhaz_h{h:g}y", f"death_risk_h{h:g}y"])
|
||||
|
||||
quantile_summary(rows, [], value_cols).to_csv(output_dir / "death_weibull_parameter_summary_overall.csv", index=False)
|
||||
quantile_summary(rows, ["landmark_age"], value_cols).to_csv(output_dir / "death_weibull_parameter_summary_by_landmark_age.csv", index=False)
|
||||
quantile_summary(rows, ["sex_label"], value_cols).to_csv(output_dir / "death_weibull_parameter_summary_by_sex.csv", index=False)
|
||||
quantile_summary(rows, ["sex_label", "landmark_age"], value_cols).to_csv(output_dir / "death_weibull_parameter_summary_by_sex_landmark_age.csv", index=False)
|
||||
|
||||
metadata = {
|
||||
"run_path": str(run_path),
|
||||
"config_path": str(config_path),
|
||||
"checkpoint_path": str(ckpt_path),
|
||||
"eval_split": str(args.eval_split),
|
||||
"model_target_mode": model_target_mode,
|
||||
"time_mode": str(cfg.get("time_mode")),
|
||||
"dist_mode_config": str(cfg.get("dist_mode")),
|
||||
"dist_mode_resolved": dist_mode,
|
||||
"extra_info_types": cfg.get("extra_info_types"),
|
||||
"death_token_id": death_idx,
|
||||
"death_label_code": dataset.label_id_to_code.get(death_idx, "Death"),
|
||||
"n_selected_patients": int(len(subset_indices)),
|
||||
"n_landmark_rows": int(len(rows)),
|
||||
"landmark_ages": [float(x) for x in landmark_ages.tolist()],
|
||||
"horizons": [float(x) for x in horizons.tolist()],
|
||||
}
|
||||
with (output_dir / "metadata.json").open("w", encoding="utf-8") as f:
|
||||
json.dump(metadata, f, indent=2)
|
||||
|
||||
if args.include_all_token_rho_summary and dist_mode == "weibull":
|
||||
export_all_token_rho_summary(
|
||||
model=model,
|
||||
hidden_all=hidden_all,
|
||||
dataset=dataset,
|
||||
device=device,
|
||||
output_dir=output_dir,
|
||||
token_chunk_size=int(args.token_chunk_size),
|
||||
row_batch_size=int(args.row_batch_size),
|
||||
use_amp=use_amp,
|
||||
horizons=horizons,
|
||||
)
|
||||
elif dist_mode == "mixed":
|
||||
pd.DataFrame([
|
||||
{
|
||||
"dist_mode": dist_mode,
|
||||
"disease_shape_available": False,
|
||||
"death_shape_available": True,
|
||||
"note": (
|
||||
"The mixed model uses Weibull rho only for Death. "
|
||||
"Non-death disease hazards are exponential, equivalent to fixed rho=1."
|
||||
),
|
||||
}
|
||||
]).to_csv(output_dir / "disease_shape_not_available_for_mixed_model.csv", index=False)
|
||||
|
||||
print(f"Wrote Weibull shape parameter statistics to: {output_dir}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
797
export_weibull_parameters.py
Normal file
797
export_weibull_parameters.py
Normal file
@@ -0,0 +1,797 @@
|
||||
"""Export individual Weibull parameters to one HDF5 file.
|
||||
|
||||
The test population is read from ``test_eid_file`` in the specified run's
|
||||
``train_config.json``. Each eligible patient is queried at ages
|
||||
40, 42, ..., 80 using only disease history observed by that age.
|
||||
|
||||
The model parameterization is
|
||||
|
||||
S(t) = exp(-rate * t ** shape)
|
||||
|
||||
and the standard Weibull parameters exported by this script are
|
||||
|
||||
shape = rho
|
||||
scale = rate ** (-1 / shape)
|
||||
|
||||
The output is one HDF5 file. It contains the token table, age summary and one
|
||||
group per landmark age. Each age group stores aligned patient metadata plus
|
||||
chunked, compressed ``shape`` and ``scale`` matrices. Chunking is internal to
|
||||
HDF5, so callers receive one file without loading the full multi-gigabyte
|
||||
export into memory.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import contextlib
|
||||
import importlib
|
||||
import json
|
||||
import math
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Sequence
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch.utils.data import DataLoader
|
||||
from tqdm.auto import tqdm
|
||||
|
||||
from dataset import (
|
||||
DISEASE_HISTORY_MODE_TIMED,
|
||||
NO_EVENT_IDX,
|
||||
PAD_IDX,
|
||||
RESERVED_IDX,
|
||||
normalize_disease_history_mode,
|
||||
)
|
||||
from eval_data import (
|
||||
build_model_from_dataset,
|
||||
load_json_config,
|
||||
load_sequence_eval_dataset,
|
||||
resolve_eval_device,
|
||||
select_indices_by_eid_file,
|
||||
validate_dataset_metadata,
|
||||
validate_training_mode_config,
|
||||
)
|
||||
from evaluate_auc_v2 import (
|
||||
LandmarkDataset,
|
||||
collate_landmark_fn,
|
||||
load_checkpoint_state_dict,
|
||||
load_model_state,
|
||||
resolve_dist_mode_for_checkpoint,
|
||||
)
|
||||
from model_architectures import resolve_model_architecture
|
||||
|
||||
|
||||
SPECIAL_TOKENS = {PAD_IDX, RESERVED_IDX, NO_EVENT_IDX}
|
||||
FORMAT_VERSION = 2
|
||||
DEFAULT_COMPRESSION_LEVEL = 4
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description=(
|
||||
"Export each test patient's complete disease/death Weibull "
|
||||
"shape and scale matrices at two-year age landmarks."
|
||||
)
|
||||
)
|
||||
parser.add_argument(
|
||||
"--run_path",
|
||||
required=True,
|
||||
help="Run directory containing train_config.json and best_model.pt.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output_path",
|
||||
default=None,
|
||||
help=(
|
||||
"Single output HDF5 file. Defaults to "
|
||||
"<run_path>/weibull_parameters_test_age40_80_step2.h5."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--test_eid_file",
|
||||
default=None,
|
||||
help=(
|
||||
"Optional override. By default use test_eid_file from the run's "
|
||||
"train_config.json, or ukb_test_eid.csv if the field is absent."
|
||||
),
|
||||
)
|
||||
parser.add_argument("--age_start", type=float, default=40.0)
|
||||
parser.add_argument("--age_stop", type=float, default=80.0)
|
||||
parser.add_argument("--age_step", type=float, default=2.0)
|
||||
parser.add_argument("--batch_size", type=int, default=128)
|
||||
parser.add_argument(
|
||||
"--rows_per_chunk",
|
||||
type=int,
|
||||
default=256,
|
||||
help=(
|
||||
"Patient rows per internal HDF5 chunk for shape/scale matrices. "
|
||||
"This does not create separate output files."
|
||||
),
|
||||
)
|
||||
parser.add_argument("--num_workers", type=int, default=4)
|
||||
parser.add_argument(
|
||||
"--device",
|
||||
default=None,
|
||||
help="For example cpu, cuda, or cuda:1. Defaults to CUDA when available.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--use_amp",
|
||||
action=argparse.BooleanOptionalAction,
|
||||
default=False,
|
||||
help="Use float16 autocast during CUDA inference.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dataset_subset_size",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Use only the first N matched test patients for a smoke test.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def build_age_grid(start: float, stop: float, step: float) -> np.ndarray:
|
||||
"""Build an inclusive age grid and reject a stop not aligned to the step."""
|
||||
if not all(math.isfinite(x) for x in (start, stop, step)):
|
||||
raise ValueError("Age start/stop/step must be finite.")
|
||||
if step <= 0:
|
||||
raise ValueError("age_step must be > 0.")
|
||||
if stop < start:
|
||||
raise ValueError("age_stop must be >= age_start.")
|
||||
count = int(math.floor((stop - start) / step + 1e-9)) + 1
|
||||
ages = start + step * np.arange(count, dtype=np.float64)
|
||||
if ages.size == 0 or not math.isclose(
|
||||
float(ages[-1]), stop, rel_tol=0.0, abs_tol=1e-7
|
||||
):
|
||||
raise ValueError(
|
||||
"age_stop must lie on the grid defined by age_start and age_step."
|
||||
)
|
||||
ages[-1] = stop
|
||||
return ages.astype(np.float32)
|
||||
|
||||
|
||||
def parse_int_list(value: Any) -> List[int] | None:
|
||||
if value is None:
|
||||
return None
|
||||
if isinstance(value, (list, tuple, np.ndarray)):
|
||||
return [int(x) for x in value]
|
||||
text = str(value).strip()
|
||||
if not text:
|
||||
return None
|
||||
if text.startswith("["):
|
||||
parsed = json.loads(text)
|
||||
if not isinstance(parsed, list):
|
||||
raise ValueError("extra_info_types must be a list of integers.")
|
||||
return [int(x) for x in parsed]
|
||||
return [int(x.strip()) for x in text.split(",") if x.strip()]
|
||||
|
||||
|
||||
def select_outcome_tokens(dataset: Any) -> List[int]:
|
||||
"""Select every real outcome token, including Death, in token-id order."""
|
||||
tokens = sorted(
|
||||
int(token)
|
||||
for token, code in dataset.label_id_to_code.items()
|
||||
if int(token) not in SPECIAL_TOKENS and not str(code).startswith("<")
|
||||
)
|
||||
if not tokens:
|
||||
raise RuntimeError("The dataset contains no disease/death outcome tokens.")
|
||||
return tokens
|
||||
|
||||
|
||||
def resolve_project_file(path_value: str | Path) -> Path:
|
||||
path = Path(path_value)
|
||||
if path.is_absolute():
|
||||
return path
|
||||
direct = Path.cwd() / path
|
||||
return direct if direct.is_file() else Path(__file__).resolve().parent / path
|
||||
|
||||
|
||||
def load_label_text(labels_file: str | Path) -> Dict[str, str]:
|
||||
result: Dict[str, str] = {}
|
||||
with resolve_project_file(labels_file).open("r", encoding="utf-8") as handle:
|
||||
for line in handle:
|
||||
text = line.strip()
|
||||
if text:
|
||||
result[text.split()[0]] = text
|
||||
return result
|
||||
|
||||
|
||||
def require_h5py() -> Any:
|
||||
"""Import h5py lazily so ``--help`` remains available without it."""
|
||||
try:
|
||||
return importlib.import_module("h5py")
|
||||
except ImportError as exc:
|
||||
raise RuntimeError(
|
||||
"This exporter writes one HDF5 file and requires h5py. Install "
|
||||
"h5py in the same Python environment used for DeepHealth."
|
||||
) from exc
|
||||
|
||||
|
||||
def age_group_name(age: float) -> str:
|
||||
text = f"{age:g}".replace("-", "minus_").replace(".", "p")
|
||||
return f"age_{text}"
|
||||
|
||||
|
||||
def write_metadata_json(
|
||||
metadata_dataset: Any,
|
||||
payload: Dict[str, Any],
|
||||
) -> None:
|
||||
metadata_dataset[()] = json.dumps(
|
||||
payload,
|
||||
ensure_ascii=False,
|
||||
separators=(",", ":"),
|
||||
)
|
||||
|
||||
|
||||
def write_empty_age_group(
|
||||
*,
|
||||
age: float,
|
||||
age_group: Any,
|
||||
token_count: int,
|
||||
) -> Dict[str, Any]:
|
||||
"""Represent an age with no eligible patients inside the HDF5 file."""
|
||||
age_group.attrs["age"] = float(age)
|
||||
age_group.attrs["n_rows"] = 0
|
||||
age_group.attrs["n_tokens"] = int(token_count)
|
||||
age_group.attrs["nonfinite_shape_values"] = 0
|
||||
age_group.attrs["nonfinite_scale_values"] = 0
|
||||
age_group.create_dataset("eid", shape=(0,), dtype=np.int64)
|
||||
age_group.create_dataset("dataset_index", shape=(0,), dtype=np.int64)
|
||||
age_group.create_dataset("sex", shape=(0,), dtype=np.int8)
|
||||
age_group.create_dataset("age", shape=(0,), dtype=np.float32)
|
||||
age_group.create_dataset(
|
||||
"shape", shape=(0, token_count), dtype=np.float32
|
||||
)
|
||||
age_group.create_dataset(
|
||||
"scale", shape=(0, token_count), dtype=np.float32
|
||||
)
|
||||
return {
|
||||
"age": float(age),
|
||||
"n_rows": 0,
|
||||
"n_tokens": int(token_count),
|
||||
"nonfinite_shape_values": 0,
|
||||
"nonfinite_scale_values": 0,
|
||||
}
|
||||
|
||||
|
||||
def validate_hdf5_export(
|
||||
output_file: Any,
|
||||
*,
|
||||
ages: np.ndarray,
|
||||
token_count: int,
|
||||
summaries: Sequence[Dict[str, Any]],
|
||||
) -> None:
|
||||
"""Validate the unified file structure before it is finalized."""
|
||||
for group_name in (
|
||||
"tokens",
|
||||
"test_population",
|
||||
"landmarks",
|
||||
"age_summary",
|
||||
):
|
||||
if group_name not in output_file:
|
||||
raise RuntimeError(f"HDF5 export is missing /{group_name}.")
|
||||
if output_file["tokens/token_id"].shape != (token_count,):
|
||||
raise RuntimeError("HDF5 token table length does not match n_tokens.")
|
||||
if len(summaries) != int(ages.size):
|
||||
raise RuntimeError("HDF5 age summary length does not match age grid.")
|
||||
|
||||
for age_value, summary in zip(ages.tolist(), summaries):
|
||||
age = float(age_value)
|
||||
group_path = f"landmarks/{age_group_name(age)}"
|
||||
if group_path not in output_file:
|
||||
raise RuntimeError(f"HDF5 export is missing /{group_path}.")
|
||||
group = output_file[group_path]
|
||||
n_rows = int(summary["n_rows"])
|
||||
matrix_shape = (n_rows, token_count)
|
||||
if group["shape"].shape != matrix_shape:
|
||||
raise RuntimeError(
|
||||
f"/{group_path}/shape has {group['shape'].shape}, expected "
|
||||
f"{matrix_shape}."
|
||||
)
|
||||
if group["scale"].shape != matrix_shape:
|
||||
raise RuntimeError(
|
||||
f"/{group_path}/scale has {group['scale'].shape}, expected "
|
||||
f"{matrix_shape}."
|
||||
)
|
||||
for dataset_name in ("eid", "dataset_index", "sex", "age"):
|
||||
if group[dataset_name].shape != (n_rows,):
|
||||
raise RuntimeError(
|
||||
f"/{group_path}/{dataset_name} length does not match "
|
||||
"the parameter matrices."
|
||||
)
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def export_age(
|
||||
*,
|
||||
age: float,
|
||||
model: Any,
|
||||
loader: DataLoader,
|
||||
tokens: Sequence[int],
|
||||
device: torch.device,
|
||||
use_amp: bool,
|
||||
subset_indices: np.ndarray,
|
||||
selected_eids: np.ndarray,
|
||||
age_group: Any,
|
||||
rows_per_chunk: int,
|
||||
) -> Dict[str, Any]:
|
||||
"""Write one age directly into its group in the unified HDF5 file."""
|
||||
token_index = torch.as_tensor(tokens, dtype=torch.long, device=device)
|
||||
amp_enabled = bool(use_amp and device.type == "cuda")
|
||||
n_rows_expected = int(len(loader.dataset))
|
||||
n_tokens = int(len(tokens))
|
||||
age_group.attrs["age"] = float(age)
|
||||
age_group.attrs["n_rows"] = n_rows_expected
|
||||
age_group.attrs["n_tokens"] = n_tokens
|
||||
|
||||
if n_rows_expected == 0:
|
||||
return write_empty_age_group(
|
||||
age=age,
|
||||
age_group=age_group,
|
||||
token_count=n_tokens,
|
||||
)
|
||||
|
||||
row_chunk = min(int(rows_per_chunk), n_rows_expected)
|
||||
vector_options = {
|
||||
"chunks": (row_chunk,),
|
||||
"compression": "gzip",
|
||||
"compression_opts": DEFAULT_COMPRESSION_LEVEL,
|
||||
"shuffle": True,
|
||||
}
|
||||
matrix_options = {
|
||||
"chunks": (row_chunk, n_tokens),
|
||||
"compression": "gzip",
|
||||
"compression_opts": DEFAULT_COMPRESSION_LEVEL,
|
||||
"shuffle": True,
|
||||
}
|
||||
eid_dataset = age_group.create_dataset(
|
||||
"eid", shape=(n_rows_expected,), dtype=np.int64, **vector_options
|
||||
)
|
||||
dataset_index_dataset = age_group.create_dataset(
|
||||
"dataset_index",
|
||||
shape=(n_rows_expected,),
|
||||
dtype=np.int64,
|
||||
**vector_options,
|
||||
)
|
||||
sex_dataset = age_group.create_dataset(
|
||||
"sex", shape=(n_rows_expected,), dtype=np.int8, **vector_options
|
||||
)
|
||||
age_dataset = age_group.create_dataset(
|
||||
"age", shape=(n_rows_expected,), dtype=np.float32, **vector_options
|
||||
)
|
||||
shape_dataset = age_group.create_dataset(
|
||||
"shape",
|
||||
shape=(n_rows_expected, n_tokens),
|
||||
dtype=np.float32,
|
||||
**matrix_options,
|
||||
)
|
||||
scale_dataset = age_group.create_dataset(
|
||||
"scale",
|
||||
shape=(n_rows_expected, n_tokens),
|
||||
dtype=np.float32,
|
||||
**matrix_options,
|
||||
)
|
||||
|
||||
row_start = 0
|
||||
nonfinite_shape = 0
|
||||
nonfinite_scale = 0
|
||||
|
||||
for batch in tqdm(loader, desc=f"Age {age:g}", dynamic_ncols=True):
|
||||
batch_device = {
|
||||
key: (
|
||||
value.to(device, non_blocking=True)
|
||||
if isinstance(value, torch.Tensor)
|
||||
else value
|
||||
)
|
||||
for key, value in batch.items()
|
||||
}
|
||||
amp_context = (
|
||||
torch.autocast(device_type="cuda", dtype=torch.float16)
|
||||
if amp_enabled
|
||||
else contextlib.nullcontext()
|
||||
)
|
||||
with amp_context:
|
||||
hidden = model(
|
||||
event_seq=batch_device["event_seq"],
|
||||
time_seq=batch_device["time_seq"],
|
||||
sex=batch_device["sex"],
|
||||
padding_mask=batch_device["padding_mask"],
|
||||
t_query=batch_device["t_query"],
|
||||
other_type=batch_device["other_type"],
|
||||
other_value=batch_device["other_value"],
|
||||
other_value_kind=batch_device["other_value_kind"],
|
||||
other_time=batch_device["other_time"],
|
||||
)
|
||||
logits = model.calc_risk(hidden).index_select(1, token_index).float()
|
||||
shape = (
|
||||
model.calc_weibull_rho(hidden)
|
||||
.index_select(1, token_index)
|
||||
.float()
|
||||
)
|
||||
rate = F.softplus(logits) + 1e-8
|
||||
scale = torch.exp(-torch.log(rate) / shape)
|
||||
|
||||
shape_np = shape.cpu().numpy().astype(np.float32, copy=False)
|
||||
scale_np = scale.cpu().numpy().astype(np.float32, copy=False)
|
||||
patient_id = batch["patient_id"].cpu().numpy().astype(np.int64)
|
||||
dataset_index = subset_indices[patient_id]
|
||||
n_rows = int(shape_np.shape[0])
|
||||
row_stop = row_start + n_rows
|
||||
if row_stop > n_rows_expected:
|
||||
raise RuntimeError(
|
||||
f"Age {age:g} produced more rows than expected: "
|
||||
f"{row_stop} > {n_rows_expected}."
|
||||
)
|
||||
eid_dataset[row_start:row_stop] = selected_eids[patient_id]
|
||||
dataset_index_dataset[row_start:row_stop] = dataset_index.astype(
|
||||
np.int64, copy=False
|
||||
)
|
||||
sex_dataset[row_start:row_stop] = (
|
||||
batch["sex"].cpu().numpy().astype(np.int8, copy=False)
|
||||
)
|
||||
age_dataset[row_start:row_stop] = (
|
||||
batch["landmark_age"]
|
||||
.cpu()
|
||||
.numpy()
|
||||
.astype(np.float32, copy=False)
|
||||
)
|
||||
shape_dataset[row_start:row_stop, :] = shape_np
|
||||
scale_dataset[row_start:row_stop, :] = scale_np
|
||||
nonfinite_shape += int((~np.isfinite(shape_np)).sum())
|
||||
nonfinite_scale += int((~np.isfinite(scale_np)).sum())
|
||||
row_start = row_stop
|
||||
|
||||
if row_start != n_rows_expected:
|
||||
raise RuntimeError(
|
||||
f"Age {age:g} row count mismatch: wrote {row_start}, "
|
||||
f"expected {n_rows_expected}."
|
||||
)
|
||||
age_group.attrs["nonfinite_shape_values"] = nonfinite_shape
|
||||
age_group.attrs["nonfinite_scale_values"] = nonfinite_scale
|
||||
return {
|
||||
"age": float(age),
|
||||
"n_rows": n_rows_expected,
|
||||
"n_tokens": n_tokens,
|
||||
"nonfinite_shape_values": nonfinite_shape,
|
||||
"nonfinite_scale_values": nonfinite_scale,
|
||||
}
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
run_path = Path(args.run_path).resolve()
|
||||
config_path = run_path / "train_config.json"
|
||||
checkpoint_path = run_path / "best_model.pt"
|
||||
if not config_path.is_file():
|
||||
raise FileNotFoundError(config_path)
|
||||
if not checkpoint_path.is_file():
|
||||
raise FileNotFoundError(checkpoint_path)
|
||||
if args.batch_size <= 0:
|
||||
raise ValueError("batch_size must be > 0.")
|
||||
if args.rows_per_chunk <= 0:
|
||||
raise ValueError("rows_per_chunk must be > 0.")
|
||||
if args.num_workers < 0:
|
||||
raise ValueError("num_workers must be >= 0.")
|
||||
if args.dataset_subset_size is not None and args.dataset_subset_size <= 0:
|
||||
raise ValueError("dataset_subset_size must be > 0.")
|
||||
|
||||
cfg = load_json_config(config_path)
|
||||
validate_training_mode_config(cfg)
|
||||
model_target_mode = str(cfg.get("model_target_mode", "next_token")).lower()
|
||||
if model_target_mode != "all_future":
|
||||
raise ValueError(
|
||||
"This exporter requires model_target_mode='all_future'; got "
|
||||
f"{model_target_mode!r}."
|
||||
)
|
||||
|
||||
ages = build_age_grid(args.age_start, args.age_stop, args.age_step)
|
||||
output_path = (
|
||||
Path(args.output_path).resolve()
|
||||
if args.output_path
|
||||
else run_path / "weibull_parameters_test_age40_80_step2.h5"
|
||||
)
|
||||
if output_path.exists():
|
||||
raise FileExistsError(
|
||||
f"Output file already exists: {output_path}. Choose a new "
|
||||
"--output_path."
|
||||
)
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
temporary_output = output_path.with_name(f".{output_path.name}.partial")
|
||||
if temporary_output.exists():
|
||||
raise FileExistsError(
|
||||
f"Partial output already exists: {temporary_output}. Remove or "
|
||||
"rename it before retrying."
|
||||
)
|
||||
h5py = require_h5py()
|
||||
|
||||
data_prefix = str(cfg.get("data_prefix", "ukb"))
|
||||
labels_file = str(cfg.get("labels_file", "labels.csv"))
|
||||
disease_history_mode = normalize_disease_history_mode(
|
||||
cfg.get("disease_history_mode", DISEASE_HISTORY_MODE_TIMED)
|
||||
)
|
||||
min_history_events = int(
|
||||
cfg.get("all_future_min_history_events", cfg.get("min_history_events", 1))
|
||||
)
|
||||
min_future_events = int(
|
||||
cfg.get("all_future_min_future_events", cfg.get("min_future_events", 1))
|
||||
)
|
||||
print("Loading dataset...")
|
||||
dataset = load_sequence_eval_dataset(
|
||||
model_target_mode=model_target_mode,
|
||||
data_prefix=data_prefix,
|
||||
labels_file=labels_file,
|
||||
no_event_interval_years=float(cfg.get("no_event_interval_years", 5.0)),
|
||||
min_history_events=min_history_events,
|
||||
min_future_events=min_future_events,
|
||||
extra_info_types=parse_int_list(cfg.get("extra_info_types")),
|
||||
disease_history_mode=disease_history_mode,
|
||||
)
|
||||
validate_dataset_metadata(dataset, cfg)
|
||||
|
||||
test_eid_file = args.test_eid_file or cfg.get(
|
||||
"test_eid_file", "ukb_test_eid.csv"
|
||||
)
|
||||
if not test_eid_file:
|
||||
raise ValueError(
|
||||
"No test_eid_file is defined. Provide --test_eid_file explicitly."
|
||||
)
|
||||
subset_indices, resolved_eid_file = select_indices_by_eid_file(
|
||||
dataset, str(test_eid_file)
|
||||
)
|
||||
if args.dataset_subset_size is not None:
|
||||
subset_indices = subset_indices[: args.dataset_subset_size]
|
||||
if subset_indices.size == 0:
|
||||
raise RuntimeError("The selected test subset is empty.")
|
||||
|
||||
state_dict = load_checkpoint_state_dict(checkpoint_path, map_location="cpu")
|
||||
dist_mode = resolve_dist_mode_for_checkpoint(
|
||||
str(cfg.get("dist_mode", "exponential")), state_dict
|
||||
)
|
||||
if dist_mode != "weibull":
|
||||
raise ValueError(
|
||||
f"The specified run uses dist_mode={dist_mode!r}, not 'weibull'."
|
||||
)
|
||||
cfg_model = dict(cfg)
|
||||
cfg_model["dist_mode"] = dist_mode
|
||||
cfg_model["model_architecture"] = resolve_model_architecture(
|
||||
cfg_model, state_dict
|
||||
)
|
||||
device = resolve_eval_device(args.device)
|
||||
model = build_model_from_dataset(
|
||||
args, cfg_model, dataset, state_dict=state_dict
|
||||
).to(device)
|
||||
load_model_state(model, state_dict)
|
||||
model.eval()
|
||||
|
||||
tokens = select_outcome_tokens(dataset)
|
||||
label_text = load_label_text(labels_file)
|
||||
death_tokens = [
|
||||
token
|
||||
for token in tokens
|
||||
if str(dataset.label_id_to_code[token]).lower() == "death"
|
||||
]
|
||||
if not death_tokens:
|
||||
raise RuntimeError("Death token was not found in the outcome vocabulary.")
|
||||
|
||||
token_codes = [str(dataset.label_id_to_code[token]) for token in tokens]
|
||||
token_text = [label_text.get(code, code) for code in token_codes]
|
||||
outcome_types = [
|
||||
"death" if code.lower() == "death" else "disease"
|
||||
for code in token_codes
|
||||
]
|
||||
selected_eids = np.asarray(
|
||||
[int(dataset.samples[int(index)]["eid"]) for index in subset_indices],
|
||||
dtype=np.int64,
|
||||
)
|
||||
metadata: Dict[str, Any] = {
|
||||
"format_version": FORMAT_VERSION,
|
||||
"complete": False,
|
||||
"run_path": str(run_path),
|
||||
"checkpoint": str(checkpoint_path),
|
||||
"test_eid_file": str(resolved_eid_file),
|
||||
"n_selected_test_patients": int(subset_indices.size),
|
||||
"ages": [float(x) for x in ages],
|
||||
"n_tokens": len(tokens),
|
||||
"matrix_dtype": "float32",
|
||||
"hdf5_layout": {
|
||||
"tokens": (
|
||||
"/tokens/{column,token_id,label_code,label_text,outcome_type}"
|
||||
),
|
||||
"test_population": "/test_population/{eid,dataset_index}",
|
||||
"age_summary": (
|
||||
"/age_summary/{age,n_rows,nonfinite_shape_values,"
|
||||
"nonfinite_scale_values}"
|
||||
),
|
||||
"landmarks": (
|
||||
"/landmarks/age_*/{eid,dataset_index,sex,age,shape,scale}"
|
||||
),
|
||||
},
|
||||
"compression": "gzip",
|
||||
"compression_level": DEFAULT_COMPRESSION_LEVEL,
|
||||
"rows_per_chunk": int(args.rows_per_chunk),
|
||||
"parameterization": {
|
||||
"survival": "S(t) = exp(-rate * t^shape)",
|
||||
"shape": "rho = softplus(rho_logit) + 1e-6",
|
||||
"rate": "softplus(risk_logit) + 1e-8",
|
||||
"scale": "rate^(-1/shape)",
|
||||
"time_unit": "years after landmark age",
|
||||
},
|
||||
"eligibility": (
|
||||
"At each age: follow-up extends beyond the landmark, the patient "
|
||||
"is alive at the landmark, and the configured minimum disease "
|
||||
"history is available. All disease and death token parameters are "
|
||||
"exported, including tokens already prevalent by the landmark."
|
||||
),
|
||||
}
|
||||
|
||||
string_dtype = h5py.string_dtype(encoding="utf-8")
|
||||
summaries: List[Dict[str, Any]] = []
|
||||
with h5py.File(temporary_output, "w") as output_file:
|
||||
output_file.attrs["format_version"] = FORMAT_VERSION
|
||||
output_file.attrs["complete"] = False
|
||||
output_file.attrs["run_path"] = str(run_path)
|
||||
output_file.attrs["checkpoint"] = str(checkpoint_path)
|
||||
output_file.attrs["test_eid_file"] = str(resolved_eid_file)
|
||||
output_file.attrs["n_selected_test_patients"] = int(
|
||||
subset_indices.size
|
||||
)
|
||||
output_file.attrs["n_tokens"] = int(len(tokens))
|
||||
output_file.attrs["matrix_dtype"] = "float32"
|
||||
output_file.attrs["compression"] = "gzip"
|
||||
output_file.attrs["compression_level"] = DEFAULT_COMPRESSION_LEVEL
|
||||
|
||||
metadata_dataset = output_file.create_dataset(
|
||||
"metadata_json", shape=(), dtype=string_dtype
|
||||
)
|
||||
write_metadata_json(metadata_dataset, metadata)
|
||||
output_file.create_dataset("ages", data=ages.astype(np.float32))
|
||||
|
||||
token_group = output_file.create_group("tokens")
|
||||
token_group.create_dataset(
|
||||
"column", data=np.arange(len(tokens), dtype=np.int64)
|
||||
)
|
||||
token_group.create_dataset(
|
||||
"token_id", data=np.asarray(tokens, dtype=np.int64)
|
||||
)
|
||||
token_group.create_dataset(
|
||||
"label_code",
|
||||
data=np.asarray(token_codes, dtype=object),
|
||||
dtype=string_dtype,
|
||||
)
|
||||
token_group.create_dataset(
|
||||
"label_text",
|
||||
data=np.asarray(token_text, dtype=object),
|
||||
dtype=string_dtype,
|
||||
)
|
||||
token_group.create_dataset(
|
||||
"outcome_type",
|
||||
data=np.asarray(outcome_types, dtype=object),
|
||||
dtype=string_dtype,
|
||||
)
|
||||
|
||||
population_group = output_file.create_group("test_population")
|
||||
population_group.create_dataset(
|
||||
"eid", data=selected_eids, compression="gzip", shuffle=True
|
||||
)
|
||||
population_group.create_dataset(
|
||||
"dataset_index",
|
||||
data=subset_indices.astype(np.int64, copy=False),
|
||||
compression="gzip",
|
||||
shuffle=True,
|
||||
)
|
||||
|
||||
landmark_root = output_file.create_group("landmarks")
|
||||
eligible_queries_by_age: Dict[str, int] = {}
|
||||
for age_value in ages.tolist():
|
||||
age = float(age_value)
|
||||
age_group = landmark_root.create_group(age_group_name(age))
|
||||
try:
|
||||
landmark_dataset = LandmarkDataset(
|
||||
dataset=dataset,
|
||||
subset_indices=subset_indices,
|
||||
landmark_ages=np.asarray([age], dtype=np.float32),
|
||||
model_target_mode=model_target_mode,
|
||||
min_history_events=min_history_events,
|
||||
first_occurrence_by_token={},
|
||||
death_token_ids=death_tokens,
|
||||
disease_history_mode=disease_history_mode,
|
||||
)
|
||||
except RuntimeError as exc:
|
||||
if "No eligible landmark query samples" not in str(exc):
|
||||
raise
|
||||
summary = write_empty_age_group(
|
||||
age=age,
|
||||
age_group=age_group,
|
||||
token_count=len(tokens),
|
||||
)
|
||||
else:
|
||||
loader = DataLoader(
|
||||
landmark_dataset,
|
||||
batch_size=int(args.batch_size),
|
||||
shuffle=False,
|
||||
collate_fn=collate_landmark_fn,
|
||||
num_workers=int(args.num_workers),
|
||||
pin_memory=device.type == "cuda",
|
||||
persistent_workers=args.num_workers > 0,
|
||||
prefetch_factor=2 if args.num_workers > 0 else None,
|
||||
)
|
||||
summary = export_age(
|
||||
age=age,
|
||||
model=model,
|
||||
loader=loader,
|
||||
tokens=tokens,
|
||||
device=device,
|
||||
use_amp=bool(args.use_amp),
|
||||
subset_indices=subset_indices,
|
||||
selected_eids=selected_eids,
|
||||
age_group=age_group,
|
||||
rows_per_chunk=int(args.rows_per_chunk),
|
||||
)
|
||||
|
||||
summaries.append(summary)
|
||||
eligible_count = int(summary["n_rows"])
|
||||
eligible_queries_by_age[f"{age:g}"] = eligible_count
|
||||
metadata["eligible_queries_by_age"] = eligible_queries_by_age
|
||||
write_metadata_json(metadata_dataset, metadata)
|
||||
output_file.flush()
|
||||
print(f"Age {age:g}: exported {eligible_count} patient rows")
|
||||
|
||||
summary_group = output_file.create_group("age_summary")
|
||||
summary_group.create_dataset(
|
||||
"age",
|
||||
data=np.asarray([row["age"] for row in summaries], dtype=np.float32),
|
||||
)
|
||||
summary_group.create_dataset(
|
||||
"n_rows",
|
||||
data=np.asarray([row["n_rows"] for row in summaries], dtype=np.int64),
|
||||
)
|
||||
summary_group.create_dataset(
|
||||
"nonfinite_shape_values",
|
||||
data=np.asarray(
|
||||
[row["nonfinite_shape_values"] for row in summaries],
|
||||
dtype=np.int64,
|
||||
),
|
||||
)
|
||||
summary_group.create_dataset(
|
||||
"nonfinite_scale_values",
|
||||
data=np.asarray(
|
||||
[row["nonfinite_scale_values"] for row in summaries],
|
||||
dtype=np.int64,
|
||||
),
|
||||
)
|
||||
|
||||
nonfinite_shape = sum(
|
||||
int(row["nonfinite_shape_values"]) for row in summaries
|
||||
)
|
||||
nonfinite_scale = sum(
|
||||
int(row["nonfinite_scale_values"]) for row in summaries
|
||||
)
|
||||
metadata["total_exported_query_rows"] = sum(
|
||||
int(row["n_rows"]) for row in summaries
|
||||
)
|
||||
metadata["nonfinite_shape_values"] = nonfinite_shape
|
||||
metadata["nonfinite_scale_values"] = nonfinite_scale
|
||||
validate_hdf5_export(
|
||||
output_file,
|
||||
ages=ages,
|
||||
token_count=len(tokens),
|
||||
summaries=summaries,
|
||||
)
|
||||
metadata["validated"] = True
|
||||
metadata["complete"] = True
|
||||
write_metadata_json(metadata_dataset, metadata)
|
||||
output_file.attrs["validated"] = True
|
||||
output_file.attrs.modify("complete", True)
|
||||
output_file.flush()
|
||||
|
||||
temporary_output.replace(output_path)
|
||||
|
||||
if nonfinite_shape or nonfinite_scale:
|
||||
print(
|
||||
"WARNING: exported non-finite values: "
|
||||
f"shape={nonfinite_shape}, scale={nonfinite_scale}."
|
||||
)
|
||||
print(f"Saved unified Weibull parameter file to: {output_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,7 +0,0 @@
|
||||
"""Compatibility entry point for Weibull shape-parameter export."""
|
||||
|
||||
from export_weibull_death_parameter_stats import main
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
161
extra_info_assessment_all_analysis_plan.md
Normal file
161
extra_info_assessment_all_analysis_plan.md
Normal file
@@ -0,0 +1,161 @@
|
||||
# assessment_only 与 all 实验及分析方案
|
||||
|
||||
## 1. 实验目的
|
||||
|
||||
补充现有四级 extra-information 证据链：
|
||||
|
||||
1. `disease_only`：疾病事件、相对患病时间和 sex；无 extra-info token。
|
||||
2. `smoking_alcohol_bmi`：疾病史、sex、smoking/alcohol/BMI token。
|
||||
3. `assessment_only`：疾病史、sex、65项常规体格、肺功能、血液、尿液和生化指标 token。
|
||||
4. `all`：疾病史、sex、全部265项体检和暴露信息 token。
|
||||
|
||||
所有配置均不使用 CHECKUP。额外信息以独立 token 注入，并使用各自的 assessment 时间。
|
||||
所有连续变量均强制使用训练子集拟合的 RobustScale；训练与评估不设置未标准化对照或兼容模式。
|
||||
|
||||
目标是区分：
|
||||
|
||||
- 疾病史本身能够提供多少未来疾病信息；
|
||||
- 常规体检在疾病史之外增加多少信息；
|
||||
- 生活方式、社会经济、心理和环境暴露在完整体检之外增加多少信息；
|
||||
- 额外信息对疾病预测和死亡预测是否具有不同作用。
|
||||
|
||||
## 2. 固定模型
|
||||
|
||||
所有新增实验固定为：
|
||||
|
||||
```text
|
||||
TrajMixer
|
||||
+ all_future
|
||||
+ relative
|
||||
+ Weibull
|
||||
+ timed disease history
|
||||
+ sex
|
||||
```
|
||||
|
||||
仅改变 `extra_info_types_file`：
|
||||
|
||||
- `extra_info_types_assessment_only.txt`
|
||||
- `extra_info_types_all.txt`
|
||||
|
||||
每种配置运行 seed 42、43、44。
|
||||
|
||||
## 3. A6000 48GB 设置
|
||||
|
||||
| 配置 | Batch size | 原因 |
|
||||
|---|---:|---|
|
||||
| assessment_only | 256 | 最多65个 extra-info token，48GB余量充足 |
|
||||
| all | 128 | 最多265个 extra-info token，relative RBF attention 显存随总序列长度平方增长 |
|
||||
|
||||
当前训练代码为 FP32/TF32，并未使用 AMP。若 `all batch=128` 在极端长序列 batch 上出现 CUDA OOM，降为64，不自动重试，避免同一配置产生多个不完整 run。
|
||||
|
||||
batch size 是 `all` 与其他模型之间的潜在训练差异。代码会按 batch size 自动缩放学习率，但最终报告仍需明确记录该差异。若 `all` 的结果处于模型选择临界区，再补 seed 42 的 batch-size sensitivity，而不是预先扩大实验矩阵。
|
||||
|
||||
## 4. 主要比较
|
||||
|
||||
### 4.1 常规体检的增量价值
|
||||
|
||||
```text
|
||||
assessment_only − disease_only
|
||||
```
|
||||
|
||||
回答常规器官功能检测指标在疾病序列之外提供多少信息。
|
||||
|
||||
### 4.2 全部信息相对常规体检
|
||||
|
||||
```text
|
||||
all − assessment_only
|
||||
```
|
||||
|
||||
回答生活方式、社会经济、心理和环境暴露是否在常规体检之后仍有增量价值。两组使用完全相同的疾病事件、Landmark 和随访边界，仅 extra-info token 集合不同。
|
||||
|
||||
### 4.3 全体检相对紧凑变量集
|
||||
|
||||
```text
|
||||
assessment_only − smoking_alcohol_bmi
|
||||
```
|
||||
|
||||
回答65项常规体检是否优于紧凑的 smoking/alcohol/BMI 输入。
|
||||
|
||||
### 4.4 全部信息相对紧凑变量集
|
||||
|
||||
```text
|
||||
all − smoking_alcohol_bmi
|
||||
```
|
||||
|
||||
衡量从当前最终模型扩展到全部 extra information 的最大增益。
|
||||
|
||||
### 4.5 既有比较
|
||||
|
||||
保留：
|
||||
|
||||
```text
|
||||
smoking_alcohol_bmi − disease_only
|
||||
```
|
||||
|
||||
与新增结果共同形成完整的信息增量路径。
|
||||
|
||||
## 5. 评估指标
|
||||
|
||||
疾病和死亡分开分析。
|
||||
|
||||
### 判别能力
|
||||
|
||||
- Landmark AUC；
|
||||
- 各 horizon AUC；
|
||||
- disease cell win rate；
|
||||
- 三个 seed 的均值、标准差和方向一致性。
|
||||
|
||||
### 概率与似然质量
|
||||
|
||||
- IPCW Brier；
|
||||
- 固定时点 IPCW NLL；
|
||||
- 连续时间 point-process NLL；
|
||||
- Expected/Observed ratio；
|
||||
- calibration-in-the-large；
|
||||
- calibration slope；
|
||||
- 校准曲线。
|
||||
|
||||
### 复杂度与稳定性
|
||||
|
||||
- 参数量；
|
||||
- 每个 epoch 运行时间；
|
||||
- 峰值显存；
|
||||
- 三个 seed 的性能波动；
|
||||
- 缺失值较多的 extra-info 类型是否造成训练不稳定。
|
||||
|
||||
## 6. 配对和汇总方法
|
||||
|
||||
- replicate unit 为 seed；
|
||||
- 按相同 `seed × label_code × sex × horizon` 配对；
|
||||
- 每个比较使用三个 seed 共同存在的 cell；
|
||||
- 先在 seed 内汇总，再计算三个 seed 的均值和样本标准差；
|
||||
- AUC 越高越好；
|
||||
- Brier、NLL及绝对校准偏差越低越好；
|
||||
- 不以单个 seed 或单个 horizon 决定模型。
|
||||
|
||||
## 7. disease_only 比较的评估边界
|
||||
|
||||
所有配置均从同一疾病/死亡事件流构造历史、查询点、随访终点和 censoring，并且都不使用 CHECKUP。不同配置只改变 extra-info token，因此可以在共同支持集上把差异解释为额外信息的增量价值。
|
||||
|
||||
## 8. 决策规则
|
||||
|
||||
1. 如果 `assessment_only` 已达到 `all` 的绝大部分性能，并且校准更稳定，优先选择 `assessment_only`，因为它与器官功能重建目标一致且解释更清楚。
|
||||
2. 如果 `all` 在疾病和死亡 AUC、Brier、NLL上均稳定优于 `assessment_only`，则将 `all` 作为性能上限模型，但不直接作为器官负担教师模型。
|
||||
3. 如果 `all` 只提高 AUC而恶化 Brier/NLL或 seed 波动明显，不升级最终模型。
|
||||
4. disease-only Timed 仍是生成纯疾病来源器官负担分数的教师模型；assessment/all 实验用于界定疾病史遗漏的信息，而不是改变该分数的纯疾病定义。
|
||||
|
||||
## 9. 运行与后续评估
|
||||
|
||||
训练：
|
||||
|
||||
```bash
|
||||
bash train_extra_info_assessment_all_multiseed_linux.sh --gpus 0
|
||||
```
|
||||
|
||||
多GPU：
|
||||
|
||||
```bash
|
||||
bash train_extra_info_assessment_all_multiseed_linux.sh --gpus 0,1,2
|
||||
```
|
||||
|
||||
训练完成后，现有扫描脚本分别生成 AUC 和 calibration/Brier/NLL。合并分析时，将新增 `assessment_only` 和 `all` 两个配置纳入主 Timed 表，不纳入 disease-history Ordered/Set 专表。
|
||||
71
extra_info_types_assessment_smoking_alcohol.txt
Normal file
71
extra_info_types_assessment_smoking_alcohol.txt
Normal file
@@ -0,0 +1,71 @@
|
||||
# Assessment/body-measurement variables plus smoking and alcohol (field_type=1 plus types 66-67)
|
||||
# Generated from field_ids_enriched.csv using prepare_data.py other-info type ordering.
|
||||
# BMI is already included as assessment type 11 and is not duplicated.
|
||||
# Format: <extra_info_type_id> # <var_name> | <full_name>
|
||||
1 # waist_circumference | Waist circumference
|
||||
2 # hip_circumference | Hip circumference
|
||||
3 # standing_height | Standing height
|
||||
4 # fasting_time | Fasting time
|
||||
5 # pulse_rate | Pulse rate automated reading
|
||||
6 # dbp | Diastolic blood pressure automated reading
|
||||
7 # sbp | Systolic blood pressure automated reading
|
||||
8 # fev1_best | Forced expiratory volume in 1-second (FEV1) Best measure
|
||||
9 # fvc_best | Forced vital capacity (FVC) Best measure
|
||||
10 # fev1_fvc_ratio | FEV1/ FVC ratio Z-score
|
||||
11 # bmi | Body mass index (BMI)
|
||||
12 # WBC | White blood cell (leukocyte) count
|
||||
13 # RBC | Red blood cell (erythrocyte) count
|
||||
14 # hemoglobin | Haemoglobin concentration
|
||||
15 # hematocrit | Haematocrit percentage
|
||||
16 # MCV | Mean corpuscular volume
|
||||
17 # MCH | Mean corpuscular haemoglobin
|
||||
18 # MCHC | Mean corpuscular haemoglobin concentration
|
||||
19 # Pc | Platelet count
|
||||
20 # MPV | Mean platelet (thrombocyte) volume
|
||||
21 # LymC | Lymphocyte count
|
||||
22 # MonC | Monocyte count
|
||||
23 # NeuC | Neutrophill count
|
||||
24 # EosC | Eosinophill count
|
||||
25 # BasC | Basophill count
|
||||
26 # nRBC | Nucleated red blood cell count
|
||||
27 # RC | Reticulocyte count
|
||||
28 # MRV | Mean reticulocyte volume
|
||||
29 # MSCV | Mean sphered cell volume
|
||||
30 # IRF | Immature reticulocyte fraction
|
||||
31 # HLSRC | High light scatter reticulocyte count
|
||||
32 # MicU | Microalbumin in urine
|
||||
33 # CreaU | Creatinine (enzymatic) in urine
|
||||
34 # PotU | Potassium in urine
|
||||
35 # SodU | Sodium in urine
|
||||
36 # Alb | Albumin
|
||||
37 # ALP | Alkaline phosphatase
|
||||
38 # Alanine | Alanine aminotransferase
|
||||
39 # ApoA | Apolipoprotein A
|
||||
40 # ApoB | Apolipoprotein B
|
||||
41 # AA | Aspartate aminotransferase
|
||||
42 # DBil | Direct bilirubin
|
||||
43 # Urea | Urea
|
||||
44 # Calcium | Calcium
|
||||
45 # Cholesterol | Cholesterol
|
||||
46 # Creatinine | Creatinine
|
||||
47 # CRP | C-reactive protein
|
||||
48 # CystatinC | Cystatin C
|
||||
49 # GGT | Gamma glutamyltransferase
|
||||
50 # Glu | Glucose
|
||||
51 # HbA1c | Glycated haemoglobin (HbA1c)
|
||||
52 # HDL | HDL cholesterol
|
||||
53 # IGF1 | IGF-1
|
||||
54 # LDL | LDL direct
|
||||
55 # LpA | Lipoprotein A
|
||||
56 # Oestradiol | Oestradiol
|
||||
57 # Phosphate | Phosphate
|
||||
58 # Rheu | Rheumatoid factor
|
||||
59 # SHBG | SHBG
|
||||
60 # TotalBil | Total bilirubin
|
||||
61 # Testosterone | Testosterone
|
||||
62 # TotalProtein | Total protein
|
||||
63 # Tri | Triglycerides
|
||||
64 # Urate | Urate
|
||||
65 # VitaminD | Vitamin D
|
||||
66 # smoking | Current tobacco smoking
|
||||
67 # alcohol | Alcohol intake frequency.
|
||||
115
future_risk.py
115
future_risk.py
@@ -1,115 +0,0 @@
|
||||
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))
|
||||
@@ -1,511 +0,0 @@
|
||||
"""Shared landmark evaluation helpers for attribution scripts."""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional, Sequence
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import torch
|
||||
from torch.nn.utils.rnn import pad_sequence
|
||||
from torch.utils.data import Dataset
|
||||
|
||||
from dataset import HealthDataset
|
||||
from eval_data import load_sequence_eval_dataset
|
||||
from evaluate_auc_v2 import (
|
||||
LandmarkDataset,
|
||||
build_model_from_dataset,
|
||||
cfg_get,
|
||||
make_eval_indices,
|
||||
)
|
||||
from models import DeepHealth
|
||||
from readouts import build_readout
|
||||
from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX
|
||||
from train_util import load_eid_file, load_extra_info_types_file
|
||||
|
||||
|
||||
SPECIAL_TOKENS = {PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX}
|
||||
|
||||
|
||||
def parse_int_list(value: Any) -> Optional[List[int]]:
|
||||
if value is None:
|
||||
return None
|
||||
if isinstance(value, (list, tuple, np.ndarray)):
|
||||
return [int(x) for x in value]
|
||||
text = str(value).strip()
|
||||
if text == "":
|
||||
return None
|
||||
if text.startswith("["):
|
||||
values = json.loads(text)
|
||||
if not isinstance(values, list):
|
||||
raise ValueError(f"Expected a JSON list, got {type(values).__name__}")
|
||||
return [int(x) for x in values]
|
||||
return [int(x.strip()) for x in text.split(",") if x.strip()]
|
||||
|
||||
|
||||
def load_extra_info_types(value: Any) -> Optional[List[int]]:
|
||||
if value is None:
|
||||
return None
|
||||
text = str(value)
|
||||
path = Path(text)
|
||||
if path.exists():
|
||||
return load_extra_info_types_file(text)
|
||||
return parse_int_list(value)
|
||||
|
||||
|
||||
def make_landmark_ages(start: float, stop: float, step: float) -> np.ndarray:
|
||||
if step <= 0:
|
||||
raise ValueError("landmark_step must be positive")
|
||||
if stop < start:
|
||||
raise ValueError("landmark_stop must be >= landmark_start")
|
||||
# Include stop when it lands on the grid, e.g. 40,45,...,80.
|
||||
return np.arange(start, stop + step * 0.5, step, dtype=np.float32)
|
||||
|
||||
|
||||
def build_first_occurrence_maps_for_landmarks(
|
||||
dataset: HealthDataset,
|
||||
subset_indices: np.ndarray,
|
||||
) -> Dict[int, tuple[np.ndarray, np.ndarray]]:
|
||||
first_lists: Dict[int, list[tuple[int, float]]] = {}
|
||||
for patient_id, dataset_index in enumerate(np.asarray(subset_indices, dtype=np.int64).tolist()):
|
||||
s = dataset.samples[int(dataset_index)]
|
||||
seq_event = np.asarray(s["event_seq"], dtype=np.int64)
|
||||
seq_time = np.asarray(s["time_seq"], dtype=np.float32)
|
||||
tgt_event = np.asarray(s["target_event_seq"], dtype=np.int64)
|
||||
tgt_time = np.asarray(s["target_time_seq"], dtype=np.float32)
|
||||
if seq_event.size == 0 or tgt_event.size == 0:
|
||||
continue
|
||||
|
||||
full_event = np.concatenate([seq_event, tgt_event[-1:]])
|
||||
full_time = np.concatenate([seq_time, tgt_time[-1:]])
|
||||
uniq_tokens, first_idx = np.unique(full_event, return_index=True)
|
||||
for token, idx in zip(uniq_tokens.tolist(), first_idx.tolist()):
|
||||
token = int(token)
|
||||
if token in SPECIAL_TOKENS:
|
||||
continue
|
||||
first_lists.setdefault(token, []).append((patient_id, float(full_time[int(idx)])))
|
||||
|
||||
return {
|
||||
int(token): (
|
||||
np.asarray([p for p, _ in pairs], dtype=np.int32),
|
||||
np.asarray([t for _, t in pairs], dtype=np.float32),
|
||||
)
|
||||
for token, pairs in first_lists.items()
|
||||
if pairs
|
||||
}
|
||||
|
||||
|
||||
def normalize_eval_split(args: argparse.Namespace, cfg: Dict[str, Any]) -> str:
|
||||
eval_split = str(cfg_get(args, cfg, "eval_split", "test")).lower()
|
||||
if eval_split in {"valid", "validation"}:
|
||||
return "val"
|
||||
if eval_split not in {"train", "val", "test", "all"}:
|
||||
raise ValueError(f"Unsupported eval_split={eval_split!r}")
|
||||
return eval_split
|
||||
|
||||
|
||||
def load_eval_sequence_dataset(
|
||||
args: argparse.Namespace,
|
||||
cfg: Dict[str, Any],
|
||||
) -> tuple[Any, np.ndarray, str, str]:
|
||||
eval_split = normalize_eval_split(args, cfg)
|
||||
model_target_mode = str(cfg.get("model_target_mode", "next_token")).lower()
|
||||
data_prefix = str(cfg.get("data_prefix", "ukb"))
|
||||
labels_file = str(cfg.get("labels_file", "labels.csv"))
|
||||
no_event_interval_years = float(cfg.get("no_event_interval_years", 5.0))
|
||||
include_no_event_in_uts_target = bool(cfg.get("include_no_event_in_uts_target", False))
|
||||
extra_info_types = load_extra_info_types(args.extra_info_types)
|
||||
if extra_info_types is None:
|
||||
extra_info_types = parse_int_list(cfg.get("extra_info_types", None))
|
||||
|
||||
print("Loading one sequence eval dataset...")
|
||||
dataset = load_sequence_eval_dataset(
|
||||
model_target_mode=model_target_mode,
|
||||
data_prefix=data_prefix,
|
||||
labels_file=labels_file,
|
||||
no_event_interval_years=no_event_interval_years,
|
||||
include_no_event_in_uts_target=include_no_event_in_uts_target,
|
||||
min_history_events=int(cfg.get("all_future_min_history_events", 1)),
|
||||
min_future_events=int(cfg.get("all_future_min_future_events", 1)),
|
||||
extra_info_types=extra_info_types,
|
||||
)
|
||||
|
||||
train_eid_file = cfg_get(args, cfg, "train_eid_file", "ukb_train_eid.csv")
|
||||
val_eid_file = cfg_get(args, cfg, "val_eid_file", "ukb_val_eid.csv")
|
||||
test_eid_file = cfg_get(args, cfg, "test_eid_file", "ukb_test_eid.csv")
|
||||
split_files_exist = all(
|
||||
Path(str(path)).exists()
|
||||
for path in (train_eid_file, val_eid_file, test_eid_file)
|
||||
)
|
||||
|
||||
if eval_split != "all" and split_files_exist:
|
||||
split_files = {
|
||||
"train": train_eid_file,
|
||||
"val": val_eid_file,
|
||||
"test": test_eid_file,
|
||||
}
|
||||
selected_eids = load_eid_file(split_files[eval_split])
|
||||
out = np.asarray(
|
||||
[
|
||||
idx
|
||||
for idx, sample in enumerate(dataset.samples)
|
||||
if int(sample["eid"]) in selected_eids
|
||||
],
|
||||
dtype=np.int64,
|
||||
)
|
||||
if out.size == 0:
|
||||
raise ValueError(
|
||||
f"No samples found for eval_split={eval_split!r} using {split_files[eval_split]}"
|
||||
)
|
||||
split_source = "eid_files"
|
||||
else:
|
||||
if eval_split == "all":
|
||||
out = np.arange(len(dataset.samples), dtype=np.int64)
|
||||
split_source = "all"
|
||||
else:
|
||||
out = make_eval_indices(dataset, args, cfg)
|
||||
split_source = "ratio_split"
|
||||
|
||||
subset_size = cfg_get(args, cfg, "dataset_subset_size", None)
|
||||
if subset_size is not None and int(subset_size) > 0:
|
||||
out = out[: int(subset_size)]
|
||||
return dataset, np.asarray(out, dtype=np.int64), eval_split, split_source
|
||||
|
||||
|
||||
def load_organ_groups(
|
||||
path: Path,
|
||||
*,
|
||||
vocab_size: int,
|
||||
) -> tuple[dict[str, list[int]], dict[str, str], dict[int, str]]:
|
||||
table = pd.read_csv(path)
|
||||
required = {"token_id", "organ_system", "organ_system_label", "is_death"}
|
||||
missing = required - set(table.columns)
|
||||
if missing:
|
||||
raise ValueError(f"{path} is missing columns: {sorted(missing)}")
|
||||
|
||||
death_idx = int(vocab_size) - 1
|
||||
groups: dict[str, list[int]] = {}
|
||||
labels: dict[str, str] = {}
|
||||
token_to_group: dict[int, str] = {}
|
||||
for row in table.itertuples(index=False):
|
||||
token = int(getattr(row, "token_id"))
|
||||
if token in SPECIAL_TOKENS or token == death_idx:
|
||||
continue
|
||||
if token < 0 or token >= int(vocab_size):
|
||||
continue
|
||||
if int(getattr(row, "is_death")) == 1:
|
||||
continue
|
||||
group = str(getattr(row, "organ_system"))
|
||||
label = str(getattr(row, "organ_system_label"))
|
||||
groups.setdefault(group, []).append(token)
|
||||
labels[group] = label
|
||||
token_to_group[token] = group
|
||||
|
||||
groups = {k: sorted(set(v)) for k, v in groups.items() if v}
|
||||
return groups, labels, token_to_group
|
||||
|
||||
|
||||
class IndexedLandmarkDataset(Dataset):
|
||||
def __init__(self, base: LandmarkDataset) -> None:
|
||||
self.base = base
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.base)
|
||||
|
||||
def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
|
||||
item = dict(self.base[idx])
|
||||
item["row_idx"] = torch.tensor(int(idx), dtype=torch.long)
|
||||
return item
|
||||
|
||||
|
||||
def collate_indexed_landmark_fn(batch: List[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]:
|
||||
event_seq = pad_sequence(
|
||||
[x["event_seq"] for x in batch], batch_first=True, padding_value=PAD_IDX
|
||||
)
|
||||
time_seq = pad_sequence(
|
||||
[x["time_seq"] for x in batch], batch_first=True, padding_value=0.0
|
||||
)
|
||||
readout_mask = pad_sequence(
|
||||
[x["readout_mask"] for x in batch], batch_first=True, padding_value=False
|
||||
)
|
||||
other_type = pad_sequence(
|
||||
[x["other_type"] for x in batch], batch_first=True, padding_value=0
|
||||
)
|
||||
other_value = pad_sequence(
|
||||
[x["other_value"] for x in batch], batch_first=True, padding_value=0.0
|
||||
)
|
||||
other_value_kind = pad_sequence(
|
||||
[x["other_value_kind"] for x in batch], batch_first=True, padding_value=0
|
||||
)
|
||||
other_time = pad_sequence(
|
||||
[x["other_time"] for x in batch], batch_first=True, padding_value=0.0
|
||||
)
|
||||
return {
|
||||
"event_seq": event_seq,
|
||||
"time_seq": time_seq,
|
||||
"padding_mask": event_seq > PAD_IDX,
|
||||
"readout_mask": readout_mask,
|
||||
"sex": torch.stack([x["sex"] for x in batch]),
|
||||
"other_type": other_type,
|
||||
"other_value": other_value,
|
||||
"other_value_kind": other_value_kind,
|
||||
"other_time": other_time,
|
||||
"landmark_pos": torch.stack([x["landmark_pos"] for x in batch]),
|
||||
"t_query": torch.stack([x["t_query"] for x in batch]),
|
||||
"patient_id": torch.stack([x["patient_id"] for x in batch]),
|
||||
"landmark_age": torch.stack([x["landmark_age"] for x in batch]),
|
||||
"followup_end_time": torch.stack([x["followup_end_time"] for x in batch]),
|
||||
"death_time": torch.stack([x["death_time"] for x in batch]),
|
||||
"row_idx": torch.stack([x["row_idx"] for x in batch]),
|
||||
}
|
||||
|
||||
|
||||
def build_group_ablated_slice(
|
||||
batch: Dict[str, torch.Tensor],
|
||||
token_ids: Sequence[int],
|
||||
row_indices: torch.Tensor,
|
||||
) -> Dict[str, torch.Tensor]:
|
||||
"""Build one fixed-width ablated slice without rebuilding variable-length rows."""
|
||||
event_seq = batch["event_seq"]
|
||||
|
||||
out: Dict[str, torch.Tensor] = {}
|
||||
out["event_seq"] = event_seq[row_indices].clone()
|
||||
out["time_seq"] = batch["time_seq"][row_indices]
|
||||
out["readout_mask"] = batch["readout_mask"][row_indices].clone()
|
||||
out["padding_mask"] = batch["padding_mask"][row_indices].bool().clone()
|
||||
out["landmark_pos"] = batch["landmark_pos"][row_indices].clone()
|
||||
|
||||
seq_len = int(event_seq.shape[1])
|
||||
positions = torch.arange(seq_len, device=event_seq.device)[None, :]
|
||||
ids = torch.as_tensor(token_ids, dtype=event_seq.dtype, device=event_seq.device)
|
||||
remove = torch.isin(out["event_seq"], ids) & out["padding_mask"]
|
||||
out["event_seq"] = torch.where(
|
||||
remove,
|
||||
torch.full_like(out["event_seq"], PAD_IDX),
|
||||
out["event_seq"],
|
||||
)
|
||||
out["padding_mask"] &= ~remove
|
||||
out["readout_mask"] &= ~remove
|
||||
|
||||
has_valid = out["padding_mask"].any(dim=1)
|
||||
if not bool(has_valid.all().item()):
|
||||
empty_rows = torch.nonzero(~has_valid, as_tuple=False).flatten()
|
||||
out["event_seq"][empty_rows, 0] = CHECKUP_IDX
|
||||
out["time_seq"][empty_rows, 0] = batch["t_query"][row_indices[empty_rows]].to(
|
||||
dtype=out["time_seq"].dtype
|
||||
)
|
||||
out["padding_mask"][empty_rows, 0] = True
|
||||
out["readout_mask"][empty_rows, 0] = True
|
||||
out["landmark_pos"][empty_rows] = 0
|
||||
|
||||
has_readout = out["readout_mask"].any(dim=1)
|
||||
if not bool(has_readout.all().item()):
|
||||
rows = torch.nonzero(~has_readout, as_tuple=False).flatten()
|
||||
local_valid = out["padding_mask"][rows]
|
||||
last_pos = torch.where(
|
||||
local_valid,
|
||||
positions.expand(local_valid.shape[0], -1),
|
||||
torch.zeros_like(positions.expand(local_valid.shape[0], -1)),
|
||||
).amax(dim=1)
|
||||
out["readout_mask"][rows] = False
|
||||
out["readout_mask"][rows, last_pos] = True
|
||||
out["landmark_pos"][rows] = last_pos.to(dtype=out["landmark_pos"].dtype)
|
||||
|
||||
repeated_keys = (
|
||||
"sex",
|
||||
"other_type",
|
||||
"other_value",
|
||||
"other_value_kind",
|
||||
"other_time",
|
||||
"t_query",
|
||||
"patient_id",
|
||||
"landmark_age",
|
||||
"followup_end_time",
|
||||
"death_time",
|
||||
"row_idx",
|
||||
)
|
||||
for key in repeated_keys:
|
||||
out[key] = batch[key][row_indices]
|
||||
return out
|
||||
|
||||
|
||||
def concat_tensor_batches(chunks: Sequence[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]:
|
||||
return {
|
||||
key: torch.cat([chunk[key] for chunk in chunks], dim=0)
|
||||
for key in chunks[0]
|
||||
}
|
||||
|
||||
|
||||
def iter_group_ablated_batches(
|
||||
batch: Dict[str, torch.Tensor],
|
||||
group_names: Sequence[str],
|
||||
organ_groups: dict[str, list[int]],
|
||||
occurred: torch.Tensor,
|
||||
max_batch_size: int,
|
||||
):
|
||||
"""Yield ablated chunks as soon as enough rows are available for a forward pass."""
|
||||
pending_batches: list[Dict[str, torch.Tensor]] = []
|
||||
pending_groups: list[str] = []
|
||||
pending_rows: list[int] = []
|
||||
pending_n = 0
|
||||
|
||||
for group in group_names:
|
||||
ids = torch.as_tensor(organ_groups[group], dtype=torch.long, device=occurred.device)
|
||||
if ids.numel() == 0:
|
||||
continue
|
||||
active_rows = torch.nonzero(occurred[:, ids].any(dim=1), as_tuple=False).flatten()
|
||||
if active_rows.numel() == 0:
|
||||
continue
|
||||
|
||||
row_offset = 0
|
||||
while row_offset < int(active_rows.numel()):
|
||||
capacity = int(max_batch_size) - pending_n
|
||||
row_stop = min(int(active_rows.numel()), row_offset + capacity)
|
||||
row_indices = active_rows[row_offset:row_stop].to(device=batch["event_seq"].device)
|
||||
chunk = build_group_ablated_slice(
|
||||
batch=batch,
|
||||
token_ids=organ_groups[group],
|
||||
row_indices=row_indices,
|
||||
)
|
||||
chunk_n = int(row_indices.numel())
|
||||
pending_batches.append(chunk)
|
||||
pending_groups.extend([group] * chunk_n)
|
||||
pending_rows.extend(int(x) for x in row_indices.detach().cpu().tolist())
|
||||
pending_n += chunk_n
|
||||
row_offset = row_stop
|
||||
|
||||
if pending_n >= int(max_batch_size):
|
||||
yield concat_tensor_batches(pending_batches), pending_groups, pending_rows
|
||||
pending_batches = []
|
||||
pending_groups = []
|
||||
pending_rows = []
|
||||
pending_n = 0
|
||||
|
||||
if pending_batches:
|
||||
yield concat_tensor_batches(pending_batches), pending_groups, pending_rows
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def infer_landmark_hidden(
|
||||
*,
|
||||
model: DeepHealth,
|
||||
batch: Dict[str, torch.Tensor],
|
||||
device: torch.device,
|
||||
model_target_mode: str,
|
||||
readout_name: str,
|
||||
readout_reduce: str,
|
||||
) -> torch.Tensor:
|
||||
batch_dev = {
|
||||
k: (v.to(device, non_blocking=True) if isinstance(v, torch.Tensor) else v)
|
||||
for k, v in batch.items()
|
||||
}
|
||||
if model_target_mode == "all_future":
|
||||
return model(
|
||||
event_seq=batch_dev["event_seq"].long(),
|
||||
time_seq=batch_dev["time_seq"].float(),
|
||||
sex=batch_dev["sex"].long(),
|
||||
padding_mask=batch_dev["padding_mask"].bool(),
|
||||
t_query=batch_dev["t_query"].float(),
|
||||
other_type=batch_dev["other_type"].long(),
|
||||
other_value=batch_dev["other_value"].float(),
|
||||
other_value_kind=batch_dev["other_value_kind"].long(),
|
||||
other_time=batch_dev["other_time"].float(),
|
||||
target_mode="all_future",
|
||||
)
|
||||
|
||||
hidden = model(
|
||||
event_seq=batch_dev["event_seq"].long(),
|
||||
time_seq=batch_dev["time_seq"].float(),
|
||||
sex=batch_dev["sex"].long(),
|
||||
padding_mask=batch_dev["padding_mask"].bool(),
|
||||
other_type=batch_dev["other_type"].long(),
|
||||
other_value=batch_dev["other_value"].float(),
|
||||
other_value_kind=batch_dev["other_value_kind"].long(),
|
||||
other_time=batch_dev["other_time"].float(),
|
||||
target_mode="next_token",
|
||||
)
|
||||
readout = build_readout(readout_name, reduce=readout_reduce)
|
||||
readout_out = readout(
|
||||
hidden=hidden,
|
||||
time_seq=batch_dev["time_seq"].float(),
|
||||
padding_mask=batch_dev["padding_mask"].bool(),
|
||||
readout_mask=batch_dev["readout_mask"].bool(),
|
||||
)
|
||||
return readout_out.hidden.gather(
|
||||
1,
|
||||
batch_dev["landmark_pos"].long()[:, None, None].expand(
|
||||
-1, 1, readout_out.hidden.shape[-1]
|
||||
),
|
||||
).squeeze(1)
|
||||
|
||||
|
||||
def make_occurred_mask(
|
||||
event_seq: torch.Tensor,
|
||||
*,
|
||||
vocab_size: int,
|
||||
device: torch.device,
|
||||
) -> torch.Tensor:
|
||||
occurred = torch.zeros(event_seq.shape[0], int(vocab_size), dtype=torch.bool, device=device)
|
||||
valid = (event_seq >= 0) & (event_seq < int(vocab_size))
|
||||
safe = event_seq.clamp(min=0, max=int(vocab_size) - 1).to(device)
|
||||
occurred.scatter_(1, safe, valid.to(device))
|
||||
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,
|
||||
*,
|
||||
death_idx: int,
|
||||
token_to_group: dict[int, str],
|
||||
group_names: Sequence[str],
|
||||
) -> tuple[int, dict[str, int]]:
|
||||
unique_tokens = {
|
||||
int(token)
|
||||
for token in np.asarray(tokens, dtype=np.int64).tolist()
|
||||
if int(token) not in SPECIAL_TOKENS and int(token) != int(death_idx)
|
||||
}
|
||||
total = len(unique_tokens)
|
||||
out = {group: 0 for group in group_names}
|
||||
for token in unique_tokens:
|
||||
group = token_to_group.get(token)
|
||||
if group in out:
|
||||
out[group] += 1
|
||||
return total, out
|
||||
320
losses.py
320
losses.py
@@ -8,7 +8,7 @@ import torch.nn.functional as F
|
||||
|
||||
|
||||
PAD_IDX = 0
|
||||
CHECKUP_IDX = 1
|
||||
RESERVED_IDX = 1
|
||||
NO_EVENT_IDX = 2
|
||||
|
||||
|
||||
@@ -37,6 +37,54 @@ def _zero_loss_like(logits: torch.Tensor) -> torch.Tensor:
|
||||
return logits.sum() * 0.0
|
||||
|
||||
|
||||
def _all_future_at_risk_mask(
|
||||
vocab_size: int,
|
||||
ignored_idx: Iterable[int],
|
||||
logits: torch.Tensor,
|
||||
history: torch.Tensor | None,
|
||||
) -> torch.Tensor:
|
||||
"""Return outcomes that have not occurred by the query time."""
|
||||
batch_size = logits.shape[0]
|
||||
at_risk = _valid_vocab_mask(
|
||||
vocab_size,
|
||||
ignored_idx,
|
||||
logits.device,
|
||||
).unsqueeze(0).expand(batch_size, -1).clone()
|
||||
if history is None or history.numel() == 0:
|
||||
return at_risk
|
||||
if history.dim() != 2 or history.shape[0] != batch_size:
|
||||
raise ValueError(
|
||||
"history must be (B, L). "
|
||||
f"Got logits={tuple(logits.shape)}, history={tuple(history.shape)}"
|
||||
)
|
||||
|
||||
history = history.to(device=logits.device, dtype=torch.long)
|
||||
history_valid = (history >= 0) & (history < vocab_size)
|
||||
for idx in ignored_idx:
|
||||
history_valid &= history != int(idx)
|
||||
safe_history = history.clamp(min=0, max=vocab_size - 1)
|
||||
prevalent = torch.zeros(
|
||||
(batch_size, vocab_size),
|
||||
dtype=torch.long,
|
||||
device=logits.device,
|
||||
)
|
||||
prevalent.scatter_add_(1, safe_history, history_valid.long())
|
||||
return at_risk & ~prevalent.bool()
|
||||
|
||||
|
||||
def _all_future_target_mask(
|
||||
targets: torch.Tensor,
|
||||
at_risk: torch.Tensor,
|
||||
ignored_idx: Iterable[int],
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
vocab_size = at_risk.shape[1]
|
||||
in_vocab = (targets >= 0) & (targets < vocab_size)
|
||||
for idx in ignored_idx:
|
||||
in_vocab &= targets != int(idx)
|
||||
safe_targets = targets.clamp(min=0, max=vocab_size - 1)
|
||||
return in_vocab & at_risk.gather(1, safe_targets), safe_targets
|
||||
|
||||
|
||||
class Delphi2MLoss(nn.Module):
|
||||
"""Next-token plus exponential time-to-next-token supervision."""
|
||||
|
||||
@@ -52,7 +100,7 @@ class Delphi2MLoss(nn.Module):
|
||||
super().__init__()
|
||||
self.t_min = float(t_min)
|
||||
self.ignored_tokens = (
|
||||
[PAD_IDX, CHECKUP_IDX]
|
||||
[PAD_IDX, RESERVED_IDX]
|
||||
if ignored_tokens is None
|
||||
else [int(x) for x in ignored_tokens]
|
||||
)
|
||||
@@ -145,101 +193,12 @@ class Delphi2MLoss(nn.Module):
|
||||
return total_loss
|
||||
|
||||
|
||||
class UniqueTimeSetExponentialLoss(nn.Module):
|
||||
"""Next distinct timestamp event-set supervision with sum reduction."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
ignored_idx: Iterable[int] = (PAD_IDX, CHECKUP_IDX),
|
||||
t_min: float = 1.0 / 365.25,
|
||||
max_exp_input: float = 60.0,
|
||||
exclude_ignored_from_intensity: bool = True,
|
||||
):
|
||||
super().__init__()
|
||||
self.ignored_idx = [int(x) for x in ignored_idx]
|
||||
self.t_min = float(t_min)
|
||||
self.max_exp_input = float(max_exp_input)
|
||||
self.exclude_ignored_from_intensity = bool(exclude_ignored_from_intensity)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
logits: torch.Tensor,
|
||||
target_multi_hot: torch.Tensor,
|
||||
target_dt_unique: torch.Tensor,
|
||||
readout_mask: torch.Tensor,
|
||||
return_components: bool = False,
|
||||
) -> torch.Tensor | tuple[torch.Tensor, dict[str, torch.Tensor]]:
|
||||
if logits.dim() != 3:
|
||||
raise ValueError(f"logits must be (B, L, K), got {tuple(logits.shape)}")
|
||||
bsz, seq_len, vocab_size = logits.shape
|
||||
|
||||
if target_multi_hot.shape != (bsz, seq_len, vocab_size):
|
||||
raise ValueError(
|
||||
"target_multi_hot must match logits shape, "
|
||||
f"got {tuple(target_multi_hot.shape)} vs {tuple(logits.shape)}"
|
||||
)
|
||||
if target_dt_unique.shape != (bsz, seq_len):
|
||||
raise ValueError(
|
||||
f"target_dt_unique must be {(bsz, seq_len)}, got {tuple(target_dt_unique.shape)}"
|
||||
)
|
||||
if readout_mask.shape != (bsz, seq_len):
|
||||
raise ValueError(f"readout_mask must be {(bsz, seq_len)}, got {tuple(readout_mask.shape)}")
|
||||
|
||||
ignore_mask = _make_ignore_mask(vocab_size, self.ignored_idx, logits.device)
|
||||
|
||||
num_targets = target_multi_hot[:, :, ~ignore_mask].sum(dim=-1)
|
||||
valid_mask = readout_mask.bool() & (num_targets > 0)
|
||||
|
||||
if not valid_mask.any():
|
||||
total_loss = _zero_loss_like(logits)
|
||||
if return_components:
|
||||
return total_loss, {
|
||||
"observed": total_loss.detach(),
|
||||
"penalty": total_loss.detach(),
|
||||
"total": total_loss.detach(),
|
||||
}
|
||||
return total_loss
|
||||
|
||||
logits_safe = torch.nan_to_num(
|
||||
logits[valid_mask],
|
||||
nan=0.0,
|
||||
posinf=self.max_exp_input,
|
||||
neginf=-self.max_exp_input,
|
||||
)
|
||||
target_valid = target_multi_hot[valid_mask].to(logits_safe.dtype)
|
||||
target_valid[:, ignore_mask] = 0.0
|
||||
|
||||
observed_term = (logits_safe * target_valid).sum(dim=-1)
|
||||
penalty_scale = target_valid.sum(dim=-1)
|
||||
|
||||
logits_for_lse = logits_safe
|
||||
if self.exclude_ignored_from_intensity:
|
||||
logits_for_lse = logits_safe.masked_fill(ignore_mask.unsqueeze(0), float("-inf"))
|
||||
|
||||
dt_clamped = torch.clamp(target_dt_unique[valid_mask], min=self.t_min)
|
||||
log_lambda_total = torch.logsumexp(logits_for_lse, dim=-1)
|
||||
log_penalty = log_lambda_total + dt_clamped.log()
|
||||
penalty = torch.exp(torch.clamp(log_penalty, max=self.max_exp_input))
|
||||
|
||||
observed_loss = -observed_term
|
||||
penalty_loss = penalty_scale * penalty
|
||||
total_loss = (observed_loss + penalty_loss).mean()
|
||||
|
||||
if return_components:
|
||||
return total_loss, {
|
||||
"observed": observed_loss.mean().detach(),
|
||||
"penalty": penalty_loss.mean().detach(),
|
||||
"total": total_loss.detach(),
|
||||
}
|
||||
return total_loss
|
||||
|
||||
|
||||
class ExponentialLoss(nn.Module):
|
||||
"""Query-conditioned all-future-event exponential point-process loss."""
|
||||
"""First-onset all-future exponential survival likelihood."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
ignored_idx: Iterable[int] = (PAD_IDX, CHECKUP_IDX),
|
||||
ignored_idx: Iterable[int] = (PAD_IDX, RESERVED_IDX),
|
||||
eps: float = 1e-8,
|
||||
):
|
||||
super().__init__()
|
||||
@@ -251,28 +210,65 @@ class ExponentialLoss(nn.Module):
|
||||
logits: torch.Tensor,
|
||||
targets: torch.Tensor,
|
||||
exposure: torch.Tensor,
|
||||
dt: torch.Tensor,
|
||||
history: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
_, vocab_size = logits.shape
|
||||
batch_size, vocab_size = logits.shape
|
||||
if targets.dim() != 2 or targets.shape[0] != batch_size:
|
||||
raise ValueError(
|
||||
"targets must be (B, M). "
|
||||
f"Got logits={tuple(logits.shape)}, targets={tuple(targets.shape)}"
|
||||
)
|
||||
if exposure.shape != (batch_size,):
|
||||
raise ValueError(
|
||||
f"exposure must be ({batch_size},), got {tuple(exposure.shape)}"
|
||||
)
|
||||
if dt.shape != targets.shape:
|
||||
raise ValueError(
|
||||
f"dt must match targets, got dt={tuple(dt.shape)}, "
|
||||
f"targets={tuple(targets.shape)}"
|
||||
)
|
||||
|
||||
rate = F.softplus(logits) + self.eps
|
||||
valid_vocab = _valid_vocab_mask(vocab_size, self.ignored_idx, logits.device)
|
||||
at_risk = _all_future_at_risk_mask(
|
||||
vocab_size,
|
||||
self.ignored_idx,
|
||||
logits,
|
||||
history,
|
||||
)
|
||||
target_valid, safe_targets = _all_future_target_mask(
|
||||
targets,
|
||||
at_risk,
|
||||
self.ignored_idx,
|
||||
)
|
||||
|
||||
penalty = exposure.to(rate.dtype) * rate[:, valid_vocab].sum(dim=-1)
|
||||
target_valid = torch.ones_like(targets, dtype=torch.bool, device=logits.device)
|
||||
for idx in self.ignored_idx:
|
||||
target_valid &= targets != idx
|
||||
censor_time = exposure.to(rate.dtype).clamp_min(self.eps)
|
||||
cumulative_at_censor = rate * censor_time.unsqueeze(1)
|
||||
penalty = (cumulative_at_censor * at_risk.to(rate.dtype)).sum(dim=-1)
|
||||
|
||||
# A first-onset outcome leaves the risk set at its event time, not at
|
||||
# the common end of follow-up.
|
||||
if targets.numel() > 0:
|
||||
event_time = dt.to(rate.dtype).clamp_min(self.eps)
|
||||
event_time = torch.minimum(event_time, censor_time.unsqueeze(1))
|
||||
event_cumulative = rate.gather(1, safe_targets) * event_time
|
||||
censor_cumulative = cumulative_at_censor.gather(1, safe_targets)
|
||||
penalty = penalty + (
|
||||
(event_cumulative - censor_cumulative)
|
||||
* target_valid.to(rate.dtype)
|
||||
).sum(dim=-1)
|
||||
|
||||
safe_targets = targets.clamp(min=0, max=vocab_size - 1)
|
||||
observed = rate.log().gather(1, safe_targets)
|
||||
observed = (observed * target_valid.to(rate.dtype)).sum(dim=-1)
|
||||
return (-observed + penalty).mean()
|
||||
|
||||
|
||||
class WeibullLoss(nn.Module):
|
||||
"""Query-conditioned all-future-event Weibull point-process loss."""
|
||||
"""First-onset all-future Weibull survival likelihood."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
ignored_idx: Iterable[int] = (PAD_IDX, CHECKUP_IDX),
|
||||
ignored_idx: Iterable[int] = (PAD_IDX, RESERVED_IDX),
|
||||
eps: float = 1e-8,
|
||||
):
|
||||
super().__init__()
|
||||
@@ -286,8 +282,9 @@ class WeibullLoss(nn.Module):
|
||||
targets: torch.Tensor,
|
||||
dt: torch.Tensor,
|
||||
exposure: torch.Tensor,
|
||||
history: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
_, vocab_size = logits.shape
|
||||
batch_size, vocab_size = logits.shape
|
||||
if weibull_rho is None:
|
||||
raise ValueError("weibull_rho is required for WeibullLoss")
|
||||
if weibull_rho.shape != logits.shape:
|
||||
@@ -295,23 +292,50 @@ class WeibullLoss(nn.Module):
|
||||
"weibull_rho must have the same shape as logits. "
|
||||
f"Got logits={tuple(logits.shape)}, weibull_rho={tuple(weibull_rho.shape)}"
|
||||
)
|
||||
if targets.dim() != 2 or targets.shape[0] != batch_size:
|
||||
raise ValueError(
|
||||
"targets must be (B, M). "
|
||||
f"Got logits={tuple(logits.shape)}, targets={tuple(targets.shape)}"
|
||||
)
|
||||
if dt.shape != targets.shape:
|
||||
raise ValueError(
|
||||
f"dt must match targets, got dt={tuple(dt.shape)}, "
|
||||
f"targets={tuple(targets.shape)}"
|
||||
)
|
||||
if exposure.shape != (batch_size,):
|
||||
raise ValueError(
|
||||
f"exposure must be ({batch_size},), got {tuple(exposure.shape)}"
|
||||
)
|
||||
|
||||
dtype = logits.dtype
|
||||
rate = F.softplus(logits) + self.eps
|
||||
rho = weibull_rho.to(device=logits.device, dtype=dtype).clamp_min(self.eps)
|
||||
valid_vocab = _valid_vocab_mask(vocab_size, self.ignored_idx, logits.device)
|
||||
at_risk = _all_future_at_risk_mask(
|
||||
vocab_size,
|
||||
self.ignored_idx,
|
||||
logits,
|
||||
history,
|
||||
)
|
||||
target_valid, safe_targets = _all_future_target_mask(
|
||||
targets,
|
||||
at_risk,
|
||||
self.ignored_idx,
|
||||
)
|
||||
|
||||
t_exp = exposure.to(dtype).clamp_min(self.eps).unsqueeze(1)
|
||||
penalty = (rate * torch.pow(t_exp, rho))[:, valid_vocab].sum(dim=-1)
|
||||
censor_time = exposure.to(dtype).clamp_min(self.eps)
|
||||
cumulative_at_censor = rate * torch.pow(censor_time.unsqueeze(1), rho)
|
||||
penalty = (cumulative_at_censor * at_risk.to(dtype)).sum(dim=-1)
|
||||
|
||||
target_valid = torch.ones_like(targets, dtype=torch.bool, device=logits.device)
|
||||
for idx in self.ignored_idx:
|
||||
target_valid &= targets != idx
|
||||
|
||||
safe_targets = targets.clamp(min=0, max=vocab_size - 1)
|
||||
target_rate = rate.gather(1, safe_targets)
|
||||
target_rho = rho.gather(1, safe_targets)
|
||||
target_dt = dt.to(dtype).clamp_min(self.eps)
|
||||
target_dt = torch.minimum(target_dt, censor_time.unsqueeze(1))
|
||||
event_cumulative = target_rate * torch.pow(target_dt, target_rho)
|
||||
censor_cumulative = cumulative_at_censor.gather(1, safe_targets)
|
||||
penalty = penalty + (
|
||||
(event_cumulative - censor_cumulative) * target_valid.to(dtype)
|
||||
).sum(dim=-1)
|
||||
|
||||
log_intensity = (
|
||||
target_rate.log()
|
||||
+ target_rho.log()
|
||||
@@ -321,80 +345,14 @@ class WeibullLoss(nn.Module):
|
||||
return (-observed + penalty).mean()
|
||||
|
||||
|
||||
class MixedLoss(nn.Module):
|
||||
"""Exponential diseases plus one Weibull death endpoint."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
death_idx: int,
|
||||
ignored_idx: Iterable[int] = (PAD_IDX, CHECKUP_IDX),
|
||||
eps: float = 1e-8,
|
||||
):
|
||||
super().__init__()
|
||||
self.death_idx = int(death_idx)
|
||||
self.ignored_idx = tuple(int(i) for i in ignored_idx)
|
||||
self.eps = eps
|
||||
|
||||
def forward(
|
||||
self,
|
||||
logits: torch.Tensor,
|
||||
death_rho: torch.Tensor,
|
||||
targets: torch.Tensor,
|
||||
dt: torch.Tensor,
|
||||
exposure: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
_, vocab_size = logits.shape
|
||||
dtype = logits.dtype
|
||||
rate = F.softplus(logits) + self.eps
|
||||
|
||||
if death_rho.dim() == 2:
|
||||
death_rho = death_rho.squeeze(-1)
|
||||
death_rho = death_rho.to(device=logits.device, dtype=dtype).clamp_min(self.eps)
|
||||
|
||||
valid_vocab = _valid_vocab_mask(vocab_size, self.ignored_idx, logits.device)
|
||||
valid_disease_vocab = valid_vocab.clone()
|
||||
valid_disease_vocab[self.death_idx] = False
|
||||
|
||||
t_exp = exposure.to(dtype).clamp_min(self.eps)
|
||||
disease_penalty = t_exp * rate[:, valid_disease_vocab].sum(dim=-1)
|
||||
death_rate = rate[:, self.death_idx]
|
||||
death_penalty = death_rate * torch.pow(t_exp, death_rho)
|
||||
penalty = disease_penalty + death_penalty
|
||||
|
||||
target_valid = torch.ones_like(targets, dtype=torch.bool, device=logits.device)
|
||||
for idx in self.ignored_idx:
|
||||
target_valid &= targets != idx
|
||||
|
||||
disease_event_mask = target_valid & (targets != self.death_idx)
|
||||
safe_targets = targets.clamp(min=0, max=vocab_size - 1)
|
||||
disease_log_rate = rate.log().gather(1, safe_targets)
|
||||
observed_disease = (disease_log_rate * disease_event_mask.to(dtype)).sum(dim=-1)
|
||||
|
||||
death_event_mask = target_valid & (targets == self.death_idx)
|
||||
death_observed = death_event_mask.any(dim=1)
|
||||
death_dt = (dt.to(dtype).clamp_min(self.eps) * death_event_mask.to(dtype)).sum(dim=1)
|
||||
death_log_intensity = (
|
||||
death_rate.log()
|
||||
+ death_rho.log()
|
||||
+ (death_rho - 1.0) * death_dt.clamp_min(self.eps).log()
|
||||
)
|
||||
observed_death = death_log_intensity * death_observed.to(dtype)
|
||||
|
||||
return (-observed_disease - observed_death + penalty).mean()
|
||||
|
||||
|
||||
def build_loss(name: str, **kwargs) -> nn.Module:
|
||||
name = name.lower()
|
||||
if name in {"delphi2m", "d2m", "next_token"}:
|
||||
if name == "delphi2m":
|
||||
return Delphi2MLoss(**kwargs)
|
||||
if name in {"uts", "unique_time_set", "unique_time_exponential"}:
|
||||
return UniqueTimeSetExponentialLoss(**kwargs)
|
||||
if name in {"exponential", "query_exponential"}:
|
||||
if name == "exponential":
|
||||
return ExponentialLoss(**kwargs)
|
||||
if name in {"weibull", "query_weibull"}:
|
||||
if name == "weibull":
|
||||
return WeibullLoss(**kwargs)
|
||||
if name in {"mixed", "query_mixed"}:
|
||||
return MixedLoss(**kwargs)
|
||||
raise ValueError(
|
||||
f"Unknown loss {name!r}. Available: delphi2m, uts, exponential, weibull, mixed."
|
||||
f"Unknown loss {name!r}. Available: delphi2m, exponential, weibull."
|
||||
)
|
||||
|
||||
131
model_architectures.py
Normal file
131
model_architectures.py
Normal file
@@ -0,0 +1,131 @@
|
||||
"""Model-architecture identifiers and checkpoint validation helpers."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from collections.abc import Mapping
|
||||
|
||||
|
||||
TRANSFORMER_FFN_ARCHITECTURE = "transformer_ffn_v1"
|
||||
TRAJ_MIXER_ARCHITECTURE = "traj_mixer_v5"
|
||||
DEFAULT_MODEL_ARCHITECTURE = TRANSFORMER_FFN_ARCHITECTURE
|
||||
SUPPORTED_MODEL_ARCHITECTURES = (
|
||||
TRANSFORMER_FFN_ARCHITECTURE,
|
||||
TRAJ_MIXER_ARCHITECTURE,
|
||||
)
|
||||
|
||||
|
||||
_FFN_STATE_KEY = re.compile(
|
||||
r"(?:^|\.)blocks\.\d+\.mlp\.w[123]\.(?:weight|bias)$"
|
||||
)
|
||||
_TRAJ_MIXER_STATE_KEY = re.compile(
|
||||
r"(?:^|\.)blocks\.\d+\.mlp\.(?:"
|
||||
r"norm\.(?:weight|bias)|"
|
||||
r"intra_gate_proj|"
|
||||
r"intra_value_proj|"
|
||||
r"intra_output_proj|"
|
||||
r"intra_gate_logits|"
|
||||
r"gate_proj|"
|
||||
r"value_proj|"
|
||||
r"output_proj"
|
||||
r")$"
|
||||
)
|
||||
|
||||
|
||||
def _validate_model_architecture(model_architecture: object) -> str:
|
||||
if not isinstance(model_architecture, str):
|
||||
raise ValueError(
|
||||
"model_architecture must be one of "
|
||||
f"{SUPPORTED_MODEL_ARCHITECTURES}, got {model_architecture!r}"
|
||||
)
|
||||
if model_architecture not in SUPPORTED_MODEL_ARCHITECTURES:
|
||||
raise ValueError(
|
||||
f"Unsupported model_architecture={model_architecture!r}; "
|
||||
f"expected one of {SUPPORTED_MODEL_ARCHITECTURES}."
|
||||
)
|
||||
return model_architecture
|
||||
|
||||
|
||||
def detect_model_architecture_from_state_dict(
|
||||
state_dict: Mapping[str, object],
|
||||
) -> str:
|
||||
"""Infer the architecture from block parameter names.
|
||||
|
||||
Detection deliberately accepts any ``blocks.<index>`` prefix rather than
|
||||
assuming that block zero is present.
|
||||
"""
|
||||
|
||||
if not isinstance(state_dict, Mapping):
|
||||
raise TypeError(
|
||||
"state_dict must be a mapping, got "
|
||||
f"{type(state_dict).__name__}"
|
||||
)
|
||||
|
||||
has_ffn = False
|
||||
has_traj_mixer = False
|
||||
for raw_key in state_dict:
|
||||
key = str(raw_key)
|
||||
has_ffn = has_ffn or _FFN_STATE_KEY.search(key) is not None
|
||||
has_traj_mixer = (
|
||||
has_traj_mixer
|
||||
or _TRAJ_MIXER_STATE_KEY.search(key) is not None
|
||||
)
|
||||
if has_ffn and has_traj_mixer:
|
||||
raise ValueError(
|
||||
"Checkpoint contains both Transformer FFN and TrajMixer "
|
||||
"block parameters; its model architecture is ambiguous."
|
||||
)
|
||||
|
||||
if has_ffn:
|
||||
return TRANSFORMER_FFN_ARCHITECTURE
|
||||
if has_traj_mixer:
|
||||
return TRAJ_MIXER_ARCHITECTURE
|
||||
raise ValueError(
|
||||
"Could not detect model architecture from checkpoint parameters. "
|
||||
"Expected a blocks.<index>.mlp FFN or TrajMixer parameter."
|
||||
)
|
||||
|
||||
|
||||
def resolve_model_architecture(
|
||||
config_or_marker: Mapping[str, object] | str | None = None,
|
||||
state_dict: Mapping[str, object] | None = None,
|
||||
) -> str:
|
||||
"""Resolve and cross-check a configured and checkpoint architecture.
|
||||
|
||||
Every saved run must provide an explicit architecture marker. Checkpoint
|
||||
parameter names are used only to verify that the marker describes the
|
||||
weights being loaded.
|
||||
"""
|
||||
|
||||
if isinstance(config_or_marker, Mapping):
|
||||
configured = config_or_marker.get("model_architecture")
|
||||
elif isinstance(config_or_marker, str) or config_or_marker is None:
|
||||
configured = config_or_marker
|
||||
else:
|
||||
raise TypeError(
|
||||
"config_or_marker must be a config mapping, string, or None, got "
|
||||
f"{type(config_or_marker).__name__}"
|
||||
)
|
||||
|
||||
resolved_config = (
|
||||
_validate_model_architecture(configured)
|
||||
if configured is not None
|
||||
else None
|
||||
)
|
||||
detected = (
|
||||
detect_model_architecture_from_state_dict(state_dict)
|
||||
if state_dict is not None
|
||||
else None
|
||||
)
|
||||
|
||||
if resolved_config is None:
|
||||
raise ValueError(
|
||||
"model_architecture is required; expected one of "
|
||||
f"{SUPPORTED_MODEL_ARCHITECTURES}."
|
||||
)
|
||||
if detected is not None and resolved_config != detected:
|
||||
raise ValueError(
|
||||
"Configured model architecture conflicts with checkpoint: "
|
||||
f"config={resolved_config!r}, checkpoint={detected!r}."
|
||||
)
|
||||
return resolved_config
|
||||
609
models.py
609
models.py
@@ -1,4 +1,3 @@
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
|
||||
import torch
|
||||
@@ -8,199 +7,14 @@ import torch.nn.functional as F
|
||||
from backbones import (
|
||||
AgeSinusoidalEncoding,
|
||||
GaussianRBFTimeBasis,
|
||||
SharedEventTrajectoryCore,
|
||||
TimeRoPE,
|
||||
TokenAutoDiscretization,
|
||||
build_backbone_block,
|
||||
)
|
||||
from model_architectures import resolve_model_architecture
|
||||
from targets import PAD_IDX
|
||||
|
||||
|
||||
EVENT_TRAJECTORY_ARCHITECTURE = "event_trajectory_shared_v2"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class EventTrajectoryModelSize:
|
||||
d_model: int
|
||||
n_trajectory: int
|
||||
|
||||
@property
|
||||
def trajectory_dim(self) -> int:
|
||||
return self.d_model // self.n_trajectory
|
||||
|
||||
@property
|
||||
def traj_hidden(self) -> int:
|
||||
return 4 * self.n_trajectory
|
||||
|
||||
|
||||
MODEL_SIZE_PRESETS = {
|
||||
"nano": EventTrajectoryModelSize(d_model=120, n_trajectory=6),
|
||||
"tiny": EventTrajectoryModelSize(d_model=256, n_trajectory=8),
|
||||
"small": EventTrajectoryModelSize(d_model=512, n_trajectory=8),
|
||||
"medium": EventTrajectoryModelSize(d_model=768, n_trajectory=12),
|
||||
"huge": EventTrajectoryModelSize(d_model=1024, n_trajectory=16),
|
||||
}
|
||||
MODEL_SIZE_NAMES = tuple(MODEL_SIZE_PRESETS)
|
||||
|
||||
|
||||
def resolve_model_size(model_size: str) -> EventTrajectoryModelSize:
|
||||
if not isinstance(model_size, str):
|
||||
raise ValueError(
|
||||
f"model_size must be a string, got {type(model_size).__name__}"
|
||||
)
|
||||
normalized = model_size.strip().lower()
|
||||
try:
|
||||
return MODEL_SIZE_PRESETS[normalized]
|
||||
except KeyError as exc:
|
||||
choices = ", ".join(MODEL_SIZE_NAMES)
|
||||
raise ValueError(
|
||||
f"Unknown model_size {model_size!r}; expected one of: {choices}"
|
||||
) from exc
|
||||
|
||||
|
||||
def _required_config_int(
|
||||
config: Mapping[str, object],
|
||||
key: str,
|
||||
) -> int:
|
||||
raw_value = config.get(key)
|
||||
if isinstance(raw_value, bool):
|
||||
raise ValueError(f"Config field {key!r} must be an integer")
|
||||
try:
|
||||
value = int(raw_value)
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise ValueError(
|
||||
f"Config field {key!r} must be present and integer-valued; "
|
||||
f"got {raw_value!r}"
|
||||
) from exc
|
||||
if isinstance(raw_value, float) and not raw_value.is_integer():
|
||||
raise ValueError(f"Config field {key!r} must be an integer")
|
||||
return value
|
||||
|
||||
|
||||
def validate_event_trajectory_config(config: Mapping[str, object]) -> None:
|
||||
actual = config.get("model_architecture")
|
||||
if actual != EVENT_TRAJECTORY_ARCHITECTURE:
|
||||
raise ValueError(
|
||||
"This branch only accepts models trained with the shared "
|
||||
"event-trajectory architecture marker "
|
||||
f"{EVENT_TRAJECTORY_ARCHITECTURE!r}; got {actual!r}."
|
||||
)
|
||||
raw_model_size = config.get("model_size")
|
||||
if not isinstance(raw_model_size, str):
|
||||
raise ValueError(
|
||||
"Config field 'model_size' must be one of: "
|
||||
+ ", ".join(MODEL_SIZE_NAMES)
|
||||
)
|
||||
model_size = raw_model_size.strip().lower()
|
||||
preset = resolve_model_size(model_size)
|
||||
d_model = _required_config_int(config, "d_model")
|
||||
n_trajectory = _required_config_int(config, "n_trajectory")
|
||||
n_reasoning_rounds = _required_config_int(
|
||||
config,
|
||||
"n_reasoning_rounds",
|
||||
)
|
||||
trajectory_dim = _required_config_int(config, "trajectory_dim")
|
||||
traj_hidden = _required_config_int(config, "traj_hidden")
|
||||
if n_reasoning_rounds <= 0:
|
||||
raise ValueError(
|
||||
"n_reasoning_rounds must be positive"
|
||||
)
|
||||
expected_values = {
|
||||
"d_model": preset.d_model,
|
||||
"n_trajectory": preset.n_trajectory,
|
||||
"trajectory_dim": preset.trajectory_dim,
|
||||
"traj_hidden": preset.traj_hidden,
|
||||
}
|
||||
actual_values = {
|
||||
"d_model": d_model,
|
||||
"n_trajectory": n_trajectory,
|
||||
"trajectory_dim": trajectory_dim,
|
||||
"traj_hidden": traj_hidden,
|
||||
}
|
||||
mismatches = [
|
||||
f"{key}: expected {expected}, got {actual_values[key]}"
|
||||
for key, expected in expected_values.items()
|
||||
if actual_values[key] != expected
|
||||
]
|
||||
if mismatches:
|
||||
raise ValueError(
|
||||
f"Config does not match model_size={model_size!r}: "
|
||||
+ "; ".join(mismatches)
|
||||
)
|
||||
|
||||
|
||||
def _checkpoint_scalar_int(
|
||||
state_dict: Mapping[str, object],
|
||||
key: str,
|
||||
) -> int:
|
||||
value = state_dict[key]
|
||||
if not isinstance(value, torch.Tensor) or value.numel() != 1:
|
||||
raise ValueError(
|
||||
f"Checkpoint architecture field {key!r} must be a scalar tensor"
|
||||
)
|
||||
return int(value.detach().cpu().item())
|
||||
|
||||
|
||||
def validate_event_trajectory_state_dict(
|
||||
state_dict: Mapping[str, object],
|
||||
*,
|
||||
expected_d_model: int | None = None,
|
||||
expected_n_trajectory: int | None = None,
|
||||
expected_n_reasoning_rounds: int | None = None,
|
||||
) -> None:
|
||||
required_keys = {
|
||||
"architecture_d_model",
|
||||
"architecture_n_trajectory",
|
||||
"architecture_n_reasoning_rounds",
|
||||
"event_projection.weight",
|
||||
"trajectory_prototypes",
|
||||
"query_projection.weight",
|
||||
"reasoning_core.cross_attention.q_proj.weight",
|
||||
"reasoning_core.cross_attention.k_proj.weight",
|
||||
"reasoning_core.cross_attention.v_proj.weight",
|
||||
"reasoning_core.traj_mixer.gate_proj",
|
||||
"reasoning_core.traj_mixer.value_proj",
|
||||
"reasoning_core.traj_mixer.output_proj",
|
||||
"reasoning_core.attn_scale",
|
||||
"reasoning_core.mixer_scale",
|
||||
}
|
||||
missing = sorted(required_keys.difference(state_dict))
|
||||
if missing:
|
||||
raise ValueError(
|
||||
"Checkpoint is not a shared event-trajectory checkpoint; "
|
||||
"missing required "
|
||||
f"parameters: {', '.join(missing)}"
|
||||
)
|
||||
checkpoint_values = {
|
||||
"d_model": _checkpoint_scalar_int(
|
||||
state_dict,
|
||||
"architecture_d_model",
|
||||
),
|
||||
"n_trajectory": _checkpoint_scalar_int(
|
||||
state_dict,
|
||||
"architecture_n_trajectory",
|
||||
),
|
||||
"n_reasoning_rounds": _checkpoint_scalar_int(
|
||||
state_dict,
|
||||
"architecture_n_reasoning_rounds",
|
||||
),
|
||||
}
|
||||
expected_values = {
|
||||
"d_model": expected_d_model,
|
||||
"n_trajectory": expected_n_trajectory,
|
||||
"n_reasoning_rounds": expected_n_reasoning_rounds,
|
||||
}
|
||||
mismatches = [
|
||||
f"{name}: checkpoint={checkpoint_values[name]}, expected={expected}"
|
||||
for name, expected in expected_values.items()
|
||||
if expected is not None and checkpoint_values[name] != expected
|
||||
]
|
||||
if mismatches:
|
||||
raise ValueError(
|
||||
"Checkpoint architecture does not match the constructed model: "
|
||||
+ "; ".join(mismatches)
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class DeepHealthOutput:
|
||||
hidden: torch.Tensor
|
||||
@@ -223,6 +37,8 @@ class OtherInfoTokenizer(nn.Module):
|
||||
cont_type_ids: list[int],
|
||||
n_value_kinds: int = 3,
|
||||
n_bins: int = 16,
|
||||
continuous_value_center: torch.Tensor | list[float] | None = None,
|
||||
continuous_value_scale: torch.Tensor | list[float] | None = None,
|
||||
):
|
||||
super().__init__()
|
||||
if len(cont_type_ids) != n_cont_types:
|
||||
@@ -240,7 +56,6 @@ class OtherInfoTokenizer(nn.Module):
|
||||
raise ValueError(
|
||||
f"n_value_kinds must be > {self.CATE_KIND}, got {n_value_kinds}"
|
||||
)
|
||||
|
||||
self.type_emb = nn.Embedding(n_types, n_embd, padding_idx=0)
|
||||
self.kind_emb = nn.Embedding(n_value_kinds, n_embd, padding_idx=0)
|
||||
self.cont_value_encoder = (
|
||||
@@ -257,6 +72,38 @@ class OtherInfoTokenizer(nn.Module):
|
||||
n_embd,
|
||||
padding_idx=0,
|
||||
)
|
||||
if n_cont_types > 0:
|
||||
if continuous_value_center is None or continuous_value_scale is None:
|
||||
raise ValueError(
|
||||
"Continuous values require train-split RobustScale center "
|
||||
"and scale statistics"
|
||||
)
|
||||
center = self._coerce_scaler_buffer(
|
||||
continuous_value_center,
|
||||
n_cont_types=n_cont_types,
|
||||
name="continuous_value_center",
|
||||
)
|
||||
scale = self._coerce_scaler_buffer(
|
||||
continuous_value_scale,
|
||||
n_cont_types=n_cont_types,
|
||||
name="continuous_value_scale",
|
||||
)
|
||||
if not torch.isfinite(center).all():
|
||||
raise ValueError(
|
||||
"continuous_value_center must contain only finite values"
|
||||
)
|
||||
if not torch.isfinite(scale).all() or torch.any(scale <= 0):
|
||||
raise ValueError(
|
||||
"continuous_value_scale must be finite and strictly positive"
|
||||
)
|
||||
self.register_buffer("continuous_value_center", center)
|
||||
self.register_buffer("continuous_value_scale", scale)
|
||||
else:
|
||||
# ``None`` buffers are omitted from state_dict. This preserves the
|
||||
# exact checkpoint schema used by models trained before continuous
|
||||
# value scaling was introduced.
|
||||
self.register_buffer("continuous_value_center", None)
|
||||
self.register_buffer("continuous_value_scale", None)
|
||||
|
||||
cont_type_index = torch.full((n_types,), -1, dtype=torch.long)
|
||||
for idx, type_id in enumerate(cont_type_ids):
|
||||
@@ -272,6 +119,22 @@ class OtherInfoTokenizer(nn.Module):
|
||||
)
|
||||
self.reset_parameters()
|
||||
|
||||
@staticmethod
|
||||
def _coerce_scaler_buffer(
|
||||
value: torch.Tensor | list[float] | None,
|
||||
*,
|
||||
n_cont_types: int,
|
||||
name: str,
|
||||
) -> torch.Tensor:
|
||||
if value is None:
|
||||
raise ValueError(f"{name} is required")
|
||||
tensor = torch.as_tensor(value, dtype=torch.float32).detach().clone()
|
||||
if tensor.shape != (n_cont_types,):
|
||||
raise ValueError(
|
||||
f"{name} must have shape ({n_cont_types},), got {tuple(tensor.shape)}"
|
||||
)
|
||||
return tensor
|
||||
|
||||
def reset_parameters(self) -> None:
|
||||
nn.init.normal_(self.type_emb.weight, mean=0.0, std=0.02)
|
||||
nn.init.zeros_(self.type_emb.weight[0])
|
||||
@@ -313,9 +176,18 @@ class OtherInfoTokenizer(nn.Module):
|
||||
f"type_id={bad_type} is marked continuous but is not in "
|
||||
"cont_type_ids"
|
||||
)
|
||||
cont_value = other_value[cont_pos].to(type_emb.dtype)
|
||||
if (
|
||||
self.continuous_value_center is None
|
||||
or self.continuous_value_scale is None
|
||||
):
|
||||
raise RuntimeError("RobustScale buffers are missing")
|
||||
center = self.continuous_value_center[cont_idx].to(type_emb.dtype)
|
||||
scale = self.continuous_value_scale[cont_idx].to(type_emb.dtype)
|
||||
cont_value = (cont_value - center) / scale
|
||||
value_emb[cont_pos] = self.cont_value_encoder(
|
||||
cont_type_idx=cont_idx,
|
||||
value=other_value[cont_pos].to(type_emb.dtype),
|
||||
value=cont_value,
|
||||
)
|
||||
|
||||
cate_pos = other_valid & (other_value_kind == self.CATE_KIND)
|
||||
@@ -332,19 +204,25 @@ class DeepHealth(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
vocab_size: int,
|
||||
model_size: str,
|
||||
n_reasoning_rounds: int,
|
||||
n_embd: int,
|
||||
n_head: int,
|
||||
n_layer: int,
|
||||
n_types: int,
|
||||
n_cont_types: int,
|
||||
n_categories: int,
|
||||
cont_type_ids: list[int],
|
||||
n_value_kinds: int = 3,
|
||||
n_bins: int = 16,
|
||||
continuous_value_center: torch.Tensor | list[float] | None = None,
|
||||
continuous_value_scale: torch.Tensor | list[float] | None = None,
|
||||
target_mode: str = "next_token", # "next_token" or "all_future"
|
||||
time_mode: str = "relative", # "relative" or "absolute"
|
||||
dist_mode: str = "exponential", # "exponential", "weibull" or "mixed"
|
||||
time_mode: str = "absolute", # next_token requires absolute
|
||||
dist_mode: str = "exponential", # "exponential" or "weibull"
|
||||
extra_pool_reduce: str = "mean",
|
||||
dropout: float = 0.0,
|
||||
model_architecture: str | None = None,
|
||||
risk_head_bias: bool = False,
|
||||
risk_head_bias_init: torch.Tensor | list[float] | None = None,
|
||||
):
|
||||
super().__init__()
|
||||
if target_mode not in ["next_token", "all_future"]:
|
||||
@@ -353,132 +231,125 @@ class DeepHealth(nn.Module):
|
||||
if time_mode not in ["relative", "absolute"]:
|
||||
raise ValueError(
|
||||
"time_mode must be either 'relative' or 'absolute'")
|
||||
if dist_mode not in ["exponential", "weibull", "mixed"]:
|
||||
if target_mode == "next_token" and time_mode != "absolute":
|
||||
raise ValueError(
|
||||
"dist_mode must be either 'exponential', 'weibull' or 'mixed'")
|
||||
"next_token is reserved for Delphi2M reproduction and "
|
||||
"requires time_mode='absolute'"
|
||||
)
|
||||
if dist_mode not in ["exponential", "weibull"]:
|
||||
raise ValueError(
|
||||
"dist_mode must be either 'exponential' or 'weibull'")
|
||||
if extra_pool_reduce not in {"mean", "sum"}:
|
||||
raise ValueError("extra_pool_reduce must be either 'mean' or 'sum'")
|
||||
if n_reasoning_rounds <= 0:
|
||||
raise ValueError(
|
||||
"n_reasoning_rounds must be positive, got "
|
||||
f"{n_reasoning_rounds}"
|
||||
)
|
||||
size_config = resolve_model_size(model_size)
|
||||
normalized_model_size = model_size.strip().lower()
|
||||
d_model = size_config.d_model
|
||||
n_trajectory = size_config.n_trajectory
|
||||
|
||||
self.token_embedding = nn.Embedding(vocab_size, d_model, padding_idx=0)
|
||||
if n_layer < 1:
|
||||
raise ValueError(f"n_layer must be >= 1, got {n_layer}")
|
||||
model_architecture = resolve_model_architecture(model_architecture)
|
||||
self.token_embedding = nn.Embedding(vocab_size, n_embd, padding_idx=0)
|
||||
self.gender_embedding = nn.Embedding(
|
||||
2, d_model) # Assuming binary gender
|
||||
2, n_embd) # Assuming binary gender
|
||||
self.tokenizer = OtherInfoTokenizer(
|
||||
n_embd=d_model,
|
||||
n_embd=n_embd,
|
||||
n_types=n_types,
|
||||
n_cont_types=n_cont_types,
|
||||
n_categories=n_categories,
|
||||
cont_type_ids=cont_type_ids,
|
||||
n_value_kinds=n_value_kinds,
|
||||
n_bins=n_bins,
|
||||
continuous_value_center=continuous_value_center,
|
||||
continuous_value_scale=continuous_value_scale,
|
||||
)
|
||||
self.target_mode = target_mode
|
||||
self.time_mode = time_mode
|
||||
self.dist_mode = dist_mode
|
||||
self.extra_pool_reduce = extra_pool_reduce
|
||||
self.model_size = normalized_model_size
|
||||
self.d_model = d_model
|
||||
self.n_trajectory = n_trajectory
|
||||
self.trajectory_dim = d_model // n_trajectory
|
||||
self.traj_hidden = 4 * n_trajectory
|
||||
self.n_reasoning_rounds = n_reasoning_rounds
|
||||
self.model_architecture = model_architecture
|
||||
self.n_layer = n_layer
|
||||
self.n_embd = n_embd
|
||||
self.vocab_size = vocab_size
|
||||
self.register_buffer(
|
||||
"architecture_d_model",
|
||||
torch.tensor(d_model, dtype=torch.int64),
|
||||
)
|
||||
self.register_buffer(
|
||||
"architecture_n_trajectory",
|
||||
torch.tensor(n_trajectory, dtype=torch.int64),
|
||||
)
|
||||
self.register_buffer(
|
||||
"architecture_n_reasoning_rounds",
|
||||
torch.tensor(n_reasoning_rounds, dtype=torch.int64),
|
||||
)
|
||||
nn.init.normal_(self.token_embedding.weight, mean=0.0, std=0.02)
|
||||
nn.init.zeros_(self.token_embedding.weight[0])
|
||||
nn.init.normal_(self.gender_embedding.weight, mean=0.0, std=0.02)
|
||||
if dist_mode == "weibull":
|
||||
self.rho_head = nn.Linear(d_model, vocab_size)
|
||||
self.rho_head = nn.Linear(n_embd, vocab_size)
|
||||
nn.init.zeros_(self.rho_head.weight)
|
||||
nn.init.constant_(self.rho_head.bias, 0.5413)
|
||||
|
||||
if dist_mode == "mixed":
|
||||
self.death_idx = vocab_size - 1
|
||||
self.rho_death_head = nn.Linear(d_model, 1)
|
||||
nn.init.zeros_(self.rho_death_head.weight)
|
||||
nn.init.constant_(self.rho_death_head.bias, 0.5413)
|
||||
|
||||
# Event and query time are encoded once before shared reasoning. In
|
||||
# relative mode, cross-attention additionally uses TimeRoPE and RBF.
|
||||
self.age_encoding = AgeSinusoidalEncoding(d_model)
|
||||
self.event_projection = nn.Linear(d_model, d_model, bias=False)
|
||||
self.event_norm = nn.LayerNorm(d_model)
|
||||
self.query_projection = nn.Linear(d_model, d_model, bias=False)
|
||||
self.trajectory_prototypes = nn.Parameter(
|
||||
torch.empty(n_trajectory, self.trajectory_dim)
|
||||
)
|
||||
self.query_token = nn.Parameter(torch.empty(d_model))
|
||||
nn.init.normal_(self.event_projection.weight, mean=0.0, std=0.02)
|
||||
nn.init.normal_(self.query_projection.weight, mean=0.0, std=0.02)
|
||||
nn.init.normal_(self.trajectory_prototypes, mean=0.0, std=0.02)
|
||||
nn.init.normal_(self.query_token, mean=0.0, std=0.02)
|
||||
|
||||
use_relative_time = time_mode == "relative"
|
||||
self.reasoning_core = SharedEventTrajectoryCore(
|
||||
d_model=d_model,
|
||||
n_trajectory=n_trajectory,
|
||||
n_reasoning_rounds=n_reasoning_rounds,
|
||||
dropout=dropout,
|
||||
n_rbf_bases=16,
|
||||
use_time_rope=use_relative_time,
|
||||
use_rbf_bias=use_relative_time,
|
||||
)
|
||||
if use_relative_time:
|
||||
self.rope = TimeRoPE(self.trajectory_dim)
|
||||
self.rbf = GaussianRBFTimeBasis(
|
||||
n_bases=16,
|
||||
max_time_diff=40.0,
|
||||
)
|
||||
else:
|
||||
if time_mode == "absolute":
|
||||
self.age_encoding = AgeSinusoidalEncoding(n_embd)
|
||||
self.blocks = nn.ModuleList([
|
||||
build_backbone_block(
|
||||
model_architecture,
|
||||
n_embd=n_embd,
|
||||
n_head=n_head,
|
||||
use_time_rope=False,
|
||||
use_rbf_bias=False,
|
||||
mlp_dropout=dropout,
|
||||
) for _ in range(n_layer)
|
||||
])
|
||||
self.rope = None
|
||||
self.rbf = None
|
||||
elif time_mode == "relative":
|
||||
self.age_encoding = None
|
||||
self.blocks = nn.ModuleList([
|
||||
build_backbone_block(
|
||||
model_architecture,
|
||||
n_embd=n_embd,
|
||||
n_head=n_head,
|
||||
use_time_rope=True,
|
||||
use_rbf_bias=True,
|
||||
mlp_dropout=dropout,
|
||||
) for _ in range(n_layer)
|
||||
])
|
||||
self.rope = TimeRoPE(n_embd // n_head)
|
||||
self.rbf = GaussianRBFTimeBasis(n_bases=16, max_time_diff=40.0)
|
||||
|
||||
self.final_ln = nn.LayerNorm(d_model)
|
||||
self.risk_head = nn.Linear(d_model, vocab_size, bias=False)
|
||||
self.final_ln = nn.LayerNorm(n_embd)
|
||||
self.risk_head = nn.Linear(
|
||||
n_embd,
|
||||
vocab_size,
|
||||
bias=bool(risk_head_bias),
|
||||
)
|
||||
if risk_head_bias_init is not None:
|
||||
if self.risk_head.bias is None:
|
||||
raise ValueError(
|
||||
"risk_head_bias_init requires risk_head_bias=True"
|
||||
)
|
||||
initial_bias = torch.as_tensor(
|
||||
risk_head_bias_init,
|
||||
dtype=self.risk_head.bias.dtype,
|
||||
device=self.risk_head.bias.device,
|
||||
)
|
||||
if initial_bias.shape != (vocab_size,):
|
||||
raise ValueError(
|
||||
"risk_head_bias_init must have shape "
|
||||
f"({vocab_size},), got {tuple(initial_bias.shape)}"
|
||||
)
|
||||
if not torch.isfinite(initial_bias).all():
|
||||
raise ValueError("risk_head_bias_init must contain only finite values")
|
||||
with torch.no_grad():
|
||||
# Start exactly at the fitted marginal baseline; covariate and
|
||||
# history effects are learned away from zero during training.
|
||||
self.risk_head.weight.zero_()
|
||||
self.risk_head.bias.copy_(initial_bias)
|
||||
if target_mode == "next_token":
|
||||
self.risk_head.weight = self.token_embedding.weight
|
||||
self.query_token = nn.Parameter(torch.zeros(n_embd))
|
||||
nn.init.normal_(self.query_token, mean=0.0, std=0.02)
|
||||
|
||||
def _make_event_invalid_mask(
|
||||
def _make_history_attn_mask(
|
||||
self,
|
||||
event_valid_mask: torch.Tensor,
|
||||
event_time: torch.Tensor,
|
||||
query_time: torch.Tensor,
|
||||
query_position: torch.Tensor | None = None,
|
||||
padding_mask: torch.Tensor,
|
||||
time_seq: torch.Tensor,
|
||||
dtype: torch.dtype,
|
||||
) -> torch.Tensor:
|
||||
valid_key = event_valid_mask[:, None, :]
|
||||
key_time = event_time[:, None, :]
|
||||
query_time = query_time[:, :, None]
|
||||
if query_position is None:
|
||||
visible_by_time = key_time <= query_time
|
||||
else:
|
||||
key_position = torch.arange(
|
||||
event_time.size(1),
|
||||
device=event_time.device,
|
||||
).view(1, 1, -1)
|
||||
visible_by_time = (key_time < query_time) | (
|
||||
(key_time == query_time)
|
||||
& (key_position <= query_position[:, :, None])
|
||||
)
|
||||
return ~(valid_key & visible_by_time)
|
||||
valid_key = padding_mask[:, None, :] # (B, 1, L)
|
||||
visible_by_time = time_seq[:, None, :] <= time_seq[:, :, None]
|
||||
valid = valid_key & visible_by_time
|
||||
return torch.zeros(
|
||||
valid.shape,
|
||||
device=valid.device,
|
||||
dtype=dtype,
|
||||
).masked_fill(~valid, -1e4)[:, None, :, :]
|
||||
|
||||
def _pool_other_by_time(
|
||||
self,
|
||||
@@ -582,8 +453,8 @@ class DeepHealth(nn.Module):
|
||||
padding_mask = padding_mask.to(device=event_seq.device, dtype=torch.bool)
|
||||
|
||||
event_len = event_seq.size(1)
|
||||
event_features = self.token_embedding(event_seq)
|
||||
event_time = time_seq
|
||||
h_disease = self.token_embedding(event_seq)
|
||||
t_disease = time_seq
|
||||
|
||||
if other_time.shape != other_type.shape:
|
||||
raise ValueError(
|
||||
@@ -591,120 +462,64 @@ class DeepHealth(nn.Module):
|
||||
f"{tuple(other_time.shape)} vs {tuple(other_type.shape)}"
|
||||
)
|
||||
other_time = other_time.to(device=event_seq.device, dtype=time_seq.dtype)
|
||||
other_features, other_mask = self.tokenizer(
|
||||
h_other, other_mask = self.tokenizer(
|
||||
other_type=other_type,
|
||||
other_value=other_value,
|
||||
other_value_kind=other_value_kind,
|
||||
)
|
||||
other_features = other_features.to(device=event_seq.device)
|
||||
h_other = h_other.to(device=event_seq.device)
|
||||
other_mask = other_mask.to(device=event_seq.device, dtype=torch.bool)
|
||||
|
||||
event_features = torch.cat([event_features, other_features], dim=1)
|
||||
event_time = torch.cat([event_time, other_time], dim=1)
|
||||
event_valid_mask = torch.cat([padding_mask, other_mask], dim=1)
|
||||
|
||||
batch_size = event_seq.size(0)
|
||||
sex_context = self.gender_embedding(sex)[:, None, :]
|
||||
event_features = (
|
||||
event_features
|
||||
+ sex_context
|
||||
+ self.age_encoding(event_time)
|
||||
)
|
||||
event_features = event_features * event_valid_mask.unsqueeze(-1).to(
|
||||
event_features.dtype
|
||||
)
|
||||
event_memory = self.event_norm(
|
||||
self.event_projection(event_features)
|
||||
)
|
||||
event_memory = event_memory * event_valid_mask.unsqueeze(-1).to(
|
||||
event_memory.dtype
|
||||
)
|
||||
h_disease = torch.cat([h_disease, h_other], dim=1)
|
||||
t_disease = torch.cat([t_disease, other_time], dim=1)
|
||||
padding_mask = torch.cat([padding_mask, other_mask], dim=1)
|
||||
h_disease = h_disease * padding_mask.unsqueeze(-1).to(h_disease.dtype)
|
||||
|
||||
if mode == "all_future":
|
||||
query_time = t_query[:, None]
|
||||
query_position = None
|
||||
query_features = (
|
||||
self.query_token.view(1, 1, -1)
|
||||
+ sex_context
|
||||
+ self.age_encoding(query_time)
|
||||
)
|
||||
query_valid_mask = torch.ones(
|
||||
batch_size = event_seq.size(0)
|
||||
query = self.query_token.view(1, 1, -1).expand(batch_size, 1, -1)
|
||||
h_disease = torch.cat([h_disease, query], dim=1)
|
||||
t_disease = torch.cat([t_disease, t_query[:, None]], dim=1)
|
||||
query_mask = torch.ones(
|
||||
batch_size,
|
||||
1,
|
||||
dtype=torch.bool,
|
||||
device=event_seq.device,
|
||||
)
|
||||
else:
|
||||
# Each event position is an independent parallel query. Including
|
||||
# its event feature preserves token-level next-step semantics.
|
||||
# Equal-time memory is additionally position-causal so a token
|
||||
# cannot read a later token that may be its Delphi2M target.
|
||||
query_time = event_time
|
||||
query_position = torch.arange(
|
||||
event_time.size(1),
|
||||
device=event_time.device,
|
||||
).view(1, -1).expand(batch_size, -1)
|
||||
query_features = event_features
|
||||
query_valid_mask = event_valid_mask
|
||||
padding_mask = torch.cat([padding_mask, query_mask], dim=1)
|
||||
|
||||
n_query = query_time.size(1)
|
||||
query_context = self.query_projection(query_features).reshape(
|
||||
batch_size,
|
||||
n_query,
|
||||
self.n_trajectory,
|
||||
self.trajectory_dim,
|
||||
)
|
||||
trajectory_state = (
|
||||
self.trajectory_prototypes.view(
|
||||
1,
|
||||
1,
|
||||
self.n_trajectory,
|
||||
self.trajectory_dim,
|
||||
)
|
||||
+ query_context
|
||||
)
|
||||
event_invalid_mask = self._make_event_invalid_mask(
|
||||
event_valid_mask=event_valid_mask,
|
||||
event_time=event_time,
|
||||
query_time=query_time,
|
||||
query_position=query_position,
|
||||
)
|
||||
sex_emb = self.gender_embedding(sex)[:, None, :]
|
||||
h_disease = h_disease + sex_emb
|
||||
h_disease = h_disease * padding_mask.unsqueeze(-1).to(h_disease.dtype)
|
||||
|
||||
event_rope_cache = None
|
||||
query_rope_cache = None
|
||||
rope_cache = None
|
||||
rbf_cache = None
|
||||
if self.time_mode == "relative":
|
||||
if self.rope is None or self.rbf is None:
|
||||
raise RuntimeError("Relative-time modules are not initialized")
|
||||
event_rope_cache = self.rope.precompute_cache(event_time)
|
||||
query_rope_cache = self.rope.precompute_cache(query_time)
|
||||
rbf_cache = self.rbf.precompute_cross_cache(
|
||||
query_time,
|
||||
event_time,
|
||||
)
|
||||
if self.time_mode == "absolute":
|
||||
h_disease = h_disease + self.age_encoding(t_disease)
|
||||
h_disease = h_disease * padding_mask.unsqueeze(-1).to(h_disease.dtype)
|
||||
elif self.time_mode == "relative":
|
||||
rope_cache = self.rope.precompute_cache(t_disease)
|
||||
rbf_cache = self.rbf.precompute_cache(t_disease)
|
||||
|
||||
event_key_value = self.reasoning_core.project_event_memory(
|
||||
event_memory,
|
||||
event_rope_cache=event_rope_cache,
|
||||
attn_mask = self._make_history_attn_mask(
|
||||
padding_mask=padding_mask,
|
||||
time_seq=t_disease,
|
||||
dtype=h_disease.dtype,
|
||||
)
|
||||
for _ in range(self.n_reasoning_rounds):
|
||||
trajectory_state = self.reasoning_core(
|
||||
trajectory_state=trajectory_state,
|
||||
event_key_value=event_key_value,
|
||||
event_invalid_mask=event_invalid_mask,
|
||||
query_rope_cache=query_rope_cache,
|
||||
for block in self.blocks:
|
||||
h_disease = block(
|
||||
h_disease,
|
||||
rope_cache=rope_cache,
|
||||
rbf_cache=rbf_cache,
|
||||
attn_mask=attn_mask,
|
||||
)
|
||||
h_disease = h_disease * padding_mask.unsqueeze(-1).to(h_disease.dtype)
|
||||
|
||||
hidden_sequence = self.final_ln(
|
||||
trajectory_state.reshape(batch_size, n_query, self.d_model)
|
||||
)
|
||||
hidden_sequence = hidden_sequence * query_valid_mask.unsqueeze(-1).to(
|
||||
hidden_sequence.dtype
|
||||
)
|
||||
h_disease = self.final_ln(h_disease)
|
||||
h_disease = h_disease * padding_mask.unsqueeze(-1).to(h_disease.dtype)
|
||||
|
||||
if mode == "all_future":
|
||||
hidden = hidden_sequence[:, 0, :]
|
||||
hidden = h_disease[:, -1, :]
|
||||
if return_output:
|
||||
return DeepHealthOutput(
|
||||
hidden=hidden,
|
||||
@@ -719,13 +534,13 @@ class DeepHealth(nn.Module):
|
||||
)
|
||||
return hidden
|
||||
if return_output:
|
||||
h_event = hidden_sequence[:, :event_len, :]
|
||||
t_event = event_time[:, :event_len]
|
||||
event_mask = event_valid_mask[:, :event_len]
|
||||
h_event = h_disease[:, :event_len, :]
|
||||
t_event = t_disease[:, :event_len]
|
||||
event_mask = padding_mask[:, :event_len]
|
||||
h_extra, t_extra, extra_mask = self._pool_other_by_time(
|
||||
h_other=hidden_sequence[:, event_len:, :],
|
||||
other_time=event_time[:, event_len:],
|
||||
other_mask=event_valid_mask[:, event_len:],
|
||||
h_other=h_disease[:, event_len:, :],
|
||||
other_time=t_disease[:, event_len:],
|
||||
other_mask=padding_mask[:, event_len:],
|
||||
)
|
||||
return DeepHealthOutput(
|
||||
hidden=torch.cat([h_event, h_extra], dim=1),
|
||||
@@ -733,17 +548,10 @@ class DeepHealth(nn.Module):
|
||||
padding_mask=torch.cat([event_mask, extra_mask], dim=1),
|
||||
event_len=event_len,
|
||||
)
|
||||
return hidden_sequence[:, :event_len, :]
|
||||
return h_disease[:, :event_len, :]
|
||||
|
||||
def forward_next_token(self, **kwargs) -> torch.Tensor:
|
||||
return self._forward_shared(mode="next_token", **kwargs)
|
||||
|
||||
def forward_all_future(self, **kwargs) -> torch.Tensor:
|
||||
return self._forward_shared(mode="all_future", **kwargs)
|
||||
|
||||
def forward(self, target_mode: str | None = None, **kwargs) -> torch.Tensor:
|
||||
mode = self.target_mode if target_mode is None else target_mode
|
||||
return self._forward_shared(mode=mode, **kwargs)
|
||||
def forward(self, **kwargs) -> torch.Tensor:
|
||||
return self._forward_shared(mode=self.target_mode, **kwargs)
|
||||
|
||||
def calc_risk(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return self.risk_head(x)
|
||||
@@ -754,10 +562,3 @@ class DeepHealth(nn.Module):
|
||||
f"calc_weibull_rho called with dist_mode={self.dist_mode!r}"
|
||||
)
|
||||
return F.softplus(self.rho_head(x)) + 1e-6
|
||||
|
||||
def calc_death_rho(self, x: torch.Tensor) -> torch.Tensor:
|
||||
if self.dist_mode != "mixed":
|
||||
raise RuntimeError(
|
||||
f"calc_death_rho called with dist_mode={self.dist_mode!r}"
|
||||
)
|
||||
return F.softplus(self.rho_death_head(x)).squeeze(-1) + 1e-6
|
||||
|
||||
@@ -1,553 +0,0 @@
|
||||
#!/usr/bin/env Rscript
|
||||
|
||||
# Paper-grade single-panel figures supporting the conclusion that fixed-landmark
|
||||
# horizon evaluation favors all_future over next_token.
|
||||
#
|
||||
# Outputs are written as separate panel files. This script intentionally does not
|
||||
# combine panels with plot_grid().
|
||||
|
||||
suppressPackageStartupMessages({
|
||||
library(cowplot)
|
||||
library(dplyr)
|
||||
library(ggplot2)
|
||||
library(jsonlite)
|
||||
library(readr)
|
||||
library(stringr)
|
||||
library(tibble)
|
||||
library(tidyr)
|
||||
})
|
||||
|
||||
root_dir <- "."
|
||||
runs_dir <- file.path(root_dir, "runs")
|
||||
out_dir <- file.path(root_dir, "figures_next_token_to_all_future_absolute_smoking_alcohol_bmi")
|
||||
dir.create(out_dir, showWarnings = FALSE, recursive = TRUE)
|
||||
|
||||
required_time_mode <- "absolute"
|
||||
required_extra_info_types <- c(11L, 66L, 67L)
|
||||
required_extra_info_signature <- paste(sort(required_extra_info_types), collapse = ",")
|
||||
|
||||
theme_set(
|
||||
theme_cowplot(font_size = 9) +
|
||||
theme(
|
||||
plot.background = element_rect(fill = "white", color = NA),
|
||||
panel.background = element_rect(fill = "white", color = NA),
|
||||
legend.background = element_rect(fill = "white", color = NA),
|
||||
legend.key = element_rect(fill = "white", color = NA)
|
||||
)
|
||||
)
|
||||
|
||||
target_cols <- c(
|
||||
"next_token" = "#B54A3A",
|
||||
"all_future" = "#2C7FB8"
|
||||
)
|
||||
|
||||
dist_shapes <- c(
|
||||
"exponential" = 16,
|
||||
"weibull" = 17,
|
||||
"mixed" = 15
|
||||
)
|
||||
|
||||
read_run_config <- function(run_path) {
|
||||
cfg_path <- file.path(run_path, "train_config.json")
|
||||
if (!file.exists(cfg_path)) return(NULL)
|
||||
cfg <- jsonlite::read_json(cfg_path, simplifyVector = TRUE)
|
||||
extra_info_types <- cfg$extra_info_types %||% integer(0)
|
||||
extra_info_signature <- paste(sort(as.integer(extra_info_types)), collapse = ",")
|
||||
tibble(
|
||||
run = basename(run_path),
|
||||
model_target_mode = as.character(cfg$model_target_mode %||% NA_character_),
|
||||
target_mode = as.character(cfg$target_mode %||% NA_character_),
|
||||
dist_mode = as.character(cfg$dist_mode %||% NA_character_),
|
||||
time_mode = as.character(cfg$time_mode %||% NA_character_),
|
||||
readout_name = as.character(cfg$readout_name %||% NA_character_),
|
||||
attn_mask_mode = as.character(cfg$attn_mask_mode %||% NA_character_),
|
||||
extra_info_signature = extra_info_signature
|
||||
)
|
||||
}
|
||||
|
||||
`%||%` <- function(x, y) {
|
||||
if (is.null(x) || length(x) == 0) y else x
|
||||
}
|
||||
|
||||
load_one_result <- function(run_path, file_name, eval_family) {
|
||||
cfg <- read_run_config(run_path)
|
||||
if (is.null(cfg)) return(NULL)
|
||||
fp <- file.path(run_path, file_name)
|
||||
if (!file.exists(fp)) return(NULL)
|
||||
|
||||
df <- suppressMessages(readr::read_csv(fp, show_col_types = FALSE))
|
||||
if (!("auc" %in% names(df)) || nrow(df) == 0) return(NULL)
|
||||
|
||||
out <- df %>%
|
||||
mutate(
|
||||
run = basename(run_path),
|
||||
eval_family = eval_family,
|
||||
auc = as.numeric(auc)
|
||||
) %>%
|
||||
left_join(cfg, by = "run", suffix = c("", "_cfg"))
|
||||
|
||||
coalesce_joined <- function(data, col) {
|
||||
cfg_col <- paste0(col, "_cfg")
|
||||
if (col %in% names(data) && cfg_col %in% names(data)) {
|
||||
dplyr::coalesce(data[[col]], data[[cfg_col]])
|
||||
} else if (col %in% names(data)) {
|
||||
data[[col]]
|
||||
} else if (cfg_col %in% names(data)) {
|
||||
data[[cfg_col]]
|
||||
} else {
|
||||
rep(NA_character_, nrow(data))
|
||||
}
|
||||
}
|
||||
|
||||
for (col in c("model_target_mode", "target_mode", "dist_mode", "time_mode", "readout_name", "attn_mask_mode")) {
|
||||
out[[col]] <- coalesce_joined(out, col)
|
||||
}
|
||||
|
||||
out %>%
|
||||
select(-any_of(c(
|
||||
"model_target_mode_cfg", "target_mode_cfg", "dist_mode_cfg",
|
||||
"time_mode_cfg", "readout_name_cfg", "attn_mask_mode_cfg"
|
||||
)))
|
||||
}
|
||||
|
||||
run_paths <- list.dirs(runs_dir, recursive = FALSE, full.names = TRUE)
|
||||
|
||||
landmark_auc <- bind_rows(lapply(
|
||||
run_paths,
|
||||
load_one_result,
|
||||
file_name = "df_auc_landmark.csv",
|
||||
eval_family = "Fixed landmark + horizon"
|
||||
)) %>%
|
||||
filter(time_mode == "absolute")
|
||||
|
||||
token_auc <- bind_rows(lapply(
|
||||
run_paths,
|
||||
load_one_result,
|
||||
file_name = "df_both.csv",
|
||||
eval_family = "Delphi2M-style token"
|
||||
)) %>%
|
||||
filter(time_mode == "absolute")
|
||||
|
||||
if (nrow(landmark_auc) == 0) {
|
||||
stop("No landmark AUC files found under runs/*/df_auc_landmark.csv")
|
||||
}
|
||||
if (nrow(token_auc) == 0) {
|
||||
stop("No token AUC files found under runs/*/df_both.csv")
|
||||
}
|
||||
|
||||
landmark_auc <- landmark_auc %>%
|
||||
filter(
|
||||
time_mode == required_time_mode,
|
||||
extra_info_signature == required_extra_info_signature
|
||||
)
|
||||
|
||||
token_auc <- token_auc %>%
|
||||
filter(
|
||||
time_mode == required_time_mode,
|
||||
extra_info_signature == required_extra_info_signature
|
||||
)
|
||||
|
||||
if (nrow(landmark_auc) == 0 || nrow(token_auc) == 0) {
|
||||
stop(
|
||||
"No AUC rows remain after filtering for time_mode='",
|
||||
required_time_mode,
|
||||
"' and extra_info_types='",
|
||||
required_extra_info_signature,
|
||||
"'."
|
||||
)
|
||||
}
|
||||
|
||||
message(
|
||||
"Using runs with time_mode='", required_time_mode,
|
||||
"' and extra_info_types='", required_extra_info_signature, "':"
|
||||
)
|
||||
print(sort(unique(landmark_auc$run)))
|
||||
|
||||
classify_endpoint <- function(data) {
|
||||
data %>%
|
||||
mutate(
|
||||
endpoint_type = if_else(
|
||||
str_to_lower(as.character(label_code)) == "death",
|
||||
"Death",
|
||||
"Non-death disease"
|
||||
),
|
||||
endpoint_type = factor(endpoint_type, levels = c("Non-death disease", "Death"))
|
||||
)
|
||||
}
|
||||
|
||||
landmark_auc <- classify_endpoint(landmark_auc)
|
||||
token_auc <- classify_endpoint(token_auc)
|
||||
|
||||
landmark_auc_disease <- landmark_auc %>% filter(endpoint_type == "Non-death disease")
|
||||
token_auc_disease <- token_auc %>% filter(endpoint_type == "Non-death disease")
|
||||
landmark_auc_death <- landmark_auc %>% filter(endpoint_type == "Death")
|
||||
token_auc_death <- token_auc %>% filter(endpoint_type == "Death")
|
||||
|
||||
if (nrow(landmark_auc_death) == 0 || nrow(token_auc_death) == 0) {
|
||||
warning("Death rows were not found in one or both AUC tables.")
|
||||
}
|
||||
|
||||
auc_all <- bind_rows(
|
||||
landmark_auc_disease %>% mutate(horizon = as.numeric(horizon), offset = NA_real_),
|
||||
token_auc_disease %>% mutate(horizon = NA_real_, offset = as.numeric(offset))
|
||||
) %>%
|
||||
mutate(
|
||||
model_target_mode = factor(model_target_mode, levels = c("next_token", "all_future")),
|
||||
eval_family = factor(eval_family, levels = c("Delphi2M-style token", "Fixed landmark + horizon")),
|
||||
dist_mode = factor(dist_mode, levels = c("exponential", "weibull", "mixed")),
|
||||
model_label = recode(
|
||||
as.character(model_target_mode),
|
||||
"next_token" = "next-token objective",
|
||||
"all_future" = "all-future objective"
|
||||
)
|
||||
)
|
||||
|
||||
mean_ci <- function(x) {
|
||||
x <- x[is.finite(x)]
|
||||
n <- length(x)
|
||||
m <- mean(x)
|
||||
se <- sd(x) / sqrt(n)
|
||||
tibble(mean = m, ymin = m - 1.96 * se, ymax = m + 1.96 * se, n = n)
|
||||
}
|
||||
|
||||
save_panel <- function(plot, name, width = 3.6, height = 3.0) {
|
||||
pdf_path <- file.path(out_dir, paste0(name, ".pdf"))
|
||||
png_path <- file.path(out_dir, paste0(name, ".png"))
|
||||
cowplot::save_plot(pdf_path, plot, base_width = width, base_height = height, bg = "white")
|
||||
cowplot::save_plot(png_path, plot, base_width = width, base_height = height, dpi = 600, bg = "white")
|
||||
message("Wrote: ", pdf_path)
|
||||
message("Wrote: ", png_path)
|
||||
}
|
||||
|
||||
# Panel 1: run-level mean AUC under the clinically aligned landmark/horizon task.
|
||||
# Death is excluded here and plotted separately below.
|
||||
landmark_run <- landmark_auc_disease %>%
|
||||
mutate(model_target_mode = factor(model_target_mode, levels = c("next_token", "all_future"))) %>%
|
||||
group_by(run, model_target_mode, dist_mode, time_mode, target_mode) %>%
|
||||
summarise(mean_auc = mean(auc, na.rm = TRUE), median_auc = median(auc, na.rm = TRUE), .groups = "drop")
|
||||
|
||||
landmark_summary <- landmark_run %>%
|
||||
group_by(model_target_mode) %>%
|
||||
summarise(mean_ci(mean_auc), .groups = "drop")
|
||||
|
||||
p1 <- ggplot(landmark_run, aes(x = model_target_mode, y = mean_auc)) +
|
||||
geom_point(
|
||||
aes(color = model_target_mode, shape = dist_mode),
|
||||
position = position_jitter(width = 0.09, height = 0, seed = 1),
|
||||
size = 2.2,
|
||||
alpha = 0.88
|
||||
) +
|
||||
geom_errorbar(
|
||||
data = landmark_summary,
|
||||
aes(x = model_target_mode, y = mean, ymin = ymin, ymax = ymax, color = model_target_mode),
|
||||
width = 0.12,
|
||||
linewidth = 0.55,
|
||||
inherit.aes = FALSE
|
||||
) +
|
||||
geom_point(
|
||||
data = landmark_summary,
|
||||
aes(x = model_target_mode, y = mean, color = model_target_mode),
|
||||
size = 3.4,
|
||||
inherit.aes = FALSE
|
||||
) +
|
||||
scale_color_manual(values = target_cols, guide = "none") +
|
||||
scale_shape_manual(values = dist_shapes, na.translate = FALSE) +
|
||||
scale_x_discrete(labels = c("next_token", "all_future")) +
|
||||
coord_cartesian(ylim = c(0.58, 0.78)) +
|
||||
labs(
|
||||
x = NULL,
|
||||
y = "Mean AUC per run",
|
||||
shape = "Risk head",
|
||||
title = "Non-death landmark AUC (absolute time)"
|
||||
) +
|
||||
theme(
|
||||
plot.title = element_text(face = "bold", size = 10),
|
||||
axis.text.x = element_text(size = 9),
|
||||
legend.position = c(0.72, 0.20),
|
||||
legend.background = element_blank()
|
||||
)
|
||||
|
||||
save_panel(p1, "panel_01_landmark_overall")
|
||||
|
||||
# Panel 2: landmark AUC by prediction horizon.
|
||||
landmark_horizon_run <- landmark_auc_disease %>%
|
||||
mutate(
|
||||
horizon = as.numeric(horizon),
|
||||
model_target_mode = factor(model_target_mode, levels = c("next_token", "all_future"))
|
||||
) %>%
|
||||
group_by(run, model_target_mode, horizon) %>%
|
||||
summarise(mean_auc = mean(auc, na.rm = TRUE), .groups = "drop")
|
||||
|
||||
landmark_horizon_summary <- landmark_horizon_run %>%
|
||||
group_by(model_target_mode, horizon) %>%
|
||||
summarise(mean_ci(mean_auc), .groups = "drop")
|
||||
|
||||
p2 <- ggplot(landmark_horizon_run, aes(x = horizon, y = mean_auc, color = model_target_mode)) +
|
||||
geom_line(aes(group = run), alpha = 0.18, linewidth = 0.35) +
|
||||
geom_point(alpha = 0.32, size = 1.1) +
|
||||
geom_ribbon(
|
||||
data = landmark_horizon_summary,
|
||||
aes(x = horizon, y = mean, ymin = ymin, ymax = ymax, fill = model_target_mode, group = model_target_mode),
|
||||
alpha = 0.13,
|
||||
color = NA,
|
||||
inherit.aes = FALSE
|
||||
) +
|
||||
geom_line(data = landmark_horizon_summary, aes(y = mean), linewidth = 0.85) +
|
||||
geom_point(data = landmark_horizon_summary, aes(y = mean), size = 2.0) +
|
||||
scale_color_manual(
|
||||
values = target_cols,
|
||||
labels = c("next_token", "all_future"),
|
||||
name = NULL
|
||||
) +
|
||||
scale_fill_manual(values = target_cols, guide = "none") +
|
||||
scale_x_continuous(breaks = c(1, 5, 10)) +
|
||||
coord_cartesian(ylim = c(0.58, 0.78)) +
|
||||
labs(
|
||||
x = "Prediction horizon, years",
|
||||
y = "Mean AUC per run",
|
||||
title = "Non-death landmark AUC across horizons"
|
||||
) +
|
||||
theme(
|
||||
plot.title = element_text(face = "bold", size = 10),
|
||||
legend.position = c(0.31, 0.20),
|
||||
legend.background = element_blank()
|
||||
)
|
||||
|
||||
save_panel(p2, "panel_02_landmark_by_horizon", width = 3.8, height = 3.0)
|
||||
|
||||
# Panel 3: Delphi2M-style token AUC by offset. This documents why the old
|
||||
# evaluation can make next_token look competitive, especially near the event.
|
||||
token_offset_run <- token_auc_disease %>%
|
||||
mutate(
|
||||
offset = as.numeric(offset),
|
||||
model_target_mode = factor(model_target_mode, levels = c("next_token", "all_future"))
|
||||
) %>%
|
||||
group_by(run, model_target_mode, offset) %>%
|
||||
summarise(mean_auc = mean(auc, na.rm = TRUE), .groups = "drop")
|
||||
|
||||
token_offset_summary <- token_offset_run %>%
|
||||
group_by(model_target_mode, offset) %>%
|
||||
summarise(mean_ci(mean_auc), .groups = "drop")
|
||||
|
||||
p3 <- ggplot(token_offset_run, aes(x = offset, y = mean_auc, color = model_target_mode)) +
|
||||
geom_line(aes(group = run), alpha = 0.18, linewidth = 0.35) +
|
||||
geom_point(alpha = 0.32, size = 1.1) +
|
||||
geom_ribbon(
|
||||
data = token_offset_summary,
|
||||
aes(x = offset, y = mean, ymin = ymin, ymax = ymax, fill = model_target_mode, group = model_target_mode),
|
||||
alpha = 0.13,
|
||||
color = NA,
|
||||
inherit.aes = FALSE
|
||||
) +
|
||||
geom_line(data = token_offset_summary, aes(y = mean), linewidth = 0.85) +
|
||||
geom_point(data = token_offset_summary, aes(y = mean), size = 2.0) +
|
||||
scale_color_manual(
|
||||
values = target_cols,
|
||||
labels = c("next_token", "all_future"),
|
||||
name = NULL
|
||||
) +
|
||||
scale_fill_manual(values = target_cols, guide = "none") +
|
||||
scale_x_continuous(breaks = c(0.1, 1, 5, 10), trans = "log10") +
|
||||
coord_cartesian(ylim = c(0.55, 0.82)) +
|
||||
labs(
|
||||
x = "Minimum offset before event, years",
|
||||
y = "Mean AUC per run",
|
||||
title = "Non-death token AUC by offset"
|
||||
) +
|
||||
theme(
|
||||
plot.title = element_text(face = "bold", size = 10),
|
||||
legend.position = c(0.31, 0.20),
|
||||
legend.background = element_blank()
|
||||
)
|
||||
|
||||
save_panel(p3, "panel_03_token_auc_by_offset", width = 3.8, height = 3.0)
|
||||
|
||||
# Panel 4: within-run contrast between old token evaluation and landmark
|
||||
# evaluation. Each run contributes one point per evaluation family.
|
||||
run_eval_contrast <- auc_all %>%
|
||||
group_by(run, model_target_mode, dist_mode, eval_family) %>%
|
||||
summarise(mean_auc = mean(auc, na.rm = TRUE), .groups = "drop")
|
||||
|
||||
p4 <- ggplot(run_eval_contrast, aes(x = eval_family, y = mean_auc, color = model_target_mode)) +
|
||||
geom_line(aes(group = run), alpha = 0.34, linewidth = 0.45) +
|
||||
geom_point(aes(shape = dist_mode), size = 2.0, alpha = 0.84) +
|
||||
stat_summary(
|
||||
aes(group = model_target_mode),
|
||||
fun = mean,
|
||||
geom = "point",
|
||||
size = 3.3,
|
||||
shape = 18,
|
||||
position = position_dodge(width = 0.16)
|
||||
) +
|
||||
scale_color_manual(
|
||||
values = target_cols,
|
||||
labels = c("next_token", "all_future"),
|
||||
name = NULL
|
||||
) +
|
||||
scale_shape_manual(values = dist_shapes, na.translate = FALSE, name = "Risk head") +
|
||||
coord_cartesian(ylim = c(0.58, 0.78)) +
|
||||
labs(
|
||||
x = NULL,
|
||||
y = "Mean AUC per run",
|
||||
title = "Evaluation choice changes the conclusion (absolute time)"
|
||||
) +
|
||||
theme(
|
||||
plot.title = element_text(face = "bold", size = 10),
|
||||
axis.text.x = element_text(angle = 18, hjust = 1),
|
||||
legend.position = "right"
|
||||
)
|
||||
|
||||
save_panel(p4, "panel_04_evaluation_contrast", width = 4.3, height = 3.1)
|
||||
|
||||
# Panel 5: disease-level distribution for the landmark task, pooled over
|
||||
# horizons and runs. This shows the shift without hiding heterogeneity.
|
||||
landmark_density <- landmark_auc_disease %>%
|
||||
mutate(model_target_mode = factor(model_target_mode, levels = c("next_token", "all_future"))) %>%
|
||||
filter(is.finite(auc))
|
||||
|
||||
p5 <- ggplot(landmark_density, aes(x = auc, fill = model_target_mode, color = model_target_mode)) +
|
||||
geom_density(alpha = 0.20, linewidth = 0.65, adjust = 1.1) +
|
||||
geom_vline(
|
||||
data = landmark_density %>%
|
||||
group_by(model_target_mode) %>%
|
||||
summarise(mean_auc = mean(auc), .groups = "drop"),
|
||||
aes(xintercept = mean_auc, color = model_target_mode),
|
||||
linewidth = 0.75,
|
||||
linetype = "22"
|
||||
) +
|
||||
scale_color_manual(values = target_cols, labels = c("next_token", "all_future"), name = NULL) +
|
||||
scale_fill_manual(values = target_cols, labels = c("next_token", "all_future"), name = NULL) +
|
||||
coord_cartesian(xlim = c(0.35, 1.0)) +
|
||||
labs(
|
||||
x = "AUC",
|
||||
y = "Density",
|
||||
title = "Non-death landmark AUC distribution"
|
||||
) +
|
||||
theme(
|
||||
plot.title = element_text(face = "bold", size = 10),
|
||||
legend.position = c(0.24, 0.82),
|
||||
legend.background = element_blank()
|
||||
)
|
||||
|
||||
save_panel(p5, "panel_05_landmark_auc_distribution", width = 3.8, height = 3.0)
|
||||
|
||||
# Panel 6: death-only fixed landmark + horizon AUC. Death has one endpoint token,
|
||||
# so each line is a run trajectory across horizons.
|
||||
death_landmark_run <- landmark_auc_death %>%
|
||||
mutate(
|
||||
horizon = as.numeric(horizon),
|
||||
model_target_mode = factor(model_target_mode, levels = c("next_token", "all_future")),
|
||||
dist_mode = factor(dist_mode, levels = c("exponential", "weibull", "mixed"))
|
||||
) %>%
|
||||
group_by(run, model_target_mode, dist_mode, horizon) %>%
|
||||
summarise(mean_auc = mean(auc, na.rm = TRUE), .groups = "drop")
|
||||
|
||||
death_landmark_summary <- death_landmark_run %>%
|
||||
group_by(model_target_mode, horizon) %>%
|
||||
summarise(mean_ci(mean_auc), .groups = "drop")
|
||||
|
||||
p6 <- ggplot(death_landmark_run, aes(x = horizon, y = mean_auc, color = model_target_mode)) +
|
||||
geom_line(aes(group = run), alpha = 0.42, linewidth = 0.45) +
|
||||
geom_point(aes(shape = dist_mode), alpha = 0.9, size = 2.0) +
|
||||
geom_line(data = death_landmark_summary, aes(y = mean, group = model_target_mode), linewidth = 0.9) +
|
||||
geom_point(data = death_landmark_summary, aes(y = mean), size = 2.2) +
|
||||
scale_color_manual(values = target_cols, labels = c("next_token", "all_future"), name = NULL) +
|
||||
scale_shape_manual(values = dist_shapes, na.translate = FALSE, name = "Risk head") +
|
||||
scale_x_continuous(breaks = c(1, 5, 10)) +
|
||||
coord_cartesian(ylim = c(0.58, 0.95)) +
|
||||
labs(
|
||||
x = "Prediction horizon, years",
|
||||
y = "AUC",
|
||||
title = "Death-only landmark AUC"
|
||||
) +
|
||||
theme(
|
||||
plot.title = element_text(face = "bold", size = 10),
|
||||
legend.position = "right"
|
||||
)
|
||||
|
||||
save_panel(p6, "panel_06_death_landmark_by_horizon", width = 3.9, height = 3.0)
|
||||
|
||||
# Panel 7: death-only Delphi2M-style token AUC by offset.
|
||||
death_token_run <- token_auc_death %>%
|
||||
mutate(
|
||||
offset = as.numeric(offset),
|
||||
model_target_mode = factor(model_target_mode, levels = c("next_token", "all_future")),
|
||||
dist_mode = factor(dist_mode, levels = c("exponential", "weibull", "mixed"))
|
||||
) %>%
|
||||
group_by(run, model_target_mode, dist_mode, offset) %>%
|
||||
summarise(mean_auc = mean(auc, na.rm = TRUE), .groups = "drop")
|
||||
|
||||
death_token_summary <- death_token_run %>%
|
||||
group_by(model_target_mode, offset) %>%
|
||||
summarise(mean_ci(mean_auc), .groups = "drop")
|
||||
|
||||
p7 <- ggplot(death_token_run, aes(x = offset, y = mean_auc, color = model_target_mode)) +
|
||||
geom_line(aes(group = run), alpha = 0.42, linewidth = 0.45) +
|
||||
geom_point(aes(shape = dist_mode), alpha = 0.9, size = 2.0) +
|
||||
geom_line(data = death_token_summary, aes(y = mean, group = model_target_mode), linewidth = 0.9) +
|
||||
geom_point(data = death_token_summary, aes(y = mean), size = 2.2) +
|
||||
scale_color_manual(values = target_cols, labels = c("next_token", "all_future"), name = NULL) +
|
||||
scale_shape_manual(values = dist_shapes, na.translate = FALSE, name = "Risk head") +
|
||||
scale_x_continuous(breaks = c(0.1, 1, 5, 10), trans = "log10") +
|
||||
coord_cartesian(ylim = c(0.58, 0.95)) +
|
||||
labs(
|
||||
x = "Minimum offset before event, years",
|
||||
y = "AUC",
|
||||
title = "Death-only token AUC"
|
||||
) +
|
||||
theme(
|
||||
plot.title = element_text(face = "bold", size = 10),
|
||||
legend.position = "right"
|
||||
)
|
||||
|
||||
save_panel(p7, "panel_07_death_token_auc_by_offset", width = 3.9, height = 3.0)
|
||||
|
||||
# Panel 8: death-only contrast between the two evaluation families.
|
||||
death_eval_contrast <- bind_rows(
|
||||
landmark_auc_death %>% mutate(horizon = as.numeric(horizon), offset = NA_real_),
|
||||
token_auc_death %>% mutate(horizon = NA_real_, offset = as.numeric(offset))
|
||||
) %>%
|
||||
mutate(
|
||||
model_target_mode = factor(model_target_mode, levels = c("next_token", "all_future")),
|
||||
eval_family = factor(eval_family, levels = c("Delphi2M-style token", "Fixed landmark + horizon")),
|
||||
dist_mode = factor(dist_mode, levels = c("exponential", "weibull", "mixed"))
|
||||
) %>%
|
||||
group_by(run, model_target_mode, dist_mode, eval_family) %>%
|
||||
summarise(mean_auc = mean(auc, na.rm = TRUE), .groups = "drop")
|
||||
|
||||
p8 <- ggplot(death_eval_contrast, aes(x = eval_family, y = mean_auc, color = model_target_mode)) +
|
||||
geom_line(aes(group = run), alpha = 0.38, linewidth = 0.5) +
|
||||
geom_point(aes(shape = dist_mode), size = 2.2, alpha = 0.9) +
|
||||
stat_summary(
|
||||
aes(group = model_target_mode),
|
||||
fun = mean,
|
||||
geom = "point",
|
||||
size = 3.4,
|
||||
shape = 18,
|
||||
position = position_dodge(width = 0.16)
|
||||
) +
|
||||
scale_color_manual(values = target_cols, labels = c("next_token", "all_future"), name = NULL) +
|
||||
scale_shape_manual(values = dist_shapes, na.translate = FALSE, name = "Risk head") +
|
||||
coord_cartesian(ylim = c(0.58, 0.95)) +
|
||||
labs(
|
||||
x = NULL,
|
||||
y = "Mean AUC per run",
|
||||
title = "Death endpoint evaluated separately"
|
||||
) +
|
||||
theme(
|
||||
plot.title = element_text(face = "bold", size = 10),
|
||||
axis.text.x = element_text(angle = 18, hjust = 1),
|
||||
legend.position = "right"
|
||||
)
|
||||
|
||||
save_panel(p8, "panel_08_death_evaluation_contrast", width = 4.3, height = 3.1)
|
||||
|
||||
# Export the exact run-level summaries used by the figures.
|
||||
readr::write_csv(landmark_run, file.path(out_dir, "landmark_run_summary.csv"))
|
||||
readr::write_csv(token_offset_run, file.path(out_dir, "token_offset_run_summary.csv"))
|
||||
readr::write_csv(run_eval_contrast, file.path(out_dir, "run_evaluation_contrast.csv"))
|
||||
readr::write_csv(death_landmark_run, file.path(out_dir, "death_landmark_run_summary.csv"))
|
||||
readr::write_csv(death_token_run, file.path(out_dir, "death_token_offset_run_summary.csv"))
|
||||
readr::write_csv(death_eval_contrast, file.path(out_dir, "death_evaluation_contrast.csv"))
|
||||
|
||||
message("Done. Panels are in: ", normalizePath(out_dir, winslash = "/"))
|
||||
@@ -4,7 +4,7 @@ This script converts raw UK Biobank CSV exports into the artefacts consumed by
|
||||
DeepHealth:
|
||||
|
||||
* ``ukb_event_data.npy``: ``(N, 3)`` uint32 array of ``(eid, days, label)``
|
||||
disease/death/checkup events sorted by patient then time.
|
||||
disease/death events sorted by patient then time.
|
||||
* ``ukb_basic_info.csv``: basic patient table indexed by ``eid`` with ``sex``.
|
||||
* ``ukb_other_info.npy``: ``(M, 5)`` float64 array of
|
||||
``(eid, type, value, value_kind, time)`` rows. ``type=0`` is reserved for
|
||||
@@ -219,7 +219,8 @@ with open(labels_file, encoding="utf-8") as f: # Open labels file
|
||||
for idx, line in enumerate(f): # Enumerate to assign incremental label IDs
|
||||
parts = line.strip().split(" ") # Split by space
|
||||
if parts and parts[0]: # Guard against empty lines
|
||||
# Start labels from 1 to reserve 0 for padding, 1 for checkup
|
||||
# Keep raw disease ids at 2+ so existing prepared data and model
|
||||
# vocabulary indices remain stable; raw id 1 is unused.
|
||||
label_dict[parts[0]] = idx + 2
|
||||
|
||||
# Pre-build lookup: ICD/Death column name -> integer label for fast per-column extraction
|
||||
@@ -327,19 +328,6 @@ for ukb_chunk in tqdm.tqdm(ukb_iterator, desc="Processing UK Biobank data"):
|
||||
if cancer_frames:
|
||||
event_list.append(np.vstack(cancer_frames))
|
||||
|
||||
# Add checkup events with label=1 using date_of_assessment (already in days from dob)
|
||||
if "date_of_assessment" in ukb_chunk.columns:
|
||||
doa_series = ukb_chunk["date_of_assessment"].dropna()
|
||||
if not doa_series.empty:
|
||||
checkup_data = np.column_stack(
|
||||
(
|
||||
doa_series.index.values,
|
||||
doa_series.values.astype(int),
|
||||
np.ones(len(doa_series), dtype=int),
|
||||
)
|
||||
)
|
||||
event_list.append(checkup_data)
|
||||
|
||||
# Combine tabular chunks
|
||||
final_tabular = pd.concat(tabular_list, axis=0, ignore_index=False)
|
||||
final_tabular.index.name = "eid" # Ensure index named consistently
|
||||
|
||||
@@ -9,12 +9,12 @@ with exactly three fields:
|
||||
token int32
|
||||
|
||||
``token`` follows the existing ``labels.csv`` convention used by
|
||||
``prepare_data.py``: padding=0, checkup=1 (not emitted here), and the first
|
||||
``prepare_data.py``: padding=0, token 1 is reserved and unused, and the first
|
||||
label in ``labels.csv`` receives token 2. Each ``(eid, token)`` is deduplicated
|
||||
to the first known event date.
|
||||
|
||||
The output is intended for calendar-indexed temperature and air-pollution
|
||||
queries. It contains no date of birth, sex, covariates, or checkup events.
|
||||
queries. It contains no date of birth, sex, covariates, or assessment events.
|
||||
|
||||
Usage
|
||||
-----
|
||||
|
||||
107
readouts.py
107
readouts.py
@@ -1,107 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
@dataclass
|
||||
class ReadoutOutput:
|
||||
hidden: torch.Tensor
|
||||
readout_mask: torch.Tensor
|
||||
|
||||
|
||||
class TokenReadout(nn.Module):
|
||||
def forward(
|
||||
self,
|
||||
hidden: torch.Tensor,
|
||||
time_seq: torch.Tensor,
|
||||
padding_mask: torch.Tensor,
|
||||
readout_mask: torch.Tensor | None = None,
|
||||
) -> ReadoutOutput:
|
||||
mask = padding_mask if readout_mask is None else readout_mask
|
||||
return ReadoutOutput(hidden=hidden, readout_mask=mask.bool())
|
||||
|
||||
|
||||
class SameTimeGroupEndReadout(nn.Module):
|
||||
def __init__(self, reduce: str = "mean"):
|
||||
super().__init__()
|
||||
if reduce not in {"mean", "sum"}:
|
||||
raise ValueError("reduce must be either 'mean' or 'sum'")
|
||||
self.reduce = reduce
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden: torch.Tensor,
|
||||
time_seq: torch.Tensor,
|
||||
padding_mask: torch.Tensor,
|
||||
readout_mask: torch.Tensor | None = None,
|
||||
) -> ReadoutOutput:
|
||||
if readout_mask is None:
|
||||
next_is_new_time = torch.ones_like(padding_mask, dtype=torch.bool)
|
||||
next_is_new_time[:, :-1] = time_seq[:, 1:] != time_seq[:, :-1]
|
||||
readout_mask = padding_mask.bool() & next_is_new_time
|
||||
else:
|
||||
readout_mask = readout_mask.bool()
|
||||
|
||||
group_start = torch.ones_like(padding_mask, dtype=torch.bool)
|
||||
group_start[:, 1:] = time_seq[:, 1:] != time_seq[:, :-1]
|
||||
group_start = group_start & padding_mask.bool()
|
||||
|
||||
group_id = group_start.long().cumsum(dim=1) - 1
|
||||
group_id = group_id.clamp_min(0)
|
||||
max_groups = hidden.size(1)
|
||||
|
||||
group_sum = hidden.new_zeros(hidden.size(0), max_groups, hidden.size(2))
|
||||
group_sum.scatter_add_(
|
||||
1,
|
||||
group_id.unsqueeze(-1).expand_as(hidden),
|
||||
hidden * padding_mask.unsqueeze(-1).to(hidden.dtype),
|
||||
)
|
||||
|
||||
if self.reduce == "mean":
|
||||
group_count = hidden.new_zeros(hidden.size(0), max_groups, 1)
|
||||
group_count.scatter_add_(
|
||||
1,
|
||||
group_id.unsqueeze(-1),
|
||||
padding_mask.unsqueeze(-1).to(hidden.dtype),
|
||||
)
|
||||
group_sum = group_sum / group_count.clamp_min(1.0)
|
||||
|
||||
out = hidden.clone()
|
||||
out[readout_mask] = group_sum.gather(
|
||||
1,
|
||||
group_id.unsqueeze(-1).expand_as(hidden),
|
||||
)[readout_mask]
|
||||
return ReadoutOutput(hidden=out, readout_mask=readout_mask)
|
||||
|
||||
|
||||
class LastValidReadout(nn.Module):
|
||||
def forward(
|
||||
self,
|
||||
hidden: torch.Tensor,
|
||||
time_seq: torch.Tensor,
|
||||
padding_mask: torch.Tensor,
|
||||
readout_mask: torch.Tensor | None = None,
|
||||
) -> ReadoutOutput:
|
||||
batch_size, seq_len = padding_mask.shape
|
||||
last_idx = padding_mask.long().sum(dim=1).clamp_min(1) - 1
|
||||
out = hidden[torch.arange(batch_size, device=hidden.device), last_idx]
|
||||
mask = torch.ones(batch_size, dtype=torch.bool, device=hidden.device)
|
||||
return ReadoutOutput(hidden=out, readout_mask=mask)
|
||||
|
||||
|
||||
def build_readout(name: str, **kwargs) -> nn.Module:
|
||||
name = name.lower()
|
||||
if name == "token":
|
||||
return TokenReadout()
|
||||
if name in {"same_time_group_end", "same_time"}:
|
||||
return SameTimeGroupEndReadout(**kwargs)
|
||||
if name == "last_valid":
|
||||
return LastValidReadout()
|
||||
raise ValueError(
|
||||
"Unknown readout {!r}. Available: token, same_time_group_end, last_valid.".format(
|
||||
name
|
||||
)
|
||||
)
|
||||
234
rerun_all_required_experiments_linux.sh
Executable file
234
rerun_all_required_experiments_linux.sh
Executable file
@@ -0,0 +1,234 @@
|
||||
#!/usr/bin/env bash
|
||||
#
|
||||
# Re-run every experiment invalidated by removing the assessment event token,
|
||||
# requiring train-split RobustScale, and deleting the mixed distribution.
|
||||
#
|
||||
# The workflow runs:
|
||||
# 1. seed-42 full factorial experiments;
|
||||
# 2. seed-43/44 key-model experiments;
|
||||
# 3. seed-42/43/44 disease-history T/O/S ablation;
|
||||
# 4. seed-42/43/44 assessment + smoking + alcohol experiment;
|
||||
# 5. AUC and calibration evaluation for all four campaigns.
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
SCRIPT_DIR="${BASH_SOURCE[0]%/*}"
|
||||
if [[ "$SCRIPT_DIR" == "${BASH_SOURCE[0]}" ]]; then
|
||||
SCRIPT_DIR="."
|
||||
fi
|
||||
SCRIPT_DIR="$(cd -- "$SCRIPT_DIR" && pwd)"
|
||||
cd "$SCRIPT_DIR"
|
||||
|
||||
NUM_GPUS=1
|
||||
NUM_GPUS_SET=0
|
||||
GPU_CSV=""
|
||||
NUM_CPUS=16
|
||||
PYTHON_BIN="${PYTHON_BIN:-python}"
|
||||
BASH_BIN="${BASH:-bash}"
|
||||
DRY_RUN=0
|
||||
|
||||
FULL_FACTORIAL_CAMPAIGN="full_factorial_smoking_alcohol_bmi"
|
||||
KEY_MODELS_CAMPAIGN="key_models_multiseed_smoking_alcohol_bmi"
|
||||
HISTORY_CAMPAIGN="disease_history_ablation_no_extra"
|
||||
EXTRA_INFO_CAMPAIGN="extra_info_assessment_smoking_alcohol_robust_multiseed"
|
||||
|
||||
usage() {
|
||||
cat <<'EOF'
|
||||
Usage:
|
||||
bash rerun_all_required_experiments_linux.sh [options]
|
||||
|
||||
Options:
|
||||
--num-gpus N Number of GPUs; uses GPU ids 0 through N-1 (default: 1).
|
||||
--gpus LIST Explicit comma-separated GPU ids, for example 0,2.
|
||||
Cannot be combined with --num-gpus.
|
||||
--num-cpus N Total CPU workers shared across active GPUs (default: 16).
|
||||
--python PATH Python executable used by every child script.
|
||||
--dry-run Print commands without training or evaluation.
|
||||
-h, --help Show this help message.
|
||||
|
||||
Examples:
|
||||
bash rerun_all_required_experiments_linux.sh --num-gpus 1 --num-cpus 16
|
||||
bash rerun_all_required_experiments_linux.sh --num-gpus 4 --num-cpus 32
|
||||
bash rerun_all_required_experiments_linux.sh --gpus 0,2 --num-cpus 16
|
||||
|
||||
CPU allocation:
|
||||
The script divides --num-cpus evenly across simultaneously active GPUs.
|
||||
For example, 4 GPUs and 32 CPUs gives each GPU job 8 workers.
|
||||
EOF
|
||||
}
|
||||
|
||||
while (($# > 0)); do
|
||||
case "$1" in
|
||||
--num-gpus)
|
||||
[[ $# -ge 2 ]] || {
|
||||
echo "ERROR: --num-gpus requires a value." >&2
|
||||
exit 2
|
||||
}
|
||||
[[ -z "$GPU_CSV" ]] || {
|
||||
echo "ERROR: --num-gpus cannot be combined with --gpus." >&2
|
||||
exit 2
|
||||
}
|
||||
NUM_GPUS="$2"
|
||||
NUM_GPUS_SET=1
|
||||
shift 2
|
||||
;;
|
||||
--gpus)
|
||||
[[ $# -ge 2 ]] || {
|
||||
echo "ERROR: --gpus requires a value." >&2
|
||||
exit 2
|
||||
}
|
||||
((NUM_GPUS_SET == 0)) || {
|
||||
echo "ERROR: --gpus cannot be combined with --num-gpus." >&2
|
||||
exit 2
|
||||
}
|
||||
GPU_CSV="$2"
|
||||
shift 2
|
||||
;;
|
||||
--num-cpus)
|
||||
[[ $# -ge 2 ]] || {
|
||||
echo "ERROR: --num-cpus requires a value." >&2
|
||||
exit 2
|
||||
}
|
||||
NUM_CPUS="$2"
|
||||
shift 2
|
||||
;;
|
||||
--python)
|
||||
[[ $# -ge 2 ]] || {
|
||||
echo "ERROR: --python requires a value." >&2
|
||||
exit 2
|
||||
}
|
||||
PYTHON_BIN="$2"
|
||||
shift 2
|
||||
;;
|
||||
--dry-run)
|
||||
DRY_RUN=1
|
||||
shift
|
||||
;;
|
||||
-h|--help)
|
||||
usage
|
||||
exit 0
|
||||
;;
|
||||
*)
|
||||
echo "ERROR: unknown argument: $1" >&2
|
||||
usage >&2
|
||||
exit 2
|
||||
;;
|
||||
esac
|
||||
done
|
||||
|
||||
[[ "$NUM_GPUS" =~ ^[1-9][0-9]*$ ]] || {
|
||||
echo "ERROR: --num-gpus must be a positive integer." >&2
|
||||
exit 2
|
||||
}
|
||||
[[ "$NUM_CPUS" =~ ^[1-9][0-9]*$ ]] || {
|
||||
echo "ERROR: --num-cpus must be a positive integer." >&2
|
||||
exit 2
|
||||
}
|
||||
|
||||
if [[ -z "$GPU_CSV" ]]; then
|
||||
for ((gpu_id = 0; gpu_id < NUM_GPUS; gpu_id++)); do
|
||||
if [[ -n "$GPU_CSV" ]]; then
|
||||
GPU_CSV+=","
|
||||
fi
|
||||
GPU_CSV+="$gpu_id"
|
||||
done
|
||||
else
|
||||
IFS=',' read -r -a GPU_IDS <<< "$GPU_CSV"
|
||||
NUM_GPUS="${#GPU_IDS[@]}"
|
||||
((NUM_GPUS > 0)) || {
|
||||
echo "ERROR: --gpus must not be empty." >&2
|
||||
exit 2
|
||||
}
|
||||
for gpu_id in "${GPU_IDS[@]}"; do
|
||||
[[ "$gpu_id" =~ ^[0-9]+$ ]] || {
|
||||
echo "ERROR: invalid GPU id: $gpu_id" >&2
|
||||
exit 2
|
||||
}
|
||||
done
|
||||
fi
|
||||
|
||||
((NUM_CPUS >= NUM_GPUS)) || {
|
||||
echo "ERROR: --num-cpus must be at least the number of GPUs." >&2
|
||||
exit 2
|
||||
}
|
||||
WORKERS_PER_GPU=$((NUM_CPUS / NUM_GPUS))
|
||||
|
||||
required_scripts=(
|
||||
train_batch_linux.sh
|
||||
train_key_models_multiseed_linux.sh
|
||||
train_disease_history_ablation_linux.sh
|
||||
train_extra_info_assessment_all_multiseed_linux.sh
|
||||
evaluate_all_runs_linux.sh
|
||||
evaluate_calibration_all_runs_linux.sh
|
||||
)
|
||||
for script in "${required_scripts[@]}"; do
|
||||
[[ -f "$SCRIPT_DIR/$script" ]] || {
|
||||
echo "ERROR: missing required script: $SCRIPT_DIR/$script" >&2
|
||||
exit 2
|
||||
}
|
||||
done
|
||||
|
||||
run_command() {
|
||||
printf '>>'
|
||||
printf ' %q' "$@"
|
||||
printf '\n'
|
||||
if ((DRY_RUN == 0)); then
|
||||
"$@"
|
||||
fi
|
||||
}
|
||||
|
||||
echo "GPUs: $GPU_CSV ($NUM_GPUS total)"
|
||||
echo "CPUs: $NUM_CPUS total; $WORKERS_PER_GPU workers per GPU job"
|
||||
echo "Python: $PYTHON_BIN"
|
||||
|
||||
run_command "$BASH_BIN" "$SCRIPT_DIR/train_batch_linux.sh" \
|
||||
--gpus "$GPU_CSV" \
|
||||
--seed 42 \
|
||||
--num-workers "$WORKERS_PER_GPU" \
|
||||
--python "$PYTHON_BIN"
|
||||
|
||||
run_command "$BASH_BIN" "$SCRIPT_DIR/train_key_models_multiseed_linux.sh" \
|
||||
--gpus "$GPU_CSV" \
|
||||
--seeds 43,44 \
|
||||
--num-workers "$WORKERS_PER_GPU" \
|
||||
--python "$PYTHON_BIN"
|
||||
|
||||
run_command "$BASH_BIN" "$SCRIPT_DIR/train_disease_history_ablation_linux.sh" \
|
||||
--gpus "$GPU_CSV" \
|
||||
--seeds 42,43,44 \
|
||||
--num-workers "$WORKERS_PER_GPU" \
|
||||
--python "$PYTHON_BIN"
|
||||
|
||||
run_command "$BASH_BIN" "$SCRIPT_DIR/train_extra_info_assessment_all_multiseed_linux.sh" \
|
||||
--gpus "$GPU_CSV" \
|
||||
--seeds 42,43,44 \
|
||||
--num-workers "$WORKERS_PER_GPU" \
|
||||
--python "$PYTHON_BIN"
|
||||
|
||||
campaigns=(
|
||||
"$FULL_FACTORIAL_CAMPAIGN"
|
||||
"$KEY_MODELS_CAMPAIGN"
|
||||
"$HISTORY_CAMPAIGN"
|
||||
"$EXTRA_INFO_CAMPAIGN"
|
||||
)
|
||||
|
||||
for campaign in "${campaigns[@]}"; do
|
||||
runs_root="$SCRIPT_DIR/runs/$campaign"
|
||||
run_command "$BASH_BIN" "$SCRIPT_DIR/evaluate_all_runs_linux.sh" \
|
||||
--gpus "$GPU_CSV" \
|
||||
--runs-root "$runs_root" \
|
||||
--log-root "$SCRIPT_DIR/batch_logs/evaluate_all_required/$campaign" \
|
||||
--python "$PYTHON_BIN" \
|
||||
--num-workers "$WORKERS_PER_GPU" \
|
||||
--num-workers-auc "$WORKERS_PER_GPU"
|
||||
|
||||
run_command "$BASH_BIN" "$SCRIPT_DIR/evaluate_calibration_all_runs_linux.sh" \
|
||||
--gpus "$GPU_CSV" \
|
||||
--runs-root "$runs_root" \
|
||||
--log-root "$SCRIPT_DIR/batch_logs/evaluate_calibration_all_required/$campaign" \
|
||||
--python "$PYTHON_BIN" \
|
||||
--num-workers "$WORKERS_PER_GPU" \
|
||||
--num-workers-calibration "$WORKERS_PER_GPU"
|
||||
done
|
||||
|
||||
echo "All required re-training and evaluation workflows completed."
|
||||
390
run_all_future_first_onset_experiments_linux.sh
Normal file
390
run_all_future_first_onset_experiments_linux.sh
Normal file
@@ -0,0 +1,390 @@
|
||||
#!/usr/bin/env bash
|
||||
#
|
||||
# Run the matched experiments required after the all-future first-onset update.
|
||||
#
|
||||
# Shared by every experiment:
|
||||
# - corrected first-onset likelihood and outcome-specific risk exposure;
|
||||
# - patient/interval/time-uniform query sampling in train/valid/test;
|
||||
# - timed disease history and smoking/alcohol/BMI extra information.
|
||||
#
|
||||
# Per seed, the script trains:
|
||||
# 1. TrajMixer + relative Weibull + fitted risk baseline (primary model);
|
||||
# 2. the primary model without the fitted risk baseline (baseline ablation);
|
||||
# 3. TrajMixer + relative exponential (time-distribution control);
|
||||
# 4. FFN + relative Weibull (architecture control).
|
||||
#
|
||||
# After training, AUC and calibration/point-process-NLL evaluation run
|
||||
# automatically unless --train-only is supplied.
|
||||
|
||||
set -uo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
|
||||
cd "$SCRIPT_DIR"
|
||||
|
||||
GPU_CSV=""
|
||||
SEED_CSV="42,43,44"
|
||||
NUM_WORKERS=4
|
||||
PYTHON_BIN="${PYTHON_BIN:-python}"
|
||||
CAMPAIGN_NAME="all_future_first_onset_v2_multiseed"
|
||||
DRY_RUN=0
|
||||
TRAIN_ONLY=0
|
||||
|
||||
BATCH_SIZE=256
|
||||
EXTRA_INFO_TYPES_FILE="$SCRIPT_DIR/extra_info_types_smoking_alcohol_bmi.txt"
|
||||
TRAIN_EID_FILE="$SCRIPT_DIR/ukb_train_eid.csv"
|
||||
VAL_EID_FILE="$SCRIPT_DIR/ukb_val_eid.csv"
|
||||
TEST_EID_FILE="$SCRIPT_DIR/ukb_test_eid.csv"
|
||||
|
||||
usage() {
|
||||
cat <<'EOF'
|
||||
Usage:
|
||||
bash run_all_future_first_onset_experiments_linux.sh --gpus LIST [options]
|
||||
|
||||
Required:
|
||||
--gpus LIST Comma-separated GPU ids, for example 0 or 0,1,2,3.
|
||||
|
||||
Options:
|
||||
--seeds LIST Comma-separated seeds (default: 42,43,44).
|
||||
--num-workers N DataLoader workers per active GPU (default: 4).
|
||||
--python PATH Python executable (default: $PYTHON_BIN or python).
|
||||
--campaign NAME Output campaign directory name.
|
||||
--train-only Skip AUC and calibration/NLL evaluation.
|
||||
--dry-run Print commands without running them.
|
||||
-h, --help Show this help message.
|
||||
|
||||
Fixed settings:
|
||||
batch_size 256
|
||||
extra information smoking + alcohol + BMI
|
||||
patient split fixed train/validation/test EID files
|
||||
disease history timed
|
||||
experiments per seed 4
|
||||
|
||||
Outputs:
|
||||
runs/<campaign>/seed_<seed>/<architecture>/...
|
||||
batch_logs/<campaign>/seed_<seed>/<experiment>.log
|
||||
EOF
|
||||
}
|
||||
|
||||
while (($# > 0)); do
|
||||
case "$1" in
|
||||
--gpus)
|
||||
[[ $# -ge 2 ]] || {
|
||||
echo "ERROR: --gpus requires a value." >&2
|
||||
exit 2
|
||||
}
|
||||
GPU_CSV="$2"
|
||||
shift 2
|
||||
;;
|
||||
--seeds)
|
||||
[[ $# -ge 2 ]] || {
|
||||
echo "ERROR: --seeds requires a value." >&2
|
||||
exit 2
|
||||
}
|
||||
SEED_CSV="$2"
|
||||
shift 2
|
||||
;;
|
||||
--num-workers)
|
||||
[[ $# -ge 2 ]] || {
|
||||
echo "ERROR: --num-workers requires a value." >&2
|
||||
exit 2
|
||||
}
|
||||
NUM_WORKERS="$2"
|
||||
shift 2
|
||||
;;
|
||||
--python)
|
||||
[[ $# -ge 2 ]] || {
|
||||
echo "ERROR: --python requires a value." >&2
|
||||
exit 2
|
||||
}
|
||||
PYTHON_BIN="$2"
|
||||
shift 2
|
||||
;;
|
||||
--campaign)
|
||||
[[ $# -ge 2 ]] || {
|
||||
echo "ERROR: --campaign requires a value." >&2
|
||||
exit 2
|
||||
}
|
||||
CAMPAIGN_NAME="$2"
|
||||
shift 2
|
||||
;;
|
||||
--train-only)
|
||||
TRAIN_ONLY=1
|
||||
shift
|
||||
;;
|
||||
--dry-run)
|
||||
DRY_RUN=1
|
||||
shift
|
||||
;;
|
||||
-h|--help)
|
||||
usage
|
||||
exit 0
|
||||
;;
|
||||
*)
|
||||
echo "ERROR: unknown argument: $1" >&2
|
||||
usage >&2
|
||||
exit 2
|
||||
;;
|
||||
esac
|
||||
done
|
||||
|
||||
[[ -n "$GPU_CSV" ]] || {
|
||||
echo "ERROR: --gpus is required." >&2
|
||||
usage >&2
|
||||
exit 2
|
||||
}
|
||||
[[ -n "$SEED_CSV" ]] || {
|
||||
echo "ERROR: --seeds must not be empty." >&2
|
||||
exit 2
|
||||
}
|
||||
[[ "$NUM_WORKERS" =~ ^[0-9]+$ ]] || {
|
||||
echo "ERROR: --num-workers must be a non-negative integer." >&2
|
||||
exit 2
|
||||
}
|
||||
[[ "$CAMPAIGN_NAME" =~ ^[A-Za-z0-9._-]+$ ]] || {
|
||||
echo "ERROR: --campaign may contain only letters, numbers, ., _, and -." >&2
|
||||
exit 2
|
||||
}
|
||||
|
||||
required_files=(
|
||||
"$SCRIPT_DIR/train_all_future.py"
|
||||
"$SCRIPT_DIR/evaluate_all_runs_linux.sh"
|
||||
"$SCRIPT_DIR/evaluate_calibration_all_runs_linux.sh"
|
||||
"$EXTRA_INFO_TYPES_FILE"
|
||||
"$TRAIN_EID_FILE"
|
||||
"$VAL_EID_FILE"
|
||||
"$TEST_EID_FILE"
|
||||
)
|
||||
for required_file in "${required_files[@]}"; do
|
||||
[[ -f "$required_file" ]] || {
|
||||
echo "ERROR: missing required file: $required_file" >&2
|
||||
exit 2
|
||||
}
|
||||
done
|
||||
command -v "$PYTHON_BIN" >/dev/null 2>&1 || {
|
||||
echo "ERROR: Python executable not found: $PYTHON_BIN" >&2
|
||||
exit 2
|
||||
}
|
||||
|
||||
IFS=',' read -r -a GPU_IDS <<< "$GPU_CSV"
|
||||
declare -A SEEN_GPUS=()
|
||||
for gpu in "${GPU_IDS[@]}"; do
|
||||
[[ -n "$gpu" && "$gpu" =~ ^[A-Za-z0-9._:-]+$ ]] || {
|
||||
echo "ERROR: invalid GPU id: $gpu" >&2
|
||||
exit 2
|
||||
}
|
||||
[[ -z "${SEEN_GPUS[$gpu]+x}" ]] || {
|
||||
echo "ERROR: duplicate GPU id: $gpu" >&2
|
||||
exit 2
|
||||
}
|
||||
SEEN_GPUS["$gpu"]=1
|
||||
done
|
||||
|
||||
IFS=',' read -r -a SEEDS <<< "$SEED_CSV"
|
||||
declare -A SEEN_SEEDS=()
|
||||
for seed in "${SEEDS[@]}"; do
|
||||
[[ "$seed" =~ ^[0-9]+$ ]] || {
|
||||
echo "ERROR: invalid seed: $seed" >&2
|
||||
exit 2
|
||||
}
|
||||
[[ -z "${SEEN_SEEDS[$seed]+x}" ]] || {
|
||||
echo "ERROR: duplicate seed: $seed" >&2
|
||||
exit 2
|
||||
}
|
||||
SEEN_SEEDS["$seed"]=1
|
||||
done
|
||||
|
||||
RUNS_ROOT="$SCRIPT_DIR/runs/$CAMPAIGN_NAME"
|
||||
LOG_ROOT="$SCRIPT_DIR/batch_logs/$CAMPAIGN_NAME"
|
||||
if ((DRY_RUN == 0)); then
|
||||
mkdir -p "$RUNS_ROOT" "$LOG_ROOT"
|
||||
fi
|
||||
|
||||
declare -a JOB_NAMES=()
|
||||
declare -a JOB_SEEDS=()
|
||||
declare -a JOB_ARCHITECTURES=()
|
||||
declare -a JOB_DIST_MODES=()
|
||||
declare -a JOB_BASELINES=()
|
||||
|
||||
add_job() {
|
||||
JOB_NAMES+=("$1")
|
||||
JOB_SEEDS+=("$2")
|
||||
JOB_ARCHITECTURES+=("$3")
|
||||
JOB_DIST_MODES+=("$4")
|
||||
JOB_BASELINES+=("$5")
|
||||
}
|
||||
|
||||
for seed in "${SEEDS[@]}"; do
|
||||
add_job \
|
||||
"traj_mixer_relative_weibull_first_onset" \
|
||||
"$seed" \
|
||||
"traj_mixer_v5" \
|
||||
"weibull" \
|
||||
"on"
|
||||
add_job \
|
||||
"traj_mixer_relative_weibull_no_rate_baseline" \
|
||||
"$seed" \
|
||||
"traj_mixer_v5" \
|
||||
"weibull" \
|
||||
"off"
|
||||
add_job \
|
||||
"traj_mixer_relative_exponential_first_onset" \
|
||||
"$seed" \
|
||||
"traj_mixer_v5" \
|
||||
"exponential" \
|
||||
"on"
|
||||
add_job \
|
||||
"ffn_relative_weibull_first_onset" \
|
||||
"$seed" \
|
||||
"transformer_ffn_v1" \
|
||||
"weibull" \
|
||||
"on"
|
||||
done
|
||||
|
||||
print_command() {
|
||||
printf '%q ' "$@"
|
||||
printf '\n'
|
||||
}
|
||||
|
||||
run_job() {
|
||||
local job_index="$1"
|
||||
local gpu="$2"
|
||||
local job_name="${JOB_NAMES[$job_index]}"
|
||||
local seed="${JOB_SEEDS[$job_index]}"
|
||||
local architecture="${JOB_ARCHITECTURES[$job_index]}"
|
||||
local dist_mode="${JOB_DIST_MODES[$job_index]}"
|
||||
local baseline="${JOB_BASELINES[$job_index]}"
|
||||
local seed_runs_root="$RUNS_ROOT/seed_$seed"
|
||||
local seed_log_root="$LOG_ROOT/seed_$seed"
|
||||
local log_file="$seed_log_root/$job_name.log"
|
||||
local -a command=(
|
||||
"$PYTHON_BIN"
|
||||
-u
|
||||
"$SCRIPT_DIR/train_all_future.py"
|
||||
--runs_root "$seed_runs_root"
|
||||
--seed "$seed"
|
||||
--batch_size "$BATCH_SIZE"
|
||||
--num_workers "$NUM_WORKERS"
|
||||
--device cuda
|
||||
--model_architecture "$architecture"
|
||||
--time_mode relative
|
||||
--dist_mode "$dist_mode"
|
||||
--disease_history_mode timed
|
||||
--extra_info_types_file "$EXTRA_INFO_TYPES_FILE"
|
||||
--train_eid_file "$TRAIN_EID_FILE"
|
||||
--val_eid_file "$VAL_EID_FILE"
|
||||
--test_eid_file "$TEST_EID_FILE"
|
||||
)
|
||||
|
||||
if [[ "$baseline" == "on" ]]; then
|
||||
command+=(--risk_head_bias)
|
||||
else
|
||||
command+=(--no-risk_head_bias)
|
||||
fi
|
||||
|
||||
if ((DRY_RUN == 0)); then
|
||||
mkdir -p "$seed_runs_root" "$seed_log_root"
|
||||
fi
|
||||
|
||||
echo "[$(date '+%F %T')] START seed=$seed job=$job_name gpu=$gpu"
|
||||
echo " log=$log_file"
|
||||
if ((DRY_RUN)); then
|
||||
printf ' CUDA_VISIBLE_DEVICES=%q ' "$gpu"
|
||||
print_command "${command[@]}"
|
||||
return 0
|
||||
fi
|
||||
|
||||
if CUDA_VISIBLE_DEVICES="$gpu" PYTHONUNBUFFERED=1 \
|
||||
"${command[@]}" >"$log_file" 2>&1; then
|
||||
echo "[$(date '+%F %T')] DONE seed=$seed job=$job_name gpu=$gpu"
|
||||
return 0
|
||||
else
|
||||
local exit_code=$?
|
||||
echo "[$(date '+%F %T')] FAIL seed=$seed job=$job_name gpu=$gpu exit=$exit_code" >&2
|
||||
echo " See: $log_file" >&2
|
||||
return "$exit_code"
|
||||
fi
|
||||
}
|
||||
|
||||
worker() {
|
||||
local slot="$1"
|
||||
local gpu="${GPU_IDS[$slot]}"
|
||||
local job_index
|
||||
local failed=0
|
||||
|
||||
for ((job_index = slot; job_index < ${#JOB_NAMES[@]}; job_index += ${#GPU_IDS[@]})); do
|
||||
run_job "$job_index" "$gpu" || failed=1
|
||||
done
|
||||
return "$failed"
|
||||
}
|
||||
|
||||
echo "Campaign: $CAMPAIGN_NAME"
|
||||
echo "Seeds: ${SEEDS[*]}"
|
||||
echo "GPUs: ${GPU_IDS[*]}"
|
||||
echo "Experiments per seed: 4"
|
||||
echo "Total training tasks: ${#JOB_NAMES[@]}"
|
||||
echo "Runs root: $RUNS_ROOT"
|
||||
echo "Log root: $LOG_ROOT"
|
||||
echo
|
||||
|
||||
declare -a WORKER_PIDS=()
|
||||
for ((slot = 0; slot < ${#GPU_IDS[@]}; slot++)); do
|
||||
worker "$slot" &
|
||||
WORKER_PIDS+=("$!")
|
||||
done
|
||||
|
||||
overall_status=0
|
||||
for pid in "${WORKER_PIDS[@]}"; do
|
||||
wait "$pid" || overall_status=1
|
||||
done
|
||||
|
||||
if ((overall_status != 0)); then
|
||||
echo "One or more training tasks failed. Inspect: $LOG_ROOT" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
if ((TRAIN_ONLY == 0)); then
|
||||
evaluation_status=0
|
||||
auc_command=(
|
||||
bash
|
||||
"$SCRIPT_DIR/evaluate_all_runs_linux.sh"
|
||||
--gpus "$GPU_CSV"
|
||||
--runs-root "$RUNS_ROOT"
|
||||
--log-root "$LOG_ROOT/evaluate_auc"
|
||||
--python "$PYTHON_BIN"
|
||||
--num-workers "$NUM_WORKERS"
|
||||
--num-workers-auc "$NUM_WORKERS"
|
||||
)
|
||||
calibration_command=(
|
||||
bash
|
||||
"$SCRIPT_DIR/evaluate_calibration_all_runs_linux.sh"
|
||||
--gpus "$GPU_CSV"
|
||||
--runs-root "$RUNS_ROOT"
|
||||
--log-root "$LOG_ROOT/evaluate_calibration"
|
||||
--python "$PYTHON_BIN"
|
||||
--num-workers "$NUM_WORKERS"
|
||||
--num-workers-calibration "$NUM_WORKERS"
|
||||
)
|
||||
|
||||
echo ">> AUC evaluation"
|
||||
print_command "${auc_command[@]}"
|
||||
if ((DRY_RUN == 0)); then
|
||||
"${auc_command[@]}" || evaluation_status=1
|
||||
fi
|
||||
|
||||
echo ">> Calibration and point-process NLL evaluation"
|
||||
print_command "${calibration_command[@]}"
|
||||
if ((DRY_RUN == 0)); then
|
||||
"${calibration_command[@]}" || evaluation_status=1
|
||||
fi
|
||||
|
||||
if ((evaluation_status != 0)); then
|
||||
echo "One or more evaluation workflows failed. Inspect: $LOG_ROOT" >&2
|
||||
exit 1
|
||||
fi
|
||||
fi
|
||||
|
||||
if ((DRY_RUN)); then
|
||||
echo "Dry run completed successfully."
|
||||
else
|
||||
echo "All required all-future experiments completed successfully."
|
||||
fi
|
||||
@@ -1,182 +0,0 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
# Run all non-wrapper evaluation scripts for every completed experiment under
|
||||
# runs/. The script is written for Linux servers with bash 4.2.
|
||||
|
||||
cd "$(dirname "${BASH_SOURCE[0]}")"
|
||||
|
||||
PYTHON_BIN="${PYTHON_BIN:-python}"
|
||||
DEVICE="${DEVICE:-cuda}"
|
||||
EVAL_SPLIT="${EVAL_SPLIT:-test}"
|
||||
NUM_WORKERS="${NUM_WORKERS:-4}"
|
||||
CPU_REDUCE_WORKERS="${CPU_REDUCE_WORKERS:-}"
|
||||
NUM_WORKERS_AUC="${NUM_WORKERS_AUC:-}"
|
||||
BATCH_SIZE="${BATCH_SIZE:-}"
|
||||
DATASET_SUBSET_SIZE="${DATASET_SUBSET_SIZE:-}"
|
||||
DRY_RUN="${DRY_RUN:-0}"
|
||||
|
||||
# These attribution jobs can be expensive, but they are part of the evaluation
|
||||
# surface in this repository. Set either variable to 0 to leave that family out.
|
||||
RUN_EXTRA_INFO_ATTRIBUTION="${RUN_EXTRA_INFO_ATTRIBUTION:-1}"
|
||||
RUN_SINGLE_DISEASE_MORTALITY_ATTRIBUTION="${RUN_SINGLE_DISEASE_MORTALITY_ATTRIBUTION:-1}"
|
||||
|
||||
common_args_base() {
|
||||
printf '%s\n' --run_path "$1" --eval_split "${EVAL_SPLIT}" --num_workers "${NUM_WORKERS}"
|
||||
if [[ -n "${BATCH_SIZE}" ]]; then
|
||||
printf '%s\n' --batch_size "${BATCH_SIZE}"
|
||||
fi
|
||||
if [[ -n "${DATASET_SUBSET_SIZE}" ]]; then
|
||||
printf '%s\n' --dataset_subset_size "${DATASET_SUBSET_SIZE}"
|
||||
fi
|
||||
}
|
||||
|
||||
common_args_with_device() {
|
||||
common_args_base "$1"
|
||||
printf '%s\n' --device "${DEVICE}"
|
||||
}
|
||||
|
||||
auc_args() {
|
||||
if [[ -n "${NUM_WORKERS_AUC}" ]]; then
|
||||
printf '%s\n' --num_workers_auc "${NUM_WORKERS_AUC}"
|
||||
fi
|
||||
}
|
||||
|
||||
cpu_reduce_args() {
|
||||
if [[ -n "${CPU_REDUCE_WORKERS}" ]]; then
|
||||
printf '%s\n' --cpu_reduce_workers "${CPU_REDUCE_WORKERS}"
|
||||
fi
|
||||
}
|
||||
|
||||
has_completed_dir() {
|
||||
local dir="$1"
|
||||
shift
|
||||
[[ -d "${dir}" ]] || return 1
|
||||
local required
|
||||
for required in "$@"; do
|
||||
[[ -s "${dir}/${required}" ]] || return 1
|
||||
done
|
||||
}
|
||||
|
||||
run_command() {
|
||||
echo " run: $*"
|
||||
if [[ "${DRY_RUN}" == "1" ]]; then
|
||||
return 0
|
||||
fi
|
||||
"$@"
|
||||
}
|
||||
|
||||
run_dir_result_if_missing() {
|
||||
local label="$1"
|
||||
local result_dir="$2"
|
||||
local required_1="$3"
|
||||
local required_2="$4"
|
||||
shift 4
|
||||
|
||||
if has_completed_dir "${result_dir}" "${required_1}" "${required_2}"; then
|
||||
echo " skip ${label}: found ${result_dir}"
|
||||
return 0
|
||||
fi
|
||||
|
||||
run_command "$@"
|
||||
}
|
||||
|
||||
run_has_extra_info() {
|
||||
"${PYTHON_BIN}" - "$1" <<'PY'
|
||||
import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
cfg_path = Path(sys.argv[1]) / "train_config.json"
|
||||
try:
|
||||
cfg = json.loads(cfg_path.read_text(encoding="utf-8"))
|
||||
except Exception:
|
||||
raise SystemExit(1)
|
||||
|
||||
extra = cfg.get("extra_info_types", [])
|
||||
raise SystemExit(0 if isinstance(extra, list) and len(extra) > 0 else 1)
|
||||
PY
|
||||
}
|
||||
|
||||
run_is_all_future() {
|
||||
"${PYTHON_BIN}" - "$1" <<'PY'
|
||||
import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
cfg_path = Path(sys.argv[1]) / "train_config.json"
|
||||
try:
|
||||
cfg = json.loads(cfg_path.read_text(encoding="utf-8"))
|
||||
except Exception:
|
||||
raise SystemExit(1)
|
||||
|
||||
mode = str(cfg.get("model_target_mode", "next_token")).lower()
|
||||
raise SystemExit(0 if mode == "all_future" else 1)
|
||||
PY
|
||||
}
|
||||
|
||||
for run_path in runs/*; do
|
||||
[[ -d "${run_path}" ]] || continue
|
||||
|
||||
echo "==> ${run_path}"
|
||||
if [[ ! -f "${run_path}/train_config.json" ]]; then
|
||||
echo " skip run: missing train_config.json"
|
||||
continue
|
||||
fi
|
||||
if [[ ! -s "${run_path}/best_model.pt" ]]; then
|
||||
echo " skip run: missing best_model.pt"
|
||||
continue
|
||||
fi
|
||||
|
||||
common=()
|
||||
while IFS= read -r arg; do common+=("${arg}"); done < <(common_args_with_device "${run_path}")
|
||||
|
||||
auc_extra=()
|
||||
while IFS= read -r arg; do auc_extra+=("${arg}"); done < <(auc_args)
|
||||
|
||||
cpu_reduce_extra=()
|
||||
while IFS= read -r arg; do cpu_reduce_extra+=("${arg}"); done < <(cpu_reduce_args)
|
||||
|
||||
run_dir_result_if_missing \
|
||||
"evaluate_auc.py" \
|
||||
"${run_path}" \
|
||||
"df_both.csv" \
|
||||
"df_auc_unpooled.csv" \
|
||||
"${PYTHON_BIN}" evaluate_auc.py "${common[@]}" "${auc_extra[@]}"
|
||||
|
||||
run_dir_result_if_missing \
|
||||
"evaluate_auc_v2.py" \
|
||||
"${run_path}" \
|
||||
"df_auc_landmark.csv" \
|
||||
"df_auc_landmark_unpooled.csv" \
|
||||
"${PYTHON_BIN}" evaluate_auc_v2.py "${common[@]}" "${auc_extra[@]}"
|
||||
|
||||
if ! run_is_all_future "${run_path}"; then
|
||||
echo " skip attribution evaluations: model_target_mode is not all_future"
|
||||
continue
|
||||
fi
|
||||
|
||||
if [[ "${RUN_EXTRA_INFO_ATTRIBUTION}" == "1" ]]; then
|
||||
if run_has_extra_info "${run_path}"; then
|
||||
run_dir_result_if_missing \
|
||||
"evaluate_extra_info_attribution.py" \
|
||||
"${run_path}/extra_info_attribution_${EVAL_SPLIT}" \
|
||||
"manifest.json" \
|
||||
"summary_extra_info_disease_parameters.csv" \
|
||||
"${PYTHON_BIN}" evaluate_extra_info_attribution.py "${common[@]}" "${cpu_reduce_extra[@]}"
|
||||
else
|
||||
echo " skip evaluate_extra_info_attribution.py: run has no extra-info types"
|
||||
fi
|
||||
fi
|
||||
|
||||
if [[ "${RUN_SINGLE_DISEASE_MORTALITY_ATTRIBUTION}" == "1" ]]; then
|
||||
run_dir_result_if_missing \
|
||||
"evaluate_single_disease_mortality_attribution.py" \
|
||||
"${run_path}/single_disease_mortality_parameters_${EVAL_SPLIT}_all_diseases" \
|
||||
"manifest.json" \
|
||||
"summary_by_disease_age_sex.csv" \
|
||||
"${PYTHON_BIN}" evaluate_single_disease_mortality_attribution.py "${common[@]}"
|
||||
fi
|
||||
done
|
||||
|
||||
echo "All missing evaluations are complete."
|
||||
@@ -1,130 +0,0 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
# Linux bash 5.2+ training-only script.
|
||||
#
|
||||
# Based on the existing runs, the objective/time/death-distribution checks are
|
||||
# already covered. The remaining gap for the current proof chain is the
|
||||
# extra-info ablation under the final candidate model:
|
||||
#
|
||||
# all_future + relative time + mixed death/risk head
|
||||
#
|
||||
# This script only launches those missing training jobs. It intentionally does
|
||||
# not call evaluate_*.py and does not add extra random seeds.
|
||||
|
||||
cd "$(dirname "${BASH_SOURCE[0]}")"
|
||||
|
||||
PYTHON_BIN="${PYTHON_BIN:-python}"
|
||||
DEVICE="${DEVICE:-cuda}"
|
||||
NUM_WORKERS="${NUM_WORKERS:-4}"
|
||||
PROGRESS_INTERVAL="${PROGRESS_INTERVAL:-20}"
|
||||
|
||||
TIME_MODE="relative"
|
||||
DIST_MODE="mixed"
|
||||
SEED="42"
|
||||
VALIDATION_QUERY_SEED="42"
|
||||
|
||||
COMMON_ARGS=(
|
||||
--data_prefix ukb
|
||||
--labels_file labels.csv
|
||||
--seed "${SEED}"
|
||||
--validation_query_seed "${VALIDATION_QUERY_SEED}"
|
||||
--train_eid_file ukb_train_eid.csv
|
||||
--val_eid_file ukb_val_eid.csv
|
||||
--test_eid_file ukb_test_eid.csv
|
||||
--min_history_events 1
|
||||
--min_future_events 1
|
||||
--n_embd 120
|
||||
--n_head 10
|
||||
--n_hist_layer 12
|
||||
--n_tab_layer 4
|
||||
--n_bins 16
|
||||
--extra_pool_reduce mean
|
||||
--dropout 0.0
|
||||
--batch_size 256
|
||||
--base_lr 0.0003
|
||||
--weight_decay 0.1
|
||||
--betas 0.9 0.99
|
||||
--grad_clip 1.0
|
||||
--max_epochs 200
|
||||
--warmup_epochs 10
|
||||
--patience 15
|
||||
--min_lr_ratio 0.1
|
||||
--num_workers "${NUM_WORKERS}"
|
||||
--device "${DEVICE}"
|
||||
--progress_interval "${PROGRESS_INTERVAL}"
|
||||
)
|
||||
|
||||
already_trained() {
|
||||
local extra_file="$1"
|
||||
"${PYTHON_BIN}" - "$TIME_MODE" "$DIST_MODE" "$extra_file" "$SEED" "$VALIDATION_QUERY_SEED" <<'PY'
|
||||
import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
time_mode, dist_mode, extra_file, seed, validation_query_seed = sys.argv[1:6]
|
||||
extra_name = Path(extra_file).name
|
||||
|
||||
for config_path in Path("runs").glob("*/train_config.json"):
|
||||
try:
|
||||
cfg = json.loads(config_path.read_text(encoding="utf-8"))
|
||||
except Exception:
|
||||
continue
|
||||
|
||||
observed_query_seed = cfg.get(
|
||||
"all_future_validation_query_seed",
|
||||
cfg.get("validation_query_seed", -1),
|
||||
)
|
||||
|
||||
if (
|
||||
cfg.get("model_target_mode") == "all_future"
|
||||
and cfg.get("time_mode") == time_mode
|
||||
and cfg.get("dist_mode") == dist_mode
|
||||
and Path(str(cfg.get("extra_info_types_file", ""))).name == extra_name
|
||||
and int(cfg.get("seed", -1)) == int(seed)
|
||||
and int(observed_query_seed) == int(validation_query_seed)
|
||||
):
|
||||
print(config_path.parent)
|
||||
raise SystemExit(0)
|
||||
|
||||
raise SystemExit(1)
|
||||
PY
|
||||
}
|
||||
|
||||
train_if_missing() {
|
||||
local label="$1"
|
||||
local extra_file="$2"
|
||||
|
||||
if [[ ! -f "${extra_file}" ]]; then
|
||||
echo "ERROR: missing extra-info type file: ${extra_file}" >&2
|
||||
return 2
|
||||
fi
|
||||
|
||||
echo "==> Checking ${label}: ${TIME_MODE} ${DIST_MODE} all_future with ${extra_file}"
|
||||
if existing_run="$(already_trained "$extra_file")"; then
|
||||
echo " skip: already trained at ${existing_run}"
|
||||
return 0
|
||||
fi
|
||||
|
||||
echo " train: ${label}"
|
||||
"${PYTHON_BIN}" train_all_future.py \
|
||||
"${COMMON_ARGS[@]}" \
|
||||
--time_mode "${TIME_MODE}" \
|
||||
--dist_mode "${DIST_MODE}" \
|
||||
--extra_info_types_file "${extra_file}"
|
||||
}
|
||||
|
||||
# Already present in runs/:
|
||||
# - next-token objective checks under SAB, plus older absolute extra ablations
|
||||
# - all-future absolute/relative x exponential/weibull/mixed under SAB
|
||||
#
|
||||
# Still needed:
|
||||
# - final all-future relative+mixed extra-info ablations beyond the existing
|
||||
# SAB baseline. These close the disease-only question without expanding seed
|
||||
# count or running downstream evaluation.
|
||||
train_if_missing "true_disease_only" "extra_info_types_none.txt"
|
||||
train_if_missing "assessment_only_extra" "extra_info_types_assessment_only.txt"
|
||||
train_if_missing "exposure_only_extra" "extra_info_types_exposure_only.txt"
|
||||
train_if_missing "all_extra_info" "extra_info_types_all.txt"
|
||||
|
||||
echo "All requested training-only missing configurations are done."
|
||||
@@ -1,76 +0,0 @@
|
||||
#!/usr/bin/env bash
|
||||
|
||||
# Export Weibull shape-parameter summaries for the all_future models trained
|
||||
# with smoking/alcohol/BMI extra information.
|
||||
#
|
||||
# Bash 4.2 compatible. Run from the DeepHealth repository root on the Linux
|
||||
# server, for example:
|
||||
#
|
||||
# bash run_weibull_shape_exports.sh
|
||||
#
|
||||
# Optional overrides:
|
||||
# PYTHON=python3 DEVICE=cuda BATCH_SIZE=128 NUM_WORKERS=0 ROW_BATCH_SIZE=512 \
|
||||
# bash run_weibull_shape_exports.sh
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
PYTHON="${PYTHON:-python}"
|
||||
DEVICE="${DEVICE:-cuda}"
|
||||
BATCH_SIZE="${BATCH_SIZE:-128}"
|
||||
NUM_WORKERS="${NUM_WORKERS:-0}"
|
||||
ROW_BATCH_SIZE="${ROW_BATCH_SIZE:-512}"
|
||||
LANDMARK_START="${LANDMARK_START:-40}"
|
||||
LANDMARK_STOP="${LANDMARK_STOP:-80}"
|
||||
LANDMARK_STEP="${LANDMARK_STEP:-5}"
|
||||
HORIZONS="${HORIZONS:-1,5,10}"
|
||||
|
||||
RUNS=(
|
||||
"runs/relative_weibull_all_future_pure_disease_20260620_095229"
|
||||
"runs/relative_mixed_all_future_pure_disease_20260620_132415"
|
||||
"runs/absolute_weibull_all_future_pure_disease_20260620_114816"
|
||||
"runs/absolute_mixed_all_future_pure_disease_20260620_161804"
|
||||
)
|
||||
|
||||
echo "Python: ${PYTHON}"
|
||||
echo "Device: ${DEVICE}"
|
||||
echo "Batch size: ${BATCH_SIZE}"
|
||||
echo "Workers: ${NUM_WORKERS}"
|
||||
echo "Row batch size: ${ROW_BATCH_SIZE}"
|
||||
echo "Horizons: ${HORIZONS}"
|
||||
echo
|
||||
|
||||
for run_path in "${RUNS[@]}"; do
|
||||
if [[ ! -d "${run_path}" ]]; then
|
||||
echo "[ERROR] Missing run directory: ${run_path}" >&2
|
||||
exit 1
|
||||
fi
|
||||
if [[ ! -f "${run_path}/best_model.pt" ]]; then
|
||||
echo "[ERROR] Missing checkpoint: ${run_path}/best_model.pt" >&2
|
||||
exit 1
|
||||
fi
|
||||
if [[ ! -f "${run_path}/train_config.json" ]]; then
|
||||
echo "[ERROR] Missing config: ${run_path}/train_config.json" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
output_path="${run_path}/weibull_shape_parameter_stats_test"
|
||||
echo "=== Exporting Weibull shape stats: ${run_path} ==="
|
||||
"${PYTHON}" export_weibull_shape_parameter_stats.py \
|
||||
--run_path "${run_path}" \
|
||||
--output_path "${output_path}" \
|
||||
--eval_split test \
|
||||
--device "${DEVICE}" \
|
||||
--batch_size "${BATCH_SIZE}" \
|
||||
--num_workers "${NUM_WORKERS}" \
|
||||
--row_batch_size "${ROW_BATCH_SIZE}" \
|
||||
--hidden_cache_dtype float32 \
|
||||
--landmark_start "${LANDMARK_START}" \
|
||||
--landmark_stop "${LANDMARK_STOP}" \
|
||||
--landmark_step "${LANDMARK_STEP}" \
|
||||
--horizons "${HORIZONS}" \
|
||||
--include_all_token_rho_summary
|
||||
echo "Wrote: ${output_path}"
|
||||
echo
|
||||
done
|
||||
|
||||
echo "All Weibull shape exports completed."
|
||||
267
targets.py
267
targets.py
@@ -2,13 +2,11 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Iterable
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
PAD_IDX = 0
|
||||
CHECKUP_IDX = 1
|
||||
RESERVED_IDX = 1
|
||||
NO_EVENT_IDX = 2
|
||||
DAYS_PER_YEAR = 365.25
|
||||
|
||||
@@ -32,38 +30,6 @@ class NextTokenTargets:
|
||||
target_times_years: np.ndarray
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class UniqueTimeSetTargets:
|
||||
"""
|
||||
Unique-time set supervision targets.
|
||||
|
||||
Shapes:
|
||||
readout_mask: (L,)
|
||||
target_dt_unique: (L,)
|
||||
target_multi_hot: (L, vocab_size)
|
||||
|
||||
where L = N - 1.
|
||||
|
||||
Only group-end positions can have readout_mask=True.
|
||||
target_dt_unique is measured in years.
|
||||
"""
|
||||
readout_mask: np.ndarray
|
||||
target_dt_unique: np.ndarray
|
||||
target_multi_hot: np.ndarray
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class TargetPack:
|
||||
"""
|
||||
Combined target package for one patient sequence.
|
||||
|
||||
Contains both next-token targets and unique-time-set targets.
|
||||
The training pipeline decides which one to use.
|
||||
"""
|
||||
next_token: NextTokenTargets
|
||||
unique_time_set: UniqueTimeSetTargets
|
||||
|
||||
|
||||
def _as_numpy_1d(
|
||||
x: np.ndarray,
|
||||
name: str,
|
||||
@@ -139,7 +105,7 @@ def build_next_token_targets(
|
||||
target_events: [x1, x2, ..., xN-1]
|
||||
target_times_years: [t1, t2, ..., tN-1] / 365.25
|
||||
|
||||
This function does not ignore PAD/CHECKUP/NO_EVENT. Ignoring belongs to
|
||||
This function does not ignore PAD/RESERVED/NO_EVENT. Ignoring belongs to
|
||||
the loss function because different objectives may use different ignore ids.
|
||||
"""
|
||||
labels = _as_numpy_1d(labels, "labels", np.int64)
|
||||
@@ -163,232 +129,3 @@ def build_next_token_targets(
|
||||
target_events=target_events,
|
||||
target_times_years=target_times_years,
|
||||
)
|
||||
|
||||
|
||||
def build_unique_time_set_targets(
|
||||
labels: np.ndarray,
|
||||
times_days: np.ndarray,
|
||||
*,
|
||||
vocab_size: int,
|
||||
ignored_target_ids: Iterable[int] = (PAD_IDX, CHECKUP_IDX),
|
||||
require_sorted: bool = True,
|
||||
) -> UniqueTimeSetTargets:
|
||||
"""
|
||||
Build next-unique-time set targets.
|
||||
|
||||
This is the target construction used by your UTS / default mode.
|
||||
|
||||
For each input position i:
|
||||
- only if i is the last token of its timestamp group;
|
||||
- find the next distinct timestamp group;
|
||||
- target is the set of valid event labels at that next timestamp.
|
||||
|
||||
Example:
|
||||
|
||||
t=49: X
|
||||
t=50: A, B, C
|
||||
t=51: D, E
|
||||
|
||||
Supervises:
|
||||
|
||||
X@49 -> {A, B, C}@50
|
||||
group_end@50 -> {D, E}@51
|
||||
|
||||
It does NOT supervise:
|
||||
|
||||
A@50 -> B@50
|
||||
B@50 -> C@50
|
||||
|
||||
Parameters
|
||||
----------
|
||||
labels:
|
||||
Full event sequence labels, shape (N,).
|
||||
|
||||
times_days:
|
||||
Full event sequence times in days, shape (N,).
|
||||
|
||||
vocab_size:
|
||||
Size of output vocabulary.
|
||||
|
||||
ignored_target_ids:
|
||||
Label ids that should not enter target_multi_hot.
|
||||
Usually:
|
||||
no no-event: {0, 1}
|
||||
with no-event: {0, 1, 2}
|
||||
For UTS, I recommend ignoring <NO_EVENT> unless explicitly testing it
|
||||
as an event target.
|
||||
|
||||
Returns
|
||||
-------
|
||||
UniqueTimeSetTargets
|
||||
"""
|
||||
labels = _as_numpy_1d(labels, "labels", np.int64)
|
||||
times_days = _as_numpy_1d(times_days, "times_days", np.float32)
|
||||
validate_event_sequence(labels, times_days, require_sorted=require_sorted)
|
||||
|
||||
if vocab_size <= 0:
|
||||
raise ValueError(f"vocab_size must be positive, got {vocab_size}")
|
||||
|
||||
if len(labels) < 2:
|
||||
raise ValueError(
|
||||
"Need at least two events to build unique-time-set targets."
|
||||
)
|
||||
|
||||
input_len = len(labels) - 1
|
||||
|
||||
readout_mask = np.zeros(input_len, dtype=bool)
|
||||
target_dt_unique = np.zeros(input_len, dtype=np.float32)
|
||||
target_multi_hot = np.zeros((input_len, vocab_size), dtype=bool)
|
||||
|
||||
ignored = {int(x) for x in ignored_target_ids}
|
||||
|
||||
unique_times = np.unique(times_days)
|
||||
time_to_group_idx = {t: i for i, t in enumerate(unique_times)}
|
||||
group_indices = np.array([time_to_group_idx[t]
|
||||
for t in times_days], dtype=np.int64)
|
||||
|
||||
for i in range(input_len):
|
||||
current_group = group_indices[i]
|
||||
|
||||
is_last_in_group = (
|
||||
i == input_len - 1
|
||||
or group_indices[i + 1] != current_group
|
||||
)
|
||||
if not is_last_in_group:
|
||||
continue
|
||||
|
||||
next_group_idx = current_group + 1
|
||||
if next_group_idx >= len(unique_times):
|
||||
continue
|
||||
|
||||
next_time = unique_times[next_group_idx]
|
||||
next_labels = labels[group_indices == next_group_idx]
|
||||
|
||||
valid_next_labels: list[int] = []
|
||||
for lab in next_labels:
|
||||
lab_int = int(lab)
|
||||
if lab_int in ignored:
|
||||
continue
|
||||
if lab_int < 0 or lab_int >= vocab_size:
|
||||
continue
|
||||
valid_next_labels.append(lab_int)
|
||||
|
||||
# If next timestamp contains only technical tokens, do not supervise UTS.
|
||||
if len(valid_next_labels) == 0:
|
||||
continue
|
||||
|
||||
readout_mask[i] = True
|
||||
target_dt_unique[i] = float(next_time - times_days[i]) / DAYS_PER_YEAR
|
||||
target_multi_hot[i, valid_next_labels] = True
|
||||
|
||||
return UniqueTimeSetTargets(
|
||||
readout_mask=readout_mask,
|
||||
target_dt_unique=target_dt_unique.astype(np.float32),
|
||||
target_multi_hot=target_multi_hot,
|
||||
)
|
||||
|
||||
|
||||
def build_all_targets(
|
||||
labels: np.ndarray,
|
||||
times_days: np.ndarray,
|
||||
*,
|
||||
vocab_size: int,
|
||||
ignored_uts_target_ids: Iterable[int] = (PAD_IDX, CHECKUP_IDX),
|
||||
require_sorted: bool = True,
|
||||
) -> TargetPack:
|
||||
"""
|
||||
Build both next-token targets and unique-time-set targets for one patient.
|
||||
|
||||
This is the function dataset.py should usually call during initialization.
|
||||
|
||||
The dataset can then store:
|
||||
event_seq = target_pack.next_token.input_events
|
||||
time_seq = target_pack.next_token.input_times_years
|
||||
|
||||
target_event_seq = target_pack.next_token.target_events
|
||||
target_time_seq = target_pack.next_token.target_times_years
|
||||
|
||||
readout_mask = target_pack.unique_time_set.readout_mask
|
||||
target_dt_unique = target_pack.unique_time_set.target_dt_unique
|
||||
target_multi_hot = target_pack.unique_time_set.target_multi_hot
|
||||
"""
|
||||
next_token = build_next_token_targets(
|
||||
labels=labels,
|
||||
times_days=times_days,
|
||||
require_sorted=require_sorted,
|
||||
)
|
||||
|
||||
unique_time_set = build_unique_time_set_targets(
|
||||
labels=labels,
|
||||
times_days=times_days,
|
||||
vocab_size=vocab_size,
|
||||
ignored_target_ids=ignored_uts_target_ids,
|
||||
require_sorted=require_sorted,
|
||||
)
|
||||
|
||||
return TargetPack(
|
||||
next_token=next_token,
|
||||
unique_time_set=unique_time_set,
|
||||
)
|
||||
|
||||
|
||||
def get_group_end_mask_from_times(
|
||||
times_days: np.ndarray,
|
||||
*,
|
||||
input_len: int | None = None,
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Convenience utility for debugging.
|
||||
|
||||
Returns a bool mask indicating the last token of each same-time group
|
||||
within the input sequence.
|
||||
|
||||
If input_len is None, uses len(times_days) - 1, matching model input length.
|
||||
"""
|
||||
times_days = _as_numpy_1d(times_days, "times_days", np.float32)
|
||||
|
||||
if input_len is None:
|
||||
input_len = len(times_days) - 1
|
||||
|
||||
if input_len < 0 or input_len > len(times_days):
|
||||
raise ValueError(
|
||||
f"Invalid input_len={input_len} for sequence length {len(times_days)}"
|
||||
)
|
||||
|
||||
out = np.zeros(input_len, dtype=bool)
|
||||
|
||||
for i in range(input_len):
|
||||
is_last_in_group = (
|
||||
i == input_len - 1
|
||||
or times_days[i + 1] != times_days[i]
|
||||
)
|
||||
out[i] = is_last_in_group
|
||||
|
||||
return out
|
||||
|
||||
|
||||
def summarize_targets(
|
||||
target_pack: TargetPack,
|
||||
) -> dict[str, int | float]:
|
||||
"""
|
||||
Small debugging helper for logging.
|
||||
"""
|
||||
nt = target_pack.next_token
|
||||
uts = target_pack.unique_time_set
|
||||
|
||||
n_tokens = int(len(nt.input_events))
|
||||
n_readout = int(uts.readout_mask.sum())
|
||||
n_positive_labels = int(uts.target_multi_hot.sum())
|
||||
|
||||
mean_set_size = (
|
||||
float(n_positive_labels / n_readout)
|
||||
if n_readout > 0
|
||||
else 0.0
|
||||
)
|
||||
|
||||
return {
|
||||
"n_input_tokens": n_tokens,
|
||||
"n_uts_readouts": n_readout,
|
||||
"n_uts_positive_labels": n_positive_labels,
|
||||
"mean_uts_set_size": mean_set_size,
|
||||
}
|
||||
|
||||
@@ -1,330 +0,0 @@
|
||||
import math
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
|
||||
from backbones import (
|
||||
SharedEventTrajectoryCore,
|
||||
SharedTrajectoryMixer,
|
||||
TrajectoryCrossAttention,
|
||||
)
|
||||
from models import (
|
||||
EVENT_TRAJECTORY_ARCHITECTURE,
|
||||
MODEL_SIZE_PRESETS,
|
||||
DeepHealth,
|
||||
resolve_model_size,
|
||||
validate_event_trajectory_config,
|
||||
validate_event_trajectory_state_dict,
|
||||
)
|
||||
from train_util import get_model_parameter_counts
|
||||
|
||||
|
||||
def build_test_model(
|
||||
*,
|
||||
target_mode: str = "next_token",
|
||||
time_mode: str = "absolute",
|
||||
n_reasoning_rounds: int = 3,
|
||||
) -> DeepHealth:
|
||||
return DeepHealth(
|
||||
vocab_size=32,
|
||||
model_size="nano",
|
||||
n_reasoning_rounds=n_reasoning_rounds,
|
||||
n_types=2,
|
||||
n_cont_types=0,
|
||||
n_categories=2,
|
||||
cont_type_ids=[],
|
||||
target_mode=target_mode,
|
||||
time_mode=time_mode,
|
||||
)
|
||||
|
||||
|
||||
def model_inputs() -> dict[str, torch.Tensor]:
|
||||
return {
|
||||
"event_seq": torch.tensor([[1, 2, 3, 4], [5, 6, 0, 0]]),
|
||||
"time_seq": torch.tensor(
|
||||
[[1.0, 2.0, 3.0, 4.0], [1.0, 2.0, 0.0, 0.0]]
|
||||
),
|
||||
"sex": torch.tensor([0, 1]),
|
||||
"padding_mask": torch.tensor(
|
||||
[[True, True, True, True], [True, True, False, False]]
|
||||
),
|
||||
"other_type": torch.zeros(2, 1, dtype=torch.long),
|
||||
"other_value": torch.zeros(2, 1),
|
||||
"other_value_kind": torch.zeros(2, 1, dtype=torch.long),
|
||||
"other_time": torch.zeros(2, 1),
|
||||
}
|
||||
|
||||
|
||||
class EventTrajectoryBackboneTest(unittest.TestCase):
|
||||
def test_model_size_presets(self) -> None:
|
||||
expected = {
|
||||
"nano": (120, 6, 20, 24),
|
||||
"tiny": (256, 8, 32, 32),
|
||||
"small": (512, 8, 64, 32),
|
||||
"medium": (768, 12, 64, 48),
|
||||
"huge": (1024, 16, 64, 64),
|
||||
}
|
||||
self.assertEqual(set(MODEL_SIZE_PRESETS), set(expected))
|
||||
for name, values in expected.items():
|
||||
preset = resolve_model_size(name)
|
||||
self.assertEqual(
|
||||
(
|
||||
preset.d_model,
|
||||
preset.n_trajectory,
|
||||
preset.trajectory_dim,
|
||||
preset.traj_hidden,
|
||||
),
|
||||
values,
|
||||
)
|
||||
|
||||
def test_default_mixer_shapes_and_parameter_count(self) -> None:
|
||||
mixer = SharedTrajectoryMixer(
|
||||
n_trajectory=8,
|
||||
trajectory_dim=32,
|
||||
)
|
||||
state = torch.randn(2, 5, 8, 32)
|
||||
self.assertEqual(mixer(state).shape, state.shape)
|
||||
self.assertEqual(mixer.traj_hidden, 32)
|
||||
self.assertEqual(tuple(mixer.gate_proj.shape), (32, 8, 32))
|
||||
self.assertEqual(tuple(mixer.value_proj.shape), (32, 8, 32))
|
||||
self.assertEqual(tuple(mixer.output_proj.shape), (32, 32, 8))
|
||||
self.assertEqual(
|
||||
get_model_parameter_counts(mixer),
|
||||
{
|
||||
"model_parameter_count": 24_576,
|
||||
"trainable_parameter_count": 24_576,
|
||||
},
|
||||
)
|
||||
|
||||
def test_reasoning_rounds_share_one_core_parameter_set(self) -> None:
|
||||
core_one = SharedEventTrajectoryCore(
|
||||
d_model=256,
|
||||
n_trajectory=8,
|
||||
n_reasoning_rounds=1,
|
||||
)
|
||||
core_twelve = SharedEventTrajectoryCore(
|
||||
d_model=256,
|
||||
n_trajectory=8,
|
||||
n_reasoning_rounds=12,
|
||||
)
|
||||
self.assertEqual(
|
||||
sum(p.numel() for p in core_one.parameters()),
|
||||
sum(p.numel() for p in core_twelve.parameters()),
|
||||
)
|
||||
self.assertAlmostEqual(core_one.attn_scale.item(), 1.0)
|
||||
self.assertAlmostEqual(
|
||||
core_twelve.attn_scale.item(),
|
||||
1.0 / math.sqrt(12),
|
||||
places=6,
|
||||
)
|
||||
self.assertAlmostEqual(
|
||||
core_twelve.mixer_scale.item(),
|
||||
1.0 / math.sqrt(12),
|
||||
places=6,
|
||||
)
|
||||
|
||||
def test_all_masked_attention_is_finite_and_zero(self) -> None:
|
||||
attention = TrajectoryCrossAttention(
|
||||
d_model=32,
|
||||
n_trajectory=4,
|
||||
)
|
||||
memory = torch.randn(2, 3, 32)
|
||||
key_value = attention.project_event_memory(memory)
|
||||
state = torch.randn(2, 2, 4, 8)
|
||||
invalid_mask = torch.ones(2, 2, 3, dtype=torch.bool)
|
||||
output = attention(
|
||||
trajectory_state=state,
|
||||
event_key_value=key_value,
|
||||
event_invalid_mask=invalid_mask,
|
||||
)
|
||||
self.assertTrue(torch.isfinite(output).all())
|
||||
torch.testing.assert_close(output, torch.zeros_like(output))
|
||||
|
||||
def test_next_token_future_events_do_not_change_earlier_query(self) -> None:
|
||||
torch.manual_seed(0)
|
||||
model = build_test_model(n_reasoning_rounds=2)
|
||||
model.eval()
|
||||
inputs = model_inputs()
|
||||
original = model(**inputs)
|
||||
changed_inputs = dict(inputs)
|
||||
changed_inputs["event_seq"] = inputs["event_seq"].clone()
|
||||
changed_inputs["event_seq"][0, 3] = 9
|
||||
changed = model(**changed_inputs)
|
||||
torch.testing.assert_close(original[0, 1], changed[0, 1])
|
||||
|
||||
def test_next_token_later_equal_time_event_is_not_visible(self) -> None:
|
||||
torch.manual_seed(0)
|
||||
model = build_test_model(n_reasoning_rounds=2)
|
||||
model.eval()
|
||||
inputs = model_inputs()
|
||||
inputs["time_seq"] = inputs["time_seq"].clone()
|
||||
inputs["time_seq"][0] = torch.tensor([1.0, 1.0, 2.0, 3.0])
|
||||
original = model(**inputs)
|
||||
changed_inputs = dict(inputs)
|
||||
changed_inputs["event_seq"] = inputs["event_seq"].clone()
|
||||
changed_inputs["event_seq"][0, 1] = 9
|
||||
changed = model(**changed_inputs)
|
||||
torch.testing.assert_close(original[0, 0], changed[0, 0])
|
||||
|
||||
def test_padding_content_does_not_change_valid_queries(self) -> None:
|
||||
torch.manual_seed(0)
|
||||
model = build_test_model(
|
||||
time_mode="relative",
|
||||
n_reasoning_rounds=2,
|
||||
)
|
||||
model.eval()
|
||||
inputs = model_inputs()
|
||||
original = model(**inputs)
|
||||
changed_inputs = dict(inputs)
|
||||
changed_inputs["event_seq"] = inputs["event_seq"].clone()
|
||||
changed_inputs["time_seq"] = inputs["time_seq"].clone()
|
||||
changed_inputs["event_seq"][1, 2:] = torch.tensor([9, 10])
|
||||
changed_inputs["time_seq"][1, 2:] = torch.tensor([30.0, 40.0])
|
||||
changed = model(**changed_inputs)
|
||||
torch.testing.assert_close(original[1, :2], changed[1, :2])
|
||||
|
||||
def test_next_token_and_all_future_output_contracts(self) -> None:
|
||||
inputs = model_inputs()
|
||||
next_model = build_test_model(target_mode="next_token")
|
||||
next_hidden = next_model(**inputs)
|
||||
self.assertEqual(tuple(next_hidden.shape), (2, 4, 120))
|
||||
next_output = next_model(**inputs, return_output=True)
|
||||
self.assertEqual(tuple(next_output.hidden.shape), (2, 4, 120))
|
||||
self.assertEqual(tuple(next_output.padding_mask.shape), (2, 4))
|
||||
|
||||
future_model = build_test_model(target_mode="all_future")
|
||||
future_hidden = future_model(
|
||||
**inputs,
|
||||
t_query=torch.tensor([5.0, 3.0]),
|
||||
)
|
||||
self.assertEqual(tuple(future_hidden.shape), (2, 120))
|
||||
|
||||
def test_model_contains_one_shared_core_and_no_block_stack(self) -> None:
|
||||
model = build_test_model(n_reasoning_rounds=12)
|
||||
self.assertFalse(hasattr(model, "blocks"))
|
||||
reasoning_keys = [
|
||||
key
|
||||
for key in model.state_dict()
|
||||
if key.startswith("reasoning_core.")
|
||||
]
|
||||
self.assertTrue(reasoning_keys)
|
||||
self.assertFalse(any("blocks." in key for key in model.state_dict()))
|
||||
self.assertFalse(any("out_proj" in key for key in reasoning_keys))
|
||||
self.assertFalse(any("group_align" in key for key in reasoning_keys))
|
||||
|
||||
def test_event_key_and_value_are_projected_once_per_forward(self) -> None:
|
||||
model = build_test_model(n_reasoning_rounds=12)
|
||||
call_counts = {"key": 0, "value": 0}
|
||||
|
||||
def count_key(*_args) -> None:
|
||||
call_counts["key"] += 1
|
||||
|
||||
def count_value(*_args) -> None:
|
||||
call_counts["value"] += 1
|
||||
|
||||
key_handle = (
|
||||
model.reasoning_core.cross_attention.k_proj
|
||||
.register_forward_hook(count_key)
|
||||
)
|
||||
value_handle = (
|
||||
model.reasoning_core.cross_attention.v_proj
|
||||
.register_forward_hook(count_value)
|
||||
)
|
||||
try:
|
||||
model(**model_inputs())
|
||||
finally:
|
||||
key_handle.remove()
|
||||
value_handle.remove()
|
||||
self.assertEqual(call_counts, {"key": 1, "value": 1})
|
||||
|
||||
def test_relative_time_forward_and_backward_are_finite(self) -> None:
|
||||
torch.manual_seed(0)
|
||||
model = build_test_model(
|
||||
target_mode="all_future",
|
||||
time_mode="relative",
|
||||
n_reasoning_rounds=2,
|
||||
)
|
||||
hidden = model(
|
||||
**model_inputs(),
|
||||
t_query=torch.tensor([5.0, 3.0]),
|
||||
)
|
||||
(hidden * torch.randn_like(hidden)).sum().backward()
|
||||
self.assertTrue(torch.isfinite(hidden).all())
|
||||
self.assertIsNotNone(model.event_projection.weight.grad)
|
||||
self.assertTrue(torch.isfinite(model.event_projection.weight.grad).all())
|
||||
time_scale = model.reasoning_core.cross_attention.time_bias_scale
|
||||
self.assertIsNotNone(time_scale)
|
||||
self.assertIsNotNone(time_scale.grad)
|
||||
self.assertGreater(abs(float(time_scale.grad)), 0.0)
|
||||
|
||||
def test_architecture_marker_and_checkpoint_are_required(self) -> None:
|
||||
validate_event_trajectory_config(
|
||||
{
|
||||
"model_architecture": EVENT_TRAJECTORY_ARCHITECTURE,
|
||||
"model_size": "nano",
|
||||
"d_model": 120,
|
||||
"n_trajectory": 6,
|
||||
"trajectory_dim": 20,
|
||||
"traj_hidden": 24,
|
||||
"n_reasoning_rounds": 3,
|
||||
}
|
||||
)
|
||||
with self.assertRaisesRegex(ValueError, "only accepts models trained"):
|
||||
validate_event_trajectory_config(
|
||||
{"model_architecture": "traj_mixer_v2"}
|
||||
)
|
||||
with self.assertRaisesRegex(ValueError, "only accepts models trained"):
|
||||
validate_event_trajectory_config(
|
||||
{"model_architecture": "event_trajectory_shared_v1"}
|
||||
)
|
||||
with self.assertRaisesRegex(ValueError, "trajectory_dim"):
|
||||
validate_event_trajectory_config(
|
||||
{
|
||||
"model_architecture": EVENT_TRAJECTORY_ARCHITECTURE,
|
||||
"model_size": "nano",
|
||||
"d_model": 120,
|
||||
"n_trajectory": 6,
|
||||
"trajectory_dim": 10,
|
||||
"traj_hidden": 24,
|
||||
"n_reasoning_rounds": 3,
|
||||
}
|
||||
)
|
||||
|
||||
model = build_test_model()
|
||||
state_dict = model.state_dict()
|
||||
validate_event_trajectory_state_dict(
|
||||
state_dict,
|
||||
expected_d_model=120,
|
||||
expected_n_trajectory=6,
|
||||
expected_n_reasoning_rounds=3,
|
||||
)
|
||||
with self.assertRaisesRegex(
|
||||
ValueError,
|
||||
"Checkpoint architecture does not match",
|
||||
):
|
||||
validate_event_trajectory_state_dict(
|
||||
state_dict,
|
||||
expected_n_reasoning_rounds=12,
|
||||
)
|
||||
state_dict.pop("reasoning_core.attn_scale")
|
||||
with self.assertRaisesRegex(
|
||||
ValueError,
|
||||
"not a shared event-trajectory checkpoint",
|
||||
):
|
||||
validate_event_trajectory_state_dict(state_dict)
|
||||
|
||||
def test_unknown_model_size_is_rejected(self) -> None:
|
||||
with self.assertRaisesRegex(ValueError, "Unknown model_size"):
|
||||
DeepHealth(
|
||||
vocab_size=32,
|
||||
model_size="giant",
|
||||
n_reasoning_rounds=2,
|
||||
n_types=2,
|
||||
n_cont_types=0,
|
||||
n_categories=2,
|
||||
cont_type_ids=[],
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
173
tests/test_all_future_likelihood.py
Normal file
173
tests/test_all_future_likelihood.py
Normal file
@@ -0,0 +1,173 @@
|
||||
import unittest
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch.utils.data import Subset
|
||||
|
||||
from dataset import AllFutureHealthDataset
|
||||
from losses import ExponentialLoss, WeibullLoss
|
||||
from models import DeepHealth
|
||||
from train_util import fit_all_future_baseline
|
||||
|
||||
|
||||
def _inverse_softplus(value: torch.Tensor) -> torch.Tensor:
|
||||
return torch.log(torch.expm1(value))
|
||||
|
||||
|
||||
class AllFutureLikelihoodTests(unittest.TestCase):
|
||||
def test_exponential_uses_event_specific_first_onset_exposure(self):
|
||||
desired_rate = torch.tensor(
|
||||
[[0.01, 0.01, 0.01, 0.20, 0.30, 0.40]],
|
||||
dtype=torch.float64,
|
||||
)
|
||||
logits = _inverse_softplus(desired_rate)
|
||||
criterion = ExponentialLoss(ignored_idx={0, 1, 2}, eps=1e-12)
|
||||
|
||||
loss = criterion(
|
||||
logits=logits,
|
||||
targets=torch.tensor([[4, 0]]),
|
||||
exposure=torch.tensor([5.0], dtype=torch.float64),
|
||||
dt=torch.tensor([[2.0, 0.0]], dtype=torch.float64),
|
||||
history=torch.tensor([[3, 0]]),
|
||||
)
|
||||
|
||||
rate = F.softplus(logits) + criterion.eps
|
||||
expected = -(rate[0, 4].log()) + rate[0, 4] * 2.0 + rate[0, 5] * 5.0
|
||||
torch.testing.assert_close(loss, expected)
|
||||
|
||||
def test_weibull_uses_event_time_and_excludes_prevalent_outcome(self):
|
||||
desired_rate = torch.tensor(
|
||||
[[0.01, 0.01, 0.01, 0.20, 0.30, 0.40]],
|
||||
dtype=torch.float64,
|
||||
)
|
||||
logits = _inverse_softplus(desired_rate)
|
||||
rho = torch.tensor(
|
||||
[[1.0, 1.0, 1.0, 1.2, 1.5, 0.8]],
|
||||
dtype=torch.float64,
|
||||
)
|
||||
criterion = WeibullLoss(ignored_idx={0, 1, 2}, eps=1e-12)
|
||||
|
||||
loss = criterion(
|
||||
logits=logits,
|
||||
weibull_rho=rho,
|
||||
targets=torch.tensor([[4, 0]]),
|
||||
dt=torch.tensor([[2.0, 0.0]], dtype=torch.float64),
|
||||
exposure=torch.tensor([5.0], dtype=torch.float64),
|
||||
history=torch.tensor([[3, 0]]),
|
||||
)
|
||||
|
||||
rate = F.softplus(logits) + criterion.eps
|
||||
log_hazard = (
|
||||
rate[0, 4].log()
|
||||
+ rho[0, 4].log()
|
||||
+ (rho[0, 4] - 1.0) * torch.tensor(2.0).log()
|
||||
)
|
||||
expected = (
|
||||
-log_hazard
|
||||
+ rate[0, 4] * torch.pow(torch.tensor(2.0), rho[0, 4])
|
||||
+ rate[0, 5] * torch.pow(torch.tensor(5.0), rho[0, 5])
|
||||
)
|
||||
torch.testing.assert_close(loss, expected)
|
||||
|
||||
|
||||
class AllFutureQueryDistributionTests(unittest.TestCase):
|
||||
@staticmethod
|
||||
def _dataset() -> AllFutureHealthDataset:
|
||||
dataset = AllFutureHealthDataset.__new__(AllFutureHealthDataset)
|
||||
dataset.min_history_events = 1
|
||||
dataset.min_future_events = 1
|
||||
return dataset
|
||||
|
||||
@staticmethod
|
||||
def _patient():
|
||||
return {
|
||||
"times": np.asarray([1.0, 2.0, 3.0, 4.0], dtype=np.float32),
|
||||
"labels": np.asarray([3, 4, 5, 6], dtype=np.int64),
|
||||
"t_obs": 4.0,
|
||||
}
|
||||
|
||||
def test_train_validation_and_test_share_one_query_sampler(self):
|
||||
dataset = self._dataset()
|
||||
patient = self._patient()
|
||||
patient["query_intervals"] = dataset._eligible_query_intervals(patient)
|
||||
self.assertEqual(len(patient["query_intervals"]), 3)
|
||||
|
||||
rng = np.random.RandomState(123)
|
||||
fixed = dataset._sample_fixed_validation_queries(patient, rng)
|
||||
self.assertEqual(len(fixed), 1)
|
||||
self.assertTrue(dataset._is_valid_query(patient, fixed[0]))
|
||||
|
||||
seen_intervals = set()
|
||||
rng = np.random.RandomState(456)
|
||||
for _ in range(200):
|
||||
query = dataset.sample_query(patient, rng)
|
||||
self.assertTrue(dataset._is_valid_query(patient, query))
|
||||
seen_intervals.add(int(np.floor(query)))
|
||||
self.assertEqual(seen_intervals, {1, 2, 3})
|
||||
|
||||
|
||||
class AllFutureBaselineTests(unittest.TestCase):
|
||||
def test_training_baseline_matches_first_onset_exposure(self):
|
||||
dataset = AllFutureHealthDataset.__new__(AllFutureHealthDataset)
|
||||
dataset.vocab_size = 7
|
||||
dataset.min_history_events = 1
|
||||
dataset.min_future_events = 2
|
||||
patient = {
|
||||
"times": np.asarray([1.0, 3.0, 5.0], dtype=np.float32),
|
||||
"labels": np.asarray([3, 4, 5], dtype=np.int64),
|
||||
"t_obs": 5.0,
|
||||
"query_intervals": [(1.0, 3.0)],
|
||||
}
|
||||
dataset.patients = [patient]
|
||||
subset = Subset(dataset, [0])
|
||||
|
||||
stats = fit_all_future_baseline(dataset, subset, seed=17)
|
||||
query = dataset.sample_query(patient, np.random.RandomState(17))
|
||||
|
||||
self.assertEqual(stats.event_count[3], 0)
|
||||
self.assertEqual(stats.at_risk_exposure[3], 0.0)
|
||||
self.assertEqual(stats.event_count[4], 1)
|
||||
self.assertEqual(stats.event_count[5], 1)
|
||||
self.assertAlmostEqual(stats.at_risk_exposure[4], 3.0 - query, places=5)
|
||||
self.assertAlmostEqual(stats.at_risk_exposure[5], 5.0 - query, places=5)
|
||||
self.assertAlmostEqual(
|
||||
float(F.softplus(torch.tensor(stats.bias[4]))),
|
||||
float(stats.rate[4]),
|
||||
places=6,
|
||||
)
|
||||
|
||||
def test_model_starts_exactly_at_fitted_output_baseline(self):
|
||||
baseline_rate = torch.tensor([0.0, 0.0, 0.0, 0.02, 0.05, 0.10])
|
||||
baseline_bias = torch.zeros_like(baseline_rate)
|
||||
baseline_bias[3:] = _inverse_softplus(baseline_rate[3:])
|
||||
model = DeepHealth(
|
||||
vocab_size=6,
|
||||
n_embd=4,
|
||||
n_head=1,
|
||||
n_layer=1,
|
||||
n_types=1,
|
||||
n_cont_types=0,
|
||||
n_categories=1,
|
||||
cont_type_ids=[],
|
||||
target_mode="all_future",
|
||||
time_mode="absolute",
|
||||
dist_mode="weibull",
|
||||
model_architecture="transformer_ffn_v1",
|
||||
risk_head_bias=True,
|
||||
risk_head_bias_init=baseline_bias,
|
||||
)
|
||||
|
||||
torch.testing.assert_close(model.risk_head.weight, torch.zeros_like(model.risk_head.weight))
|
||||
output_rate = F.softplus(model.risk_head(torch.randn(3, 4)))
|
||||
torch.testing.assert_close(output_rate[:, 3:], baseline_rate[None, 3:].expand(3, -1))
|
||||
torch.testing.assert_close(
|
||||
F.softplus(model.rho_head.bias),
|
||||
torch.ones_like(model.rho_head.bias),
|
||||
atol=2e-5,
|
||||
rtol=0.0,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
162
tests/test_auc_eid_split.py
Normal file
162
tests/test_auc_eid_split.py
Normal file
@@ -0,0 +1,162 @@
|
||||
import argparse
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
import numpy as np
|
||||
|
||||
from eval_data import split_indices
|
||||
from evaluate_auc import make_eval_subset
|
||||
from evaluate_auc_v2 import make_eval_indices
|
||||
|
||||
|
||||
class _DummyDataset:
|
||||
def __init__(self, eids: list[int]) -> None:
|
||||
self.samples = [{"eid": eid} for eid in eids]
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.samples)
|
||||
|
||||
def __getitem__(self, index: int) -> dict[str, int]:
|
||||
return self.samples[index]
|
||||
|
||||
|
||||
class AUCEidSplitTests(unittest.TestCase):
|
||||
def test_both_auc_evaluators_default_to_ukb_test_eid_file(self) -> None:
|
||||
dataset = _DummyDataset([101, 102, 103, 104, 105])
|
||||
args = argparse.Namespace(
|
||||
eval_split="test",
|
||||
dataset_subset_size=None,
|
||||
test_eid_file=None,
|
||||
)
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
root = Path(tmp_dir)
|
||||
(root / "ukb_test_eid.csv").write_text(
|
||||
"eid\n104\n102\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
with patch.object(Path, "cwd", return_value=root):
|
||||
subset, legacy_indices = make_eval_subset(dataset, args, {})
|
||||
landmark_indices = make_eval_indices(dataset, args, {})
|
||||
|
||||
expected = np.asarray([1, 3], dtype=np.int64)
|
||||
np.testing.assert_array_equal(legacy_indices, expected)
|
||||
np.testing.assert_array_equal(landmark_indices, expected)
|
||||
self.assertEqual(subset.indices, expected.tolist())
|
||||
|
||||
def test_subset_size_is_applied_after_eid_selection(self) -> None:
|
||||
dataset = _DummyDataset([201, 202, 203, 204])
|
||||
args = argparse.Namespace(
|
||||
eval_split="test",
|
||||
dataset_subset_size=1,
|
||||
test_eid_file=None,
|
||||
)
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
eid_path = Path(tmp_dir) / "test.csv"
|
||||
eid_path.write_text("eid\n202\n204\n", encoding="utf-8")
|
||||
cfg = {"test_eid_file": str(eid_path)}
|
||||
_, legacy_indices = make_eval_subset(dataset, args, cfg)
|
||||
landmark_indices = make_eval_indices(dataset, args, cfg)
|
||||
|
||||
expected = np.asarray([1], dtype=np.int64)
|
||||
np.testing.assert_array_equal(legacy_indices, expected)
|
||||
np.testing.assert_array_equal(landmark_indices, expected)
|
||||
|
||||
def test_cli_test_eid_file_overrides_config(self) -> None:
|
||||
dataset = _DummyDataset([301, 302, 303])
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
root = Path(tmp_dir)
|
||||
config_path = root / "config.csv"
|
||||
cli_path = root / "cli.csv"
|
||||
config_path.write_text("eid\n301\n", encoding="utf-8")
|
||||
cli_path.write_text("eid\n303\n", encoding="utf-8")
|
||||
args = argparse.Namespace(
|
||||
eval_split="test",
|
||||
dataset_subset_size=None,
|
||||
test_eid_file=str(cli_path),
|
||||
)
|
||||
cfg = {"test_eid_file": str(config_path)}
|
||||
|
||||
_, legacy_indices = make_eval_subset(dataset, args, cfg)
|
||||
landmark_indices = make_eval_indices(dataset, args, cfg)
|
||||
|
||||
expected = np.asarray([2], dtype=np.int64)
|
||||
np.testing.assert_array_equal(legacy_indices, expected)
|
||||
np.testing.assert_array_equal(landmark_indices, expected)
|
||||
|
||||
def test_empty_test_eid_file_explicitly_uses_ratio_split(self) -> None:
|
||||
dataset = _DummyDataset(list(range(20)))
|
||||
args = argparse.Namespace(
|
||||
eval_split="test",
|
||||
dataset_subset_size=None,
|
||||
test_eid_file="",
|
||||
)
|
||||
cfg = {
|
||||
"train_ratio": 0.7,
|
||||
"val_ratio": 0.15,
|
||||
"test_ratio": 0.15,
|
||||
"seed": 7,
|
||||
}
|
||||
|
||||
expected = split_indices(20, 0.7, 0.15, 0.15, 7)[2]
|
||||
_, legacy_indices = make_eval_subset(dataset, args, cfg)
|
||||
landmark_indices = make_eval_indices(dataset, args, cfg)
|
||||
|
||||
np.testing.assert_array_equal(legacy_indices, expected)
|
||||
np.testing.assert_array_equal(landmark_indices, expected)
|
||||
|
||||
def test_non_test_split_does_not_read_test_eid_file(self) -> None:
|
||||
dataset = _DummyDataset(list(range(20)))
|
||||
args = argparse.Namespace(
|
||||
eval_split="val",
|
||||
dataset_subset_size=None,
|
||||
test_eid_file="missing.csv",
|
||||
)
|
||||
cfg = {
|
||||
"train_ratio": 0.7,
|
||||
"val_ratio": 0.15,
|
||||
"test_ratio": 0.15,
|
||||
"seed": 11,
|
||||
}
|
||||
|
||||
expected = split_indices(20, 0.7, 0.15, 0.15, 11)[1]
|
||||
_, legacy_indices = make_eval_subset(dataset, args, cfg)
|
||||
landmark_indices = make_eval_indices(dataset, args, cfg)
|
||||
|
||||
np.testing.assert_array_equal(legacy_indices, expected)
|
||||
np.testing.assert_array_equal(landmark_indices, expected)
|
||||
|
||||
def test_missing_or_nonmatching_eid_file_fails_closed(self) -> None:
|
||||
dataset = _DummyDataset([401, 402])
|
||||
args = argparse.Namespace(
|
||||
eval_split="test",
|
||||
dataset_subset_size=None,
|
||||
test_eid_file=None,
|
||||
)
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
root = Path(tmp_dir)
|
||||
missing = root / "missing.csv"
|
||||
with self.assertRaisesRegex(FileNotFoundError, "EID split file"):
|
||||
make_eval_indices(
|
||||
dataset,
|
||||
args,
|
||||
{"test_eid_file": str(missing)},
|
||||
)
|
||||
|
||||
nonmatching = root / "nonmatching.csv"
|
||||
nonmatching.write_text("eid\n999\n", encoding="utf-8")
|
||||
with self.assertRaisesRegex(ValueError, "No dataset patients"):
|
||||
make_eval_indices(
|
||||
dataset,
|
||||
args,
|
||||
{"test_eid_file": str(nonmatching)},
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
457
tests/test_calibration_metrics.py
Normal file
457
tests/test_calibration_metrics.py
Normal file
@@ -0,0 +1,457 @@
|
||||
import math
|
||||
import unittest
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import torch
|
||||
|
||||
from evaluate_calibration import (
|
||||
_censoring_km,
|
||||
_evaluate_calibration_token,
|
||||
_risk_probability_matrix,
|
||||
aggregate_metric_rows,
|
||||
compute_ipcw_cell,
|
||||
compute_ipcw_horizons,
|
||||
evaluate_landmark_calibration,
|
||||
fit_weighted_logistic_calibration,
|
||||
fit_weighted_logistic_calibration_batch,
|
||||
)
|
||||
from evaluate_auc_v2 import _score_to_probability
|
||||
|
||||
|
||||
class IPCWCalibrationMetricTests(unittest.TestCase):
|
||||
@staticmethod
|
||||
def _naive_censoring_km(observed_times, censor_events):
|
||||
observed_times = np.asarray(observed_times, dtype=np.float64)
|
||||
censor_events = np.asarray(censor_events, dtype=bool)
|
||||
event_times = np.unique(observed_times[censor_events])
|
||||
survival = 1.0
|
||||
survival_after = []
|
||||
for time_value in event_times:
|
||||
at_risk = np.sum(observed_times >= time_value)
|
||||
censored = np.sum(
|
||||
censor_events & (observed_times == time_value)
|
||||
)
|
||||
survival *= 1.0 - float(censored) / float(at_risk)
|
||||
survival_after.append(survival)
|
||||
return event_times, np.asarray(survival_after)
|
||||
|
||||
def test_sorted_censoring_km_matches_naive_reference(self):
|
||||
rng = np.random.RandomState(12)
|
||||
observed_times = rng.randint(1, 20, size=500).astype(np.float64)
|
||||
censor_events = rng.uniform(size=500) < 0.4
|
||||
|
||||
expected_times, expected_survival = self._naive_censoring_km(
|
||||
observed_times,
|
||||
censor_events,
|
||||
)
|
||||
actual_times, actual_survival = _censoring_km(
|
||||
observed_times,
|
||||
censor_events,
|
||||
)
|
||||
|
||||
np.testing.assert_array_equal(actual_times, expected_times)
|
||||
np.testing.assert_allclose(
|
||||
actual_survival,
|
||||
expected_survival,
|
||||
rtol=1e-14,
|
||||
atol=1e-14,
|
||||
)
|
||||
|
||||
def test_no_censoring_matches_binary_metrics(self):
|
||||
result = compute_ipcw_cell(
|
||||
probabilities=np.asarray([0.2, 0.8]),
|
||||
event_times=np.asarray([np.inf, 0.5]),
|
||||
censor_times=np.asarray([2.0, 2.0]),
|
||||
horizon=1.0,
|
||||
min_cases=1,
|
||||
min_controls=1,
|
||||
max_ipcw_weight=0.0,
|
||||
)
|
||||
|
||||
self.assertIsNotNone(result)
|
||||
row, arrays = result
|
||||
self.assertEqual(row["n_events"], 1)
|
||||
self.assertEqual(row["n_controls"], 1)
|
||||
self.assertAlmostEqual(row["brier_ipcw"], 0.04)
|
||||
self.assertAlmostEqual(row["nll_ipcw"], -math.log(0.8))
|
||||
self.assertAlmostEqual(row["predicted_mean"], 0.5)
|
||||
self.assertAlmostEqual(row["observed_rate_ipcw"], 0.5)
|
||||
np.testing.assert_allclose(arrays["metric_weights"], [1.0, 1.0])
|
||||
|
||||
def test_censored_before_horizon_gets_zero_outcome_weight(self):
|
||||
result = compute_ipcw_cell(
|
||||
probabilities=np.asarray([0.8, 0.2, 0.4]),
|
||||
event_times=np.asarray([0.5, np.inf, np.inf]),
|
||||
censor_times=np.asarray([2.0, 2.0, 0.5]),
|
||||
horizon=1.0,
|
||||
min_cases=1,
|
||||
min_controls=1,
|
||||
max_ipcw_weight=0.0,
|
||||
)
|
||||
|
||||
self.assertIsNotNone(result)
|
||||
row, arrays = result
|
||||
self.assertEqual(row["n_censored_before_horizon"], 1)
|
||||
self.assertAlmostEqual(row["known_fraction"], 2.0 / 3.0)
|
||||
np.testing.assert_allclose(
|
||||
arrays["metric_weights"],
|
||||
[1.0, 1.5, 0.0],
|
||||
)
|
||||
self.assertAlmostEqual(row["brier_ipcw"], 0.1 / 3.0)
|
||||
self.assertAlmostEqual(
|
||||
row["nll_ipcw"],
|
||||
-2.5 * math.log(0.8) / 3.0,
|
||||
)
|
||||
self.assertAlmostEqual(row["observed_rate_ipcw"], 1.0 / 3.0)
|
||||
|
||||
def test_calibration_intercept_and_slope_recover_identity(self):
|
||||
probabilities = np.repeat([0.1, 0.3, 0.7, 0.9], 100)
|
||||
outcomes = np.concatenate(
|
||||
[
|
||||
np.r_[np.ones(10), np.zeros(90)],
|
||||
np.r_[np.ones(30), np.zeros(70)],
|
||||
np.r_[np.ones(70), np.zeros(30)],
|
||||
np.r_[np.ones(90), np.zeros(10)],
|
||||
]
|
||||
)
|
||||
weights = np.ones_like(probabilities)
|
||||
|
||||
calibration_in_large, intercept, slope = (
|
||||
fit_weighted_logistic_calibration(
|
||||
probabilities,
|
||||
outcomes,
|
||||
weights,
|
||||
)
|
||||
)
|
||||
|
||||
self.assertAlmostEqual(calibration_in_large, 0.0, places=7)
|
||||
self.assertAlmostEqual(intercept, 0.0, places=7)
|
||||
self.assertAlmostEqual(slope, 1.0, places=7)
|
||||
|
||||
def test_batched_calibration_fits_multiple_horizons(self):
|
||||
probabilities = np.vstack(
|
||||
[
|
||||
np.repeat([0.1, 0.3, 0.7, 0.9], 100),
|
||||
np.repeat([0.2, 0.4, 0.6, 0.8], 100),
|
||||
]
|
||||
)
|
||||
outcomes = np.vstack(
|
||||
[
|
||||
np.concatenate(
|
||||
[
|
||||
np.r_[np.ones(10), np.zeros(90)],
|
||||
np.r_[np.ones(30), np.zeros(70)],
|
||||
np.r_[np.ones(70), np.zeros(30)],
|
||||
np.r_[np.ones(90), np.zeros(10)],
|
||||
]
|
||||
),
|
||||
np.concatenate(
|
||||
[
|
||||
np.r_[np.ones(20), np.zeros(80)],
|
||||
np.r_[np.ones(40), np.zeros(60)],
|
||||
np.r_[np.ones(60), np.zeros(40)],
|
||||
np.r_[np.ones(80), np.zeros(20)],
|
||||
]
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
calibration_in_large, intercept, slope = (
|
||||
fit_weighted_logistic_calibration_batch(
|
||||
probabilities,
|
||||
outcomes,
|
||||
np.ones_like(probabilities),
|
||||
)
|
||||
)
|
||||
|
||||
np.testing.assert_allclose(calibration_in_large, 0.0, atol=1e-7)
|
||||
np.testing.assert_allclose(intercept, 0.0, atol=1e-7)
|
||||
np.testing.assert_allclose(slope, 1.0, atol=1e-7)
|
||||
|
||||
def test_all_horizons_reuse_one_censoring_km(self):
|
||||
probabilities = np.asarray(
|
||||
[
|
||||
[0.05, 0.10, 0.15, 0.20, 0.25],
|
||||
[0.10, 0.20, 0.30, 0.40, 0.50],
|
||||
[0.20, 0.35, 0.50, 0.65, 0.80],
|
||||
]
|
||||
)
|
||||
event_times = np.asarray([0.5, 1.5, 4.0, np.inf, np.inf])
|
||||
censor_times = np.asarray([5.0, 5.0, 5.0, 2.5, 5.0])
|
||||
|
||||
with patch(
|
||||
"evaluate_calibration._censoring_km",
|
||||
wraps=_censoring_km,
|
||||
) as km:
|
||||
results = compute_ipcw_horizons(
|
||||
probabilities=probabilities,
|
||||
event_times=event_times,
|
||||
censor_times=censor_times,
|
||||
horizons=np.asarray([1.0, 2.0, 5.0]),
|
||||
min_cases=1,
|
||||
min_controls=1,
|
||||
max_ipcw_weight=0.0,
|
||||
)
|
||||
|
||||
self.assertEqual(km.call_count, 1)
|
||||
self.assertEqual(len(results), 3)
|
||||
self.assertTrue(all(result is not None for result in results))
|
||||
self.assertEqual([result[0]["n_events"] for result in results], [1, 2, 3])
|
||||
|
||||
def test_batched_risk_probabilities_match_scalar_reference(self):
|
||||
logits = np.asarray([-2.0, -0.5, 0.2, 1.5], dtype=np.float32)
|
||||
rho = np.asarray([0.8, 1.0, 1.2, 1.5], dtype=np.float32)
|
||||
horizons = np.asarray([0.1, 1.0, 5.0], dtype=np.float32)
|
||||
|
||||
for dist_mode, selected_rho in (
|
||||
("exponential", None),
|
||||
("weibull", rho),
|
||||
):
|
||||
actual = _risk_probability_matrix(
|
||||
logits=logits,
|
||||
rho=selected_rho,
|
||||
horizons=horizons,
|
||||
dist_mode=dist_mode,
|
||||
)
|
||||
expected = np.vstack(
|
||||
[
|
||||
_score_to_probability(
|
||||
logits,
|
||||
selected_rho,
|
||||
score_mode="risk",
|
||||
horizon=float(horizon),
|
||||
dist_mode=dist_mode,
|
||||
)
|
||||
for horizon in horizons
|
||||
]
|
||||
)
|
||||
np.testing.assert_allclose(actual, expected, rtol=1e-6, atol=1e-7)
|
||||
|
||||
def test_per_disease_worker_is_thread_safe(self):
|
||||
logits_chunk = np.asarray(
|
||||
[
|
||||
[-2.0, -1.5],
|
||||
[-1.0, -0.5],
|
||||
[0.0, 0.5],
|
||||
[0.5, 1.0],
|
||||
[1.0, 1.5],
|
||||
[1.5, 2.0],
|
||||
],
|
||||
dtype=np.float32,
|
||||
)
|
||||
common = {
|
||||
"logits_chunk": logits_chunk,
|
||||
"rho_chunk": None,
|
||||
"strata": [("Female", 50.0, np.arange(6, dtype=np.int64))],
|
||||
"row_patient_id": np.arange(6, dtype=np.int32),
|
||||
"row_followup_end": np.full(6, 65.0, dtype=np.float32),
|
||||
"row_death_time": np.full(6, np.inf, dtype=np.float32),
|
||||
"first_occurrence_by_token": {
|
||||
4: (
|
||||
np.asarray([0, 1], dtype=np.int32),
|
||||
np.asarray([50.5, 52.0], dtype=np.float32),
|
||||
),
|
||||
5: (
|
||||
np.asarray([2, 3], dtype=np.int32),
|
||||
np.asarray([50.7, 53.0], dtype=np.float32),
|
||||
),
|
||||
},
|
||||
"patient_count": 6,
|
||||
"death_tokens": set(),
|
||||
"label_id_to_code": {4: "D4", 5: "D5"},
|
||||
"dist_mode": "exponential",
|
||||
"horizons": np.asarray([1.0, 5.0], dtype=np.float32),
|
||||
"min_cases": 1,
|
||||
"min_controls": 1,
|
||||
"max_ipcw_weight": 0.0,
|
||||
"exclude_death_competing": True,
|
||||
"probability_bins": np.asarray([0.0, 0.5, 1.0]),
|
||||
}
|
||||
tasks = [
|
||||
{"column_index": 0, "token": 4, **common},
|
||||
{"column_index": 1, "token": 5, **common},
|
||||
]
|
||||
serial = [
|
||||
_evaluate_calibration_token(**task)
|
||||
for task in tasks
|
||||
]
|
||||
with ThreadPoolExecutor(max_workers=2) as executor:
|
||||
parallel = list(
|
||||
executor.map(
|
||||
lambda task: _evaluate_calibration_token(**task),
|
||||
tasks,
|
||||
)
|
||||
)
|
||||
|
||||
for (serial_rows, serial_curve), (
|
||||
parallel_rows,
|
||||
parallel_curve,
|
||||
) in zip(serial, parallel):
|
||||
pd.testing.assert_frame_equal(
|
||||
pd.DataFrame(serial_rows),
|
||||
pd.DataFrame(parallel_rows),
|
||||
)
|
||||
self.assertEqual(set(serial_curve), set(parallel_curve))
|
||||
for key in serial_curve:
|
||||
self.assertEqual(
|
||||
set(serial_curve[key]),
|
||||
set(parallel_curve[key]),
|
||||
)
|
||||
np.testing.assert_allclose(
|
||||
list(serial_curve[key].values()),
|
||||
list(parallel_curve[key].values()),
|
||||
equal_nan=True,
|
||||
)
|
||||
|
||||
def test_landmark_evaluation_parallel_matches_serial(self):
|
||||
logits_chunk = np.asarray(
|
||||
[
|
||||
[-2.0, -1.5],
|
||||
[-1.0, -0.5],
|
||||
[0.0, 0.5],
|
||||
[0.5, 1.0],
|
||||
[1.0, 1.5],
|
||||
[1.5, 2.0],
|
||||
],
|
||||
dtype=np.float32,
|
||||
)
|
||||
row_arrays = {
|
||||
"patient_id": np.arange(6, dtype=np.int32),
|
||||
"sex": np.zeros(6, dtype=np.int8),
|
||||
"landmark_age": np.full(6, 50.0, dtype=np.float32),
|
||||
"followup_end_time": np.full(6, 65.0, dtype=np.float32),
|
||||
"death_time": np.full(6, np.inf, dtype=np.float32),
|
||||
}
|
||||
landmark_dataset = SimpleNamespace(
|
||||
subset_indices=np.arange(6, dtype=np.int64),
|
||||
death_token_ids=[],
|
||||
first_occurrence_by_token={
|
||||
4: (
|
||||
np.asarray([0, 1], dtype=np.int32),
|
||||
np.asarray([50.5, 52.0], dtype=np.float32),
|
||||
),
|
||||
5: (
|
||||
np.asarray([2, 3], dtype=np.int32),
|
||||
np.asarray([50.7, 53.0], dtype=np.float32),
|
||||
),
|
||||
},
|
||||
dataset=SimpleNamespace(
|
||||
label_id_to_code={4: "D4", 5: "D5"}
|
||||
),
|
||||
)
|
||||
|
||||
class FakeModel:
|
||||
death_idx = 9
|
||||
vocab_size = 10
|
||||
|
||||
def eval(self):
|
||||
return self
|
||||
|
||||
def to(self, _device):
|
||||
return self
|
||||
|
||||
common = {
|
||||
"model": FakeModel(),
|
||||
"loader": [],
|
||||
"landmark_dataset": landmark_dataset,
|
||||
"disease_ids": [4, 5],
|
||||
"dist_mode": "exponential",
|
||||
"horizons": np.asarray([1.0, 5.0], dtype=np.float32),
|
||||
"device": torch.device("cpu"),
|
||||
"use_amp": False,
|
||||
"hidden_cache_dtype": "float16",
|
||||
"logit_batch_size": 8,
|
||||
"disease_chunk_size": 2,
|
||||
"min_cases": 1,
|
||||
"min_controls": 1,
|
||||
"max_ipcw_weight": 0.0,
|
||||
"exclude_death_competing": True,
|
||||
"probability_bins": np.asarray([0.0, 0.5, 1.0]),
|
||||
}
|
||||
with (
|
||||
patch(
|
||||
"evaluate_calibration.infer_landmark_hidden",
|
||||
return_value=(
|
||||
np.zeros((6, 4), dtype=np.float16),
|
||||
row_arrays,
|
||||
),
|
||||
),
|
||||
patch(
|
||||
"evaluate_calibration.project_distribution_chunk",
|
||||
return_value=(logits_chunk, None),
|
||||
),
|
||||
):
|
||||
serial_metrics, serial_curve = evaluate_landmark_calibration(
|
||||
**common,
|
||||
num_workers_calibration=1,
|
||||
)
|
||||
parallel_metrics, parallel_curve = evaluate_landmark_calibration(
|
||||
**common,
|
||||
num_workers_calibration=2,
|
||||
)
|
||||
|
||||
pd.testing.assert_frame_equal(serial_metrics, parallel_metrics)
|
||||
pd.testing.assert_frame_equal(serial_curve, parallel_curve)
|
||||
|
||||
def test_metric_aggregation_uses_contribution_sums(self):
|
||||
metrics = pd.DataFrame(
|
||||
[
|
||||
{
|
||||
"outcome": "Disease",
|
||||
"sex": "Female",
|
||||
"horizon": 5.0,
|
||||
"n_at_risk": 10,
|
||||
"n_events": 2,
|
||||
"n_controls": 7,
|
||||
"n_censored_before_horizon": 1,
|
||||
"prediction_sum": 2.0,
|
||||
"event_weight_sum": 2.0,
|
||||
"brier_ipcw_sum": 1.0,
|
||||
"nll_ipcw_sum": 3.0,
|
||||
"calibration_in_the_large": 0.1,
|
||||
"calibration_intercept": 0.2,
|
||||
"calibration_slope": 0.9,
|
||||
"ipcw_weight_max": 1.2,
|
||||
"ipcw_weights_clipped": 0,
|
||||
},
|
||||
{
|
||||
"outcome": "Disease",
|
||||
"sex": "Female",
|
||||
"horizon": 5.0,
|
||||
"n_at_risk": 10,
|
||||
"n_events": 3,
|
||||
"n_controls": 6,
|
||||
"n_censored_before_horizon": 1,
|
||||
"prediction_sum": 3.0,
|
||||
"event_weight_sum": 3.0,
|
||||
"brier_ipcw_sum": 2.0,
|
||||
"nll_ipcw_sum": 4.0,
|
||||
"calibration_in_the_large": -0.1,
|
||||
"calibration_intercept": -0.2,
|
||||
"calibration_slope": 1.1,
|
||||
"ipcw_weight_max": 1.4,
|
||||
"ipcw_weights_clipped": 1,
|
||||
},
|
||||
]
|
||||
)
|
||||
|
||||
aggregated = aggregate_metric_rows(
|
||||
metrics,
|
||||
group_columns=["outcome", "sex", "horizon"],
|
||||
).iloc[0]
|
||||
|
||||
self.assertEqual(aggregated["n_at_risk"], 20)
|
||||
self.assertAlmostEqual(aggregated["predicted_mean"], 0.25)
|
||||
self.assertAlmostEqual(aggregated["observed_rate_ipcw"], 0.25)
|
||||
self.assertAlmostEqual(aggregated["brier_ipcw"], 0.15)
|
||||
self.assertAlmostEqual(aggregated["nll_ipcw"], 0.35)
|
||||
self.assertAlmostEqual(aggregated["calibration_slope_median"], 1.0)
|
||||
self.assertEqual(aggregated["ipcw_weights_clipped"], 1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
225
tests/test_continuous_value_scaling.py
Normal file
225
tests/test_continuous_value_scaling.py
Normal file
@@ -0,0 +1,225 @@
|
||||
import unittest
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.utils.data import Subset
|
||||
|
||||
from eval_data import build_model_from_dataset
|
||||
from models import DeepHealth, OtherInfoTokenizer
|
||||
from train_util import fit_continuous_robust_scaler
|
||||
|
||||
|
||||
class _ToyAllFutureDataset:
|
||||
def __init__(self):
|
||||
self.cont_type_ids = [1, 3]
|
||||
self.n_types = 4
|
||||
self.patients = [
|
||||
self._patient([1, 3], [0.0, 10.0]),
|
||||
self._patient([1, 3], [1.0, 10.0]),
|
||||
self._patient([1, 3], [2.0, 10.0]),
|
||||
self._patient([1, 3], [3.0, 10.0]),
|
||||
self._patient([1, 3], [4.0, 10.0]),
|
||||
self._patient([1, 3], [1000.0, 999.0]),
|
||||
]
|
||||
|
||||
@staticmethod
|
||||
def _patient(types, values):
|
||||
return {
|
||||
"other_type": np.asarray(types, dtype=np.int64),
|
||||
"other_value": np.asarray(values, dtype=np.float32),
|
||||
"other_value_kind": np.ones(len(types), dtype=np.int64),
|
||||
}
|
||||
|
||||
def __len__(self):
|
||||
return len(self.patients)
|
||||
|
||||
|
||||
class _CaptureContinuousEncoder(nn.Module):
|
||||
def __init__(self, n_embd):
|
||||
super().__init__()
|
||||
self.n_embd = n_embd
|
||||
self.last_type = None
|
||||
self.last_value = None
|
||||
|
||||
def forward(self, cont_type_idx, value):
|
||||
self.last_type = cont_type_idx.detach().clone()
|
||||
self.last_value = value.detach().clone()
|
||||
return value[:, None].expand(-1, self.n_embd)
|
||||
|
||||
|
||||
class ContinuousValueScalingTests(unittest.TestCase):
|
||||
def test_fit_uses_only_training_subset_and_handles_constant_features(self):
|
||||
dataset = _ToyAllFutureDataset()
|
||||
train_subset = Subset(dataset, np.asarray([0, 1, 2, 3, 4]))
|
||||
|
||||
stats = fit_continuous_robust_scaler(dataset, train_subset)
|
||||
|
||||
self.assertEqual(stats.cont_type_ids, (1, 3))
|
||||
np.testing.assert_array_equal(stats.observation_count, np.asarray([5, 5]))
|
||||
np.testing.assert_allclose(stats.center, np.asarray([2.0, 10.0]))
|
||||
np.testing.assert_allclose(stats.scale, np.asarray([2.0, 1.0]))
|
||||
|
||||
def test_fit_supports_next_step_sample_storage(self):
|
||||
dataset = _ToyAllFutureDataset()
|
||||
dataset.samples = dataset.patients
|
||||
del dataset.patients
|
||||
train_subset = Subset(dataset, np.asarray([0, 1, 2, 3, 4]))
|
||||
|
||||
stats = fit_continuous_robust_scaler(dataset, train_subset)
|
||||
|
||||
np.testing.assert_allclose(stats.center, np.asarray([2.0, 10.0]))
|
||||
np.testing.assert_allclose(stats.scale, np.asarray([2.0, 1.0]))
|
||||
|
||||
def test_tokenizer_standardizes_only_continuous_values(self):
|
||||
tokenizer = OtherInfoTokenizer(
|
||||
n_embd=4,
|
||||
n_types=4,
|
||||
n_cont_types=2,
|
||||
n_categories=3,
|
||||
cont_type_ids=[1, 3],
|
||||
continuous_value_center=[10.0, 100.0],
|
||||
continuous_value_scale=[2.0, 20.0],
|
||||
)
|
||||
capture = _CaptureContinuousEncoder(n_embd=4)
|
||||
tokenizer.cont_value_encoder = capture
|
||||
|
||||
tokenizer(
|
||||
other_type=torch.tensor([[1, 2, 3]], dtype=torch.long),
|
||||
other_value=torch.tensor([[14.0, 1.0, 80.0]]),
|
||||
other_value_kind=torch.tensor([[1, 2, 1]], dtype=torch.long),
|
||||
)
|
||||
|
||||
torch.testing.assert_close(capture.last_type, torch.tensor([0, 1]))
|
||||
torch.testing.assert_close(capture.last_value, torch.tensor([2.0, -1.0]))
|
||||
|
||||
def test_scaler_buffers_round_trip_in_new_checkpoint(self):
|
||||
tokenizer = OtherInfoTokenizer(
|
||||
n_embd=4,
|
||||
n_types=4,
|
||||
n_cont_types=2,
|
||||
n_categories=2,
|
||||
cont_type_ids=[1, 3],
|
||||
continuous_value_center=[2.0, 10.0],
|
||||
continuous_value_scale=[1.5, 4.0],
|
||||
)
|
||||
state = tokenizer.state_dict()
|
||||
self.assertIn("continuous_value_center", state)
|
||||
self.assertIn("continuous_value_scale", state)
|
||||
|
||||
restored = OtherInfoTokenizer(
|
||||
n_embd=4,
|
||||
n_types=4,
|
||||
n_cont_types=2,
|
||||
n_categories=2,
|
||||
cont_type_ids=[1, 3],
|
||||
continuous_value_center=[0.0, 0.0],
|
||||
continuous_value_scale=[1.0, 1.0],
|
||||
)
|
||||
restored.load_state_dict(state, strict=True)
|
||||
|
||||
torch.testing.assert_close(
|
||||
restored.continuous_value_center,
|
||||
torch.tensor([2.0, 10.0]),
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
restored.continuous_value_scale,
|
||||
torch.tensor([1.5, 4.0]),
|
||||
)
|
||||
|
||||
def test_continuous_tokenizer_rejects_missing_scaler_statistics(self):
|
||||
with self.assertRaisesRegex(ValueError, "require train-split RobustScale"):
|
||||
OtherInfoTokenizer(
|
||||
n_embd=4,
|
||||
n_types=4,
|
||||
n_cont_types=2,
|
||||
n_categories=2,
|
||||
cont_type_ids=[1, 3],
|
||||
)
|
||||
|
||||
def test_evaluation_rejects_unscaled_continuous_checkpoint(self):
|
||||
dataset = type(
|
||||
"DatasetMetadata",
|
||||
(),
|
||||
{
|
||||
"vocab_size": 8,
|
||||
"n_types": 4,
|
||||
"n_cont_types": 2,
|
||||
"n_categories": 2,
|
||||
"cont_type_ids": [1, 3],
|
||||
},
|
||||
)()
|
||||
cfg = {
|
||||
"model_target_mode": "all_future",
|
||||
"target_mode": "all_future",
|
||||
"model_architecture": "transformer_ffn_v1",
|
||||
"n_layer": 1,
|
||||
"time_mode": "absolute",
|
||||
"dist_mode": "exponential",
|
||||
}
|
||||
|
||||
with self.assertRaisesRegex(RuntimeError, "unscaled checkpoints are not supported"):
|
||||
build_model_from_dataset(
|
||||
None,
|
||||
cfg,
|
||||
dataset,
|
||||
state_dict={"blocks.0.mlp.w1.weight": torch.zeros(1)},
|
||||
)
|
||||
|
||||
def test_evaluation_restores_required_scaler_buffers(self):
|
||||
dataset = type(
|
||||
"DatasetMetadata",
|
||||
(),
|
||||
{
|
||||
"vocab_size": 8,
|
||||
"n_types": 4,
|
||||
"n_cont_types": 2,
|
||||
"n_categories": 2,
|
||||
"cont_type_ids": [1, 3],
|
||||
},
|
||||
)()
|
||||
source = DeepHealth(
|
||||
vocab_size=8,
|
||||
n_embd=4,
|
||||
n_head=1,
|
||||
n_layer=1,
|
||||
n_types=4,
|
||||
n_cont_types=2,
|
||||
n_categories=2,
|
||||
cont_type_ids=[1, 3],
|
||||
continuous_value_center=[2.0, 10.0],
|
||||
continuous_value_scale=[1.5, 4.0],
|
||||
target_mode="all_future",
|
||||
time_mode="absolute",
|
||||
dist_mode="exponential",
|
||||
model_architecture="transformer_ffn_v1",
|
||||
)
|
||||
state = source.state_dict()
|
||||
cfg = {
|
||||
"model_target_mode": "all_future",
|
||||
"target_mode": "all_future",
|
||||
"model_architecture": "transformer_ffn_v1",
|
||||
"n_embd": 4,
|
||||
"n_head": 1,
|
||||
"n_layer": 1,
|
||||
"n_bins": 16,
|
||||
"time_mode": "absolute",
|
||||
"dist_mode": "exponential",
|
||||
"continuous_value_scaling": "robust",
|
||||
}
|
||||
|
||||
restored = build_model_from_dataset(None, cfg, dataset, state_dict=state)
|
||||
restored.load_state_dict(state, strict=True)
|
||||
|
||||
torch.testing.assert_close(
|
||||
restored.tokenizer.continuous_value_center,
|
||||
torch.tensor([2.0, 10.0]),
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
restored.tokenizer.continuous_value_scale,
|
||||
torch.tensor([1.5, 4.0]),
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
44
tests/test_dataset_reserved_event.py
Normal file
44
tests/test_dataset_reserved_event.py
Normal file
@@ -0,0 +1,44 @@
|
||||
import unittest
|
||||
|
||||
import numpy as np
|
||||
|
||||
from dataset import _ExpoBaseDataset
|
||||
from targets import RESERVED_IDX
|
||||
|
||||
|
||||
class ReservedEventFilteringTests(unittest.TestCase):
|
||||
@staticmethod
|
||||
def _base(extra_info_types):
|
||||
dataset = _ExpoBaseDataset.__new__(_ExpoBaseDataset)
|
||||
dataset.extra_info_types = list(extra_info_types)
|
||||
dataset.event_data = np.asarray(
|
||||
[
|
||||
[101, 10, RESERVED_IDX],
|
||||
[101, 20, 2],
|
||||
[101, 30, 3],
|
||||
],
|
||||
dtype=np.float64,
|
||||
)
|
||||
return dataset
|
||||
|
||||
def _assert_reserved_event_removed(self, extra_info_types):
|
||||
dataset = self._base(extra_info_types)
|
||||
|
||||
rows = list(dataset._iter_patient_events(impute_no_event_gaps=False))
|
||||
|
||||
self.assertEqual(len(rows), 1)
|
||||
eid, times, labels = rows[0]
|
||||
self.assertEqual(eid, 101)
|
||||
np.testing.assert_array_equal(times, np.asarray([20, 30], dtype=np.float32))
|
||||
np.testing.assert_array_equal(labels, np.asarray([3, 4], dtype=np.int64))
|
||||
self.assertNotIn(RESERVED_IDX, labels.tolist())
|
||||
|
||||
def test_empty_extra_info_removes_legacy_reserved_event(self):
|
||||
self._assert_reserved_event_removed([])
|
||||
|
||||
def test_selected_extra_info_removes_legacy_reserved_event(self):
|
||||
self._assert_reserved_event_removed([11])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
381
tests/test_disease_history_ablation.py
Normal file
381
tests/test_disease_history_ablation.py
Normal file
@@ -0,0 +1,381 @@
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from dataset import (
|
||||
AllFutureHealthDataset,
|
||||
all_future_collate_fn,
|
||||
transform_disease_history,
|
||||
transform_disease_history_batch_at_position,
|
||||
)
|
||||
from eval_data import validate_training_mode_config
|
||||
from losses import build_loss
|
||||
from models import DeepHealth
|
||||
from train_all_future import parse_args
|
||||
|
||||
|
||||
class DiseaseHistoryTransformTests(unittest.TestCase):
|
||||
def test_timed_ordered_and_set_representations(self):
|
||||
events = np.asarray([9, 4, 7], dtype=np.int64)
|
||||
times = np.asarray([50.0, 60.0, 65.0], dtype=np.float32)
|
||||
|
||||
timed_events, timed_times, timed_query = transform_disease_history(
|
||||
events, times, 70.0, "timed"
|
||||
)
|
||||
np.testing.assert_array_equal(timed_events, events)
|
||||
np.testing.assert_array_equal(timed_times, times)
|
||||
self.assertEqual(float(timed_query), 70.0)
|
||||
|
||||
ordered_events, ordered_times, ordered_query = transform_disease_history(
|
||||
events, times, 70.0, "ordered"
|
||||
)
|
||||
np.testing.assert_array_equal(ordered_events, events)
|
||||
np.testing.assert_array_equal(
|
||||
ordered_times,
|
||||
np.asarray([0.0, 1.0, 2.0], dtype=np.float32),
|
||||
)
|
||||
self.assertEqual(float(ordered_query), 3.0)
|
||||
|
||||
set_events, set_times, set_query = transform_disease_history(
|
||||
events, times, 70.0, "set"
|
||||
)
|
||||
np.testing.assert_array_equal(
|
||||
set_events,
|
||||
np.asarray([4, 7, 9], dtype=np.int64),
|
||||
)
|
||||
np.testing.assert_array_equal(
|
||||
set_times,
|
||||
np.zeros(3, dtype=np.float32),
|
||||
)
|
||||
self.assertEqual(float(set_query), 0.0)
|
||||
|
||||
def test_ordered_removes_calendar_time_but_keeps_order(self):
|
||||
events = np.asarray([9, 4, 7], dtype=np.int64)
|
||||
first = transform_disease_history(
|
||||
events,
|
||||
np.asarray([20.0, 21.0, 70.0], dtype=np.float32),
|
||||
75.0,
|
||||
"ordered",
|
||||
)
|
||||
second = transform_disease_history(
|
||||
events,
|
||||
np.asarray([50.0, 60.0, 65.0], dtype=np.float32),
|
||||
70.0,
|
||||
"ordered",
|
||||
)
|
||||
np.testing.assert_array_equal(first[0], second[0])
|
||||
np.testing.assert_array_equal(first[1], second[1])
|
||||
self.assertEqual(float(first[2]), float(second[2]))
|
||||
|
||||
reversed_events = transform_disease_history(
|
||||
events[::-1],
|
||||
np.asarray([50.0, 60.0, 65.0], dtype=np.float32),
|
||||
70.0,
|
||||
"ordered",
|
||||
)[0]
|
||||
self.assertFalse(np.array_equal(first[0], reversed_events))
|
||||
|
||||
def test_ordered_keeps_same_day_diseases_in_one_order_group(self):
|
||||
events, model_times, model_query = transform_disease_history(
|
||||
np.asarray([9, 4, 7], dtype=np.int64),
|
||||
np.asarray([50.0, 50.0, 65.0], dtype=np.float32),
|
||||
70.0,
|
||||
"ordered",
|
||||
)
|
||||
np.testing.assert_array_equal(
|
||||
events,
|
||||
np.asarray([9, 4, 7], dtype=np.int64),
|
||||
)
|
||||
np.testing.assert_array_equal(
|
||||
model_times,
|
||||
np.asarray([0.0, 0.0, 1.0], dtype=np.float32),
|
||||
)
|
||||
self.assertEqual(float(model_query), 2.0)
|
||||
|
||||
def test_set_removes_order_and_calendar_time(self):
|
||||
first = transform_disease_history(
|
||||
np.asarray([9, 4, 7], dtype=np.int64),
|
||||
np.asarray([50.0, 60.0, 65.0], dtype=np.float32),
|
||||
70.0,
|
||||
"set",
|
||||
)
|
||||
second = transform_disease_history(
|
||||
np.asarray([7, 9, 4], dtype=np.int64),
|
||||
np.asarray([20.0, 21.0, 70.0], dtype=np.float32),
|
||||
75.0,
|
||||
"set",
|
||||
)
|
||||
np.testing.assert_array_equal(first[0], second[0])
|
||||
np.testing.assert_array_equal(first[1], second[1])
|
||||
self.assertEqual(float(first[2]), float(second[2]))
|
||||
|
||||
def test_all_future_targets_stay_on_actual_time(self):
|
||||
patient = {
|
||||
"times": np.asarray([50.0, 60.0, 65.0, 75.0], dtype=np.float32),
|
||||
"labels": np.asarray([9, 4, 7, 12], dtype=np.int64),
|
||||
"t_obs": 75.0,
|
||||
"sex": 0,
|
||||
"other_type": np.zeros(0, dtype=np.int64),
|
||||
"other_value": np.zeros(0, dtype=np.float32),
|
||||
"other_value_kind": np.zeros(0, dtype=np.int64),
|
||||
"other_time": np.zeros(0, dtype=np.float32),
|
||||
}
|
||||
|
||||
items = {}
|
||||
for mode in ("timed", "ordered", "set"):
|
||||
dataset = AllFutureHealthDataset.__new__(AllFutureHealthDataset)
|
||||
dataset.disease_history_mode = mode
|
||||
items[mode] = dataset._build_item(patient, 70.0)
|
||||
|
||||
for item in items.values():
|
||||
torch.testing.assert_close(
|
||||
item["future_targets"],
|
||||
torch.tensor([12], dtype=torch.long),
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
item["future_dt"],
|
||||
torch.tensor([5.0]),
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
item["exposure"],
|
||||
torch.tensor(5.0),
|
||||
)
|
||||
|
||||
torch.testing.assert_close(
|
||||
items["timed"]["time_seq"],
|
||||
torch.tensor([50.0, 60.0, 65.0]),
|
||||
)
|
||||
self.assertEqual(float(items["timed"]["t_query"]), 70.0)
|
||||
torch.testing.assert_close(
|
||||
items["ordered"]["event_seq"],
|
||||
torch.tensor([9, 4, 7]),
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
items["ordered"]["time_seq"],
|
||||
torch.tensor([0.0, 1.0, 2.0]),
|
||||
)
|
||||
self.assertEqual(float(items["ordered"]["t_query"]), 3.0)
|
||||
torch.testing.assert_close(
|
||||
items["set"]["event_seq"],
|
||||
torch.tensor([4, 7, 9]),
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
items["set"]["time_seq"],
|
||||
torch.zeros(3),
|
||||
)
|
||||
self.assertEqual(float(items["set"]["t_query"]), 0.0)
|
||||
|
||||
def test_batch_prefix_transform_masks_future_events(self):
|
||||
events = torch.tensor(
|
||||
[
|
||||
[9, 4, 7, 12],
|
||||
[8, 5, 11, 0],
|
||||
],
|
||||
dtype=torch.long,
|
||||
)
|
||||
actual_times = torch.tensor(
|
||||
[
|
||||
[50.0, 60.0, 65.0, 75.0],
|
||||
[45.0, 55.0, 80.0, 0.0],
|
||||
]
|
||||
)
|
||||
mask = events > 0
|
||||
|
||||
timed = transform_disease_history_batch_at_position(
|
||||
events, actual_times, mask, 1, "timed", vocab_size=20
|
||||
)
|
||||
torch.testing.assert_close(timed[0], events)
|
||||
torch.testing.assert_close(timed[1], actual_times)
|
||||
torch.testing.assert_close(timed[2], mask)
|
||||
torch.testing.assert_close(timed[3], torch.tensor([60.0, 55.0]))
|
||||
|
||||
ordered = transform_disease_history_batch_at_position(
|
||||
events, actual_times, mask, 1, "ordered", vocab_size=20
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
ordered[2],
|
||||
torch.tensor(
|
||||
[
|
||||
[True, True, False, False],
|
||||
[True, True, False, False],
|
||||
]
|
||||
),
|
||||
)
|
||||
torch.testing.assert_close(ordered[3], torch.tensor([2.0, 2.0]))
|
||||
|
||||
disease_set = transform_disease_history_batch_at_position(
|
||||
events, actual_times, mask, 1, "set", vocab_size=20
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
disease_set[0],
|
||||
torch.tensor(
|
||||
[
|
||||
[4, 9, 0, 0],
|
||||
[5, 8, 0, 0],
|
||||
]
|
||||
),
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
disease_set[2],
|
||||
torch.tensor(
|
||||
[
|
||||
[True, True, False, False],
|
||||
[True, True, False, False],
|
||||
]
|
||||
),
|
||||
)
|
||||
torch.testing.assert_close(disease_set[3], torch.zeros(2))
|
||||
|
||||
|
||||
class DiseaseSetModelTests(unittest.TestCase):
|
||||
@staticmethod
|
||||
def _model():
|
||||
return DeepHealth(
|
||||
vocab_size=16,
|
||||
n_embd=24,
|
||||
n_head=4,
|
||||
n_layer=2,
|
||||
n_types=1,
|
||||
n_cont_types=0,
|
||||
n_categories=1,
|
||||
cont_type_ids=[],
|
||||
n_bins=4,
|
||||
target_mode="all_future",
|
||||
time_mode="relative",
|
||||
dist_mode="weibull",
|
||||
dropout=0.0,
|
||||
model_architecture="traj_mixer_v5",
|
||||
)
|
||||
|
||||
def test_equal_time_query_is_permutation_invariant(self):
|
||||
torch.manual_seed(7)
|
||||
model = self._model().eval()
|
||||
event_seq = torch.tensor(
|
||||
[
|
||||
[4, 7, 9],
|
||||
[9, 4, 7],
|
||||
],
|
||||
dtype=torch.long,
|
||||
)
|
||||
time_seq = torch.zeros(2, 3)
|
||||
empty_long = torch.zeros(2, 0, dtype=torch.long)
|
||||
empty_float = torch.zeros(2, 0)
|
||||
|
||||
with torch.inference_mode():
|
||||
hidden = model(
|
||||
event_seq=event_seq,
|
||||
time_seq=time_seq,
|
||||
sex=torch.zeros(2, dtype=torch.long),
|
||||
padding_mask=torch.ones(2, 3, dtype=torch.bool),
|
||||
t_query=torch.zeros(2),
|
||||
other_type=empty_long,
|
||||
other_value=empty_float,
|
||||
other_value_kind=empty_long,
|
||||
other_time=empty_float,
|
||||
)
|
||||
|
||||
torch.testing.assert_close(hidden[0], hidden[1], atol=1e-6, rtol=1e-6)
|
||||
|
||||
def test_ordered_and_set_support_finite_weibull_loss(self):
|
||||
patient = {
|
||||
"times": np.asarray([50.0, 60.0, 65.0, 75.0], dtype=np.float32),
|
||||
"labels": np.asarray([9, 4, 7, 12], dtype=np.int64),
|
||||
"t_obs": 75.0,
|
||||
"sex": 0,
|
||||
"other_type": np.zeros(0, dtype=np.int64),
|
||||
"other_value": np.zeros(0, dtype=np.float32),
|
||||
"other_value_kind": np.zeros(0, dtype=np.int64),
|
||||
"other_time": np.zeros(0, dtype=np.float32),
|
||||
}
|
||||
criterion = build_loss("weibull", ignored_idx={0, 1})
|
||||
|
||||
for mode in ("ordered", "set"):
|
||||
dataset = AllFutureHealthDataset.__new__(AllFutureHealthDataset)
|
||||
dataset.disease_history_mode = mode
|
||||
item = dataset._build_item(patient, 70.0)
|
||||
batch = all_future_collate_fn([item, item])
|
||||
model = self._model()
|
||||
hidden = model(
|
||||
event_seq=batch["event_seq"],
|
||||
time_seq=batch["time_seq"],
|
||||
sex=batch["sex"],
|
||||
padding_mask=batch["padding_mask"],
|
||||
t_query=batch["t_query"],
|
||||
other_type=batch["other_type"],
|
||||
other_value=batch["other_value"],
|
||||
other_value_kind=batch["other_value_kind"],
|
||||
other_time=batch["other_time"],
|
||||
)
|
||||
loss = criterion(
|
||||
logits=model.calc_risk(hidden),
|
||||
weibull_rho=model.calc_weibull_rho(hidden),
|
||||
targets=batch["future_targets"],
|
||||
dt=batch["future_dt"],
|
||||
exposure=batch["exposure"],
|
||||
)
|
||||
self.assertTrue(torch.isfinite(loss), msg=f"{mode} loss={loss}")
|
||||
|
||||
|
||||
class DiseaseHistoryConfigTests(unittest.TestCase):
|
||||
def test_training_cli_accepts_ordered_ablation(self):
|
||||
project_root = Path(__file__).resolve().parents[1]
|
||||
with patch(
|
||||
"sys.argv",
|
||||
[
|
||||
"train_all_future.py",
|
||||
"--disease_history_mode",
|
||||
"ordered",
|
||||
"--model_architecture",
|
||||
"traj_mixer_v5",
|
||||
"--time_mode",
|
||||
"relative",
|
||||
"--dist_mode",
|
||||
"weibull",
|
||||
"--extra_info_types_file",
|
||||
str(project_root / "extra_info_types_none.txt"),
|
||||
],
|
||||
):
|
||||
args = parse_args()
|
||||
self.assertEqual(args.disease_history_mode, "ordered")
|
||||
self.assertEqual(args.extra_info_types, [])
|
||||
|
||||
def test_legacy_config_defaults_to_timed(self):
|
||||
validate_training_mode_config(
|
||||
{
|
||||
"model_target_mode": "all_future",
|
||||
"time_mode": "relative",
|
||||
"dist_mode": "weibull",
|
||||
"model_architecture": "traj_mixer_v5",
|
||||
"extra_info_types": [11, 66, 67],
|
||||
}
|
||||
)
|
||||
|
||||
def test_ordered_config_requires_exact_ablation_setup(self):
|
||||
validate_training_mode_config(
|
||||
{
|
||||
"model_target_mode": "all_future",
|
||||
"time_mode": "relative",
|
||||
"dist_mode": "weibull",
|
||||
"model_architecture": "traj_mixer_v5",
|
||||
"extra_info_types": [],
|
||||
"disease_history_mode": "ordered",
|
||||
}
|
||||
)
|
||||
with self.assertRaises(ValueError):
|
||||
validate_training_mode_config(
|
||||
{
|
||||
"model_target_mode": "all_future",
|
||||
"time_mode": "relative",
|
||||
"dist_mode": "weibull",
|
||||
"model_architecture": "traj_mixer_v5",
|
||||
"extra_info_types": [11],
|
||||
"disease_history_mode": "ordered",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -5,8 +5,8 @@ Training samples are patient-level. For each patient and each __getitem__ call,
|
||||
AllFutureHealthDataset randomly samples a query time t_query, uses events at or
|
||||
before t_query as history, and uses events after t_query as the future target set.
|
||||
|
||||
Validation/test samples are deterministic query points built from future event
|
||||
times, then split by patient.
|
||||
All splits use the same patient/interval/time-uniform query distribution.
|
||||
Validation/test keep one deterministic query draw per patient.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -25,21 +25,31 @@ from torch.optim import AdamW
|
||||
from torch.utils.data import DataLoader, RandomSampler
|
||||
from tqdm.auto import tqdm
|
||||
|
||||
from dataset import AllFutureHealthDataset, all_future_collate_fn
|
||||
from losses import build_loss
|
||||
from models import (
|
||||
EVENT_TRAJECTORY_ARCHITECTURE,
|
||||
MODEL_SIZE_NAMES,
|
||||
DeepHealth,
|
||||
resolve_model_size,
|
||||
from dataset import (
|
||||
DISEASE_HISTORY_MODES,
|
||||
DISEASE_HISTORY_MODE_TIMED,
|
||||
AllFutureHealthDataset,
|
||||
all_future_collate_fn,
|
||||
)
|
||||
from targets import CHECKUP_IDX, PAD_IDX
|
||||
from losses import build_loss
|
||||
from model_architectures import (
|
||||
DEFAULT_MODEL_ARCHITECTURE,
|
||||
SUPPORTED_MODEL_ARCHITECTURES,
|
||||
)
|
||||
from models import DeepHealth
|
||||
from targets import NO_EVENT_IDX, PAD_IDX, RESERVED_IDX
|
||||
from train_util import (
|
||||
AllFutureBaselineStats,
|
||||
ContinuousRobustScalerStats,
|
||||
configure_torch_for_training,
|
||||
create_unique_run_dir,
|
||||
format_extra_info_types,
|
||||
fit_continuous_robust_scaler,
|
||||
fit_all_future_baseline,
|
||||
get_lr,
|
||||
get_model_parameter_counts,
|
||||
load_extra_info_types_file,
|
||||
move_batch_to_device,
|
||||
resolve_device,
|
||||
save_checkpoint,
|
||||
save_config,
|
||||
@@ -70,8 +80,19 @@ def parse_args() -> argparse.Namespace:
|
||||
|
||||
parser.add_argument("--data_prefix", type=str, default="ukb")
|
||||
parser.add_argument("--labels_file", type=str, default="labels.csv")
|
||||
parser.add_argument("--runs_root", type=str, default="runs")
|
||||
parser.add_argument("--seed", type=int, default=42)
|
||||
parser.add_argument("--extra_info_types_file", type=str, default=None)
|
||||
parser.add_argument(
|
||||
"--disease_history_mode",
|
||||
type=str,
|
||||
default=DISEASE_HISTORY_MODE_TIMED,
|
||||
choices=DISEASE_HISTORY_MODES,
|
||||
help=(
|
||||
"timed=real disease times; ordered=chronological disease order with "
|
||||
"ordinal positions; set=unordered disease set with no disease time"
|
||||
),
|
||||
)
|
||||
|
||||
parser.add_argument("--train_ratio", type=float, default=0.7)
|
||||
parser.add_argument("--val_ratio", type=float, default=0.15)
|
||||
@@ -83,21 +104,32 @@ def parse_args() -> argparse.Namespace:
|
||||
parser.add_argument("--min_future_events", type=int, default=1)
|
||||
parser.add_argument("--validation_query_seed", type=int, default=None)
|
||||
|
||||
parser.add_argument(
|
||||
"--model_size",
|
||||
type=str,
|
||||
default="nano",
|
||||
choices=MODEL_SIZE_NAMES,
|
||||
)
|
||||
parser.add_argument("--n_reasoning_rounds", type=int, default=12)
|
||||
parser.add_argument("--n_embd", type=int, default=120)
|
||||
parser.add_argument("--n_head", type=int, default=10)
|
||||
parser.add_argument("--n_layer", type=int, default=12)
|
||||
parser.add_argument("--n_bins", type=int, default=16)
|
||||
parser.add_argument("--extra_pool_reduce", type=str, default="mean",
|
||||
choices=["mean", "sum"])
|
||||
parser.add_argument("--time_mode", type=str, default="relative",
|
||||
choices=["relative", "absolute"])
|
||||
parser.add_argument("--dist_mode", type=str, default="exponential",
|
||||
choices=["exponential", "weibull", "mixed"])
|
||||
choices=["exponential", "weibull"])
|
||||
parser.add_argument("--dropout", type=float, default=0.0)
|
||||
parser.add_argument(
|
||||
"--risk_head_bias",
|
||||
action=argparse.BooleanOptionalAction,
|
||||
default=True,
|
||||
help=(
|
||||
"Initialize a learnable all-future output bias from training-only "
|
||||
"marginal first-onset rates"
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model_architecture",
|
||||
type=str,
|
||||
default=DEFAULT_MODEL_ARCHITECTURE,
|
||||
choices=SUPPORTED_MODEL_ARCHITECTURES,
|
||||
)
|
||||
|
||||
parser.add_argument("--batch_size", type=int, default=128)
|
||||
parser.add_argument("--base_lr", type=float, default=3e-4)
|
||||
@@ -130,57 +162,77 @@ def parse_args() -> argparse.Namespace:
|
||||
if args.extra_info_types_file is not None
|
||||
else None
|
||||
)
|
||||
if args.disease_history_mode != DISEASE_HISTORY_MODE_TIMED:
|
||||
expected = {
|
||||
"model_architecture": "traj_mixer_v5",
|
||||
"time_mode": "relative",
|
||||
"dist_mode": "weibull",
|
||||
}
|
||||
mismatches = [
|
||||
f"{name}={getattr(args, name)!r} (expected {value!r})"
|
||||
for name, value in expected.items()
|
||||
if getattr(args, name) != value
|
||||
]
|
||||
if args.extra_info_types != []:
|
||||
mismatches.append(
|
||||
"extra_info_types must be [] via extra_info_types_none.txt"
|
||||
)
|
||||
if mismatches:
|
||||
raise ValueError(
|
||||
f"disease_history_mode={args.disease_history_mode!r} is reserved "
|
||||
"for the no-extra TrajMixer + all_future + relative + Weibull "
|
||||
"ablation; " + "; ".join(mismatches)
|
||||
)
|
||||
return args
|
||||
|
||||
|
||||
def get_lr(epoch: int, args: argparse.Namespace, adaptive_lr: float) -> float:
|
||||
if epoch < args.warmup_epochs:
|
||||
return adaptive_lr * (epoch + 1) / args.warmup_epochs
|
||||
progress = (epoch - args.warmup_epochs) / max(1, args.max_epochs - args.warmup_epochs)
|
||||
cosine = 0.5 * (1 + math.cos(math.pi * progress))
|
||||
return adaptive_lr * (args.min_lr_ratio + cosine * (1 - args.min_lr_ratio))
|
||||
|
||||
|
||||
def move_batch_to_device(batch: Dict[str, torch.Tensor], device: torch.device) -> Dict[str, torch.Tensor]:
|
||||
non_blocking = device.type == "cuda"
|
||||
return {
|
||||
key: value.to(device, non_blocking=non_blocking)
|
||||
if isinstance(value, torch.Tensor)
|
||||
else value
|
||||
for key, value in batch.items()
|
||||
}
|
||||
|
||||
|
||||
def build_model(args: argparse.Namespace, dataset: AllFutureHealthDataset) -> DeepHealth:
|
||||
def build_model(
|
||||
args: argparse.Namespace,
|
||||
dataset: AllFutureHealthDataset,
|
||||
scaler_stats: ContinuousRobustScalerStats,
|
||||
baseline_stats: AllFutureBaselineStats | None,
|
||||
) -> DeepHealth:
|
||||
if tuple(int(x) for x in dataset.cont_type_ids) != scaler_stats.cont_type_ids:
|
||||
raise ValueError(
|
||||
"RobustScale statistics are not aligned with dataset.cont_type_ids"
|
||||
)
|
||||
center = scaler_stats.center if dataset.n_cont_types > 0 else None
|
||||
scale = scaler_stats.scale if dataset.n_cont_types > 0 else None
|
||||
if args.risk_head_bias and baseline_stats is None:
|
||||
raise ValueError("baseline_stats is required when risk_head_bias is enabled")
|
||||
return DeepHealth(
|
||||
vocab_size=dataset.vocab_size,
|
||||
model_size=args.model_size,
|
||||
n_reasoning_rounds=args.n_reasoning_rounds,
|
||||
n_embd=args.n_embd,
|
||||
n_head=args.n_head,
|
||||
n_layer=args.n_layer,
|
||||
n_types=dataset.n_types,
|
||||
n_cont_types=dataset.n_cont_types,
|
||||
n_categories=dataset.n_categories,
|
||||
cont_type_ids=dataset.cont_type_ids,
|
||||
n_bins=args.n_bins,
|
||||
continuous_value_center=center,
|
||||
continuous_value_scale=scale,
|
||||
extra_pool_reduce=args.extra_pool_reduce,
|
||||
target_mode="all_future",
|
||||
time_mode=args.time_mode,
|
||||
dist_mode=args.dist_mode,
|
||||
dropout=args.dropout,
|
||||
model_architecture=args.model_architecture,
|
||||
risk_head_bias=args.risk_head_bias,
|
||||
risk_head_bias_init=(
|
||||
baseline_stats.bias
|
||||
if args.risk_head_bias and baseline_stats is not None
|
||||
else None
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def build_criterion(args: argparse.Namespace, dataset: AllFutureHealthDataset):
|
||||
ignored_idx = {PAD_IDX, CHECKUP_IDX}
|
||||
def build_criterion(args: argparse.Namespace):
|
||||
ignored_idx = {PAD_IDX, RESERVED_IDX, NO_EVENT_IDX}
|
||||
if args.dist_mode == "exponential":
|
||||
return build_loss("exponential", ignored_idx=ignored_idx)
|
||||
if args.dist_mode == "weibull":
|
||||
return build_loss("weibull", ignored_idx=ignored_idx)
|
||||
if args.dist_mode == "mixed":
|
||||
return build_loss(
|
||||
"mixed",
|
||||
death_idx=dataset.vocab_size - 1,
|
||||
ignored_idx=ignored_idx,
|
||||
)
|
||||
raise ValueError(f"Unknown dist_mode: {args.dist_mode}")
|
||||
|
||||
|
||||
@@ -192,9 +244,7 @@ def compute_all_future_loss(
|
||||
device: torch.device,
|
||||
) -> torch.Tensor:
|
||||
required_keys = set(MODEL_INPUT_KEYS)
|
||||
required_keys.update(("future_targets", "exposure"))
|
||||
if args.dist_mode in {"weibull", "mixed"}:
|
||||
required_keys.add("future_dt")
|
||||
required_keys.update(("future_targets", "future_dt", "exposure"))
|
||||
batch = move_batch_to_device(
|
||||
{key: batch[key] for key in required_keys},
|
||||
device,
|
||||
@@ -210,7 +260,6 @@ def compute_all_future_loss(
|
||||
other_value=batch["other_value"],
|
||||
other_value_kind=batch["other_value_kind"],
|
||||
other_time=batch["other_time"],
|
||||
target_mode="all_future",
|
||||
)
|
||||
logits = model.calc_risk(hidden)
|
||||
|
||||
@@ -219,6 +268,8 @@ def compute_all_future_loss(
|
||||
logits=logits,
|
||||
targets=batch["future_targets"],
|
||||
exposure=batch["exposure"],
|
||||
dt=batch["future_dt"],
|
||||
history=batch["event_seq"],
|
||||
)
|
||||
elif args.dist_mode == "weibull":
|
||||
loss = criterion(
|
||||
@@ -227,15 +278,10 @@ def compute_all_future_loss(
|
||||
targets=batch["future_targets"],
|
||||
dt=batch["future_dt"],
|
||||
exposure=batch["exposure"],
|
||||
history=batch["event_seq"],
|
||||
)
|
||||
else:
|
||||
loss = criterion(
|
||||
logits=logits,
|
||||
death_rho=model.calc_death_rho(hidden),
|
||||
targets=batch["future_targets"],
|
||||
dt=batch["future_dt"],
|
||||
exposure=batch["exposure"],
|
||||
)
|
||||
raise ValueError(f"Unknown dist_mode: {args.dist_mode}")
|
||||
|
||||
if not torch.isfinite(loss):
|
||||
raise RuntimeError(f"Loss is not finite: {float(loss.detach().cpu())}")
|
||||
@@ -299,30 +345,44 @@ def build_metadata(
|
||||
train_subset,
|
||||
val_subset,
|
||||
test_subset,
|
||||
scaler_stats: ContinuousRobustScalerStats,
|
||||
baseline_stats: AllFutureBaselineStats | None,
|
||||
) -> Dict[str, Any]:
|
||||
size_config = resolve_model_size(args.model_size)
|
||||
scaler_metadata = scaler_stats.as_metadata()
|
||||
return {
|
||||
"run_name": run_name,
|
||||
"dataset_class": "AllFutureHealthDataset",
|
||||
"collate_fn": "all_future_collate_fn",
|
||||
"model_class": "DeepHealth",
|
||||
"model_architecture": EVENT_TRAJECTORY_ARCHITECTURE,
|
||||
"d_model": size_config.d_model,
|
||||
"n_trajectory": size_config.n_trajectory,
|
||||
"trajectory_dim": size_config.trajectory_dim,
|
||||
"traj_hidden": size_config.traj_hidden,
|
||||
"model_architecture": args.model_architecture,
|
||||
"model_target_mode": "all_future",
|
||||
"target_mode": "all_future",
|
||||
"event_stream_version": "disease_death_only_v1",
|
||||
"uses_assessment_event_token": False,
|
||||
"dist_mode": args.dist_mode,
|
||||
"disease_history_mode": args.disease_history_mode,
|
||||
"all_future_min_history_events": int(args.min_history_events),
|
||||
"all_future_min_future_events": int(args.min_future_events),
|
||||
"all_future_validation_query_seed": int(args.validation_query_seed),
|
||||
"all_future_query_distribution": "patient_interval_time_uniform",
|
||||
"all_future_queries_per_validation_patient": 1,
|
||||
"all_future_likelihood": "first_onset_survival_v2",
|
||||
"all_future_prevalent_outcomes_excluded": True,
|
||||
"risk_head_bias": bool(args.risk_head_bias),
|
||||
"risk_head_bias_weight_decay": 0.0 if args.risk_head_bias else None,
|
||||
"risk_head_baseline": (
|
||||
baseline_stats.as_metadata()
|
||||
if args.risk_head_bias and baseline_stats is not None
|
||||
else {"method": "none"}
|
||||
),
|
||||
"extra_info_types_file": (
|
||||
Path(args.extra_info_types_file).name
|
||||
if args.extra_info_types_file is not None
|
||||
else None
|
||||
),
|
||||
"extra_info_types": [int(x) for x in dataset.extra_info_types],
|
||||
"continuous_value_scaling": "robust",
|
||||
"continuous_value_scaler": scaler_metadata,
|
||||
"dataset_metadata": {
|
||||
"vocab_size": int(dataset.vocab_size),
|
||||
"n_types": int(dataset.n_types),
|
||||
@@ -330,6 +390,8 @@ def build_metadata(
|
||||
"n_categories": int(dataset.n_categories),
|
||||
"cont_type_ids": [int(x) for x in dataset.cont_type_ids],
|
||||
"extra_info_types": [int(x) for x in dataset.extra_info_types],
|
||||
"event_stream_version": "disease_death_only_v1",
|
||||
"uses_assessment_event_token": False,
|
||||
},
|
||||
"split_sizes": {
|
||||
"train": int(len(train_subset)),
|
||||
@@ -349,26 +411,24 @@ def main() -> None:
|
||||
|
||||
run_dir, run_name = create_unique_run_dir(
|
||||
lambda timestamp: (
|
||||
f"{args.model_size}_r{args.n_reasoning_rounds}_"
|
||||
f"{args.time_mode}_{args.dist_mode}_"
|
||||
f"all_future_pure_disease_{timestamp}"
|
||||
(
|
||||
""
|
||||
if args.disease_history_mode == DISEASE_HISTORY_MODE_TIMED
|
||||
else f"{args.disease_history_mode}_"
|
||||
)
|
||||
+ f"{args.time_mode}_{args.dist_mode}_"
|
||||
f"all_future_pure_disease_{timestamp}"
|
||||
),
|
||||
runs_root=Path(args.runs_root) / args.model_architecture,
|
||||
)
|
||||
logger = setup_logging(run_dir)
|
||||
|
||||
logger.info(f"Starting all-future training run: {run_name}")
|
||||
logger.info(f"Device: {device}")
|
||||
size_config = resolve_model_size(args.model_size)
|
||||
logger.info(
|
||||
"Model size: "
|
||||
f"{args.model_size} "
|
||||
f"(d_model={size_config.d_model}, "
|
||||
f"n_trajectory={size_config.n_trajectory}, "
|
||||
f"trajectory_dim={size_config.trajectory_dim}, "
|
||||
f"traj_hidden={size_config.traj_hidden}); "
|
||||
f"reasoning_rounds={args.n_reasoning_rounds}"
|
||||
)
|
||||
logger.info(f"Model architecture: {args.model_architecture}")
|
||||
logger.info(f"Disease history mode: {args.disease_history_mode}")
|
||||
logger.info(f"extra_info_types: {format_extra_info_types(args.extra_info_types)}")
|
||||
logger.info("Continuous value scaling: RobustScale (required)")
|
||||
|
||||
logger.info("Loading all-future datasets...")
|
||||
train_dataset = AllFutureHealthDataset(
|
||||
@@ -379,6 +439,7 @@ def main() -> None:
|
||||
min_future_events=args.min_future_events,
|
||||
validation_query_seed=args.validation_query_seed,
|
||||
extra_info_types=args.extra_info_types,
|
||||
disease_history_mode=args.disease_history_mode,
|
||||
)
|
||||
val_dataset = AllFutureHealthDataset(
|
||||
data_prefix=args.data_prefix,
|
||||
@@ -388,6 +449,7 @@ def main() -> None:
|
||||
min_future_events=args.min_future_events,
|
||||
validation_query_seed=args.validation_query_seed,
|
||||
extra_info_types=args.extra_info_types,
|
||||
disease_history_mode=args.disease_history_mode,
|
||||
)
|
||||
test_dataset = AllFutureHealthDataset(
|
||||
data_prefix=args.data_prefix,
|
||||
@@ -397,6 +459,7 @@ def main() -> None:
|
||||
min_future_events=args.min_future_events,
|
||||
validation_query_seed=args.validation_query_seed,
|
||||
extra_info_types=args.extra_info_types,
|
||||
disease_history_mode=args.disease_history_mode,
|
||||
)
|
||||
if args.train_eid_file and args.val_eid_file and args.test_eid_file:
|
||||
train_subset, val_subset, test_subset = split_all_future_datasets_by_eid_files(
|
||||
@@ -429,6 +492,44 @@ def main() -> None:
|
||||
f"Patients/queries: train={len(train_subset)}, val={len(val_subset)}, test={len(test_subset)}"
|
||||
)
|
||||
|
||||
if train_dataset.n_cont_types > 0:
|
||||
logger.info(
|
||||
"Fitting continuous RobustScaler on the complete training subset: "
|
||||
f"patients={len(train_subset):,}, features={train_dataset.n_cont_types}"
|
||||
)
|
||||
scaler_stats = fit_continuous_robust_scaler(
|
||||
train_dataset,
|
||||
train_subset,
|
||||
)
|
||||
if train_dataset.n_cont_types > 0:
|
||||
logger.info(
|
||||
"Continuous RobustScaler fitted: "
|
||||
f"observations={int(scaler_stats.observation_count.sum()):,}, "
|
||||
f"min_per_feature={int(scaler_stats.observation_count.min()):,}, "
|
||||
f"max_per_feature={int(scaler_stats.observation_count.max()):,}"
|
||||
)
|
||||
|
||||
baseline_stats = None
|
||||
if args.risk_head_bias:
|
||||
logger.info(
|
||||
"Fitting all-future output baseline on one seeded query per "
|
||||
f"training patient: patients={len(train_subset):,}"
|
||||
)
|
||||
baseline_stats = fit_all_future_baseline(
|
||||
train_dataset,
|
||||
train_subset,
|
||||
seed=args.seed,
|
||||
)
|
||||
baseline_summary = baseline_stats.as_metadata()
|
||||
rate_summary = baseline_summary["rate_per_year"]
|
||||
logger.info(
|
||||
"All-future baseline rates/year: "
|
||||
f"min={rate_summary['min']:.8g}, "
|
||||
f"median={rate_summary['median']:.8g}, "
|
||||
f"max={rate_summary['max']:.8g}, "
|
||||
f"first_onsets={baseline_summary['observed_first_onsets']:,}"
|
||||
)
|
||||
|
||||
train_loader = DataLoader(
|
||||
train_subset,
|
||||
batch_size=args.batch_size,
|
||||
@@ -460,24 +561,53 @@ def main() -> None:
|
||||
prefetch_factor=2 if args.num_workers > 0 else None,
|
||||
)
|
||||
|
||||
model = build_model(args, train_dataset).to(device)
|
||||
model = build_model(
|
||||
args,
|
||||
train_dataset,
|
||||
scaler_stats=scaler_stats,
|
||||
baseline_stats=baseline_stats,
|
||||
).to(device)
|
||||
parameter_counts = get_model_parameter_counts(model)
|
||||
logger.info(
|
||||
"Model parameters: "
|
||||
f"total={parameter_counts['model_parameter_count']:,}, "
|
||||
f"trainable={parameter_counts['trainable_parameter_count']:,}"
|
||||
)
|
||||
if model.risk_head.bias is None:
|
||||
optimizer_parameters = model.parameters()
|
||||
else:
|
||||
optimizer_parameters = [
|
||||
{
|
||||
"params": [
|
||||
parameter
|
||||
for parameter in model.parameters()
|
||||
if parameter is not model.risk_head.bias
|
||||
],
|
||||
"weight_decay": args.weight_decay,
|
||||
},
|
||||
{
|
||||
"params": [model.risk_head.bias],
|
||||
"weight_decay": 0.0,
|
||||
},
|
||||
]
|
||||
optimizer = AdamW(
|
||||
model.parameters(),
|
||||
optimizer_parameters,
|
||||
lr=args.base_lr,
|
||||
betas=tuple(args.betas),
|
||||
weight_decay=args.weight_decay,
|
||||
)
|
||||
criterion = build_criterion(args, train_dataset)
|
||||
criterion = build_criterion(args)
|
||||
adaptive_lr = args.base_lr * math.sqrt(args.batch_size / 128)
|
||||
|
||||
train_metadata = build_metadata(
|
||||
args, train_dataset, run_name, train_subset, val_subset, test_subset
|
||||
args,
|
||||
train_dataset,
|
||||
run_name,
|
||||
train_subset,
|
||||
val_subset,
|
||||
test_subset,
|
||||
scaler_stats,
|
||||
baseline_stats,
|
||||
)
|
||||
train_metadata.update(parameter_counts)
|
||||
save_config(
|
||||
|
||||
289
train_batch_linux.sh
Executable file
289
train_batch_linux.sh
Executable file
@@ -0,0 +1,289 @@
|
||||
#!/usr/bin/env bash
|
||||
#
|
||||
# Run the complete DeepHealth training matrix on a Linux GPU server.
|
||||
#
|
||||
# The matrix contains:
|
||||
# 1. FFN + Delphi2M next-token reproduction
|
||||
# 2. FFN/TrajMixer x absolute/relative x exponential/Weibull
|
||||
#
|
||||
# Every task uses one GPU. Each selected GPU runs its assigned tasks
|
||||
# sequentially, while different GPUs run in parallel.
|
||||
#
|
||||
# Example:
|
||||
# bash train_batch_linux.sh --gpus 0,1,2,3
|
||||
#
|
||||
|
||||
set -uo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
|
||||
cd "$SCRIPT_DIR"
|
||||
|
||||
GPU_CSV=""
|
||||
SEED=42
|
||||
NUM_WORKERS=4
|
||||
PYTHON_BIN="${PYTHON_BIN:-python}"
|
||||
CAMPAIGN_NAME="full_factorial_smoking_alcohol_bmi"
|
||||
DRY_RUN=0
|
||||
|
||||
BATCH_SIZE=256
|
||||
EXTRA_INFO_TYPES_FILE="$SCRIPT_DIR/extra_info_types_smoking_alcohol_bmi.txt"
|
||||
|
||||
usage() {
|
||||
cat <<'EOF'
|
||||
Usage:
|
||||
bash train_batch_linux.sh --gpus GPU_LIST [options]
|
||||
|
||||
Required:
|
||||
--gpus LIST Comma-separated GPU ids, for example 0,1,2,3.
|
||||
|
||||
Options:
|
||||
--seed N Training seed for every task (default: 42).
|
||||
--num-workers N DataLoader workers per task (default: 4).
|
||||
--python PATH Python executable (default: $PYTHON_BIN or python).
|
||||
--campaign NAME Output campaign name.
|
||||
--dry-run Print all commands without running them.
|
||||
-h, --help Show this help message.
|
||||
|
||||
Fixed experiment settings:
|
||||
batch_size 256
|
||||
extra_info_types extra_info_types_smoking_alcohol_bmi.txt
|
||||
continuous scaling required train-split RobustScale
|
||||
model size Defaults from the individual training entrypoints
|
||||
|
||||
Outputs:
|
||||
runs/<campaign>/<architecture>/...
|
||||
batch_logs/<campaign>/seed_<seed>/<job>.log
|
||||
EOF
|
||||
}
|
||||
|
||||
while (($# > 0)); do
|
||||
case "$1" in
|
||||
--gpus)
|
||||
[[ $# -ge 2 ]] || {
|
||||
echo "ERROR: --gpus requires a value." >&2
|
||||
exit 2
|
||||
}
|
||||
GPU_CSV="$2"
|
||||
shift 2
|
||||
;;
|
||||
--seed)
|
||||
[[ $# -ge 2 ]] || {
|
||||
echo "ERROR: --seed requires a value." >&2
|
||||
exit 2
|
||||
}
|
||||
SEED="$2"
|
||||
shift 2
|
||||
;;
|
||||
--num-workers)
|
||||
[[ $# -ge 2 ]] || {
|
||||
echo "ERROR: --num-workers requires a value." >&2
|
||||
exit 2
|
||||
}
|
||||
NUM_WORKERS="$2"
|
||||
shift 2
|
||||
;;
|
||||
--python)
|
||||
[[ $# -ge 2 ]] || {
|
||||
echo "ERROR: --python requires a value." >&2
|
||||
exit 2
|
||||
}
|
||||
PYTHON_BIN="$2"
|
||||
shift 2
|
||||
;;
|
||||
--campaign)
|
||||
[[ $# -ge 2 ]] || {
|
||||
echo "ERROR: --campaign requires a value." >&2
|
||||
exit 2
|
||||
}
|
||||
CAMPAIGN_NAME="$2"
|
||||
shift 2
|
||||
;;
|
||||
--dry-run)
|
||||
DRY_RUN=1
|
||||
shift
|
||||
;;
|
||||
-h|--help)
|
||||
usage
|
||||
exit 0
|
||||
;;
|
||||
*)
|
||||
echo "ERROR: unknown argument: $1" >&2
|
||||
usage >&2
|
||||
exit 2
|
||||
;;
|
||||
esac
|
||||
done
|
||||
|
||||
[[ -n "$GPU_CSV" ]] || {
|
||||
echo "ERROR: --gpus is required." >&2
|
||||
usage >&2
|
||||
exit 2
|
||||
}
|
||||
[[ "$SEED" =~ ^[0-9]+$ ]] || {
|
||||
echo "ERROR: --seed must be a non-negative integer." >&2
|
||||
exit 2
|
||||
}
|
||||
[[ "$NUM_WORKERS" =~ ^[0-9]+$ ]] || {
|
||||
echo "ERROR: --num-workers must be a non-negative integer." >&2
|
||||
exit 2
|
||||
}
|
||||
[[ "$CAMPAIGN_NAME" =~ ^[A-Za-z0-9._-]+$ ]] || {
|
||||
echo "ERROR: --campaign may contain only letters, numbers, ., _, and -." >&2
|
||||
exit 2
|
||||
}
|
||||
[[ -f "$EXTRA_INFO_TYPES_FILE" ]] || {
|
||||
echo "ERROR: missing extra-info file: $EXTRA_INFO_TYPES_FILE" >&2
|
||||
exit 2
|
||||
}
|
||||
command -v "$PYTHON_BIN" >/dev/null 2>&1 || {
|
||||
echo "ERROR: Python executable not found: $PYTHON_BIN" >&2
|
||||
exit 2
|
||||
}
|
||||
|
||||
IFS=',' read -r -a GPU_IDS <<< "$GPU_CSV"
|
||||
declare -A SEEN_GPUS=()
|
||||
for gpu in "${GPU_IDS[@]}"; do
|
||||
[[ -n "$gpu" && "$gpu" =~ ^[A-Za-z0-9._:-]+$ ]] || {
|
||||
echo "ERROR: invalid GPU id: $gpu" >&2
|
||||
exit 2
|
||||
}
|
||||
[[ -z "${SEEN_GPUS[$gpu]+x}" ]] || {
|
||||
echo "ERROR: duplicate GPU id: $gpu" >&2
|
||||
exit 2
|
||||
}
|
||||
SEEN_GPUS["$gpu"]=1
|
||||
done
|
||||
|
||||
RUNS_ROOT="$SCRIPT_DIR/runs/$CAMPAIGN_NAME"
|
||||
LOG_ROOT="$SCRIPT_DIR/batch_logs/$CAMPAIGN_NAME/seed_$SEED"
|
||||
mkdir -p "$RUNS_ROOT" "$LOG_ROOT"
|
||||
|
||||
declare -a JOB_NAMES=()
|
||||
declare -a JOB_ENTRYPOINTS=()
|
||||
declare -a JOB_ARCHITECTURES=()
|
||||
declare -a JOB_TIME_MODES=()
|
||||
declare -a JOB_DIST_MODES=()
|
||||
|
||||
add_job() {
|
||||
JOB_NAMES+=("$1")
|
||||
JOB_ENTRYPOINTS+=("$2")
|
||||
JOB_ARCHITECTURES+=("$3")
|
||||
JOB_TIME_MODES+=("$4")
|
||||
JOB_DIST_MODES+=("$5")
|
||||
}
|
||||
|
||||
add_job \
|
||||
"ffn_next_token_absolute_delphi2m" \
|
||||
"train_next_step.py" \
|
||||
"transformer_ffn_v1" \
|
||||
"absolute" \
|
||||
"exponential"
|
||||
|
||||
for architecture in transformer_ffn_v1 traj_mixer_v5; do
|
||||
for time_mode in absolute relative; do
|
||||
for dist_mode in exponential weibull; do
|
||||
add_job \
|
||||
"${architecture}_all_future_${time_mode}_${dist_mode}" \
|
||||
"train_all_future.py" \
|
||||
"$architecture" \
|
||||
"$time_mode" \
|
||||
"$dist_mode"
|
||||
done
|
||||
done
|
||||
done
|
||||
|
||||
print_command() {
|
||||
printf '%q ' "$@"
|
||||
printf '\n'
|
||||
}
|
||||
|
||||
run_job() {
|
||||
local job_index="$1"
|
||||
local gpu="$2"
|
||||
local job_name="${JOB_NAMES[$job_index]}"
|
||||
local entrypoint="${JOB_ENTRYPOINTS[$job_index]}"
|
||||
local architecture="${JOB_ARCHITECTURES[$job_index]}"
|
||||
local time_mode="${JOB_TIME_MODES[$job_index]}"
|
||||
local dist_mode="${JOB_DIST_MODES[$job_index]}"
|
||||
local log_file="$LOG_ROOT/$job_name.log"
|
||||
local -a command=(
|
||||
"$PYTHON_BIN"
|
||||
-u
|
||||
"$SCRIPT_DIR/$entrypoint"
|
||||
--runs_root "$RUNS_ROOT"
|
||||
--seed "$SEED"
|
||||
--batch_size "$BATCH_SIZE"
|
||||
--extra_info_types_file "$EXTRA_INFO_TYPES_FILE"
|
||||
--model_architecture "$architecture"
|
||||
--num_workers "$NUM_WORKERS"
|
||||
--device cuda
|
||||
)
|
||||
|
||||
if [[ "$entrypoint" == "train_all_future.py" ]]; then
|
||||
command+=(
|
||||
--time_mode "$time_mode"
|
||||
--dist_mode "$dist_mode"
|
||||
)
|
||||
fi
|
||||
|
||||
echo "[$(date '+%F %T')] START job=$job_name gpu=$gpu"
|
||||
echo " log=$log_file"
|
||||
if ((DRY_RUN)); then
|
||||
printf ' CUDA_VISIBLE_DEVICES=%q ' "$gpu"
|
||||
print_command "${command[@]}"
|
||||
return 0
|
||||
fi
|
||||
|
||||
if CUDA_VISIBLE_DEVICES="$gpu" PYTHONUNBUFFERED=1 \
|
||||
"${command[@]}" >"$log_file" 2>&1; then
|
||||
echo "[$(date '+%F %T')] DONE job=$job_name gpu=$gpu"
|
||||
return 0
|
||||
else
|
||||
local exit_code=$?
|
||||
echo "[$(date '+%F %T')] FAIL job=$job_name gpu=$gpu exit=$exit_code" >&2
|
||||
echo " See: $log_file" >&2
|
||||
return "$exit_code"
|
||||
fi
|
||||
}
|
||||
|
||||
worker() {
|
||||
local slot="$1"
|
||||
local gpu="${GPU_IDS[$slot]}"
|
||||
local job_index
|
||||
local failed=0
|
||||
|
||||
for ((job_index = slot; job_index < ${#JOB_NAMES[@]}; job_index += ${#GPU_IDS[@]})); do
|
||||
run_job "$job_index" "$gpu" || failed=1
|
||||
done
|
||||
return "$failed"
|
||||
}
|
||||
|
||||
echo "Campaign: $CAMPAIGN_NAME"
|
||||
echo "Seed: $SEED"
|
||||
echo "GPUs: ${GPU_IDS[*]}"
|
||||
echo "Tasks: ${#JOB_NAMES[@]}"
|
||||
echo "Runs root: $RUNS_ROOT"
|
||||
echo "Log root: $LOG_ROOT"
|
||||
echo
|
||||
|
||||
declare -a WORKER_PIDS=()
|
||||
for ((slot = 0; slot < ${#GPU_IDS[@]}; slot++)); do
|
||||
worker "$slot" &
|
||||
WORKER_PIDS+=("$!")
|
||||
done
|
||||
|
||||
overall_status=0
|
||||
for pid in "${WORKER_PIDS[@]}"; do
|
||||
wait "$pid" || overall_status=1
|
||||
done
|
||||
|
||||
if ((overall_status != 0)); then
|
||||
echo "One or more training tasks failed. Inspect logs under: $LOG_ROOT" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
if ((DRY_RUN)); then
|
||||
echo "Dry run completed successfully."
|
||||
else
|
||||
echo "All training tasks completed successfully."
|
||||
fi
|
||||
337
train_disease_history_ablation_linux.sh
Normal file
337
train_disease_history_ablation_linux.sh
Normal file
@@ -0,0 +1,337 @@
|
||||
#!/usr/bin/env bash
|
||||
#
|
||||
# Train the no-extra-information T/O/S disease-history ablation.
|
||||
#
|
||||
# Fixed model:
|
||||
# TrajMixer + all_future + relative + Weibull
|
||||
#
|
||||
# History modes:
|
||||
# timed (T): disease identities, order, and real first-onset times
|
||||
# ordered (O): disease identities and chronological order only
|
||||
# set (S): unordered disease set only
|
||||
#
|
||||
# Jobs assigned to one GPU run sequentially. Different GPUs run in parallel.
|
||||
#
|
||||
# Examples:
|
||||
# bash train_disease_history_ablation_linux.sh --gpus 0,1,2
|
||||
# bash train_disease_history_ablation_linux.sh --gpus 0 --seeds 42 --modes ordered,set
|
||||
# bash train_disease_history_ablation_linux.sh --gpus 0,1 --dry-run
|
||||
#
|
||||
|
||||
set -uo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
|
||||
cd "$SCRIPT_DIR"
|
||||
|
||||
GPU_CSV=""
|
||||
SEED_CSV="42,43,44"
|
||||
MODE_CSV="timed,ordered,set"
|
||||
NUM_WORKERS=4
|
||||
BATCH_SIZE=256
|
||||
PYTHON_BIN="${PYTHON_BIN:-python}"
|
||||
CAMPAIGN_NAME="disease_history_ablation_no_extra"
|
||||
DRY_RUN=0
|
||||
|
||||
EXTRA_INFO_TYPES_FILE="$SCRIPT_DIR/extra_info_types_none.txt"
|
||||
ENTRYPOINT="$SCRIPT_DIR/train_all_future.py"
|
||||
|
||||
usage() {
|
||||
cat <<'EOF'
|
||||
Usage:
|
||||
bash train_disease_history_ablation_linux.sh --gpus GPU_LIST [options]
|
||||
|
||||
Required:
|
||||
--gpus LIST Comma-separated GPU ids, for example 0,1,2.
|
||||
|
||||
Options:
|
||||
--seeds LIST Comma-separated seeds (default: 42,43,44).
|
||||
--modes LIST Subset of timed,ordered,set (default: all three).
|
||||
--batch-size N Batch size per task (default: 256).
|
||||
--num-workers N DataLoader workers per task (default: 4).
|
||||
--python PATH Python executable (default: $PYTHON_BIN or python).
|
||||
--campaign NAME Output campaign name.
|
||||
--dry-run Print commands without creating files or training.
|
||||
-h, --help Show this help message.
|
||||
|
||||
Fixed experiment settings:
|
||||
architecture traj_mixer_v5
|
||||
target all_future
|
||||
time mode relative
|
||||
distribution weibull
|
||||
extra information extra_info_types_none.txt
|
||||
|
||||
Outputs:
|
||||
runs/<campaign>/seed_<seed>/traj_mixer_v5/...
|
||||
batch_logs/<campaign>/seed_<seed>/<mode>.log
|
||||
EOF
|
||||
}
|
||||
|
||||
while (($# > 0)); do
|
||||
case "$1" in
|
||||
--gpus)
|
||||
[[ $# -ge 2 ]] || {
|
||||
echo "ERROR: --gpus requires a value." >&2
|
||||
exit 2
|
||||
}
|
||||
GPU_CSV="$2"
|
||||
shift 2
|
||||
;;
|
||||
--seeds)
|
||||
[[ $# -ge 2 ]] || {
|
||||
echo "ERROR: --seeds requires a value." >&2
|
||||
exit 2
|
||||
}
|
||||
SEED_CSV="$2"
|
||||
shift 2
|
||||
;;
|
||||
--modes)
|
||||
[[ $# -ge 2 ]] || {
|
||||
echo "ERROR: --modes requires a value." >&2
|
||||
exit 2
|
||||
}
|
||||
MODE_CSV="$2"
|
||||
shift 2
|
||||
;;
|
||||
--batch-size)
|
||||
[[ $# -ge 2 ]] || {
|
||||
echo "ERROR: --batch-size requires a value." >&2
|
||||
exit 2
|
||||
}
|
||||
BATCH_SIZE="$2"
|
||||
shift 2
|
||||
;;
|
||||
--num-workers)
|
||||
[[ $# -ge 2 ]] || {
|
||||
echo "ERROR: --num-workers requires a value." >&2
|
||||
exit 2
|
||||
}
|
||||
NUM_WORKERS="$2"
|
||||
shift 2
|
||||
;;
|
||||
--python)
|
||||
[[ $# -ge 2 ]] || {
|
||||
echo "ERROR: --python requires a value." >&2
|
||||
exit 2
|
||||
}
|
||||
PYTHON_BIN="$2"
|
||||
shift 2
|
||||
;;
|
||||
--campaign)
|
||||
[[ $# -ge 2 ]] || {
|
||||
echo "ERROR: --campaign requires a value." >&2
|
||||
exit 2
|
||||
}
|
||||
CAMPAIGN_NAME="$2"
|
||||
shift 2
|
||||
;;
|
||||
--dry-run)
|
||||
DRY_RUN=1
|
||||
shift
|
||||
;;
|
||||
-h|--help)
|
||||
usage
|
||||
exit 0
|
||||
;;
|
||||
*)
|
||||
echo "ERROR: unknown argument: $1" >&2
|
||||
usage >&2
|
||||
exit 2
|
||||
;;
|
||||
esac
|
||||
done
|
||||
|
||||
[[ -n "$GPU_CSV" ]] || {
|
||||
echo "ERROR: --gpus is required." >&2
|
||||
usage >&2
|
||||
exit 2
|
||||
}
|
||||
[[ -n "$SEED_CSV" ]] || {
|
||||
echo "ERROR: --seeds must not be empty." >&2
|
||||
exit 2
|
||||
}
|
||||
[[ -n "$MODE_CSV" ]] || {
|
||||
echo "ERROR: --modes must not be empty." >&2
|
||||
exit 2
|
||||
}
|
||||
[[ "$BATCH_SIZE" =~ ^[1-9][0-9]*$ ]] || {
|
||||
echo "ERROR: --batch-size must be a positive integer." >&2
|
||||
exit 2
|
||||
}
|
||||
[[ "$NUM_WORKERS" =~ ^[0-9]+$ ]] || {
|
||||
echo "ERROR: --num-workers must be a non-negative integer." >&2
|
||||
exit 2
|
||||
}
|
||||
[[ "$CAMPAIGN_NAME" =~ ^[A-Za-z0-9._-]+$ ]] || {
|
||||
echo "ERROR: --campaign may contain only letters, numbers, ., _, and -." >&2
|
||||
exit 2
|
||||
}
|
||||
[[ -f "$EXTRA_INFO_TYPES_FILE" ]] || {
|
||||
echo "ERROR: missing extra-info file: $EXTRA_INFO_TYPES_FILE" >&2
|
||||
exit 2
|
||||
}
|
||||
[[ -f "$ENTRYPOINT" ]] || {
|
||||
echo "ERROR: missing training entrypoint: $ENTRYPOINT" >&2
|
||||
exit 2
|
||||
}
|
||||
command -v "$PYTHON_BIN" >/dev/null 2>&1 || {
|
||||
echo "ERROR: Python executable not found: $PYTHON_BIN" >&2
|
||||
exit 2
|
||||
}
|
||||
|
||||
IFS=',' read -r -a GPU_IDS <<< "$GPU_CSV"
|
||||
declare -A SEEN_GPUS=()
|
||||
for gpu in "${GPU_IDS[@]}"; do
|
||||
[[ -n "$gpu" && "$gpu" =~ ^[A-Za-z0-9._:-]+$ ]] || {
|
||||
echo "ERROR: invalid GPU id: $gpu" >&2
|
||||
exit 2
|
||||
}
|
||||
[[ -z "${SEEN_GPUS[$gpu]+x}" ]] || {
|
||||
echo "ERROR: duplicate GPU id: $gpu" >&2
|
||||
exit 2
|
||||
}
|
||||
SEEN_GPUS["$gpu"]=1
|
||||
done
|
||||
|
||||
IFS=',' read -r -a SEEDS <<< "$SEED_CSV"
|
||||
declare -A SEEN_SEEDS=()
|
||||
for seed in "${SEEDS[@]}"; do
|
||||
[[ "$seed" =~ ^[0-9]+$ ]] || {
|
||||
echo "ERROR: invalid seed: $seed" >&2
|
||||
exit 2
|
||||
}
|
||||
[[ -z "${SEEN_SEEDS[$seed]+x}" ]] || {
|
||||
echo "ERROR: duplicate seed: $seed" >&2
|
||||
exit 2
|
||||
}
|
||||
SEEN_SEEDS["$seed"]=1
|
||||
done
|
||||
|
||||
IFS=',' read -r -a MODES <<< "$MODE_CSV"
|
||||
declare -A SEEN_MODES=()
|
||||
for mode in "${MODES[@]}"; do
|
||||
case "$mode" in
|
||||
timed|ordered|set)
|
||||
;;
|
||||
*)
|
||||
echo "ERROR: invalid mode: $mode (expected timed, ordered, or set)." >&2
|
||||
exit 2
|
||||
;;
|
||||
esac
|
||||
[[ -z "${SEEN_MODES[$mode]+x}" ]] || {
|
||||
echo "ERROR: duplicate mode: $mode" >&2
|
||||
exit 2
|
||||
}
|
||||
SEEN_MODES["$mode"]=1
|
||||
done
|
||||
|
||||
RUNS_ROOT="$SCRIPT_DIR/runs/$CAMPAIGN_NAME"
|
||||
LOG_ROOT="$SCRIPT_DIR/batch_logs/$CAMPAIGN_NAME"
|
||||
if ((!DRY_RUN)); then
|
||||
mkdir -p "$RUNS_ROOT" "$LOG_ROOT"
|
||||
fi
|
||||
|
||||
declare -a JOB_SEEDS=()
|
||||
declare -a JOB_MODES=()
|
||||
for seed in "${SEEDS[@]}"; do
|
||||
for mode in "${MODES[@]}"; do
|
||||
JOB_SEEDS+=("$seed")
|
||||
JOB_MODES+=("$mode")
|
||||
done
|
||||
done
|
||||
|
||||
print_command() {
|
||||
printf '%q ' "$@"
|
||||
printf '\n'
|
||||
}
|
||||
|
||||
run_job() {
|
||||
local job_index="$1"
|
||||
local gpu="$2"
|
||||
local seed="${JOB_SEEDS[$job_index]}"
|
||||
local mode="${JOB_MODES[$job_index]}"
|
||||
local seed_runs_root="$RUNS_ROOT/seed_$seed"
|
||||
local seed_log_root="$LOG_ROOT/seed_$seed"
|
||||
local log_file="$seed_log_root/$mode.log"
|
||||
local -a command=(
|
||||
"$PYTHON_BIN"
|
||||
-u
|
||||
"$ENTRYPOINT"
|
||||
--runs_root "$seed_runs_root"
|
||||
--seed "$seed"
|
||||
--batch_size "$BATCH_SIZE"
|
||||
--num_workers "$NUM_WORKERS"
|
||||
--device cuda
|
||||
--model_architecture traj_mixer_v5
|
||||
--time_mode relative
|
||||
--dist_mode weibull
|
||||
--extra_info_types_file "$EXTRA_INFO_TYPES_FILE"
|
||||
--disease_history_mode "$mode"
|
||||
)
|
||||
|
||||
if ((!DRY_RUN)); then
|
||||
mkdir -p "$seed_runs_root" "$seed_log_root"
|
||||
fi
|
||||
|
||||
echo "[$(date '+%F %T')] START seed=$seed mode=$mode gpu=$gpu"
|
||||
echo " log=$log_file"
|
||||
if ((DRY_RUN)); then
|
||||
printf ' CUDA_VISIBLE_DEVICES=%q ' "$gpu"
|
||||
print_command "${command[@]}"
|
||||
return 0
|
||||
fi
|
||||
|
||||
if CUDA_VISIBLE_DEVICES="$gpu" PYTHONUNBUFFERED=1 \
|
||||
"${command[@]}" >"$log_file" 2>&1; then
|
||||
echo "[$(date '+%F %T')] DONE seed=$seed mode=$mode gpu=$gpu"
|
||||
return 0
|
||||
else
|
||||
local exit_code=$?
|
||||
echo "[$(date '+%F %T')] FAIL seed=$seed mode=$mode gpu=$gpu exit=$exit_code" >&2
|
||||
echo " See: $log_file" >&2
|
||||
return "$exit_code"
|
||||
fi
|
||||
}
|
||||
|
||||
worker() {
|
||||
local slot="$1"
|
||||
local gpu="${GPU_IDS[$slot]}"
|
||||
local job_index
|
||||
local failed=0
|
||||
|
||||
for ((job_index = slot; job_index < ${#JOB_SEEDS[@]}; job_index += ${#GPU_IDS[@]})); do
|
||||
run_job "$job_index" "$gpu" || failed=1
|
||||
done
|
||||
return "$failed"
|
||||
}
|
||||
|
||||
echo "Campaign: $CAMPAIGN_NAME"
|
||||
echo "Seeds: ${SEEDS[*]}"
|
||||
echo "Modes: ${MODES[*]}"
|
||||
echo "GPUs: ${GPU_IDS[*]}"
|
||||
echo "Configurations per seed: ${#MODES[@]}"
|
||||
echo "Total tasks: ${#JOB_SEEDS[@]}"
|
||||
echo "Runs root: $RUNS_ROOT"
|
||||
echo "Log root: $LOG_ROOT"
|
||||
echo
|
||||
|
||||
declare -a WORKER_PIDS=()
|
||||
for ((slot = 0; slot < ${#GPU_IDS[@]}; slot++)); do
|
||||
worker "$slot" &
|
||||
WORKER_PIDS+=("$!")
|
||||
done
|
||||
|
||||
overall_status=0
|
||||
for pid in "${WORKER_PIDS[@]}"; do
|
||||
wait "$pid" || overall_status=1
|
||||
done
|
||||
|
||||
if ((overall_status != 0)); then
|
||||
echo "One or more training tasks failed. Inspect logs under: $LOG_ROOT" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
if ((DRY_RUN)); then
|
||||
echo "Dry run completed successfully."
|
||||
else
|
||||
echo "All disease-history ablation tasks completed successfully."
|
||||
fi
|
||||
318
train_extra_info_assessment_all_multiseed_linux.sh
Executable file
318
train_extra_info_assessment_all_multiseed_linux.sh
Executable file
@@ -0,0 +1,318 @@
|
||||
#!/usr/bin/env bash
|
||||
#
|
||||
# Train the assessment + smoking + alcohol extra-information experiment.
|
||||
#
|
||||
# Fixed model:
|
||||
# TrajMixer + all_future + relative + Weibull + timed disease history + sex
|
||||
#
|
||||
# Extra information:
|
||||
# - 65 routine assessment/body/laboratory variables
|
||||
# - smoking
|
||||
# - alcohol
|
||||
# - BMI is already included in the assessment variables
|
||||
# - continuous values use train-split RobustScaler statistics
|
||||
#
|
||||
# A6000 48 GB default:
|
||||
# batch_size=256
|
||||
#
|
||||
# Each task uses one GPU. Tasks assigned to the same GPU run sequentially;
|
||||
# different GPUs run in parallel.
|
||||
#
|
||||
# Examples:
|
||||
# bash train_extra_info_assessment_all_multiseed_linux.sh --gpus 0
|
||||
# bash train_extra_info_assessment_all_multiseed_linux.sh --gpus 0,1,2
|
||||
# bash train_extra_info_assessment_all_multiseed_linux.sh \
|
||||
# --gpus 0 --seeds 42 --dry-run
|
||||
#
|
||||
|
||||
set -uo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
|
||||
cd "$SCRIPT_DIR"
|
||||
|
||||
GPU_CSV=""
|
||||
SEED_CSV="42,43,44"
|
||||
NUM_WORKERS=4
|
||||
BATCH_SIZE=256
|
||||
PYTHON_BIN="${PYTHON_BIN:-python}"
|
||||
CAMPAIGN_NAME="extra_info_assessment_smoking_alcohol_robust_multiseed"
|
||||
DRY_RUN=0
|
||||
|
||||
ENTRYPOINT="$SCRIPT_DIR/train_all_future.py"
|
||||
EXTRA_INFO_TYPES_FILE="$SCRIPT_DIR/extra_info_types_assessment_smoking_alcohol.txt"
|
||||
|
||||
usage() {
|
||||
cat <<'EOF'
|
||||
Usage:
|
||||
bash train_extra_info_assessment_all_multiseed_linux.sh \
|
||||
--gpus GPU_LIST [options]
|
||||
|
||||
Required:
|
||||
--gpus LIST Comma-separated GPU ids, for example 0,1,2.
|
||||
|
||||
Options:
|
||||
--seeds LIST Comma-separated seeds (default: 42,43,44).
|
||||
--batch-size N Batch size per task (default: 256).
|
||||
--num-workers N DataLoader workers per task (default: 4).
|
||||
--python PATH Python executable (default: $PYTHON_BIN or python).
|
||||
--campaign NAME Output campaign name.
|
||||
--dry-run Print commands without creating files or training.
|
||||
-h, --help Show this help message.
|
||||
|
||||
Fixed experiment settings:
|
||||
architecture traj_mixer_v5
|
||||
target all_future
|
||||
time mode relative
|
||||
distribution weibull
|
||||
disease history timed
|
||||
sex enabled by the model
|
||||
extra information extra_info_types_assessment_smoking_alcohol.txt
|
||||
continuous scaling robust (training-subset median/IQR)
|
||||
|
||||
A6000 48 GB default:
|
||||
batch_size 256
|
||||
|
||||
Outputs:
|
||||
runs/<campaign>/seed_<seed>/traj_mixer_v5/...
|
||||
batch_logs/<campaign>/seed_<seed>/assessment_smoking_alcohol.log
|
||||
EOF
|
||||
}
|
||||
|
||||
while (($# > 0)); do
|
||||
case "$1" in
|
||||
--gpus)
|
||||
[[ $# -ge 2 ]] || {
|
||||
echo "ERROR: --gpus requires a value." >&2
|
||||
exit 2
|
||||
}
|
||||
GPU_CSV="$2"
|
||||
shift 2
|
||||
;;
|
||||
--seeds)
|
||||
[[ $# -ge 2 ]] || {
|
||||
echo "ERROR: --seeds requires a value." >&2
|
||||
exit 2
|
||||
}
|
||||
SEED_CSV="$2"
|
||||
shift 2
|
||||
;;
|
||||
--batch-size)
|
||||
[[ $# -ge 2 ]] || {
|
||||
echo "ERROR: --batch-size requires a value." >&2
|
||||
exit 2
|
||||
}
|
||||
BATCH_SIZE="$2"
|
||||
shift 2
|
||||
;;
|
||||
--num-workers)
|
||||
[[ $# -ge 2 ]] || {
|
||||
echo "ERROR: --num-workers requires a value." >&2
|
||||
exit 2
|
||||
}
|
||||
NUM_WORKERS="$2"
|
||||
shift 2
|
||||
;;
|
||||
--python)
|
||||
[[ $# -ge 2 ]] || {
|
||||
echo "ERROR: --python requires a value." >&2
|
||||
exit 2
|
||||
}
|
||||
PYTHON_BIN="$2"
|
||||
shift 2
|
||||
;;
|
||||
--campaign)
|
||||
[[ $# -ge 2 ]] || {
|
||||
echo "ERROR: --campaign requires a value." >&2
|
||||
exit 2
|
||||
}
|
||||
CAMPAIGN_NAME="$2"
|
||||
shift 2
|
||||
;;
|
||||
--dry-run)
|
||||
DRY_RUN=1
|
||||
shift
|
||||
;;
|
||||
-h|--help)
|
||||
usage
|
||||
exit 0
|
||||
;;
|
||||
*)
|
||||
echo "ERROR: unknown argument: $1" >&2
|
||||
usage >&2
|
||||
exit 2
|
||||
;;
|
||||
esac
|
||||
done
|
||||
|
||||
[[ -n "$GPU_CSV" ]] || {
|
||||
echo "ERROR: --gpus is required." >&2
|
||||
usage >&2
|
||||
exit 2
|
||||
}
|
||||
[[ -n "$SEED_CSV" ]] || {
|
||||
echo "ERROR: --seeds must not be empty." >&2
|
||||
exit 2
|
||||
}
|
||||
[[ "$BATCH_SIZE" =~ ^[1-9][0-9]*$ ]] || {
|
||||
echo "ERROR: --batch-size must be a positive integer." >&2
|
||||
exit 2
|
||||
}
|
||||
[[ "$NUM_WORKERS" =~ ^[0-9]+$ ]] || {
|
||||
echo "ERROR: --num-workers must be a non-negative integer." >&2
|
||||
exit 2
|
||||
}
|
||||
[[ "$CAMPAIGN_NAME" =~ ^[A-Za-z0-9._-]+$ ]] || {
|
||||
echo "ERROR: --campaign may contain only letters, numbers, ., _, and -." >&2
|
||||
exit 2
|
||||
}
|
||||
[[ -f "$ENTRYPOINT" ]] || {
|
||||
echo "ERROR: missing training entrypoint: $ENTRYPOINT" >&2
|
||||
exit 2
|
||||
}
|
||||
[[ -f "$EXTRA_INFO_TYPES_FILE" ]] || {
|
||||
echo "ERROR: missing extra-info file: $EXTRA_INFO_TYPES_FILE" >&2
|
||||
exit 2
|
||||
}
|
||||
command -v "$PYTHON_BIN" >/dev/null 2>&1 || {
|
||||
echo "ERROR: Python executable not found: $PYTHON_BIN" >&2
|
||||
exit 2
|
||||
}
|
||||
|
||||
IFS=',' read -r -a GPU_IDS <<< "$GPU_CSV"
|
||||
declare -A SEEN_GPUS=()
|
||||
for gpu in "${GPU_IDS[@]}"; do
|
||||
[[ -n "$gpu" && "$gpu" =~ ^[A-Za-z0-9._:-]+$ ]] || {
|
||||
echo "ERROR: invalid GPU id: $gpu" >&2
|
||||
exit 2
|
||||
}
|
||||
[[ -z "${SEEN_GPUS[$gpu]+x}" ]] || {
|
||||
echo "ERROR: duplicate GPU id: $gpu" >&2
|
||||
exit 2
|
||||
}
|
||||
SEEN_GPUS["$gpu"]=1
|
||||
done
|
||||
|
||||
IFS=',' read -r -a SEEDS <<< "$SEED_CSV"
|
||||
declare -A SEEN_SEEDS=()
|
||||
for seed in "${SEEDS[@]}"; do
|
||||
[[ "$seed" =~ ^[0-9]+$ ]] || {
|
||||
echo "ERROR: invalid seed: $seed" >&2
|
||||
exit 2
|
||||
}
|
||||
[[ -z "${SEEN_SEEDS[$seed]+x}" ]] || {
|
||||
echo "ERROR: duplicate seed: $seed" >&2
|
||||
exit 2
|
||||
}
|
||||
SEEN_SEEDS["$seed"]=1
|
||||
done
|
||||
|
||||
RUNS_ROOT="$SCRIPT_DIR/runs/$CAMPAIGN_NAME"
|
||||
LOG_ROOT="$SCRIPT_DIR/batch_logs/$CAMPAIGN_NAME"
|
||||
if ((!DRY_RUN)); then
|
||||
mkdir -p "$RUNS_ROOT" "$LOG_ROOT"
|
||||
fi
|
||||
|
||||
print_command() {
|
||||
printf '%q ' "$@"
|
||||
printf '\n'
|
||||
}
|
||||
|
||||
run_job() {
|
||||
local seed="$1"
|
||||
local gpu="$2"
|
||||
local seed_runs_root="$RUNS_ROOT/seed_$seed"
|
||||
local seed_log_root="$LOG_ROOT/seed_$seed"
|
||||
local log_file="$seed_log_root/assessment_smoking_alcohol.log"
|
||||
local -a command=(
|
||||
"$PYTHON_BIN"
|
||||
-u
|
||||
"$ENTRYPOINT"
|
||||
--runs_root "$seed_runs_root"
|
||||
--seed "$seed"
|
||||
--batch_size "$BATCH_SIZE"
|
||||
--num_workers "$NUM_WORKERS"
|
||||
--device cuda
|
||||
--model_architecture traj_mixer_v5
|
||||
--time_mode relative
|
||||
--dist_mode weibull
|
||||
--disease_history_mode timed
|
||||
--extra_info_types_file "$EXTRA_INFO_TYPES_FILE"
|
||||
)
|
||||
|
||||
if ((!DRY_RUN)); then
|
||||
mkdir -p "$seed_runs_root" "$seed_log_root"
|
||||
fi
|
||||
|
||||
echo "[$(date '+%F %T')] START seed=$seed gpu=$gpu batch=$BATCH_SIZE"
|
||||
echo " log=$log_file"
|
||||
if ((DRY_RUN)); then
|
||||
printf ' CUDA_VISIBLE_DEVICES=%q ' "$gpu"
|
||||
print_command "${command[@]}"
|
||||
return 0
|
||||
fi
|
||||
|
||||
if CUDA_VISIBLE_DEVICES="$gpu" PYTHONUNBUFFERED=1 \
|
||||
"${command[@]}" >"$log_file" 2>&1; then
|
||||
echo "[$(date '+%F %T')] DONE seed=$seed gpu=$gpu"
|
||||
return 0
|
||||
else
|
||||
local exit_code=$?
|
||||
echo "[$(date '+%F %T')] FAIL seed=$seed gpu=$gpu exit=$exit_code" >&2
|
||||
echo " See: $log_file" >&2
|
||||
return "$exit_code"
|
||||
fi
|
||||
}
|
||||
|
||||
worker() {
|
||||
local slot="$1"
|
||||
local gpu="${GPU_IDS[$slot]}"
|
||||
local seed_index
|
||||
local failed=0
|
||||
|
||||
for ((seed_index = slot; seed_index < ${#SEEDS[@]}; seed_index += ${#GPU_IDS[@]})); do
|
||||
run_job "${SEEDS[$seed_index]}" "$gpu" || failed=1
|
||||
done
|
||||
return "$failed"
|
||||
}
|
||||
|
||||
echo "Campaign: $CAMPAIGN_NAME"
|
||||
echo "Seeds: ${SEEDS[*]}"
|
||||
echo "GPUs: ${GPU_IDS[*]}"
|
||||
echo "Batch size: $BATCH_SIZE"
|
||||
echo "Extra-info file: $EXTRA_INFO_TYPES_FILE"
|
||||
echo "Total tasks: ${#SEEDS[@]}"
|
||||
echo "Runs root: $RUNS_ROOT"
|
||||
echo "Log root: $LOG_ROOT"
|
||||
if command -v nvidia-smi >/dev/null 2>&1; then
|
||||
echo "Selected GPU inventory:"
|
||||
for gpu in "${GPU_IDS[@]}"; do
|
||||
nvidia-smi \
|
||||
--id="$gpu" \
|
||||
--query-gpu=index,name,memory.total \
|
||||
--format=csv,noheader \
|
||||
2>/dev/null || true
|
||||
done
|
||||
fi
|
||||
echo
|
||||
|
||||
declare -a WORKER_PIDS=()
|
||||
for ((slot = 0; slot < ${#GPU_IDS[@]}; slot++)); do
|
||||
worker "$slot" &
|
||||
WORKER_PIDS+=("$!")
|
||||
done
|
||||
|
||||
overall_status=0
|
||||
for pid in "${WORKER_PIDS[@]}"; do
|
||||
wait "$pid" || overall_status=1
|
||||
done
|
||||
|
||||
if ((overall_status != 0)); then
|
||||
echo "One or more training tasks failed. Inspect logs under: $LOG_ROOT" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
if ((DRY_RUN)); then
|
||||
echo "Dry run completed successfully."
|
||||
else
|
||||
echo "All assessment + smoking + alcohol tasks completed successfully."
|
||||
fi
|
||||
358
train_key_models_multiseed_linux.sh
Normal file
358
train_key_models_multiseed_linux.sh
Normal file
@@ -0,0 +1,358 @@
|
||||
#!/usr/bin/env bash
|
||||
#
|
||||
# Re-run the key DeepHealth experiments with additional random seeds.
|
||||
#
|
||||
# The selected configurations form a compact evidence chain:
|
||||
# 1. FFN + next_token + absolute + exponential
|
||||
# Delphi2M reproduction baseline.
|
||||
# 2. FFN + all_future + absolute + exponential
|
||||
# Isolates the target change from next_token to all_future.
|
||||
# 3. TrajMixer + all_future + absolute + exponential
|
||||
# Isolates the architecture change from FFN to TrajMixer.
|
||||
# 4. TrajMixer + all_future + relative + exponential
|
||||
# Isolates the time-mode change and is the current disease-best candidate.
|
||||
# 5. TrajMixer + all_future + relative + Weibull
|
||||
# Unified disease/death candidate.
|
||||
# 6. FFN + all_future + relative + Weibull
|
||||
# Current mortality-best control and matched architecture comparison for (5).
|
||||
#
|
||||
# Each task uses one GPU. Tasks assigned to the same GPU run sequentially,
|
||||
# while different GPUs run in parallel.
|
||||
#
|
||||
# Examples:
|
||||
# bash train_key_models_multiseed_linux.sh --gpus 0,1,2,3
|
||||
# bash train_key_models_multiseed_linux.sh --gpus 0,1 --seeds 43,44
|
||||
#
|
||||
|
||||
set -uo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
|
||||
cd "$SCRIPT_DIR"
|
||||
|
||||
GPU_CSV=""
|
||||
SEED_CSV="43,44,45"
|
||||
NUM_WORKERS=4
|
||||
PYTHON_BIN="${PYTHON_BIN:-python}"
|
||||
CAMPAIGN_NAME="key_models_multiseed_smoking_alcohol_bmi"
|
||||
DRY_RUN=0
|
||||
|
||||
BATCH_SIZE=256
|
||||
EXTRA_INFO_TYPES_FILE="$SCRIPT_DIR/extra_info_types_smoking_alcohol_bmi.txt"
|
||||
|
||||
usage() {
|
||||
cat <<'EOF'
|
||||
Usage:
|
||||
bash train_key_models_multiseed_linux.sh --gpus GPU_LIST [options]
|
||||
|
||||
Required:
|
||||
--gpus LIST Comma-separated GPU ids, for example 0,1,2,3.
|
||||
|
||||
Options:
|
||||
--seeds LIST Additional seeds (default: 43,44,45).
|
||||
Use 43,44 for two additional seeds.
|
||||
--num-workers N DataLoader workers per task (default: 4).
|
||||
--python PATH Python executable (default: $PYTHON_BIN or python).
|
||||
--campaign NAME Output campaign name.
|
||||
--dry-run Print all commands without running them.
|
||||
-h, --help Show this help message.
|
||||
|
||||
Fixed experiment settings:
|
||||
batch_size 256
|
||||
extra_info_types extra_info_types_smoking_alcohol_bmi.txt
|
||||
continuous scaling required train-split RobustScale
|
||||
model size Defaults from the training entrypoints
|
||||
tasks per seed 6
|
||||
|
||||
Outputs:
|
||||
runs/<campaign>/seed_<seed>/<architecture>/...
|
||||
batch_logs/<campaign>/seed_<seed>/<job>.log
|
||||
EOF
|
||||
}
|
||||
|
||||
while (($# > 0)); do
|
||||
case "$1" in
|
||||
--gpus)
|
||||
[[ $# -ge 2 ]] || {
|
||||
echo "ERROR: --gpus requires a value." >&2
|
||||
exit 2
|
||||
}
|
||||
GPU_CSV="$2"
|
||||
shift 2
|
||||
;;
|
||||
--seeds)
|
||||
[[ $# -ge 2 ]] || {
|
||||
echo "ERROR: --seeds requires a value." >&2
|
||||
exit 2
|
||||
}
|
||||
SEED_CSV="$2"
|
||||
shift 2
|
||||
;;
|
||||
--num-workers)
|
||||
[[ $# -ge 2 ]] || {
|
||||
echo "ERROR: --num-workers requires a value." >&2
|
||||
exit 2
|
||||
}
|
||||
NUM_WORKERS="$2"
|
||||
shift 2
|
||||
;;
|
||||
--python)
|
||||
[[ $# -ge 2 ]] || {
|
||||
echo "ERROR: --python requires a value." >&2
|
||||
exit 2
|
||||
}
|
||||
PYTHON_BIN="$2"
|
||||
shift 2
|
||||
;;
|
||||
--campaign)
|
||||
[[ $# -ge 2 ]] || {
|
||||
echo "ERROR: --campaign requires a value." >&2
|
||||
exit 2
|
||||
}
|
||||
CAMPAIGN_NAME="$2"
|
||||
shift 2
|
||||
;;
|
||||
--dry-run)
|
||||
DRY_RUN=1
|
||||
shift
|
||||
;;
|
||||
-h|--help)
|
||||
usage
|
||||
exit 0
|
||||
;;
|
||||
*)
|
||||
echo "ERROR: unknown argument: $1" >&2
|
||||
usage >&2
|
||||
exit 2
|
||||
;;
|
||||
esac
|
||||
done
|
||||
|
||||
[[ -n "$GPU_CSV" ]] || {
|
||||
echo "ERROR: --gpus is required." >&2
|
||||
usage >&2
|
||||
exit 2
|
||||
}
|
||||
[[ -n "$SEED_CSV" ]] || {
|
||||
echo "ERROR: --seeds must not be empty." >&2
|
||||
exit 2
|
||||
}
|
||||
[[ "$NUM_WORKERS" =~ ^[0-9]+$ ]] || {
|
||||
echo "ERROR: --num-workers must be a non-negative integer." >&2
|
||||
exit 2
|
||||
}
|
||||
[[ "$CAMPAIGN_NAME" =~ ^[A-Za-z0-9._-]+$ ]] || {
|
||||
echo "ERROR: --campaign may contain only letters, numbers, ., _, and -." >&2
|
||||
exit 2
|
||||
}
|
||||
[[ -f "$EXTRA_INFO_TYPES_FILE" ]] || {
|
||||
echo "ERROR: missing extra-info file: $EXTRA_INFO_TYPES_FILE" >&2
|
||||
exit 2
|
||||
}
|
||||
command -v "$PYTHON_BIN" >/dev/null 2>&1 || {
|
||||
echo "ERROR: Python executable not found: $PYTHON_BIN" >&2
|
||||
exit 2
|
||||
}
|
||||
|
||||
IFS=',' read -r -a GPU_IDS <<< "$GPU_CSV"
|
||||
declare -A SEEN_GPUS=()
|
||||
for gpu in "${GPU_IDS[@]}"; do
|
||||
[[ -n "$gpu" && "$gpu" =~ ^[A-Za-z0-9._:-]+$ ]] || {
|
||||
echo "ERROR: invalid GPU id: $gpu" >&2
|
||||
exit 2
|
||||
}
|
||||
[[ -z "${SEEN_GPUS[$gpu]+x}" ]] || {
|
||||
echo "ERROR: duplicate GPU id: $gpu" >&2
|
||||
exit 2
|
||||
}
|
||||
SEEN_GPUS["$gpu"]=1
|
||||
done
|
||||
|
||||
IFS=',' read -r -a SEEDS <<< "$SEED_CSV"
|
||||
declare -A SEEN_SEEDS=()
|
||||
for seed in "${SEEDS[@]}"; do
|
||||
[[ "$seed" =~ ^[0-9]+$ ]] || {
|
||||
echo "ERROR: invalid seed: $seed" >&2
|
||||
exit 2
|
||||
}
|
||||
[[ -z "${SEEN_SEEDS[$seed]+x}" ]] || {
|
||||
echo "ERROR: duplicate seed: $seed" >&2
|
||||
exit 2
|
||||
}
|
||||
SEEN_SEEDS["$seed"]=1
|
||||
done
|
||||
|
||||
RUNS_ROOT="$SCRIPT_DIR/runs/$CAMPAIGN_NAME"
|
||||
LOG_ROOT="$SCRIPT_DIR/batch_logs/$CAMPAIGN_NAME"
|
||||
if ((!DRY_RUN)); then
|
||||
mkdir -p "$RUNS_ROOT" "$LOG_ROOT"
|
||||
fi
|
||||
|
||||
declare -a JOB_NAMES=()
|
||||
declare -a JOB_SEEDS=()
|
||||
declare -a JOB_ENTRYPOINTS=()
|
||||
declare -a JOB_ARCHITECTURES=()
|
||||
declare -a JOB_TIME_MODES=()
|
||||
declare -a JOB_DIST_MODES=()
|
||||
|
||||
add_job() {
|
||||
JOB_NAMES+=("$1")
|
||||
JOB_SEEDS+=("$2")
|
||||
JOB_ENTRYPOINTS+=("$3")
|
||||
JOB_ARCHITECTURES+=("$4")
|
||||
JOB_TIME_MODES+=("$5")
|
||||
JOB_DIST_MODES+=("$6")
|
||||
}
|
||||
|
||||
for seed in "${SEEDS[@]}"; do
|
||||
add_job \
|
||||
"ffn_next_token_absolute_exponential" \
|
||||
"$seed" \
|
||||
"train_next_step.py" \
|
||||
"transformer_ffn_v1" \
|
||||
"absolute" \
|
||||
"exponential"
|
||||
|
||||
add_job \
|
||||
"ffn_all_future_absolute_exponential" \
|
||||
"$seed" \
|
||||
"train_all_future.py" \
|
||||
"transformer_ffn_v1" \
|
||||
"absolute" \
|
||||
"exponential"
|
||||
|
||||
add_job \
|
||||
"traj_mixer_all_future_absolute_exponential" \
|
||||
"$seed" \
|
||||
"train_all_future.py" \
|
||||
"traj_mixer_v5" \
|
||||
"absolute" \
|
||||
"exponential"
|
||||
|
||||
add_job \
|
||||
"traj_mixer_all_future_relative_exponential" \
|
||||
"$seed" \
|
||||
"train_all_future.py" \
|
||||
"traj_mixer_v5" \
|
||||
"relative" \
|
||||
"exponential"
|
||||
|
||||
add_job \
|
||||
"traj_mixer_all_future_relative_weibull" \
|
||||
"$seed" \
|
||||
"train_all_future.py" \
|
||||
"traj_mixer_v5" \
|
||||
"relative" \
|
||||
"weibull"
|
||||
|
||||
add_job \
|
||||
"ffn_all_future_relative_weibull" \
|
||||
"$seed" \
|
||||
"train_all_future.py" \
|
||||
"transformer_ffn_v1" \
|
||||
"relative" \
|
||||
"weibull"
|
||||
done
|
||||
|
||||
print_command() {
|
||||
printf '%q ' "$@"
|
||||
printf '\n'
|
||||
}
|
||||
|
||||
run_job() {
|
||||
local job_index="$1"
|
||||
local gpu="$2"
|
||||
local job_name="${JOB_NAMES[$job_index]}"
|
||||
local seed="${JOB_SEEDS[$job_index]}"
|
||||
local entrypoint="${JOB_ENTRYPOINTS[$job_index]}"
|
||||
local architecture="${JOB_ARCHITECTURES[$job_index]}"
|
||||
local time_mode="${JOB_TIME_MODES[$job_index]}"
|
||||
local dist_mode="${JOB_DIST_MODES[$job_index]}"
|
||||
local seed_runs_root="$RUNS_ROOT/seed_$seed"
|
||||
local seed_log_root="$LOG_ROOT/seed_$seed"
|
||||
local log_file="$seed_log_root/$job_name.log"
|
||||
local -a command=(
|
||||
"$PYTHON_BIN"
|
||||
-u
|
||||
"$SCRIPT_DIR/$entrypoint"
|
||||
--runs_root "$seed_runs_root"
|
||||
--seed "$seed"
|
||||
--batch_size "$BATCH_SIZE"
|
||||
--extra_info_types_file "$EXTRA_INFO_TYPES_FILE"
|
||||
--model_architecture "$architecture"
|
||||
--num_workers "$NUM_WORKERS"
|
||||
--device cuda
|
||||
)
|
||||
|
||||
if ((!DRY_RUN)); then
|
||||
mkdir -p "$seed_runs_root" "$seed_log_root"
|
||||
fi
|
||||
|
||||
if [[ "$entrypoint" == "train_all_future.py" ]]; then
|
||||
command+=(
|
||||
--time_mode "$time_mode"
|
||||
--dist_mode "$dist_mode"
|
||||
)
|
||||
fi
|
||||
|
||||
echo "[$(date '+%F %T')] START seed=$seed job=$job_name gpu=$gpu"
|
||||
echo " log=$log_file"
|
||||
if ((DRY_RUN)); then
|
||||
printf ' CUDA_VISIBLE_DEVICES=%q ' "$gpu"
|
||||
print_command "${command[@]}"
|
||||
return 0
|
||||
fi
|
||||
|
||||
if CUDA_VISIBLE_DEVICES="$gpu" PYTHONUNBUFFERED=1 \
|
||||
"${command[@]}" >"$log_file" 2>&1; then
|
||||
echo "[$(date '+%F %T')] DONE seed=$seed job=$job_name gpu=$gpu"
|
||||
return 0
|
||||
else
|
||||
local exit_code=$?
|
||||
echo "[$(date '+%F %T')] FAIL seed=$seed job=$job_name gpu=$gpu exit=$exit_code" >&2
|
||||
echo " See: $log_file" >&2
|
||||
return "$exit_code"
|
||||
fi
|
||||
}
|
||||
|
||||
worker() {
|
||||
local slot="$1"
|
||||
local gpu="${GPU_IDS[$slot]}"
|
||||
local job_index
|
||||
local failed=0
|
||||
|
||||
for ((job_index = slot; job_index < ${#JOB_NAMES[@]}; job_index += ${#GPU_IDS[@]})); do
|
||||
run_job "$job_index" "$gpu" || failed=1
|
||||
done
|
||||
return "$failed"
|
||||
}
|
||||
|
||||
echo "Campaign: $CAMPAIGN_NAME"
|
||||
echo "Additional seeds: ${SEEDS[*]}"
|
||||
echo "GPUs: ${GPU_IDS[*]}"
|
||||
echo "Configurations per seed: 6"
|
||||
echo "Total tasks: ${#JOB_NAMES[@]}"
|
||||
echo "Runs root: $RUNS_ROOT"
|
||||
echo "Log root: $LOG_ROOT"
|
||||
echo
|
||||
|
||||
declare -a WORKER_PIDS=()
|
||||
for ((slot = 0; slot < ${#GPU_IDS[@]}; slot++)); do
|
||||
worker "$slot" &
|
||||
WORKER_PIDS+=("$!")
|
||||
done
|
||||
|
||||
overall_status=0
|
||||
for pid in "${WORKER_PIDS[@]}"; do
|
||||
wait "$pid" || overall_status=1
|
||||
done
|
||||
|
||||
if ((overall_status != 0)); then
|
||||
echo "One or more training tasks failed. Inspect logs under: $LOG_ROOT" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
if ((DRY_RUN)); then
|
||||
echo "Dry run completed successfully."
|
||||
else
|
||||
echo "All training tasks completed successfully."
|
||||
fi
|
||||
@@ -1,10 +1,4 @@
|
||||
"""
|
||||
Train DeepHealth with next-token / next-time-point supervision.
|
||||
|
||||
The next-step dataset uses observed event histories, including CHECKUP state
|
||||
tokens, plus optional gap <NO_EVENT> imputation. UTS training reads out only
|
||||
same-time group ends.
|
||||
"""
|
||||
"""Reproduce Delphi2M with absolute-time next-token supervision."""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
@@ -24,21 +18,22 @@ from tqdm.auto import tqdm
|
||||
|
||||
from dataset import HealthDataset, collate_fn
|
||||
from losses import build_loss
|
||||
from models import (
|
||||
EVENT_TRAJECTORY_ARCHITECTURE,
|
||||
MODEL_SIZE_NAMES,
|
||||
DeepHealth,
|
||||
DeepHealthOutput,
|
||||
resolve_model_size,
|
||||
from model_architectures import (
|
||||
DEFAULT_MODEL_ARCHITECTURE,
|
||||
SUPPORTED_MODEL_ARCHITECTURES,
|
||||
)
|
||||
from readouts import build_readout
|
||||
from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX
|
||||
from models import DeepHealth, DeepHealthOutput
|
||||
from targets import PAD_IDX, RESERVED_IDX
|
||||
from train_util import (
|
||||
ContinuousRobustScalerStats,
|
||||
configure_torch_for_training,
|
||||
create_unique_run_dir,
|
||||
fit_continuous_robust_scaler,
|
||||
format_extra_info_types,
|
||||
get_lr,
|
||||
get_model_parameter_counts,
|
||||
load_extra_info_types_file,
|
||||
move_batch_to_device,
|
||||
resolve_device,
|
||||
save_checkpoint,
|
||||
save_config,
|
||||
@@ -68,10 +63,10 @@ def parse_args() -> argparse.Namespace:
|
||||
|
||||
parser.add_argument("--data_prefix", type=str, default="ukb")
|
||||
parser.add_argument("--labels_file", type=str, default="labels.csv")
|
||||
parser.add_argument("--runs_root", type=str, default="runs")
|
||||
parser.add_argument("--seed", type=int, default=42)
|
||||
parser.add_argument("--extra_info_types_file", type=str, default=None)
|
||||
parser.add_argument("--no_event_interval_years", type=float, default=5.0)
|
||||
parser.add_argument("--include_no_event_in_uts_target", action="store_true")
|
||||
|
||||
parser.add_argument("--train_ratio", type=float, default=0.7)
|
||||
parser.add_argument("--val_ratio", type=float, default=0.15)
|
||||
@@ -80,31 +75,24 @@ def parse_args() -> argparse.Namespace:
|
||||
parser.add_argument("--val_eid_file", type=str, default="ukb_val_eid.csv")
|
||||
parser.add_argument("--test_eid_file", type=str, default="ukb_test_eid.csv")
|
||||
|
||||
parser.add_argument(
|
||||
"--model_size",
|
||||
type=str,
|
||||
default="nano",
|
||||
choices=MODEL_SIZE_NAMES,
|
||||
)
|
||||
parser.add_argument("--n_reasoning_rounds", type=int, default=12)
|
||||
parser.add_argument("--n_embd", type=int, default=120)
|
||||
parser.add_argument("--n_head", type=int, default=10)
|
||||
parser.add_argument("--n_layer", type=int, default=12)
|
||||
parser.add_argument("--n_bins", type=int, default=16)
|
||||
parser.add_argument("--extra_pool_reduce", type=str, default="mean",
|
||||
choices=["mean", "sum"])
|
||||
parser.add_argument("--time_mode", type=str, default="relative",
|
||||
choices=["relative", "absolute"])
|
||||
parser.add_argument("--dropout", type=float, default=0.0)
|
||||
parser.add_argument(
|
||||
"--model_architecture",
|
||||
type=str,
|
||||
default=DEFAULT_MODEL_ARCHITECTURE,
|
||||
choices=SUPPORTED_MODEL_ARCHITECTURES,
|
||||
)
|
||||
|
||||
parser.add_argument("--target_mode", type=str, default="uts",
|
||||
choices=["delphi2m", "uts"])
|
||||
parser.add_argument("--readout_name", type=str, default=None,
|
||||
choices=["token", "same_time_group_end", "last_valid"])
|
||||
parser.add_argument("--readout_reduce", type=str, default="mean",
|
||||
choices=["mean", "sum"])
|
||||
parser.add_argument("--t_min", type=float, default=0.0027378507871321013)
|
||||
parser.add_argument("--max_exp_input", type=float, default=60.0)
|
||||
parser.add_argument("--ce_weight", type=float, default=1.0)
|
||||
parser.add_argument("--time_weight", type=float, default=1.0)
|
||||
parser.add_argument("--ignore_no_event_in_delphi2m", action="store_true")
|
||||
|
||||
parser.add_argument("--batch_size", type=int, default=128)
|
||||
parser.add_argument("--base_lr", type=float, default=3e-4)
|
||||
@@ -126,11 +114,6 @@ def parse_args() -> argparse.Namespace:
|
||||
)
|
||||
if not use_eid_split and not np.isclose(args.train_ratio + args.val_ratio + args.test_ratio, 1.0):
|
||||
raise ValueError("train_ratio + val_ratio + test_ratio must equal 1.0")
|
||||
if args.target_mode == "uts":
|
||||
args.readout_name = args.readout_name or "same_time_group_end"
|
||||
args.include_no_event_in_uts_target = True
|
||||
else:
|
||||
args.readout_name = args.readout_name or "token"
|
||||
args.extra_info_types = (
|
||||
load_extra_info_types_file(args.extra_info_types_file)
|
||||
if args.extra_info_types_file is not None
|
||||
@@ -139,73 +122,52 @@ def parse_args() -> argparse.Namespace:
|
||||
return args
|
||||
|
||||
|
||||
def get_lr(epoch: int, args: argparse.Namespace, adaptive_lr: float) -> float:
|
||||
if epoch < args.warmup_epochs:
|
||||
return adaptive_lr * (epoch + 1) / args.warmup_epochs
|
||||
progress = (epoch - args.warmup_epochs) / max(1, args.max_epochs - args.warmup_epochs)
|
||||
cosine = 0.5 * (1 + math.cos(math.pi * progress))
|
||||
return adaptive_lr * (args.min_lr_ratio + cosine * (1 - args.min_lr_ratio))
|
||||
|
||||
|
||||
def move_batch_to_device(batch: Dict[str, torch.Tensor], device: torch.device) -> Dict[str, torch.Tensor]:
|
||||
non_blocking = device.type == "cuda"
|
||||
return {
|
||||
key: value.to(device, non_blocking=non_blocking)
|
||||
if isinstance(value, torch.Tensor)
|
||||
else value
|
||||
for key, value in batch.items()
|
||||
}
|
||||
|
||||
|
||||
def build_model(args: argparse.Namespace, dataset: HealthDataset) -> DeepHealth:
|
||||
def build_model(
|
||||
args: argparse.Namespace,
|
||||
dataset: HealthDataset,
|
||||
scaler_stats: ContinuousRobustScalerStats,
|
||||
) -> DeepHealth:
|
||||
if tuple(int(x) for x in dataset.cont_type_ids) != scaler_stats.cont_type_ids:
|
||||
raise ValueError(
|
||||
"RobustScale statistics are not aligned with dataset.cont_type_ids"
|
||||
)
|
||||
center = scaler_stats.center if dataset.n_cont_types > 0 else None
|
||||
scale = scaler_stats.scale if dataset.n_cont_types > 0 else None
|
||||
return DeepHealth(
|
||||
vocab_size=dataset.vocab_size,
|
||||
model_size=args.model_size,
|
||||
n_reasoning_rounds=args.n_reasoning_rounds,
|
||||
n_embd=args.n_embd,
|
||||
n_head=args.n_head,
|
||||
n_layer=args.n_layer,
|
||||
n_types=dataset.n_types,
|
||||
n_cont_types=dataset.n_cont_types,
|
||||
n_categories=dataset.n_categories,
|
||||
cont_type_ids=dataset.cont_type_ids,
|
||||
n_bins=args.n_bins,
|
||||
continuous_value_center=center,
|
||||
continuous_value_scale=scale,
|
||||
extra_pool_reduce=args.extra_pool_reduce,
|
||||
target_mode="next_token",
|
||||
time_mode=args.time_mode,
|
||||
time_mode="absolute",
|
||||
dist_mode="exponential",
|
||||
dropout=args.dropout,
|
||||
model_architecture=args.model_architecture,
|
||||
)
|
||||
|
||||
|
||||
def build_next_step_readout(args: argparse.Namespace):
|
||||
if args.readout_name == "same_time_group_end":
|
||||
return build_readout("same_time_group_end", reduce=args.readout_reduce)
|
||||
return build_readout(args.readout_name)
|
||||
|
||||
|
||||
def build_next_step_loss(args: argparse.Namespace):
|
||||
if args.target_mode == "delphi2m":
|
||||
ignored_tokens = {PAD_IDX, CHECKUP_IDX}
|
||||
if args.ignore_no_event_in_delphi2m:
|
||||
ignored_tokens.add(NO_EVENT_IDX)
|
||||
return build_loss(
|
||||
"delphi2m",
|
||||
ignored_tokens=ignored_tokens,
|
||||
ignored_tokens={PAD_IDX, RESERVED_IDX},
|
||||
t_min=args.t_min,
|
||||
max_exp_input=args.max_exp_input,
|
||||
ce_weight=args.ce_weight,
|
||||
time_weight=args.time_weight,
|
||||
)
|
||||
return build_loss(
|
||||
"uts",
|
||||
ignored_idx={PAD_IDX, CHECKUP_IDX},
|
||||
t_min=args.t_min,
|
||||
max_exp_input=args.max_exp_input,
|
||||
)
|
||||
|
||||
|
||||
def build_augmented_next_step_targets(
|
||||
batch_cpu: Dict[str, torch.Tensor],
|
||||
model_out: DeepHealthOutput,
|
||||
include_uts_targets: bool,
|
||||
) -> Dict[str, torch.Tensor]:
|
||||
hidden_len = model_out.hidden.size(1)
|
||||
event_len = int(model_out.event_len)
|
||||
@@ -213,26 +175,12 @@ def build_augmented_next_step_targets(
|
||||
device = model_out.hidden.device
|
||||
non_blocking = device.type == "cuda"
|
||||
if extra_len <= 0:
|
||||
targets = {
|
||||
return {
|
||||
"target_event_seq": batch_cpu["target_event_seq"].to(device, non_blocking=non_blocking),
|
||||
"target_time_seq": batch_cpu["target_time_seq"].to(device, non_blocking=non_blocking),
|
||||
"readout_mask": batch_cpu["readout_mask"].to(device, non_blocking=non_blocking),
|
||||
}
|
||||
if include_uts_targets:
|
||||
targets["target_dt_unique"] = batch_cpu["target_dt_unique"].to(
|
||||
device, non_blocking=non_blocking
|
||||
)
|
||||
targets["target_multi_hot"] = batch_cpu["target_multi_hot"].to(
|
||||
device, non_blocking=non_blocking
|
||||
)
|
||||
return targets
|
||||
|
||||
bsz = batch_cpu["target_event_seq"].size(0)
|
||||
vocab_size = (
|
||||
batch_cpu["target_multi_hot"].size(2)
|
||||
if include_uts_targets
|
||||
else None
|
||||
)
|
||||
other_valid = batch_cpu["other_type"] > 0
|
||||
extra_time = batch_cpu["other_time"].new_zeros(bsz, extra_len)
|
||||
extra_mask = torch.zeros(bsz, extra_len, dtype=torch.bool)
|
||||
@@ -265,34 +213,6 @@ def build_augmented_next_step_targets(
|
||||
],
|
||||
dim=1,
|
||||
)
|
||||
readout_mask = torch.cat([batch_cpu["readout_mask"], extra_mask], dim=1)
|
||||
target_dt_unique = None
|
||||
target_multi_hot = None
|
||||
if include_uts_targets:
|
||||
target_dt_unique = torch.cat(
|
||||
[
|
||||
batch_cpu["target_dt_unique"],
|
||||
torch.zeros(
|
||||
bsz,
|
||||
extra_len,
|
||||
dtype=batch_cpu["target_dt_unique"].dtype,
|
||||
),
|
||||
],
|
||||
dim=1,
|
||||
)
|
||||
target_multi_hot = torch.cat(
|
||||
[
|
||||
batch_cpu["target_multi_hot"],
|
||||
torch.zeros(
|
||||
bsz,
|
||||
extra_len,
|
||||
vocab_size,
|
||||
dtype=batch_cpu["target_multi_hot"].dtype,
|
||||
),
|
||||
],
|
||||
dim=1,
|
||||
)
|
||||
|
||||
for b in range(bsz):
|
||||
valid_event = batch_cpu["padding_mask"][b].bool()
|
||||
if not valid_event.any():
|
||||
@@ -323,7 +243,6 @@ def build_augmented_next_step_targets(
|
||||
t = extra_time[b, j]
|
||||
future = times > t
|
||||
if not future.any():
|
||||
readout_mask[b, pos] = False
|
||||
continue
|
||||
|
||||
first_idx = int(torch.nonzero(future, as_tuple=False)[0].item())
|
||||
@@ -332,35 +251,14 @@ def build_augmented_next_step_targets(
|
||||
target_event_seq[b, pos] = next_event
|
||||
target_time_seq[b, pos] = next_time
|
||||
|
||||
if not include_uts_targets:
|
||||
continue
|
||||
|
||||
same_next_time = times == next_time
|
||||
next_events = events[same_next_time]
|
||||
valid_next_events = next_events[
|
||||
(next_events > PAD_IDX) & (next_events < vocab_size)
|
||||
].long()
|
||||
if valid_next_events.numel() == 0:
|
||||
readout_mask[b, pos] = False
|
||||
continue
|
||||
target_multi_hot[b, pos, valid_next_events] = True
|
||||
target_dt_unique[b, pos] = next_time - t
|
||||
|
||||
targets = {
|
||||
return {
|
||||
"target_event_seq": target_event_seq.to(device, non_blocking=non_blocking),
|
||||
"target_time_seq": target_time_seq.to(device, non_blocking=non_blocking),
|
||||
"readout_mask": readout_mask.to(device, non_blocking=non_blocking),
|
||||
}
|
||||
if include_uts_targets:
|
||||
targets["target_dt_unique"] = target_dt_unique.to(device, non_blocking=non_blocking)
|
||||
targets["target_multi_hot"] = target_multi_hot.to(device, non_blocking=non_blocking)
|
||||
return targets
|
||||
|
||||
|
||||
def compute_next_step_loss(
|
||||
args: argparse.Namespace,
|
||||
model: DeepHealth,
|
||||
readout,
|
||||
criterion,
|
||||
batch: Dict[str, torch.Tensor],
|
||||
device: torch.device,
|
||||
@@ -379,7 +277,6 @@ def compute_next_step_loss(
|
||||
other_value=batch["other_value"],
|
||||
other_value_kind=batch["other_value_kind"],
|
||||
other_time=batch["other_time"],
|
||||
target_mode="next_token",
|
||||
return_output=True,
|
||||
)
|
||||
if not isinstance(model_out, DeepHealthOutput):
|
||||
@@ -387,33 +284,15 @@ def compute_next_step_loss(
|
||||
targets = build_augmented_next_step_targets(
|
||||
batch_cpu=batch_cpu,
|
||||
model_out=model_out,
|
||||
include_uts_targets=args.target_mode == "uts",
|
||||
)
|
||||
readout_out = readout(
|
||||
hidden=model_out.hidden,
|
||||
time_seq=model_out.time_seq,
|
||||
padding_mask=model_out.padding_mask,
|
||||
readout_mask=targets["readout_mask"]
|
||||
if args.readout_name == "same_time_group_end"
|
||||
else None,
|
||||
)
|
||||
logits = model.calc_risk(readout_out.hidden)
|
||||
logits = model.calc_risk(model_out.hidden)
|
||||
|
||||
if args.target_mode == "delphi2m":
|
||||
loss, parts = criterion(
|
||||
logits=logits,
|
||||
target_events=targets["target_event_seq"],
|
||||
target_times=targets["target_time_seq"],
|
||||
current_times=model_out.time_seq,
|
||||
padding_mask=readout_out.readout_mask,
|
||||
return_components=True,
|
||||
)
|
||||
else:
|
||||
loss, parts = criterion(
|
||||
logits=logits,
|
||||
target_multi_hot=targets["target_multi_hot"],
|
||||
target_dt_unique=targets["target_dt_unique"],
|
||||
readout_mask=readout_out.readout_mask,
|
||||
padding_mask=model_out.padding_mask,
|
||||
return_components=True,
|
||||
)
|
||||
if not torch.isfinite(loss):
|
||||
@@ -425,7 +304,6 @@ def run_epoch(
|
||||
logger: logging.Logger,
|
||||
args: argparse.Namespace,
|
||||
model: DeepHealth,
|
||||
readout,
|
||||
criterion,
|
||||
loader: DataLoader,
|
||||
optimizer: AdamW | None,
|
||||
@@ -433,7 +311,6 @@ def run_epoch(
|
||||
is_train: bool,
|
||||
) -> float:
|
||||
model.train(is_train)
|
||||
readout.train(is_train)
|
||||
total = torch.zeros((), device=device)
|
||||
n_batches = 0
|
||||
skipped = 0
|
||||
@@ -444,7 +321,9 @@ def run_epoch(
|
||||
progress = tqdm(loader, desc=desc, leave=False, dynamic_ncols=True)
|
||||
for batch_idx, batch in enumerate(progress):
|
||||
try:
|
||||
loss, parts = compute_next_step_loss(args, model, readout, criterion, batch, device)
|
||||
loss, parts = compute_next_step_loss(
|
||||
model, criterion, batch, device
|
||||
)
|
||||
if is_train:
|
||||
if optimizer is None:
|
||||
raise ValueError("optimizer is required for training")
|
||||
@@ -486,20 +365,19 @@ def build_metadata(
|
||||
train_subset,
|
||||
val_subset,
|
||||
test_subset,
|
||||
scaler_stats: ContinuousRobustScalerStats,
|
||||
) -> Dict[str, Any]:
|
||||
size_config = resolve_model_size(args.model_size)
|
||||
return {
|
||||
"run_name": run_name,
|
||||
"dataset_class": "NextStepHealthDataset",
|
||||
"collate_fn": "next_step_collate_fn",
|
||||
"model_class": "DeepHealth",
|
||||
"model_architecture": EVENT_TRAJECTORY_ARCHITECTURE,
|
||||
"d_model": size_config.d_model,
|
||||
"n_trajectory": size_config.n_trajectory,
|
||||
"trajectory_dim": size_config.trajectory_dim,
|
||||
"traj_hidden": size_config.traj_hidden,
|
||||
"model_architecture": args.model_architecture,
|
||||
"model_target_mode": "next_token",
|
||||
"target_mode": args.target_mode,
|
||||
"target_mode": "delphi2m",
|
||||
"event_stream_version": "disease_death_only_v1",
|
||||
"uses_assessment_event_token": False,
|
||||
"time_mode": "absolute",
|
||||
"dist_mode": "exponential",
|
||||
"extra_info_types_file": (
|
||||
Path(args.extra_info_types_file).name
|
||||
@@ -507,6 +385,8 @@ def build_metadata(
|
||||
else None
|
||||
),
|
||||
"extra_info_types": [int(x) for x in dataset.extra_info_types],
|
||||
"continuous_value_scaling": "robust",
|
||||
"continuous_value_scaler": scaler_stats.as_metadata(),
|
||||
"dataset_metadata": {
|
||||
"vocab_size": int(dataset.vocab_size),
|
||||
"n_types": int(dataset.n_types),
|
||||
@@ -514,14 +394,16 @@ def build_metadata(
|
||||
"n_categories": int(dataset.n_categories),
|
||||
"cont_type_ids": [int(x) for x in dataset.cont_type_ids],
|
||||
"extra_info_types": [int(x) for x in dataset.extra_info_types],
|
||||
"event_stream_version": "disease_death_only_v1",
|
||||
"uses_assessment_event_token": False,
|
||||
},
|
||||
"split_sizes": {
|
||||
"train": int(len(train_subset)),
|
||||
"val": int(len(val_subset)),
|
||||
"test": int(len(test_subset)),
|
||||
},
|
||||
"resolved_readout_name": args.readout_name,
|
||||
"resolved_loss_name": args.target_mode,
|
||||
"resolved_readout_name": "token",
|
||||
"resolved_loss_name": "delphi2m",
|
||||
}
|
||||
|
||||
|
||||
@@ -533,34 +415,24 @@ def main() -> None:
|
||||
|
||||
run_dir, run_name = create_unique_run_dir(
|
||||
lambda timestamp: (
|
||||
f"{args.model_size}_r{args.n_reasoning_rounds}_"
|
||||
f"{args.time_mode}_exponential_"
|
||||
f"next_token_{args.target_mode}_"
|
||||
"absolute_exponential_next_token_delphi2m_"
|
||||
f"gap_{args.no_event_interval_years:g}y_{timestamp}"
|
||||
)
|
||||
),
|
||||
runs_root=Path(args.runs_root) / args.model_architecture,
|
||||
)
|
||||
logger = setup_logging(run_dir)
|
||||
|
||||
logger.info(f"Starting next-step training run: {run_name}")
|
||||
logger.info(f"Device: {device}")
|
||||
size_config = resolve_model_size(args.model_size)
|
||||
logger.info(
|
||||
"Model size: "
|
||||
f"{args.model_size} "
|
||||
f"(d_model={size_config.d_model}, "
|
||||
f"n_trajectory={size_config.n_trajectory}, "
|
||||
f"trajectory_dim={size_config.trajectory_dim}, "
|
||||
f"traj_hidden={size_config.traj_hidden}); "
|
||||
f"reasoning_rounds={args.n_reasoning_rounds}"
|
||||
)
|
||||
logger.info(f"Model architecture: {args.model_architecture}")
|
||||
logger.info(f"extra_info_types: {format_extra_info_types(args.extra_info_types)}")
|
||||
logger.info(f"readout={args.readout_name}, target_mode={args.target_mode}")
|
||||
logger.info("Continuous value scaling: RobustScale (required)")
|
||||
logger.info("time_mode=absolute, readout=token, target_mode=delphi2m")
|
||||
|
||||
dataset = HealthDataset(
|
||||
data_prefix=args.data_prefix,
|
||||
labels_file=args.labels_file,
|
||||
no_event_interval_years=args.no_event_interval_years,
|
||||
include_no_event_in_uts_target=args.include_no_event_in_uts_target,
|
||||
extra_info_types=args.extra_info_types,
|
||||
)
|
||||
if args.train_eid_file and args.val_eid_file and args.test_eid_file:
|
||||
@@ -590,6 +462,20 @@ def main() -> None:
|
||||
f"Samples: train={len(train_subset)}, val={len(val_subset)}, test={len(test_subset)}"
|
||||
)
|
||||
|
||||
if dataset.n_cont_types > 0:
|
||||
logger.info(
|
||||
"Fitting continuous RobustScaler on the complete training subset: "
|
||||
f"patients={len(train_subset):,}, features={dataset.n_cont_types}"
|
||||
)
|
||||
scaler_stats = fit_continuous_robust_scaler(dataset, train_subset)
|
||||
if dataset.n_cont_types > 0:
|
||||
logger.info(
|
||||
"Continuous RobustScaler fitted: "
|
||||
f"observations={int(scaler_stats.observation_count.sum()):,}, "
|
||||
f"min_per_feature={int(scaler_stats.observation_count.min()):,}, "
|
||||
f"max_per_feature={int(scaler_stats.observation_count.max()):,}"
|
||||
)
|
||||
|
||||
train_loader = DataLoader(
|
||||
train_subset,
|
||||
batch_size=args.batch_size,
|
||||
@@ -621,14 +507,13 @@ def main() -> None:
|
||||
prefetch_factor=2 if args.num_workers > 0 else None,
|
||||
)
|
||||
|
||||
model = build_model(args, dataset).to(device)
|
||||
model = build_model(args, dataset, scaler_stats).to(device)
|
||||
parameter_counts = get_model_parameter_counts(model)
|
||||
logger.info(
|
||||
"Model parameters: "
|
||||
f"total={parameter_counts['model_parameter_count']:,}, "
|
||||
f"trainable={parameter_counts['trainable_parameter_count']:,}"
|
||||
)
|
||||
readout = build_next_step_readout(args).to(device)
|
||||
criterion = build_next_step_loss(args)
|
||||
optimizer = AdamW(
|
||||
model.parameters(),
|
||||
@@ -639,7 +524,13 @@ def main() -> None:
|
||||
adaptive_lr = args.base_lr * math.sqrt(args.batch_size / 128)
|
||||
|
||||
train_metadata = build_metadata(
|
||||
args, dataset, run_name, train_subset, val_subset, test_subset
|
||||
args,
|
||||
dataset,
|
||||
run_name,
|
||||
train_subset,
|
||||
val_subset,
|
||||
test_subset,
|
||||
scaler_stats,
|
||||
)
|
||||
train_metadata.update(parameter_counts)
|
||||
save_config(
|
||||
@@ -658,9 +549,13 @@ def main() -> None:
|
||||
lr = get_lr(epoch, args, adaptive_lr)
|
||||
set_optimizer_lr(optimizer, lr)
|
||||
|
||||
train_loss = run_epoch(logger, args, model, readout, criterion, train_loader, optimizer, device, True)
|
||||
train_loss = run_epoch(
|
||||
logger, args, model, criterion, train_loader, optimizer, device, True
|
||||
)
|
||||
with torch.no_grad():
|
||||
val_loss = run_epoch(logger, args, model, readout, criterion, val_loader, None, device, False)
|
||||
val_loss = run_epoch(
|
||||
logger, args, model, criterion, val_loader, None, device, False
|
||||
)
|
||||
|
||||
is_best = val_loss < best_val
|
||||
if is_best:
|
||||
@@ -694,7 +589,9 @@ def main() -> None:
|
||||
logger.info("Evaluating best model on next-step test split...")
|
||||
model.load_state_dict(torch.load(best_model_path, map_location=device))
|
||||
with torch.no_grad():
|
||||
test_loss = run_epoch(logger, args, model, readout, criterion, test_loader, None, device, False)
|
||||
test_loss = run_epoch(
|
||||
logger, args, model, criterion, test_loader, None, device, False
|
||||
)
|
||||
logger.info(f"Test loss: {test_loss:.6f}")
|
||||
logger.info(f"Best checkpoint: {best_model_path}")
|
||||
|
||||
|
||||
304
train_util.py
304
train_util.py
@@ -5,7 +5,9 @@ import logging
|
||||
import sys
|
||||
import time
|
||||
import csv
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
import math
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Iterable, Tuple
|
||||
|
||||
@@ -16,6 +18,272 @@ from torch.utils.data import Subset
|
||||
|
||||
from dataset import AllFutureHealthDataset, HealthDataset
|
||||
from models import DeepHealth
|
||||
from targets import NO_EVENT_IDX, PAD_IDX, RESERVED_IDX
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ContinuousRobustScalerStats:
|
||||
"""Train-split robust scaling statistics aligned to ``cont_type_ids``."""
|
||||
|
||||
cont_type_ids: tuple[int, ...]
|
||||
center: np.ndarray
|
||||
scale: np.ndarray
|
||||
observation_count: np.ndarray
|
||||
quantile_range: tuple[float, float] = (25.0, 75.0)
|
||||
|
||||
def as_metadata(self) -> Dict[str, Any]:
|
||||
return {
|
||||
"method": "robust",
|
||||
"fitted_on": "train_subset",
|
||||
"quantile_range": [float(x) for x in self.quantile_range],
|
||||
"cont_type_ids": [int(x) for x in self.cont_type_ids],
|
||||
"observation_count": [int(x) for x in self.observation_count.tolist()],
|
||||
"center_buffer": "tokenizer.continuous_value_center",
|
||||
"scale_buffer": "tokenizer.continuous_value_scale",
|
||||
}
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class AllFutureBaselineStats:
|
||||
"""Training-only marginal first-onset rates for output initialization."""
|
||||
|
||||
event_count: np.ndarray
|
||||
at_risk_exposure: np.ndarray
|
||||
rate: np.ndarray
|
||||
bias: np.ndarray
|
||||
query_count: int
|
||||
rate_floor: float
|
||||
rate_ceiling: float
|
||||
|
||||
def as_metadata(self) -> Dict[str, Any]:
|
||||
positive = self.rate[self.at_risk_exposure > 0]
|
||||
if positive.size:
|
||||
rate_summary = {
|
||||
"min": float(positive.min()),
|
||||
"median": float(np.median(positive)),
|
||||
"max": float(positive.max()),
|
||||
}
|
||||
else:
|
||||
rate_summary = {"min": 0.0, "median": 0.0, "max": 0.0}
|
||||
return {
|
||||
"method": "training_marginal_first_onset_rate",
|
||||
"fitted_on": "one_seeded_query_per_training_patient",
|
||||
"query_distribution": "patient_interval_time_uniform",
|
||||
"query_count": int(self.query_count),
|
||||
"observed_first_onsets": int(self.event_count.sum()),
|
||||
"rate_floor": float(self.rate_floor),
|
||||
"rate_ceiling": float(self.rate_ceiling),
|
||||
"rate_per_year": rate_summary,
|
||||
"checkpoint_parameter": "risk_head.bias",
|
||||
}
|
||||
|
||||
|
||||
def fit_all_future_baseline(
|
||||
dataset: AllFutureHealthDataset,
|
||||
subset: Subset,
|
||||
*,
|
||||
seed: int,
|
||||
ignored_idx: Iterable[int] = (PAD_IDX, RESERVED_IDX, NO_EVENT_IDX),
|
||||
rate_floor: float = 1e-6,
|
||||
rate_ceiling: float = 10.0,
|
||||
) -> AllFutureBaselineStats:
|
||||
"""Fit marginal per-outcome rates under the training query distribution."""
|
||||
if subset.dataset is not dataset:
|
||||
raise ValueError("subset must reference the all-future training dataset")
|
||||
if rate_floor <= 0 or rate_ceiling <= rate_floor:
|
||||
raise ValueError("Require 0 < rate_floor < rate_ceiling")
|
||||
|
||||
subset_indices = np.asarray(subset.indices, dtype=np.int64)
|
||||
if subset_indices.ndim != 1 or subset_indices.size == 0:
|
||||
raise ValueError("training subset must contain at least one patient")
|
||||
|
||||
vocab_size = int(dataset.vocab_size)
|
||||
ignored = {
|
||||
int(idx)
|
||||
for idx in ignored_idx
|
||||
if 0 <= int(idx) < vocab_size
|
||||
}
|
||||
event_count = np.zeros(vocab_size, dtype=np.int64)
|
||||
exposure_adjustment = np.zeros(vocab_size, dtype=np.float64)
|
||||
common_exposure = 0.0
|
||||
rng = np.random.RandomState(int(seed))
|
||||
|
||||
for patient_index in subset_indices.tolist():
|
||||
patient = dataset.patients[int(patient_index)]
|
||||
t_query = dataset.sample_query(patient, rng)
|
||||
times = np.asarray(patient["times"], dtype=np.float64)
|
||||
labels = np.asarray(patient["labels"], dtype=np.int64)
|
||||
censor_time = max(float(patient["t_obs"]) - t_query, 0.0)
|
||||
common_exposure += censor_time
|
||||
|
||||
prevalent = {
|
||||
int(label)
|
||||
for label in labels[times <= t_query].tolist()
|
||||
if 0 <= int(label) < vocab_size and int(label) not in ignored
|
||||
}
|
||||
for label in prevalent:
|
||||
exposure_adjustment[label] -= censor_time
|
||||
|
||||
# Keep the earliest future occurrence defensively; prepared disease
|
||||
# events are already deduplicated to their first occurrence.
|
||||
first_future_dt: Dict[int, float] = {}
|
||||
for label, event_time in zip(labels[times > t_query], times[times > t_query]):
|
||||
label = int(label)
|
||||
if label in ignored or label in prevalent or not (0 <= label < vocab_size):
|
||||
continue
|
||||
event_dt = max(float(event_time) - t_query, 0.0)
|
||||
previous = first_future_dt.get(label)
|
||||
if previous is None or event_dt < previous:
|
||||
first_future_dt[label] = event_dt
|
||||
for label, event_dt in first_future_dt.items():
|
||||
event_count[label] += 1
|
||||
exposure_adjustment[label] += event_dt - censor_time
|
||||
|
||||
at_risk_exposure = common_exposure + exposure_adjustment
|
||||
at_risk_exposure = np.maximum(at_risk_exposure, 0.0)
|
||||
rate = np.full(vocab_size, float(rate_floor), dtype=np.float64)
|
||||
has_exposure = at_risk_exposure > 0
|
||||
rate[has_exposure] = (
|
||||
event_count[has_exposure].astype(np.float64)
|
||||
/ at_risk_exposure[has_exposure]
|
||||
)
|
||||
rate = np.clip(rate, float(rate_floor), float(rate_ceiling))
|
||||
|
||||
for idx in ignored:
|
||||
at_risk_exposure[idx] = 0.0
|
||||
rate[idx] = 0.0
|
||||
bias = np.zeros(vocab_size, dtype=np.float64)
|
||||
valid = np.ones(vocab_size, dtype=bool)
|
||||
if ignored:
|
||||
valid[np.asarray(sorted(ignored), dtype=np.int64)] = False
|
||||
bias[valid] = np.log(np.expm1(rate[valid]))
|
||||
|
||||
if not np.isfinite(bias).all():
|
||||
raise RuntimeError("All-future baseline initialization is non-finite")
|
||||
return AllFutureBaselineStats(
|
||||
event_count=event_count,
|
||||
at_risk_exposure=at_risk_exposure.astype(np.float32),
|
||||
rate=rate.astype(np.float32),
|
||||
bias=bias.astype(np.float32),
|
||||
query_count=int(subset_indices.size),
|
||||
rate_floor=float(rate_floor),
|
||||
rate_ceiling=float(rate_ceiling),
|
||||
)
|
||||
|
||||
|
||||
def fit_continuous_robust_scaler(
|
||||
dataset: HealthDataset | AllFutureHealthDataset,
|
||||
subset: Subset,
|
||||
*,
|
||||
quantile_range: tuple[float, float] = (25.0, 75.0),
|
||||
scale_epsilon: float = 1e-6,
|
||||
) -> ContinuousRobustScalerStats:
|
||||
"""Fit median/IQR statistics using only patients in the training subset.
|
||||
|
||||
The prepared arrays remain unchanged. Scaling is performed later inside the
|
||||
model tokenizer so the fitted center and scale can live in the checkpoint.
|
||||
"""
|
||||
|
||||
if subset.dataset is not dataset:
|
||||
raise ValueError("subset must reference the dataset used to fit the scaler")
|
||||
low, high = (float(quantile_range[0]), float(quantile_range[1]))
|
||||
if not (0.0 <= low < high <= 100.0):
|
||||
raise ValueError(
|
||||
"quantile_range must satisfy 0 <= low < high <= 100, got "
|
||||
f"{quantile_range!r}"
|
||||
)
|
||||
if scale_epsilon <= 0:
|
||||
raise ValueError("scale_epsilon must be > 0")
|
||||
|
||||
cont_type_ids = tuple(int(x) for x in dataset.cont_type_ids)
|
||||
n_cont_types = len(cont_type_ids)
|
||||
if n_cont_types == 0:
|
||||
empty = np.zeros(0, dtype=np.float32)
|
||||
return ContinuousRobustScalerStats(
|
||||
cont_type_ids=cont_type_ids,
|
||||
center=empty.copy(),
|
||||
scale=empty.copy(),
|
||||
observation_count=np.zeros(0, dtype=np.int64),
|
||||
quantile_range=(low, high),
|
||||
)
|
||||
|
||||
subset_indices = np.asarray(subset.indices, dtype=np.int64)
|
||||
if subset_indices.ndim != 1 or subset_indices.size == 0:
|
||||
raise ValueError("training subset must contain at least one patient")
|
||||
|
||||
type_to_column = np.full(int(dataset.n_types), -1, dtype=np.int64)
|
||||
for column, type_id in enumerate(cont_type_ids):
|
||||
if type_id <= 0 or type_id >= len(type_to_column):
|
||||
raise ValueError(
|
||||
f"continuous type id {type_id} is outside [1, {len(type_to_column)})"
|
||||
)
|
||||
type_to_column[type_id] = column
|
||||
|
||||
if hasattr(dataset, "patients"):
|
||||
records = dataset.patients
|
||||
elif hasattr(dataset, "samples"):
|
||||
records = dataset.samples
|
||||
else:
|
||||
raise TypeError(
|
||||
"dataset must expose patient records through .patients or .samples"
|
||||
)
|
||||
|
||||
values = np.full(
|
||||
(int(subset_indices.size), n_cont_types),
|
||||
np.nan,
|
||||
dtype=np.float32,
|
||||
)
|
||||
for row, patient_index in enumerate(subset_indices.tolist()):
|
||||
patient = records[int(patient_index)]
|
||||
other_type = np.asarray(patient["other_type"], dtype=np.int64)
|
||||
other_value = np.asarray(patient["other_value"], dtype=np.float32)
|
||||
other_kind = np.asarray(patient["other_value_kind"], dtype=np.int64)
|
||||
continuous = other_kind == 1
|
||||
if not np.any(continuous):
|
||||
continue
|
||||
selected_type = other_type[continuous]
|
||||
selected_value = other_value[continuous]
|
||||
valid_type = (selected_type > 0) & (selected_type < len(type_to_column))
|
||||
columns = np.full(selected_type.shape, -1, dtype=np.int64)
|
||||
columns[valid_type] = type_to_column[selected_type[valid_type]]
|
||||
valid = (columns >= 0) & np.isfinite(selected_value)
|
||||
values[row, columns[valid]] = selected_value[valid]
|
||||
|
||||
observation_count = np.isfinite(values).sum(axis=0).astype(np.int64)
|
||||
missing_types = [
|
||||
type_id
|
||||
for type_id, count in zip(cont_type_ids, observation_count.tolist())
|
||||
if count == 0
|
||||
]
|
||||
if missing_types:
|
||||
raise ValueError(
|
||||
"Training subset has no finite observations for continuous type ids: "
|
||||
f"{missing_types}"
|
||||
)
|
||||
|
||||
low_value, center, high_value = np.nanpercentile(
|
||||
values,
|
||||
[low, 50.0, high],
|
||||
axis=0,
|
||||
)
|
||||
scale = high_value - low_value
|
||||
near_constant = (~np.isfinite(scale)) | (np.abs(scale) <= float(scale_epsilon))
|
||||
scale[near_constant] = 1.0
|
||||
|
||||
center = np.asarray(center, dtype=np.float32)
|
||||
scale = np.asarray(scale, dtype=np.float32)
|
||||
if not np.isfinite(center).all():
|
||||
raise RuntimeError("Robust scaler produced non-finite center values")
|
||||
if not np.isfinite(scale).all() or np.any(scale <= 0):
|
||||
raise RuntimeError("Robust scaler produced invalid scale values")
|
||||
|
||||
return ContinuousRobustScalerStats(
|
||||
cont_type_ids=cont_type_ids,
|
||||
center=center,
|
||||
scale=scale,
|
||||
observation_count=observation_count,
|
||||
quantile_range=(low, high),
|
||||
)
|
||||
|
||||
|
||||
def create_unique_run_dir(name_fn, runs_root: Path = Path("runs")) -> tuple[Path, str]:
|
||||
@@ -141,6 +409,31 @@ def resolve_device(device_arg: str) -> torch.device:
|
||||
raise ValueError(f"Unsupported device: {device_arg}")
|
||||
|
||||
|
||||
def get_lr(epoch: int, args: Any, adaptive_lr: float) -> float:
|
||||
if epoch < args.warmup_epochs:
|
||||
return adaptive_lr * (epoch + 1) / args.warmup_epochs
|
||||
progress = (epoch - args.warmup_epochs) / max(
|
||||
1, args.max_epochs - args.warmup_epochs
|
||||
)
|
||||
cosine = 0.5 * (1 + math.cos(math.pi * progress))
|
||||
return adaptive_lr * (
|
||||
args.min_lr_ratio + cosine * (1 - args.min_lr_ratio)
|
||||
)
|
||||
|
||||
|
||||
def move_batch_to_device(
|
||||
batch: Dict[str, torch.Tensor],
|
||||
device: torch.device,
|
||||
) -> Dict[str, torch.Tensor]:
|
||||
non_blocking = device.type == "cuda"
|
||||
return {
|
||||
key: value.to(device, non_blocking=non_blocking)
|
||||
if isinstance(value, torch.Tensor)
|
||||
else value
|
||||
for key, value in batch.items()
|
||||
}
|
||||
|
||||
|
||||
def split_dataset(
|
||||
dataset: HealthDataset,
|
||||
train_ratio: float,
|
||||
@@ -291,17 +584,8 @@ def split_all_future_datasets_by_eid_files(
|
||||
)
|
||||
|
||||
|
||||
def build_optimizer(args: Any, model: DeepHealth) -> AdamW:
|
||||
return AdamW(
|
||||
model.parameters(),
|
||||
lr=args.base_lr,
|
||||
betas=tuple(args.betas),
|
||||
weight_decay=args.weight_decay,
|
||||
)
|
||||
|
||||
|
||||
def get_model_parameter_counts(model: torch.nn.Module) -> Dict[str, int]:
|
||||
"""Return stable parameter-count fields for logs and train_config.json."""
|
||||
"""Return stable total and trainable parameter counts."""
|
||||
return {
|
||||
"model_parameter_count": sum(
|
||||
parameter.numel() for parameter in model.parameters()
|
||||
|
||||
Reference in New Issue
Block a user