Improve all-future first-onset training

This reverts commit bf5cae8758.

AGENTS.md

@@ -0,0 +1,7 @@
+# 项目协作原则
+
+- 不回答或执行超出用户问题或请求范围的事项。
+- 如果确有必要超出范围，必须先停止并询问用户；得到明确同意后才能继续。
+- 回答应简单直接、逻辑清晰。
+- 永远不要使用 Node.js 读取、处理、分析或转换数据；涉及数据任务时必须使用其他工具。
+- Python 与数据处理统一使用本机 Miniconda：`C:\ProgramData\miniconda3`；不要使用 Codex 自带或其他 Python/Conda 运行时。

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

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

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,

evaluate_all_runs_linux.sh

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

evaluate_auc.py

@@ -19,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
 
@@ -39,16 +39,32 @@ import torch
 from torch.utils.data import DataLoader, Subset
 from tqdm.auto import tqdm
 
-from dataset import HealthDataset
+from dataset import (
+    DISEASE_HISTORY_MODE_TIMED,
+    HealthDataset,
+    normalize_disease_history_mode,
+    transform_disease_history_batch_at_position,
+)
 from delphi2m_auc_report import (
     DEFAULT_DELPHI2M_PERIODS_YEARS,
     build_delphi2m_auc_report,
 )
-from eval_data import load_sequence_eval_dataset, sequence_eval_collate_fn
+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 readouts import build_readout
-from targets import PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX
+from targets import NO_EVENT_IDX, PAD_IDX, RESERVED_IDX
 
 
 # ---------------------------------------------------------------------------
@@ -153,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]:
@@ -264,88 +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,
-    state_dict: Optional[Dict[str, Any]] = None,
-) -> DeepHealth:
-    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}"
-        )
-    model_architecture = resolve_model_architecture(cfg, state_dict)
-    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)),
-        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)),
-        model_architecture=model_architecture,
-    )
-
-
 def _extract_state_dict(ckpt: Any) -> Dict[str, Any]:
     if isinstance(ckpt, dict) and "model" in ckpt:
         return ckpt["model"]
@@ -365,27 +250,20 @@ 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
 
 
@@ -402,7 +280,11 @@ def load_model_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))
@@ -410,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)
 
-    train_idx, val_idx, test_idx = split_indices(
-        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:
+    if eval_split in {"valid", "validation"}:
+        eval_split = "val"
+    if eval_split not in {"train", "val", "test", "all"}:
         raise ValueError(
-            f"eval_split must be one of {sorted(split_map)}, got {eval_split!r}")
+            "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
+        )
+        split_map = {
+            "train": train_idx,
+            "val": val_idx,
+            "test": test_idx,
+            "all": np.arange(len(dataset)),
+        }
+        indices = split_map[eval_split]
 
-    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
 # ---------------------------------------------------------------------------
@@ -480,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]]:
@@ -491,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]] = {
@@ -543,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()
@@ -567,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)
@@ -1065,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,
@@ -1115,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
@@ -1125,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,
     )
@@ -1275,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,
@@ -1318,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)
@@ -1326,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(
@@ -1336,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":
@@ -1350,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)
 
@@ -1384,6 +1268,7 @@ def main() -> None:
     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."
@@ -1406,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)
@@ -1427,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)),

evaluate_auc_v2.py

