# 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-step 训练入口，使用：

```python
HealthDataset = NextStepHealthDataset
```

示例：

```bash
python train.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
```

选择额外信息变量：

```bash
python train.py --extra_info_types 1 3 7
```

如果不传 `--extra_info_types`，默认使用全部 other-info type。

## 主要文件

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