# DeepHealthNew

这个目录包含 DeepHealth 的数据准备、数据集、模型、readout 和 loss 代码。当前版本的核心设计是：

```text
疾病序列 stream + 统一的额外信息 token stream
```

疾病、死亡、checkup 事件仍然保存在事件序列里；性别单独保存在 `basic_info`；其他体检、暴露、生活方式等信息统一整理成 `(type, value, value_kind, time)` token。

## 数据准备

运行：

```bash
python prepare_data.py
```

输入文件：

- `ukb_data.csv`
- `field_ids_enriched.csv`
- `icd10_codes_mod.tsv`
- `labels.csv`

输出文件：

- `ukb_event_data.npy`
  - 形状为 `(N, 3)`
  - 每行是 `(eid, days, label)`
  - 包含疾病、死亡、checkup 事件

- `ukb_basic_info.csv`
  - index 为 `eid`
  - 当前只保留 `sex`

- `ukb_other_info.npy`
  - 形状为 `(M, 5)`
  - 每行是：
    ```text
    (eid, type, value, value_kind, time)
    ```
  - `type=0` 预留给 padding
  - `value_kind=1` 表示连续变量
  - `value_kind=2` 表示分类变量
  - 缺失值不会生成 token
  - 当前 UKB 额外信息只有一个时间点，所以 `time` 暂时都是 `date_of_assessment`

- `cate_types.csv`
  - 分类变量元信息
  - 字段：
    ```text
    type,name,n_categories
    ```
  - `ukb_other_info.npy` 里的分类 value 是变量内部的局部 id；global category id 在 dataset 中根据实验选择的变量动态计算。

## Dataset

`dataset.py` 提供两个 dataset：

- `NextStepHealthDataset`
  - 用于 next-token / next-time-point 监督
  - 对应 `Delphi2MLoss` 和 `UniqueTimeSetExponentialLoss`

- `AllFutureHealthDataset`
  - 用于 query-conditioned all-future 监督
  - 对应 `ExponentialLoss`、`WeibullLoss`、`MixedLoss`

为了兼容旧训练入口：

```python
HealthDataset = NextStepHealthDataset
collate_fn = next_step_collate_fn
```

dataset 会输出：

```python
event_seq
time_seq
padding_mask
sex
other_type
other_value
other_value_kind
other_time
```

其中 `other_type=0` 表示 padding，不额外传 other-token mask。

可以通过 `extra_info_types` 选择纳入哪些额外信息变量：

```python
dataset = NextStepHealthDataset(extra_info_types=[1, 3, 7])
```

如果不传，则使用全部可用 other-info type。

dataset 会暴露模型初始化需要的元信息：

```python
dataset.n_types
dataset.n_cont_types
dataset.n_categories
dataset.cont_type_ids
dataset.vocab_size
```

## 模型

模型主体定义在 `models.py`，通用网络模块定义在 `backbones.py`。

### BaselineEncoder

`BaselineEncoder` 编码统一的 other-info token：

```python
other_type        # (B, K)
other_value       # (B, K)
other_value_kind  # (B, K)
```

它暂时不直接使用 `other_time`。时间信息保留给后续 `CrossAttention`，用于建模疾病/query 与 other-info token 的相对时间关系。

连续值使用 `TokenAutoDiscretization`：

```text
type_id -> continuous type index -> soft bin embedding
```

分类值使用 dataset 动态计算后的 global category id：

```text
selected type offsets + local category id
```

### CrossAttention

`CrossAttention` 让 disease-side hidden state 注意到 other-info token：

```python
h_disease  # (B, L, D)
t_disease  # (B, L)
h_token    # (B, K, D)
t_token    # (B, K)
```

时间信息通过两种方式进入注意力：

- `TimeRoPE`
  - 使用 query time 和 key time 旋转 q/k
  - 让 q-k 相似度带有时间位置信息

- `GaussianRBFTimeBasis`
  - 对 `t_disease - t_token` 做 RBF 编码
  - 投影成每个 attention head 的时间 bias

注意力是时间因果的：

```text
other_info_time <= disease_or_query_time
```

如果某个 disease/query 位置没有任何可见 other-info token，则该位置保持原 hidden 不变。

### DeepHealth

`DeepHealth` 的统一路径是：

```text
disease-side sequence
-> disease temporal backbone
-> CrossAttention 到 other-info tokens
-> risk head
```

两种目标模式共用同一套语义：