@@ -1,7 +1,7 @@
 """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.
@@ -30,44 +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 dataset import (
+    DISEASE_HISTORY_MODE_TIMED,
+    HealthDataset,
+    normalize_disease_history_mode,
+    transform_disease_history,
+)
 from delphi2m_auc_report import (
     DEFAULT_DELPHI2M_PERIODS_YEARS,
     build_delphi2m_auc_report,
 )
-from eval_data import load_sequence_eval_dataset
+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 readouts import build_readout
-from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX
+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]]:
@@ -108,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))
@@ -128,18 +110,32 @@ def make_eval_indices(dataset: HealthDataset, args: argparse.Namespace, cfg: Dic
     if eval_split in {"valid", "validation"}:
         eval_split = "val"
 
-    train_idx, val_idx, test_idx = split_indices(
-        len(dataset), train_ratio, val_ratio, test_ratio, seed
-    )
-    split_map = {
-        "train": train_idx,
-        "val": val_idx,
-        "test": test_idx,
-        "all": np.arange(len(dataset), dtype=np.int64),
-    }
-    if eval_split not in split_map:
+    if eval_split not in {"train", "val", "test", "all"}:
         raise ValueError(f"Unsupported eval_split={eval_split!r}")
-    indices = split_map[eval_split]
+
+    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
+        )
+        split_map = {
+            "train": train_idx,
+            "val": val_idx,
+            "test": test_idx,
+            "all": np.arange(len(dataset), dtype=np.int64),
+        }
+        indices = split_map[eval_split]
 
     subset_size = cfg_get(args, cfg, "dataset_subset_size", None)
     if subset_size is not None and int(subset_size) > 0:
@@ -160,61 +156,19 @@ 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,
-    state_dict: Optional[Dict[str, Any]] = None,
-) -> DeepHealth:
-    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}"
+    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."
         )
-    model_architecture = resolve_model_architecture(cfg, state_dict)
-    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)),
-        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)),
-        model_architecture=model_architecture,
-    )
+    return mode
 
 
 def load_model_state(model: DeepHealth, state_dict: Dict[str, Any]) -> None:
@@ -222,32 +176,6 @@ def load_model_state(model: DeepHealth, state_dict: Dict[str, Any]) -> None:
     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:
-        raise RuntimeError(
-            "Current dataset metadata does not match train_config.json. "
-            "Use the same prepared data and extra_info_types as training. "
-            + "; ".join(mismatches)
-        )
-
-
 # ---------------------------------------------------------------------------
 # DeLong AUC utilities
 # ---------------------------------------------------------------------------
@@ -371,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],
@@ -468,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
@@ -538,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
 
@@ -555,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
-                        )
+                    time_seq_landmark[-1] = np.nextafter(
+                        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)
 
@@ -579,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,
@@ -671,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]]:
@@ -682,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": [],
@@ -727,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(
@@ -739,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)
 
@@ -799,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] = []
@@ -831,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))
@@ -875,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")
@@ -900,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": {},
         }
     )
@@ -926,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)
@@ -939,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)
 
 
@@ -960,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"]
@@ -1019,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]],
@@ -1028,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)
@@ -1085,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,
@@ -1102,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,
     )
@@ -1151,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)]
@@ -1180,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(
@@ -1235,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)
@@ -1282,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(
@@ -1297,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)
@@ -1325,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)
 
@@ -1347,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))
@@ -1389,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(
@@ -1438,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))
@@ -1489,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"
@@ -1516,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,

evaluate_calibration_all_runs_linux.sh

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

evaluate_event_free_survival.py

@@ -1,816 +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, state_dict=state_dict
-    ).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()

evaluate_extra_info_attribution.py

@@ -1,952 +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, state_dict=state_dict
-    ).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()

evaluate_landmark_auc.py

@@ -1,7 +0,0 @@
-from __future__ import annotations
-
-from evaluate_auc_v2 import main
-
-
-if __name__ == "__main__":
-    main()

evaluate_single_disease_mortality_attribution.py

@@ -1,819 +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, state_dict=state_dict
-    ).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()

evaluate_token_auc.py

@@ -1,7 +0,0 @@
-from __future__ import annotations
-
-from evaluate_auc import main
-
-
-if __name__ == "__main__":
-    main()

event_date_utils.py

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

export_tquery_logits_hidden.py

@@ -1,327 +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, state_dict=state_dict
-    ).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(getattr(model, "n_embd", cfg_get(args, cfg_model, "n_embd", 120)))
-    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()

export_weibull_death_parameter_stats.py

@@ -1,519 +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, state_dict=state_dict
-    ).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()

export_weibull_parameters.py

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

export_weibull_shape_parameter_stats.py

@@ -1,7 +0,0 @@
-"""Compatibility entry point for Weibull shape-parameter export."""
-
-from export_weibull_death_parameter_stats import main
-
-
-if __name__ == "__main__":
-    main()

extra_info_assessment_all_analysis_plan.md

@@ -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 专表。

extra_info_types_assessment_smoking_alcohol.txt

@@ -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.

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

landmark_eval_utils.py

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

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

models.py

@@ -37,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:
@@ -54,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 = (
@@ -71,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):
@@ -86,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])
@@ -127,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)
@@ -155,12 +213,16 @@ class DeepHealth(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,
         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"]:
@@ -169,9 +231,14 @@ 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_layer < 1:
@@ -188,6 +255,8 @@ class DeepHealth(nn.Module):
             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
@@ -205,12 +274,6 @@ class DeepHealth(nn.Module):
             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(n_embd, 1)
-            nn.init.zeros_(self.rho_death_head.weight)
-            nn.init.constant_(self.rho_death_head.bias, 0.5413)
-
         if time_mode == "absolute":
             self.age_encoding = AgeSinusoidalEncoding(n_embd)
             self.blocks = nn.ModuleList([
@@ -241,7 +304,33 @@ class DeepHealth(nn.Module):
             self.rbf = GaussianRBFTimeBasis(n_bases=16, max_time_diff=40.0)
 
         self.final_ln = nn.LayerNorm(n_embd)
-        self.risk_head = nn.Linear(n_embd, vocab_size, bias=False)
+        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))
@@ -461,15 +550,8 @@ class DeepHealth(nn.Module):
             )
         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)
@@ -480,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

plot_next_token_to_all_future_auc.R

@@ -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 = "/"))

prepare_data.py

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

prepare_event_dates.py

@@ -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
-label in ``labels.csv`` receives token 2.  Each ``(eid, token)`` is deduplicated
+``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
 -----

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

rerun_all_required_experiments_linux.sh

@@ -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."

run_all_future_first_onset_experiments_linux.sh

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

run_missing_evaluations.sh

@@ -1,221 +0,0 @@
-#!/usr/bin/env bash
-set -euo pipefail
-
-# Run all non-wrapper evaluation scripts for every completed current-format
-# experiment under runs/. The script is written for Linux servers with bash 4.2.
-
-cd "$(dirname "${BASH_SOURCE[0]}")"
-shopt -s globstar nullglob
-
-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_file_result_if_missing() {
-  local label="$1"
-  local result_dir="$2"
-  local required="$3"
-  shift 3
-
-  if [[ -s "${result_dir}/${required}" ]]; then
-    echo "    skip ${label}: found ${result_dir}/${required}"
-    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
-}
-
-run_has_current_model_config() {
-  "${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"))
-    n_layer = int(cfg.get("n_layer", 0))
-except Exception:
-    raise SystemExit(1)
-
-supported = {"transformer_ffn_v1", "traj_mixer_v5"}
-raise SystemExit(
-    0
-    if cfg.get("model_architecture") in supported and n_layer >= 1
-    else 1
-)
-PY
-}
-
-for config_path in runs/**/train_config.json; do
-  run_path="${config_path%/train_config.json}"
-
-  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
-  if ! run_has_current_model_config "${run_path}"; then
-    echo "    skip run: config lacks current model_architecture/n_layer fields"
-    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_file_result_if_missing \
-    "evaluate_auc.py" \
-    "${run_path}" \
-    "df_auc_delphi2m_report.csv" \
-    "${PYTHON_BIN}" evaluate_auc.py "${common[@]}" "${auc_extra[@]}"
-
-  run_file_result_if_missing \
-    "evaluate_auc_v2.py" \
-    "${run_path}" \
-    "df_auc_landmark_delphi2m_report.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."

run_missing_training_runs.sh

@@ -1,143 +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. Set
-# MODEL_ARCHITECTURE=traj_mixer_v5 to run the TrajMixer variant.
-
-cd "$(dirname "${BASH_SOURCE[0]}")"
-
-PYTHON_BIN="${PYTHON_BIN:-python}"
-DEVICE="${DEVICE:-cuda}"
-NUM_WORKERS="${NUM_WORKERS:-4}"
-PROGRESS_INTERVAL="${PROGRESS_INTERVAL:-20}"
-MODEL_ARCHITECTURE="${MODEL_ARCHITECTURE:-transformer_ffn_v1}"
-N_LAYER="${N_LAYER:-12}"
-
-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_layer "${N_LAYER}"
-  --model_architecture "${MODEL_ARCHITECTURE}"
-  --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" "$MODEL_ARCHITECTURE" "$N_LAYER" <<'PY'
-import json
-import sys
-from pathlib import Path
-
-(
-    time_mode,
-    dist_mode,
-    extra_file,
-    seed,
-    validation_query_seed,
-    model_architecture,
-    n_layer,
-) = sys.argv[1:8]
-extra_name = Path(extra_file).name
-
-for config_path in Path("runs").rglob("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("model_architecture") == model_architecture
-        and int(cfg.get("n_layer", -1)) == int(n_layer)
-        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}: ${MODEL_ARCHITECTURE} n_layer=${N_LAYER} ${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."

run_weibull_shape_exports.sh

@@ -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."

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,
-    }

test_model_architectures.py

@@ -1,247 +0,0 @@
-import unittest
-
-import torch
-
-from backbones import (
-    SwiGLU,
-    TrajMixer,
-    TrajMixerBlock,
-    TransformerFFNBlock,
-    build_backbone_block,
-)
-from model_architectures import (
-    TRAJ_MIXER_ARCHITECTURE,
-    TRANSFORMER_FFN_ARCHITECTURE,
-    detect_model_architecture_from_state_dict,
-    resolve_model_architecture,
-)
-from models import DeepHealth
-
-
-def _build_block(model_architecture: str):
-    return build_backbone_block(
-        model_architecture,
-        n_embd=12,
-        n_head=3,
-        use_time_rope=False,
-        use_rbf_bias=False,
-        mlp_dropout=0.0,
-    )
-
-
-def _as_model_state_dict(block: torch.nn.Module) -> dict[str, torch.Tensor]:
-    return {
-        f"blocks.0.{name}": value.detach().clone()
-        for name, value in block.state_dict().items()
-    }
-
-
-def _build_model(
-    model_architecture: str | None,
-    *,
-    n_layer: int = 1,
-) -> DeepHealth:
-    return DeepHealth(
-        vocab_size=8,
-        n_embd=12,
-        n_head=3,
-        n_layer=n_layer,
-        n_types=2,
-        n_cont_types=0,
-        n_categories=2,
-        cont_type_ids=[],
-        time_mode="absolute",
-        model_architecture=model_architecture,
-    )
-
-
-class ModelArchitectureFactoryTest(unittest.TestCase):
-    def test_factory_builds_both_architectures_with_expected_topology(self) -> None:
-        ffn_block = _build_block(TRANSFORMER_FFN_ARCHITECTURE)
-        self.assertIsInstance(ffn_block, TransformerFFNBlock)
-        self.assertIsInstance(ffn_block.mlp, SwiGLU)
-        self.assertTrue(hasattr(ffn_block, "ln1"))
-        self.assertTrue(hasattr(ffn_block, "ln2"))
-
-        traj_block = _build_block(TRAJ_MIXER_ARCHITECTURE)
-        self.assertIsInstance(traj_block, TrajMixerBlock)
-        self.assertIsInstance(traj_block.mlp, TrajMixer)
-        self.assertTrue(hasattr(traj_block, "ln1"))
-        self.assertFalse(hasattr(traj_block, "ln2"))
-
-    def test_both_architectures_forward_and_backward(self) -> None:
-        for architecture in (
-            TRANSFORMER_FFN_ARCHITECTURE,
-            TRAJ_MIXER_ARCHITECTURE,
-        ):
-            with self.subTest(architecture=architecture):
-                torch.manual_seed(0)
-                block = _build_block(architecture)
-                x = torch.randn(2, 5, 12, requires_grad=True)
-
-                output = block(x)
-                self.assertEqual(output.shape, x.shape)
-                output.square().mean().backward()
-
-                self.assertIsNotNone(x.grad)
-                self.assertTrue(torch.isfinite(x.grad).all())
-                self.assertGreater(x.grad.abs().sum().item(), 0.0)
-                self.assertIsNotNone(block.attn.qkv.weight.grad)
-                self.assertGreater(
-                    block.attn.qkv.weight.grad.abs().sum().item(),
-                    0.0,
-                )
-
-                if architecture == TRANSFORMER_FFN_ARCHITECTURE:
-                    branch_parameters = (
-                        block.mlp.w1.weight,
-                        block.mlp.w2.weight,
-                        block.mlp.w3.weight,
-                    )
-                else:
-                    branch_parameters = (
-                        block.mlp.intra_gate_proj,
-                        block.mlp.intra_value_proj,
-                        block.mlp.output_proj,
-                    )
-                for parameter in branch_parameters:
-                    self.assertIsNotNone(parameter.grad)
-                    self.assertTrue(torch.isfinite(parameter.grad).all())
-                    self.assertGreater(parameter.grad.abs().sum().item(), 0.0)
-
-    def test_unknown_architecture_is_rejected(self) -> None:
-        with self.assertRaises(ValueError):
-            _build_block("unknown_architecture")
-        with self.assertRaisesRegex(ValueError, "model_architecture is required"):
-            _build_model(None)
-
-    def test_deephealth_rejects_fewer_than_one_layer(self) -> None:
-        for n_layer in (0, -1):
-            with self.subTest(n_layer=n_layer):
-                with self.assertRaisesRegex(ValueError, "n_layer must be >= 1"):
-                    _build_model(
-                        TRANSFORMER_FFN_ARCHITECTURE,
-                        n_layer=n_layer,
-                    )
-
-    def test_deephealth_uses_factory_and_strictly_reloads_both_models(self) -> None:
-        for architecture, block_class in (
-            (TRANSFORMER_FFN_ARCHITECTURE, TransformerFFNBlock),
-            (TRAJ_MIXER_ARCHITECTURE, TrajMixerBlock),
-        ):
-            with self.subTest(architecture=architecture):
-                model = _build_model(architecture)
-                self.assertEqual(model.model_architecture, architecture)
-                self.assertIsInstance(model.blocks[0], block_class)
-                self.assertEqual(
-                    detect_model_architecture_from_state_dict(
-                        model.state_dict()
-                    ),
-                    architecture,
-                )
-
-                reloaded = _build_model(architecture)
-                incompatible = reloaded.load_state_dict(
-                    model.state_dict(),
-                    strict=True,
-                )
-                self.assertEqual(incompatible.missing_keys, [])
-                self.assertEqual(incompatible.unexpected_keys, [])
-
-
-class ModelArchitectureResolutionTest(unittest.TestCase):
-    def setUp(self) -> None:
-        self.ffn_block = _build_block(TRANSFORMER_FFN_ARCHITECTURE)
-        self.traj_block = _build_block(TRAJ_MIXER_ARCHITECTURE)
-        self.ffn_state = _as_model_state_dict(self.ffn_block)
-        self.traj_state = _as_model_state_dict(self.traj_block)
-
-    def test_state_dict_detection_recognizes_both_architectures(self) -> None:
-        self.assertEqual(
-            detect_model_architecture_from_state_dict(self.ffn_state),
-            TRANSFORMER_FFN_ARCHITECTURE,
-        )
-        self.assertEqual(
-            detect_model_architecture_from_state_dict(self.traj_state),
-            TRAJ_MIXER_ARCHITECTURE,
-        )
-
-    def test_explicit_markers_resolve_when_checkpoint_matches(self) -> None:
-        for architecture, state_dict in (
-            (TRANSFORMER_FFN_ARCHITECTURE, self.ffn_state),
-            (TRAJ_MIXER_ARCHITECTURE, self.traj_state),
-        ):
-            with self.subTest(architecture=architecture):
-                self.assertEqual(
-                    resolve_model_architecture(
-                        {"model_architecture": architecture},
-                        state_dict,
-                    ),
-                    architecture,
-                )
-
-    def test_architecture_marker_is_required_for_checkpoint_loading(self) -> None:
-        with self.assertRaisesRegex(ValueError, "model_architecture is required"):
-            resolve_model_architecture({}, self.ffn_state)
-        with self.assertRaisesRegex(ValueError, "model_architecture is required"):
-            resolve_model_architecture(None, self.traj_state)
-
-    def test_explicit_marker_conflicting_with_state_dict_is_rejected(self) -> None:
-        conflicts = (
-            (TRANSFORMER_FFN_ARCHITECTURE, self.traj_state),
-            (TRAJ_MIXER_ARCHITECTURE, self.ffn_state),
-        )
-        for architecture, state_dict in conflicts:
-            with self.subTest(architecture=architecture):
-                with self.assertRaises(ValueError):
-                    resolve_model_architecture(
-                        {"model_architecture": architecture},
-                        state_dict,
-                    )
-
-    def test_unknown_marker_and_ambiguous_state_dict_are_rejected(self) -> None:
-        with self.assertRaises(ValueError):
-            resolve_model_architecture(
-                {"model_architecture": "traj_mixer_v4"}
-            )
-
-        ambiguous_state = dict(self.ffn_state)
-        ambiguous_state.update(self.traj_state)
-        with self.assertRaises(ValueError):
-            detect_model_architecture_from_state_dict(ambiguous_state)
-
-        with self.assertRaises(ValueError):
-            detect_model_architecture_from_state_dict(
-                {"token_embedding.weight": torch.empty(2, 2)}
-            )
-
-    def test_ffn_block_schema_is_stable_and_strictly_loadable(self) -> None:
-        expected_keys = {
-            "attn.time_bias_scale",
-            "attn.qkv.weight",
-            "attn.out_proj.weight",
-            "attn.rbf_proj.weight",
-            "mlp.w1.weight",
-            "mlp.w1.bias",
-            "mlp.w2.weight",
-            "mlp.w2.bias",
-            "mlp.w3.weight",
-            "mlp.w3.bias",
-            "ln1.weight",
-            "ln1.bias",
-            "ln2.weight",
-            "ln2.bias",
-        }
-        state = self.ffn_block.state_dict()
-        self.assertSetEqual(set(state), expected_keys)
-        self.assertEqual(tuple(state["mlp.w1.weight"].shape), (30, 12))
-        self.assertEqual(tuple(state["mlp.w3.weight"].shape), (12, 30))
-
-        reloaded = _build_block(TRANSFORMER_FFN_ARCHITECTURE)
-        incompatible = reloaded.load_state_dict(state, strict=True)
-        self.assertEqual(incompatible.missing_keys, [])
-        self.assertEqual(incompatible.unexpected_keys, [])
-
-
-if __name__ == "__main__":
-    unittest.main()

test_temporal_attention.py

@@ -1,46 +0,0 @@
-import unittest
-
-import torch
-
-from backbones import TemporalAttention
-
-
-class TemporalAttentionTest(unittest.TestCase):
-    def test_zero_rbf_bias_has_live_projection_gradient(self) -> None:
-        torch.manual_seed(0)
-        attention = TemporalAttention(
-            n_embd=12,
-            n_head=3,
-            use_time_rope=False,
-            use_rbf_bias=True,
-        )
-        features = torch.randn(2, 4, 4, 16)
-        target = torch.randn(2, 4, 4, 3)
-
-        initial_bias = (
-            attention.time_bias_scale.tanh()
-            * attention.rbf_proj(features)
-        )
-        torch.testing.assert_close(initial_bias, torch.zeros_like(initial_bias))
-
-        (initial_bias * target).sum().backward()
-        projection_grad = attention.rbf_proj.weight.grad
-        self.assertIsNotNone(projection_grad)
-        self.assertGreater(projection_grad.abs().sum().item(), 0.0)
-
-        with torch.no_grad():
-            attention.rbf_proj.weight.add_(projection_grad, alpha=-1e-3)
-        attention.zero_grad(set_to_none=True)
-        updated_bias = (
-            attention.time_bias_scale.tanh()
-            * attention.rbf_proj(features)
-        )
-        (updated_bias * target).sum().backward()
-
-        scale_grad = attention.time_bias_scale.grad
-        self.assertIsNotNone(scale_grad)
-        self.assertGreater(scale_grad.abs().item(), 0.0)
-
-
-if __name__ == "__main__":
-    unittest.main()

test_traj_mixer.py

@@ -1,166 +0,0 @@
-import unittest
-
-import torch
-
-from backbones import TrajMixer
-
-
-class TrajMixerTest(unittest.TestCase):
-    def test_default_shape_parameters_and_initialization(self) -> None:
-        mixer = TrajMixer(
-            n_embd=120,
-            n_head=10,
-            dropout=0.0,
-        )
-
-        x = torch.randn(2, 7, 120)
-        self.assertEqual(mixer(x).shape, x.shape)
-        self.assertEqual(sum(p.numel() for p in mixer.parameters()), 32_040)
-        self.assertEqual(tuple(mixer.norm.normalized_shape), (120,))
-        self.assertEqual(tuple(mixer.intra_gate_logits.shape), (10, 12))
-        torch.testing.assert_close(
-            torch.sigmoid(mixer.intra_gate_logits.detach()),
-            torch.full((10, 12), 0.1),
-        )
-        self.assertEqual(mixer.intra_hidden, 48)
-        self.assertEqual(
-            tuple(mixer.intra_gate_proj.shape),
-            (10, 12, 48),
-        )
-        self.assertEqual(
-            tuple(mixer.intra_value_proj.shape),
-            (10, 12, 48),
-        )
-        self.assertEqual(
-            tuple(mixer.intra_output_proj.shape),
-            (10, 48, 12),
-        )
-        self.assertEqual(mixer.hidden_group, 40)
-        self.assertEqual(tuple(mixer.gate_proj.shape), (12, 10, 40))
-        self.assertEqual(tuple(mixer.value_proj.shape), (12, 10, 40))
-        self.assertEqual(tuple(mixer.output_proj.shape), (12, 40, 10))
-
-    def test_zero_final_output_projection_makes_mixer_identity(self) -> None:
-        torch.manual_seed(0)
-        mixer = TrajMixer(12, n_head=3, dropout=0.0)
-        with torch.no_grad():
-            mixer.output_proj.zero_()
-        x = torch.randn(2, 5, 12)
-        torch.testing.assert_close(mixer(x), x)
-
-    def test_forward_matches_single_outer_residual_formula(self) -> None:
-        torch.manual_seed(0)
-        mixer = TrajMixer(12, n_head=3, dropout=0.0)
-        mixer.eval()
-        x = torch.randn(2, 5, 12)
-
-        grouped = mixer.norm(x).reshape(2, 5, 3, 4)
-        intra_output = mixer._intra_mix(grouped)
-        static_gate = torch.sigmoid(mixer.intra_gate_logits).view(
-            1, 1, 3, 4
-        )
-        mixed_input = grouped + static_gate * intra_output
-        update = mixer._cross_mix(mixed_input).reshape(2, 5, 12)
-
-        torch.testing.assert_close(mixer(x), x + update)
-
-    def test_intra_stage_is_independent_across_groups(self) -> None:
-        torch.manual_seed(0)
-        mixer = TrajMixer(12, n_head=3, dropout=0.0)
-        mixer.eval()
-
-        grouped = torch.randn(2, 4, 3, 4)
-        changed = grouped.clone()
-        changed[:, :, 1, :] += torch.randn_like(changed[:, :, 1, :])
-
-        original_out = mixer._intra_mix(grouped)
-        changed_out = mixer._intra_mix(changed)
-        unchanged_groups = torch.tensor([0, 2])
-        torch.testing.assert_close(
-            original_out.index_select(2, unchanged_groups),
-            changed_out.index_select(2, unchanged_groups),
-        )
-
-    def test_cross_stage_mixes_groups_without_mixing_coordinates(self) -> None:
-        mixer = TrajMixer(6, n_head=3, dropout=0.0)
-        mixer.eval()
-        with torch.no_grad():
-            mixer.gate_proj.zero_()
-            mixer.value_proj.zero_()
-            mixer.output_proj.zero_()
-
-            # Coordinate 0 reads group 0 through hidden unit 0 and writes it
-            # into group 1. Coordinate 1 must remain independent.
-            mixer.gate_proj[0, 0, 0] = 1.0
-            mixer.value_proj[0, 0, 0] = 1.0
-            mixer.output_proj[0, 0, 1] = 1.0
-
-        grouped = torch.tensor(
-            [[[
-                [-1.0, 4.0],
-                [0.0, 5.0],
-                [1.0, 6.0],
-            ]]]
-        )
-        changed = grouped.clone()
-        changed[0, 0, 0, 0] = 2.0
-
-        original_out = mixer._cross_mix(grouped)
-        changed_out = mixer._cross_mix(changed)
-
-        self.assertNotEqual(
-            original_out[0, 0, 1, 0].item(),
-            changed_out[0, 0, 1, 0].item(),
-        )
-        torch.testing.assert_close(
-            original_out[..., 1],
-            changed_out[..., 1],
-        )
-
-    def test_mixer_does_not_mix_sequence_positions(self) -> None:
-        torch.manual_seed(0)
-        mixer = TrajMixer(12, n_head=3, dropout=0.0)
-        mixer.eval()
-        x = torch.randn(2, 5, 12)
-        changed = x.clone()
-        changed[:, 3, :] += torch.randn_like(changed[:, 3, :])
-
-        original_out = mixer(x)
-        changed_out = mixer(changed)
-        unchanged_positions = torch.tensor([0, 1, 2, 4])
-        torch.testing.assert_close(
-            original_out.index_select(1, unchanged_positions),
-            changed_out.index_select(1, unchanged_positions),
-        )
-
-    def test_gradients_reach_every_projection_family(self) -> None:
-        torch.manual_seed(1)
-        mixer = TrajMixer(12, n_head=3, dropout=0.0)
-        x = torch.randn(2, 4, 12, requires_grad=True)
-
-        mixer(x).square().mean().backward()
-
-        self.assertIsNotNone(x.grad)
-        self.assertTrue(torch.isfinite(x.grad).all())
-        self.assertGreater(x.grad.abs().sum().item(), 0.0)
-        projection_names = (
-            "intra_gate_proj",
-            "intra_value_proj",
-            "intra_output_proj",
-            "gate_proj",
-            "value_proj",
-            "output_proj",
-        )
-        for name in projection_names:
-            parameter = getattr(mixer, name)
-            self.assertIsNotNone(parameter.grad, name)
-            self.assertTrue(torch.isfinite(parameter.grad).all(), name)
-            self.assertGreater(parameter.grad.abs().sum().item(), 0.0, name)
-
-    def test_invalid_group_partition_is_rejected(self) -> None:
-        with self.assertRaisesRegex(ValueError, "divisible"):
-            TrajMixer(n_embd=121, n_head=10)
-
-
-if __name__ == "__main__":
-    unittest.main()

tests/test_all_future_likelihood.py

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

tests/test_auc_eid_split.py

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

tests/test_calibration_metrics.py

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

tests/test_continuous_value_scaling.py

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

tests/test_dataset_reserved_event.py

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

tests/test_disease_history_ablation.py

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

train_all_future.py

@@ -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,20 +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 dataset import (
+    DISEASE_HISTORY_MODES,
+    DISEASE_HISTORY_MODE_TIMED,
+    AllFutureHealthDataset,
+    all_future_collate_fn,
+)
 from losses import build_loss
 from model_architectures import (
     DEFAULT_MODEL_ARCHITECTURE,
     SUPPORTED_MODEL_ARCHITECTURES,
 )
 from models import DeepHealth
-from targets import CHECKUP_IDX, PAD_IDX
+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,
@@ -72,6 +83,16 @@ def parse_args() -> argparse.Namespace:
     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)
@@ -92,8 +113,17 @@ def parse_args() -> argparse.Namespace:
     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,
@@ -132,28 +162,44 @@ 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,
         n_embd=args.n_embd,
@@ -164,27 +210,29 @@ def build_model(args: argparse.Namespace, dataset: AllFutureHealthDataset) -> De
         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}")
 
 
@@ -196,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,
@@ -214,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)
 
@@ -223,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(
@@ -231,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())}")
@@ -303,7 +345,10 @@ def build_metadata(
     train_subset,
     val_subset,
     test_subset,
+    scaler_stats: ContinuousRobustScalerStats,
+    baseline_stats: AllFutureBaselineStats | None,
 ) -> Dict[str, Any]:
+    scaler_metadata = scaler_stats.as_metadata()
     return {
         "run_name": run_name,
         "dataset_class": "AllFutureHealthDataset",
@@ -312,16 +357,32 @@ def build_metadata(
         "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),
@@ -329,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)),
@@ -347,7 +410,15 @@ def main() -> None:
     configure_torch_for_training(device)
 
     run_dir, run_name = create_unique_run_dir(
-        lambda timestamp: f"{args.time_mode}_{args.dist_mode}_all_future_pure_disease_{timestamp}",
+        lambda 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)
@@ -355,7 +426,9 @@ def main() -> None:
     logger.info(f"Starting all-future training run: {run_name}")
     logger.info(f"Device: {device}")
     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(
@@ -366,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,
@@ -375,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,
@@ -384,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(
@@ -416,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,
@@ -447,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(

train_batch_linux.sh

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

train_disease_history_ablation_linux.sh

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

train_extra_info_assessment_all_multiseed_linux.sh

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

train_key_models_multiseed_linux.sh

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

train_next_step.py

@@ -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
@@ -29,14 +23,17 @@ from model_architectures import (
     SUPPORTED_MODEL_ARCHITECTURES,
 )
 from models import DeepHealth, DeepHealthOutput
-from readouts import build_readout
-from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX
+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,
@@ -70,7 +67,6 @@ def parse_args() -> argparse.Namespace:
     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)
@@ -85,8 +81,6 @@ def parse_args() -> argparse.Namespace:
     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",
@@ -95,17 +89,10 @@ def parse_args() -> argparse.Namespace:
         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)
@@ -127,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
@@ -140,25 +122,17 @@ 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,
         n_embd=args.n_embd,
@@ -169,46 +143,31 @@ def build_model(args: argparse.Namespace, dataset: HealthDataset) -> DeepHealth:
         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,
-            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},
+        "delphi2m",
+        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,
     )
 
 
 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)
@@ -216,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)
@@ -268,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():
@@ -326,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())
@@ -335,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,
@@ -382,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):
@@ -390,35 +284,17 @@ 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,
-            return_components=True,
-        )
+    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=model_out.padding_mask,
+        return_components=True,
+    )
     if not torch.isfinite(loss):
         raise RuntimeError(f"Loss is not finite: {float(loss.detach().cpu())}")
     return loss, parts
@@ -428,7 +304,6 @@ def run_epoch(
     logger: logging.Logger,
     args: argparse.Namespace,
     model: DeepHealth,
-    readout,
     criterion,
     loader: DataLoader,
     optimizer: AdamW | None,
@@ -436,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
@@ -447,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")
@@ -489,6 +365,7 @@ def build_metadata(
     train_subset,
     val_subset,
     test_subset,
+    scaler_stats: ContinuousRobustScalerStats,
 ) -> Dict[str, Any]:
     return {
         "run_name": run_name,
@@ -497,7 +374,10 @@ def build_metadata(
         "model_class": "DeepHealth",
         "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
@@ -505,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),
@@ -512,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",
     }
 
 
@@ -531,7 +415,7 @@ def main() -> None:
 
     run_dir, run_name = create_unique_run_dir(
         lambda timestamp: (
-            f"{args.time_mode}_exponential_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,
@@ -542,13 +426,13 @@ def main() -> None:
     logger.info(f"Device: {device}")
     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:
@@ -578,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,
@@ -609,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(),
@@ -627,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(
@@ -646,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:
@@ -682,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}")
 

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,15 +584,6 @@ 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 total and trainable parameter counts."""
     return {