# 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
```

all-future 训练入口会显式使用 `AllFutureHealthDataset` 和 `all_future_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`。

### OtherInfoTokenizer

`DeepHealth` 内部的 `OtherInfoTokenizer` 编码统一的 other-info token：

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

连续值使用 `TokenAutoDiscretization`：

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

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

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

主序列内部保留展开的单个 extra-info token；疾病 token 因此可以直接关注到每个额外信息变量。只有在 next-token 训练需要把 extra-info 位置接入 readout/loss 时，同一患者、同一 `other_time` 的 extra hidden 才会被池化成一个监督用 token。默认池化方式是平均池化：

```bash
--extra_pool_reduce mean
```

也可以设为 `sum`。池化后，每个 extra-info 时间点最多产生一个用于预测和 loss 的 readout token。

### DeepHealth

`DeepHealth` 的统一路径是：

```text
disease tokens + expanded extra-info tokens
-> temporal backbone
-> risk head
```

extra-info 不再通过独立的 `BaselineEncoder` 或 `CrossAttention` 注入；它们作为主序列 token 直接参与同一个 temporal transformer。相对时间模式下，主序列内所有 token 共用 `TimeRoPE` 和 `GaussianRBFTimeBasis`。

两种目标模式的读出语义不同：

- `next_token`
  - 模型内部序列包含 disease tokens 和展开的单个 extra-info tokens
  - 默认 Tensor 返回值仍只返回 disease token hidden，形状为 `(B, L, D)`，兼容评估和旧调用
  - 训练时使用 `return_output=True`；此时同一 `other_time` 的 extra hidden 会在读出端池化，pooled extra readout tokens 也会产生 logits，并在有未来 disease target 时参与 prediction/loss 监督
  - next-token 模式下 `risk_head.weight` 与 `token_embedding.weight` 使用 weight tying

- `all_future`
  - 在合并后的主序列末尾拼接一个 query token
  - query token 的时间是 `t_query`
  - 只读出最后一个 query hidden，形状为 `(B, D)`
  - 展开的 extra-info tokens 只作为 query 的上下文输入，不单独读出、不参与 loss 监督
  - all-future 模式不使用 weight tying

如果需要拿到完整 next-token 输出，可使用结构化返回：

```python
out = model(..., target_mode="next_token", return_output=True)
out.hidden        # disease tokens + pooled extra-info readout tokens
out.time_seq      # 与 hidden 对齐的时间
out.padding_mask  # 与 hidden 对齐的有效位置
out.event_len     # 原 disease token 长度
```

模型初始化示例：

```python
model = DeepHealth(
    vocab_size=dataset.vocab_size,
    n_embd=120,
    n_head=10,
    n_hist_layer=12,
    n_tab_layer=4,  # 兼容旧配置；当前不再创建独立 tabular transformer
    n_types=dataset.n_types,
    n_cont_types=dataset.n_cont_types,
    n_categories=dataset.n_categories,
    cont_type_ids=dataset.cont_type_ids,
    extra_pool_reduce="mean",
)
```

## Loss

`losses.py` 中保留：

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。

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
  - 展开的 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`
  - 展开的 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` |
| `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 |

示例：

```bash
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 \
  --extra_pool_reduce mean
```

all-future 示例：

```bash
python train_all_future.py \
  --data_prefix ukb \
  --labels_file labels.csv \
  --dist_mode weibull \
  --time_mode relative \
  --extra_pool_reduce mean
```

选择额外信息变量：

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

`train_next_step.py` / `train_all_future.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 列表，用于评估脚本复现变量选择
- `extra_pool_reduce`：同一 `other_time` 的 extra-info tokens 池化方式，默认为 `mean`
- `model_target_mode`、`time_mode`、`dist_mode`、`dataset_class`、`collate_fn`、`resolved_loss_name`：用于评估脚本重建模型和输入方式

## 评估 AUC

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

评估脚本默认使用常规 Tensor 返回值：next-token checkpoint 只缓存 disease token/readout 位置的 hidden；all-future checkpoint 只缓存 `t_query` query hidden。训练中额外加入监督的 pooled extra-info readout tokens 不作为 AUC 评估位置单独输出。

### `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`
  - `DeepHealthOutput`
  - `OtherInfoTokenizer`

- `backbones.py`
  - `TimeRoPE`
  - `GaussianRBFTimeBasis`
  - `TemporalAttention`
  - `GPTBlock`
  - `TokenAutoDiscretization`
  - `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 风险