- `next_token`
  - `h_disease` 长度为 `L`
  - 输出 `(B, L, D)`

- `all_future`
  - 在 disease-side 序列末尾拼接一个 query token
  - query token 的时间是 `t_query`
  - `h_disease` 长度为 `L + 1`
  - 输出最后一个 query hidden，形状为 `(B, D)`

模型初始化示例：

```python
model = DeepHealth(
    vocab_size=dataset.vocab_size,
    n_embd=120,
    n_head=10,
    n_hist_layer=12,
    n_tab_layer=4,
    n_types=dataset.n_types,
    n_cont_types=dataset.n_cont_types,
    n_categories=dataset.n_categories,
    cont_type_ids=dataset.cont_type_ids,
)
```

## Loss

`losses.py` 中保留：

next-token 监督：

- `Delphi2MLoss`
- `UniqueTimeSetExponentialLoss`

all-future / query-conditioned 监督：

- `ExponentialLoss`
- `WeibullLoss`
- `MixedLoss`

`UniqueTimeSetExponentialLoss` 的 observed term 固定使用 sum reduction，不再暴露旧的 `observed_reduction` 参数。

## 训练

当前 `train.py` 支持 next-token 和 all-future 两类训练入口：

- `--model_target_mode next_token`
  - 使用 `NextStepHealthDataset`
  - `--target_mode delphi2m` 默认搭配 `Delphi2MLoss` + `token` readout
  - `--target_mode uts` 默认搭配 `UniqueTimeSetExponentialLoss` + `same_time_group_end` readout
- `--model_target_mode all_future`
  - 使用 `AllFutureHealthDataset`
  - 不使用 readout，直接对 query hidden 计算风险
  - `--dist_mode exponential/weibull/mixed` 分别搭配 `ExponentialLoss`、`WeibullLoss`、`MixedLoss`

模型结构组合由 `model_target_mode × time_mode × dist_mode` 决定：

| 维度 | 可选项 |
| --- | --- |
| `model_target_mode` | `next_token`, `all_future` |
| `time_mode` | `relative`, `absolute` |
| `dist_mode` | `exponential`, `weibull`, `mixed` |

示例：

```bash
python train.py \
  --data_prefix ukb \
  --labels_file labels.csv \
  --model_target_mode next_token \
  --target_mode uts \
  --n_embd 120 \
  --n_head 10 \
  --n_hist_layer 12 \
  --n_tab_layer 4
```

all-future 示例：

```bash
python train.py \
  --data_prefix ukb \
  --labels_file labels.csv \
  --model_target_mode all_future \
  --dist_mode weibull \
  --time_mode relative
```

选择额外信息变量：

```bash
python train.py --extra_info_types_file extra_info_types_smoking_alcohol_bmi.txt
```

`train.py` 只接受 `--extra_info_types_file` 指定变量列表，不接受在 CLI 里直接输入 type id。文件可以每行一个 type id，也可以带 `#` 注释；如果不传 `--extra_info_types_file`，默认使用全部 other-info type。

训练输出的 `train_config.json` 会记录：

- `extra_info_types_file`：训练时使用的列表文件名
- `extra_info_types`：解析后的实际 type id 列表，用于评估脚本复现变量选择
- `model_target_mode`、`time_mode`、`dist_mode`、`dataset_class`、`collate_fn`、`resolved_loss_name`：用于评估脚本重建模型和输入方式

## 评估 AUC

当前提供两个 AUC 评估入口，二者都已适配新的 `DeepHealth` 模型和统一的 other-info token 输入；AUC 的 DeLong 计算、病例/对照筛选和分层聚合逻辑保持原评估脚本口径。

### `evaluate_auc.py`

`evaluate_auc.py` 评估的是 **next-step / token-level 预测位置上的疾病 AUC**。

查询方式由 `train_config.json` 中的 `model_target_mode` 决定：

- `model_target_mode="next_token"`：使用训练 readout 对应的历史 token hidden 作为预测点。
- `model_target_mode="all_future"`：不使用 readout token，直接把每个预测点年龄作为 `t_query` 传入模型，取 query hidden 作为预测点。

核心流程：

- 按训练配置重新构建 `HealthDataset` 和 `DeepHealth`。
- 对评估 split 中的患者做模型推理，缓存每个预测点的 hidden。
- 对疾病 token 分块投影到 `risk_head`，避免一次性保存全词表 logits。
- AUC score 使用疾病对应的 eta/logit 排序分数；`dist_mode` 只用于正确构建模型，不会把分数转换成 horizon-specific risk probability。
- 对每个疾病、性别、年龄段、prediction offset 分别计算 AUC。
- 输出未池化分层结果和按疾病汇总后的结果。

典型用法：

```bash
python evaluate_auc.py \
  --run_path runs/your_run_dir \
  --eval_split test \
  --offsets 0.1,1,5,10
```

主要输出：

- `df_auc_unpooled.csv`
  - 疾病 token 在 sex、age bracket、offset 分层下的 AUC。
- `df_both.csv`
  - 按疾病 token 和 offset 聚合后的 AUC。

适合回答的问题：

- “模型在历史序列中的某个预测 token 上，提前 offset 年预测未来疾病的区分能力如何？”
- “不同年龄段、性别、提前量下，next-step 训练模型的疾病预测 AUC 如何？”

### `evaluate_auc_v2.py`

`evaluate_auc_v2.py` 评估的是 **landmark fixed-horizon incident disease AUC**。

它不是使用已有序列中的普通 readout 位置，而是在指定 landmark age 构造一个 landmark query，然后评估该 landmark 后固定 horizon 内是否发生 incident disease。

查询方式由 `train_config.json` 中的 `model_target_mode` 决定：

- `model_target_mode="next_token"`：在 landmark age 人工插入一个 `<NO_EVENT>` token，取该 token 的 hidden 做风险分数。
- `model_target_mode="all_future"`：不插入 `<NO_EVENT>`，直接把 landmark age 作为 `t_query` 传入模型，取 query hidden 做风险分数。

核心流程：

- 为每个患者和 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。
- `score_mode="eta"` 是诊断用排序分数，不使用 `rho`，因此不区分不同分布的风险曲线。
- 按疾病、性别、landmark age、horizon 计算 incident disease AUC。
- 可选择排除 horizon 内先于目标疾病发生的死亡竞争风险。

典型用法：

```bash
python evaluate_auc_v2.py \
  --run_path runs/your_run_dir \
  --eval_split test \
  --landmark_start 40 \
  --landmark_stop 80 \
  --landmark_step 5 \
  --horizons 1,5,10
```

主要输出：

- `df_auc_landmark_unpooled.csv`
  - 疾病 token 在 sex、landmark age、horizon 分层下的 AUC。
- `df_auc_landmark.csv`
  - 按疾病 token 和 horizon 聚合后的 landmark AUC。

适合回答的问题：

- “一个人在 40/45/50/... 岁这个固定年龄点，如果此前未患某病，未来 1/5/10 年内发生该病的风险区分能力如何？”
- “模型能否作为 landmark risk prediction 模型使用？”

### 两者区别

| 项目 | `evaluate_auc.py` | `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"` 不区分分布 |
| 病例定义 | target table 中出现目标疾病的患者/事件 | landmark 后 horizon 内首次发生目标疾病 |
| 对照定义 | 从未出现该疾病的患者的 eligible target occurrence | landmark 时未患病，且 horizon 内未发病并有足够随访 |
| 分层 | sex + age bracket + offset | sex + landmark age + horizon |
| 输出文件 | `df_auc_unpooled.csv`, `df_both.csv` | `df_auc_landmark_unpooled.csv`, `df_auc_landmark.csv` |
| 适用问题 | 提前若干年预测未来目标事件的 token-level AUC | 固定年龄点未来固定年限 incident disease risk AUC |

简单选择：

- 想复现/延续旧的 next-token Delphi 风格 AUC：用 `evaluate_auc.py`。
- 想做临床上更像 “某年龄点未来 N 年发病风险” 的 landmark AUC：用 `evaluate_auc_v2.py`。

## 主要文件

- `prepare_data.py`
  - UKB 原始数据到模型输入文件的 ETL

- `dataset.py`
  - next-step 和 all-future dataset
  - 动态选择 other-info type
  - 动态计算 categorical global id

- `models.py`
  - `DeepHealth`

- `backbones.py`
  - `TimeRoPE`
  - `GaussianRBFTimeBasis`
  - `TemporalAttention`
  - `GPTBlock`
  - `TokenAutoDiscretization`
  - `BaselineEncoder`
  - `CrossAttention`
  - `AgeSinusoidalEncoding`

- `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 分层
  - next-token 模型使用 readout hidden；all-future 模型使用 `t_query` hidden
  - score 是疾病 eta/logit，不按分布转换为固定 horizon 风险概率

- `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 风险