From de6f9b75b9660600fec140a53e23324f4fe7f7ce Mon Sep 17 00:00:00 2001 From: Jiarui Li Date: Sat, 1 Aug 2026 14:23:18 +0800 Subject: [PATCH] Remove legacy event and mixed distribution paths --- README.md | 18 +-- dataset.py | 24 ++-- eval_data.py | 38 +++++- evaluate_auc.py | 37 +++--- evaluate_auc_v2.py | 89 +++----------- evaluate_calibration.py | 52 ++------ extra_info_assessment_all_analysis_plan.md | 23 ++-- losses.py | 78 ++---------- models.py | 60 +++------ prepare_data.py | 18 +-- prepare_event_dates.py | 6 +- targets.py | 4 +- tests/test_calibration_metrics.py | 13 +- tests/test_continuous_value_scaling.py | 115 +++++++++++++++--- tests/test_dataset_checkup.py | 57 --------- tests/test_dataset_reserved_event.py | 44 +++++++ train_all_future.py | 82 ++++--------- train_batch_linux.sh | 5 +- ...tra_info_assessment_all_multiseed_linux.sh | 1 - train_key_models_multiseed_linux.sh | 1 + train_next_step.py | 55 +++++++-- train_util.py | 13 +- 22 files changed, 370 insertions(+), 463 deletions(-) delete mode 100644 tests/test_dataset_checkup.py create mode 100644 tests/test_dataset_reserved_event.py diff --git a/README.md b/README.md index bd9aa53..5aeb7ab 100644 --- a/README.md +++ b/README.md @@ -6,7 +6,9 @@ 疾病序列 stream + 统一的额外信息 token stream ``` -疾病、死亡、checkup 事件保存在预处理事件文件中；性别单独保存在 `basic_info`；其他体检、暴露、生活方式等信息统一整理成 `(type, value, value_kind, time)` token。dataset 仅在实验选择了至少一种 extra-info type 时保留 checkup；显式传入空列表时，模型输入是没有 checkup 的纯疾病历史。 +疾病和死亡事件保存在预处理事件文件中；性别单独保存在 `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` @@ -64,7 +66,7 @@ python prepare_data.py - `AllFutureHealthDataset` - 用于 query-conditioned all-future 监督 - - 对应 `ExponentialLoss`、`WeibullLoss`、`MixedLoss` + - 对应 `ExponentialLoss`、`WeibullLoss` 为了兼容旧训练入口： @@ -215,7 +217,6 @@ all-future / query-conditioned 监督： - `ExponentialLoss` - `WeibullLoss` -- `MixedLoss` all-future 训练只读出 `t_query` 对应的 query hidden。展开的 extra-info tokens 作为主序列上下文输入，但不会被单独读出，也不会被纳入 loss 监督。 @@ -232,7 +233,7 @@ all-future 训练只读出 `t_query` 对应的 query hidden。展开的 extra-in - `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` 支持所有已有训练目标定义的组合： @@ -242,7 +243,6 @@ all-future 训练只读出 `t_query` 对应的 query hidden。展开的 extra-in | `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 | 示例： @@ -366,7 +366,7 @@ python evaluate_auc.py \ - 为每个患者和 landmark age 构造 landmark query 样本。 - 根据模型模式插入 `` 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 内先于目标疾病发生的死亡竞争风险。 @@ -402,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 用人工插入的 `` 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 | @@ -450,4 +450,4 @@ python evaluate_auc_v2.py \ - landmark fixed-horizon incident disease AUC 评估 - next-token 模型通过插入 `` landmark token 查询固定年龄点风险 - all-future 模型直接通过 `t_query` 查询固定年龄点风险 - - `score_mode="risk"` 按 exponential / Weibull / mixed 分布计算固定 horizon 风险 + - `score_mode="risk"` 按 exponential / Weibull 分布计算固定 horizon 风险 diff --git a/dataset.py b/dataset.py index e8c66ca..40b278c 100644 --- a/dataset.py +++ b/dataset.py @@ -10,10 +10,10 @@ 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, + RESERVED_IDX, build_next_token_targets, ) @@ -188,12 +188,12 @@ def load_label_vocab( ) -> Tuple[Dict[str, int], Dict[int, str]]: label_id_to_code: Dict[int, str] = { PAD_IDX: "", - CHECKUP_IDX: "", + RESERVED_IDX: "", } if include_no_event: label_id_to_code[NO_EVENT_IDX] = "" - 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): @@ -416,14 +416,12 @@ class _ExpoBaseDataset(Dataset): times_days_raw = rows[:, 1].astype(np.float32) labels_raw = rows[:, 2].astype(np.int64) - # CHECKUP is the assessment landmark for selected extra-info tokens. - # An explicitly empty selection represents a disease-only history, - # so retaining CHECKUP in that case would introduce an empty - # landmark token that is not part of the disease sequence. - if not self.extra_info_types: - keep = labels_raw != CHECKUP_IDX - times_days_raw = times_days_raw[keep] - labels_raw = labels_raw[keep] + # 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 @@ -615,7 +613,7 @@ 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_future = int(((times > t_query) & real_event_mask).sum()) @@ -634,7 +632,7 @@ class AllFutureHealthDataset(_ExpoBaseDataset): 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), + np.array([PAD_IDX, RESERVED_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) diff --git a/eval_data.py b/eval_data.py index 7b76249..c671e97 100644 --- a/eval_data.py +++ b/eval_data.py @@ -151,6 +151,33 @@ def build_model_from_dataset( f"{model_target_mode!r}" ) 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)), @@ -161,9 +188,8 @@ def build_model_from_dataset( n_categories=dataset.n_categories, cont_type_ids=dataset.cont_type_ids, n_bins=int(cfg_get(args, cfg, "n_bins", 16)), - continuous_value_scaling=str( - cfg_get(args, cfg, "continuous_value_scaling", "none") - ), + continuous_value_center=continuous_value_center, + continuous_value_scale=continuous_value_scale, extra_pool_reduce=str( cfg_get(args, cfg, "extra_pool_reduce", "mean") ), @@ -251,9 +277,9 @@ class AllFutureSequenceEvalDataset: Eval-only sequence view for all-future checkpoints. All-future training uses the observed history without reusing the - next-step view that contains imputed gap tokens. CHECKUP is - retained only when the experiment selects at least one extra-info type; - an explicitly empty selection is a disease-only history. + next-step view that contains imputed gap tokens. Legacy label-1 + assessment events are removed by the shared base dataset for every + extra-info selection. """ def __init__( diff --git a/evaluate_auc.py b/evaluate_auc.py index a3b36f5..5e18b4d 100644 --- a/evaluate_auc.py +++ b/evaluate_auc.py @@ -63,7 +63,7 @@ from eval_data import ( ) from model_architectures import resolve_model_architecture from models import DeepHealth -from targets import PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX +from targets import NO_EVENT_IDX, PAD_IDX, RESERVED_IDX # --------------------------------------------------------------------------- @@ -168,7 +168,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} +SPECIAL_TOKENS = {PAD_IDX, RESERVED_IDX, NO_EVENT_IDX} def _get_death_token_ids(dataset: HealthDataset) -> List[int]: @@ -249,27 +249,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 @@ -977,7 +970,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) == "": valid_target_min_id = NO_EVENT_IDX diff --git a/evaluate_auc_v2.py b/evaluate_auc_v2.py index 05b7386..1a1ece1 100644 --- a/evaluate_auc_v2.py +++ b/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. @@ -53,9 +53,9 @@ from eval_data import ( ) from model_architectures import resolve_model_architecture from models import DeepHealth -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} +SPECIAL_TOKENS = {PAD_IDX, RESERVED_IDX, NO_EVENT_IDX} def parse_int_list(value: Any) -> Optional[List[int]]: @@ -137,29 +137,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" + if mode == "weibull" and not has_rho_head: + raise RuntimeError("Weibull checkpoint is missing rho_head parameters.") + if mode == "exponential" and has_rho_head: + raise RuntimeError( + "Exponential checkpoint unexpectedly contains rho_head parameters." + ) + return mode def load_model_state(model: DeepHealth, state_dict: Dict[str, Any]) -> None: @@ -290,7 +280,7 @@ 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] @@ -418,7 +408,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 @@ -671,24 +661,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] = [] @@ -703,14 +680,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)) @@ -747,7 +716,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") @@ -772,7 +740,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": {}, } ) @@ -798,8 +765,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) @@ -811,11 +776,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) @@ -832,7 +792,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"] @@ -891,8 +850,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]], @@ -900,8 +857,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) @@ -1019,8 +974,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)] @@ -1048,8 +1001,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( @@ -1253,10 +1204,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( @@ -1302,7 +1249,7 @@ 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, diff --git a/evaluate_calibration.py b/evaluate_calibration.py index d07423d..6422703 100644 --- a/evaluate_calibration.py +++ b/evaluate_calibration.py @@ -23,12 +23,11 @@ from __future__ import annotations import argparse import contextlib import json -import math import os from collections import defaultdict from concurrent.futures import ThreadPoolExecutor, as_completed from pathlib import Path -from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple +from typing import Any, Dict, List, Optional, Sequence, Tuple import numpy as np import pandas as pd @@ -67,7 +66,7 @@ from evaluate_auc_v2 import ( ) from losses import build_loss from model_architectures import resolve_model_architecture -from targets import CHECKUP_IDX, PAD_IDX +from targets import PAD_IDX, RESERVED_IDX from train_util import load_eid_file @@ -971,8 +970,6 @@ def _risk_probability_matrix( rho: Optional[np.ndarray], horizons: np.ndarray, dist_mode: str, - token: int, - death_idx: int, ) -> np.ndarray: """Convert one token's logits to all horizon risks at once.""" logits = np.asarray(logits, dtype=np.float32) @@ -982,13 +979,7 @@ def _risk_probability_matrix( + np.maximum(logits, np.float32(0.0)) + np.float32(1e-8) ) - use_weibull = ( - str(dist_mode).lower() == "weibull" - or ( - str(dist_mode).lower() == "mixed" - and int(token) == int(death_idx) - ) - ) + use_weibull = str(dist_mode).lower() == "weibull" if use_weibull: if rho is None: raise RuntimeError( @@ -1047,7 +1038,6 @@ def _evaluate_calibration_token( label_id_to_code: Dict[int, str], dist_mode: str, horizons: np.ndarray, - death_index: int, min_cases: int, min_controls: int, max_ipcw_weight: float, @@ -1114,8 +1104,6 @@ def _evaluate_calibration_token( ), horizons=horizons, dist_mode=dist_mode, - token=token, - death_idx=death_index, ) results = compute_ipcw_horizons( probabilities=probabilities, @@ -1200,9 +1188,6 @@ def evaluate_landmark_calibration( patient_count = len(landmark_dataset.subset_indices) death_tokens = set(int(value) for value in landmark_dataset.death_token_ids) - death_index = int( - getattr(model, "death_idx", getattr(model, "vocab_size", 1) - 1) - ) strata = _build_calibration_strata( row_arrays["sex"], row_arrays["landmark_age"], @@ -1265,7 +1250,6 @@ def evaluate_landmark_calibration( ), "dist_mode": dist_mode, "horizons": horizons, - "death_index": death_index, "min_cases": min_cases, "min_controls": min_controls, "max_ipcw_weight": max_ipcw_weight, @@ -1498,19 +1482,12 @@ def build_calibration_summary(metrics: pd.DataFrame) -> pd.DataFrame: def _build_point_process_criterion( dist_mode: str, - death_index: int, ) -> Any: - ignored = {PAD_IDX, CHECKUP_IDX} + ignored = {PAD_IDX, RESERVED_IDX} if dist_mode == "exponential": return build_loss("exponential", ignored_idx=ignored) if dist_mode == "weibull": return build_loss("weibull", ignored_idx=ignored) - if dist_mode == "mixed": - return build_loss( - "mixed", - death_idx=death_index, - ignored_idx=ignored, - ) raise ValueError(f"Unsupported dist_mode: {dist_mode!r}") @@ -1523,16 +1500,7 @@ def evaluate_point_process_nll( device: torch.device, use_amp: bool, ) -> Dict[str, Any]: - criterion = _build_point_process_criterion( - dist_mode, - death_index=int( - getattr( - model, - "death_idx", - int(getattr(model, "vocab_size", 1)) - 1, - ) - ), - ) + criterion = _build_point_process_criterion(dist_mode) model.eval().to(device) total_nll = 0.0 query_count = 0 @@ -1587,13 +1555,7 @@ def evaluate_point_process_nll( exposure=batch_device["exposure"], ) else: - loss = criterion( - logits=logits, - death_rho=model.calc_death_rho(hidden), - targets=batch_device["future_targets"], - dt=batch_device["future_dt"], - exposure=batch_device["exposure"], - ) + raise ValueError(f"Unsupported dist_mode: {dist_mode!r}") if not torch.isfinite(loss): raise RuntimeError("Non-finite point-process NLL encountered.") @@ -1601,7 +1563,7 @@ def evaluate_point_process_nll( total_nll += float(loss.detach().cpu()) * batch_size query_count += batch_size valid_targets = batch["future_targets"] > PAD_IDX - valid_targets &= batch["future_targets"] != CHECKUP_IDX + valid_targets &= batch["future_targets"] != RESERVED_IDX future_event_count += int(valid_targets.sum().item()) exposure_sum += float(batch["exposure"].sum().item()) diff --git a/extra_info_assessment_all_analysis_plan.md b/extra_info_assessment_all_analysis_plan.md index bcc8b69..7a14091 100644 --- a/extra_info_assessment_all_analysis_plan.md +++ b/extra_info_assessment_all_analysis_plan.md @@ -4,10 +4,13 @@ 补充现有四级 extra-information 证据链： -1. `disease_only`：疾病事件、相对患病时间和 sex；无 CHECKUP、无 extra-info token。 -2. `smoking_alcohol_bmi`：疾病史、sex、CHECKUP、smoking/alcohol/BMI。 -3. `assessment_only`：疾病史、sex、CHECKUP、65项常规体格、肺功能、血液、尿液和生化指标。 -4. `all`：疾病史、sex、CHECKUP、全部265项体检和暴露信息。 +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；训练与评估不设置未标准化对照或兼容模式。 目标是区分： @@ -63,7 +66,7 @@ assessment_only − disease_only all − assessment_only ``` -回答生活方式、社会经济、心理和环境暴露是否在常规体检之后仍有增量价值。这是新增实验中最干净的主要比较，因为两组都保留 CHECKUP，Landmark 和随访边界应一致。 +回答生活方式、社会经济、心理和环境暴露是否在常规体检之后仍有增量价值。两组使用完全相同的疾病事件、Landmark 和随访边界，仅 extra-info token 集合不同。 ### 4.3 全体检相对紧凑变量集 @@ -130,15 +133,9 @@ smoking_alcohol_bmi − disease_only - Brier、NLL及绝对校准偏差越低越好； - 不以单个 seed 或单个 horizon 决定模型。 -## 7. disease_only 比较的评估限制 +## 7. disease_only 比较的评估边界 -`disease_only` 按设计删除 CHECKUP，其他三组保留 CHECKUP。当前评估实现会使两类模型的随访终点和 `n_at_risk` 略有差异。 - -因此： - -- `assessment_only`、`smoking_alcohol_bmi`、`all` 三者之间可以直接比较； -- 它们与 `disease_only` 的比较应使用固定的原始随访终点、Landmark 和 censoring； -- 在共享风险集评估完成前，不能把与 `disease_only` 的全部差异严格归因于 extra-info 数值。 +所有配置均从同一疾病/死亡事件流构造历史、查询点、随访终点和 censoring，并且都不使用 CHECKUP。不同配置只改变 extra-info token，因此可以在共同支持集上把差异解释为额外信息的增量价值。 ## 8. 决策规则 diff --git a/losses.py b/losses.py index ee44c60..dfbc122 100644 --- a/losses.py +++ b/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 @@ -52,7 +52,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] ) @@ -150,7 +150,7 @@ class ExponentialLoss(nn.Module): 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__() @@ -183,7 +183,7 @@ class WeibullLoss(nn.Module): 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__() @@ -232,78 +232,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 == "delphi2m": return Delphi2MLoss(**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, exponential, weibull, mixed." + f"Unknown loss {name!r}. Available: delphi2m, exponential, weibull." ) diff --git a/models.py b/models.py index 68c3b7e..c2c1b30 100644 --- a/models.py +++ b/models.py @@ -37,7 +37,6 @@ class OtherInfoTokenizer(nn.Module): cont_type_ids: list[int], n_value_kinds: int = 3, n_bins: int = 16, - continuous_value_scaling: str = "none", continuous_value_center: torch.Tensor | list[float] | None = None, continuous_value_scale: torch.Tensor | list[float] | None = None, ): @@ -57,13 +56,6 @@ class OtherInfoTokenizer(nn.Module): raise ValueError( f"n_value_kinds must be > {self.CATE_KIND}, got {n_value_kinds}" ) - continuous_value_scaling = str(continuous_value_scaling).lower() - if continuous_value_scaling not in {"none", "robust"}: - raise ValueError( - "continuous_value_scaling must be either 'none' or 'robust', " - f"got {continuous_value_scaling!r}" - ) - 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 = ( @@ -80,18 +72,20 @@ class OtherInfoTokenizer(nn.Module): n_embd, padding_idx=0, ) - self.continuous_value_scaling = continuous_value_scaling - if continuous_value_scaling == "robust" and n_cont_types > 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, - default=0.0, name="continuous_value_center", ) scale = self._coerce_scaler_buffer( continuous_value_scale, n_cont_types=n_cont_types, - default=1.0, name="continuous_value_scale", ) if not torch.isfinite(center).all(): @@ -130,11 +124,10 @@ class OtherInfoTokenizer(nn.Module): value: torch.Tensor | list[float] | None, *, n_cont_types: int, - default: float, name: str, ) -> torch.Tensor: if value is None: - return torch.full((n_cont_types,), float(default), dtype=torch.float32) + raise ValueError(f"{name} is required") tensor = torch.as_tensor(value, dtype=torch.float32).detach().clone() if tensor.shape != (n_cont_types,): raise ValueError( @@ -184,15 +177,14 @@ class OtherInfoTokenizer(nn.Module): "cont_type_ids" ) cont_value = other_value[cont_pos].to(type_emb.dtype) - if self.continuous_value_scaling == "robust": - if ( - self.continuous_value_center is None - or self.continuous_value_scale is None - ): - raise RuntimeError("Robust continuous-value scaler 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 + 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=cont_value, @@ -221,12 +213,11 @@ class DeepHealth(nn.Module): cont_type_ids: list[int], n_value_kinds: int = 3, n_bins: int = 16, - continuous_value_scaling: str = "none", 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 = "absolute", # next_token requires absolute - dist_mode: str = "exponential", # "exponential", "weibull" or "mixed" + dist_mode: str = "exponential", # "exponential" or "weibull" extra_pool_reduce: str = "mean", dropout: float = 0.0, model_architecture: str | None = None, @@ -243,9 +234,9 @@ class DeepHealth(nn.Module): "next_token is reserved for Delphi2M reproduction and " "requires time_mode='absolute'" ) - if dist_mode not in ["exponential", "weibull", "mixed"]: + if dist_mode not in ["exponential", "weibull"]: raise ValueError( - "dist_mode must be either 'exponential', 'weibull' or 'mixed'") + "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: @@ -262,7 +253,6 @@ class DeepHealth(nn.Module): cont_type_ids=cont_type_ids, n_value_kinds=n_value_kinds, n_bins=n_bins, - continuous_value_scaling=continuous_value_scaling, continuous_value_center=continuous_value_center, continuous_value_scale=continuous_value_scale, ) @@ -270,7 +260,6 @@ class DeepHealth(nn.Module): self.time_mode = time_mode self.dist_mode = dist_mode self.extra_pool_reduce = extra_pool_reduce - self.continuous_value_scaling = str(continuous_value_scaling).lower() self.model_architecture = model_architecture self.n_layer = n_layer self.n_embd = n_embd @@ -283,12 +272,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([ @@ -551,10 +534,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 diff --git a/prepare_data.py b/prepare_data.py index 2838a70..f06c496 100644 --- a/prepare_data.py +++ b/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 diff --git a/prepare_event_dates.py b/prepare_event_dates.py index 4d84ad3..4c1d237 100644 --- a/prepare_event_dates.py +++ b/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 ----- diff --git a/targets.py b/targets.py index 818c4e2..44fb6ae 100644 --- a/targets.py +++ b/targets.py @@ -6,7 +6,7 @@ import numpy as np PAD_IDX = 0 -CHECKUP_IDX = 1 +RESERVED_IDX = 1 NO_EVENT_IDX = 2 DAYS_PER_YEAR = 365.25 @@ -105,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) diff --git a/tests/test_calibration_metrics.py b/tests/test_calibration_metrics.py index 8570345..52a0bc4 100644 --- a/tests/test_calibration_metrics.py +++ b/tests/test_calibration_metrics.py @@ -207,19 +207,15 @@ class IPCWCalibrationMetricTests(unittest.TestCase): 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, token, death_idx, selected_rho in ( - ("exponential", 4, 9, None), - ("weibull", 4, 9, rho), - ("mixed", 9, 9, rho), - ("mixed", 4, 9, None), + 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, - token=token, - death_idx=death_idx, ) expected = np.vstack( [ @@ -229,8 +225,6 @@ class IPCWCalibrationMetricTests(unittest.TestCase): score_mode="risk", horizon=float(horizon), dist_mode=dist_mode, - token=token, - death_idx=death_idx, ) for horizon in horizons ] @@ -271,7 +265,6 @@ class IPCWCalibrationMetricTests(unittest.TestCase): "label_id_to_code": {4: "D4", 5: "D5"}, "dist_mode": "exponential", "horizons": np.asarray([1.0, 5.0], dtype=np.float32), - "death_index": 9, "min_cases": 1, "min_controls": 1, "max_ipcw_weight": 0.0, diff --git a/tests/test_continuous_value_scaling.py b/tests/test_continuous_value_scaling.py index c688a18..f0fe2ab 100644 --- a/tests/test_continuous_value_scaling.py +++ b/tests/test_continuous_value_scaling.py @@ -5,7 +5,8 @@ import torch import torch.nn as nn from torch.utils.data import Subset -from models import OtherInfoTokenizer +from eval_data import build_model_from_dataset +from models import DeepHealth, OtherInfoTokenizer from train_util import fit_continuous_robust_scaler @@ -59,6 +60,17 @@ class ContinuousValueScalingTests(unittest.TestCase): 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, @@ -66,7 +78,6 @@ class ContinuousValueScalingTests(unittest.TestCase): n_cont_types=2, n_categories=3, cont_type_ids=[1, 3], - continuous_value_scaling="robust", continuous_value_center=[10.0, 100.0], continuous_value_scale=[2.0, 20.0], ) @@ -89,7 +100,6 @@ class ContinuousValueScalingTests(unittest.TestCase): n_cont_types=2, n_categories=2, cont_type_ids=[1, 3], - continuous_value_scaling="robust", continuous_value_center=[2.0, 10.0], continuous_value_scale=[1.5, 4.0], ) @@ -103,7 +113,8 @@ class ContinuousValueScalingTests(unittest.TestCase): n_cont_types=2, n_categories=2, cont_type_ids=[1, 3], - continuous_value_scaling="robust", + continuous_value_center=[0.0, 0.0], + continuous_value_scale=[1.0, 1.0], ) restored.load_state_dict(state, strict=True) @@ -116,27 +127,99 @@ class ContinuousValueScalingTests(unittest.TestCase): torch.tensor([1.5, 4.0]), ) - def test_legacy_none_mode_keeps_old_state_dict_schema(self): - tokenizer = OtherInfoTokenizer( - n_embd=4, - n_types=4, - n_cont_types=2, - n_categories=2, - cont_type_ids=[1, 3], - ) - state = tokenizer.state_dict() + 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], + ) - self.assertNotIn("continuous_value_center", state) - self.assertNotIn("continuous_value_scale", state) - restored = OtherInfoTokenizer( + 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() diff --git a/tests/test_dataset_checkup.py b/tests/test_dataset_checkup.py deleted file mode 100644 index 4bea4d0..0000000 --- a/tests/test_dataset_checkup.py +++ /dev/null @@ -1,57 +0,0 @@ -import unittest -from pathlib import Path - -import numpy as np - -from dataset import _ExpoBaseDataset -from targets import CHECKUP_IDX -from train_util import load_extra_info_types_file - - -class CheckupSelectionTests(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, CHECKUP_IDX], - [101, 20, 2], - [101, 30, 3], - ], - dtype=np.float64, - ) - return dataset - - def test_explicit_empty_extra_info_removes_checkup(self): - project_root = Path(__file__).resolve().parents[1] - selected_types = load_extra_info_types_file( - str(project_root / "extra_info_types_none.txt") - ) - self.assertEqual(selected_types, []) - dataset = self._base(selected_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)) - self.assertNotIn(CHECKUP_IDX, labels.tolist()) - - def test_selected_extra_info_keeps_checkup(self): - dataset = self._base([11]) - - rows = list(dataset._iter_patient_events(impute_no_event_gaps=False)) - - self.assertEqual(len(rows), 1) - _, times, labels = rows[0] - np.testing.assert_array_equal( - times, - np.asarray([10, 20, 30], dtype=np.float32), - ) - self.assertEqual(labels[0], CHECKUP_IDX) - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/test_dataset_reserved_event.py b/tests/test_dataset_reserved_event.py new file mode 100644 index 0000000..2f9b9b8 --- /dev/null +++ b/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() diff --git a/train_all_future.py b/train_all_future.py index 34b97a1..ab7ac3c 100644 --- a/train_all_future.py +++ b/train_all_future.py @@ -37,7 +37,7 @@ from model_architectures import ( SUPPORTED_MODEL_ARCHITECTURES, ) from models import DeepHealth -from targets import CHECKUP_IDX, PAD_IDX +from targets import PAD_IDX, RESERVED_IDX from train_util import ( ContinuousRobustScalerStats, configure_torch_for_training, @@ -106,22 +106,12 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--n_head", type=int, default=10) parser.add_argument("--n_layer", type=int, default=12) parser.add_argument("--n_bins", type=int, default=16) - parser.add_argument( - "--continuous_value_scaling", - type=str, - default="robust", - choices=["none", "robust"], - help=( - "Continuous extra-info scaling. 'robust' fits the median and IQR " - "on the complete training subset and stores them in the checkpoint." - ), - ) parser.add_argument("--extra_pool_reduce", type=str, default="mean", choices=["mean", "sum"]) parser.add_argument("--time_mode", type=str, default="relative", choices=["relative", "absolute"]) parser.add_argument("--dist_mode", type=str, default="exponential", - choices=["exponential", "weibull", "mixed"]) + choices=["exponential", "weibull"]) parser.add_argument("--dropout", type=float, default=0.0) parser.add_argument( "--model_architecture", @@ -188,19 +178,14 @@ def parse_args() -> argparse.Namespace: def build_model( args: argparse.Namespace, dataset: AllFutureHealthDataset, - scaler_stats: ContinuousRobustScalerStats | None = None, + scaler_stats: ContinuousRobustScalerStats, ) -> DeepHealth: - if ( - args.continuous_value_scaling == "robust" - and dataset.n_cont_types > 0 - and scaler_stats is None - ): + if tuple(int(x) for x in dataset.cont_type_ids) != scaler_stats.cont_type_ids: raise ValueError( - "Robust continuous-value scaling requires statistics fitted on the " - "training subset" + "RobustScale statistics are not aligned with dataset.cont_type_ids" ) - center = None if scaler_stats is None else scaler_stats.center - scale = None if scaler_stats is None else scaler_stats.scale + 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, @@ -211,7 +196,6 @@ def build_model( n_categories=dataset.n_categories, cont_type_ids=dataset.cont_type_ids, n_bins=args.n_bins, - continuous_value_scaling=args.continuous_value_scaling, continuous_value_center=center, continuous_value_scale=scale, extra_pool_reduce=args.extra_pool_reduce, @@ -223,18 +207,12 @@ def build_model( ) -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} 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}") @@ -247,7 +225,7 @@ def compute_all_future_loss( ) -> torch.Tensor: required_keys = set(MODEL_INPUT_KEYS) required_keys.update(("future_targets", "exposure")) - if args.dist_mode in {"weibull", "mixed"}: + if args.dist_mode == "weibull": required_keys.add("future_dt") batch = move_batch_to_device( {key: batch[key] for key in required_keys}, @@ -282,13 +260,7 @@ def compute_all_future_loss( exposure=batch["exposure"], ) 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())}") @@ -352,16 +324,9 @@ def build_metadata( train_subset, val_subset, test_subset, - scaler_stats: ContinuousRobustScalerStats | None, + scaler_stats: ContinuousRobustScalerStats, ) -> Dict[str, Any]: - scaler_metadata: Dict[str, Any] - if scaler_stats is None: - scaler_metadata = { - "method": "none", - "fitted_on": None, - } - else: - scaler_metadata = scaler_stats.as_metadata() + scaler_metadata = scaler_stats.as_metadata() return { "run_name": run_name, "dataset_class": "AllFutureHealthDataset", @@ -370,6 +335,8 @@ 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), @@ -381,6 +348,7 @@ 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_metadata, "dataset_metadata": { "vocab_size": int(dataset.vocab_size), @@ -389,6 +357,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)), @@ -425,7 +395,7 @@ def main() -> None: 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(f"Continuous value scaling: {args.continuous_value_scaling}") + logger.info("Continuous value scaling: RobustScale (required)") logger.info("Loading all-future datasets...") train_dataset = AllFutureHealthDataset( @@ -489,16 +459,16 @@ def main() -> None: f"Patients/queries: train={len(train_subset)}, val={len(val_subset)}, test={len(test_subset)}" ) - scaler_stats = None - if args.continuous_value_scaling == "robust" and train_dataset.n_cont_types > 0: + 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, - ) + 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()):,}, " @@ -550,7 +520,7 @@ def main() -> None: 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( diff --git a/train_batch_linux.sh b/train_batch_linux.sh index d85a44b..a7880ce 100755 --- a/train_batch_linux.sh +++ b/train_batch_linux.sh @@ -4,7 +4,7 @@ # # The matrix contains: # 1. FFN + Delphi2M next-token reproduction -# 2. FFN/TrajMixer x absolute/relative x exponential/Weibull/mixed +# 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. @@ -47,6 +47,7 @@ Options: 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: @@ -180,7 +181,7 @@ add_job \ for architecture in transformer_ffn_v1 traj_mixer_v5; do for time_mode in absolute relative; do - for dist_mode in exponential weibull mixed; do + for dist_mode in exponential weibull; do add_job \ "${architecture}_all_future_${time_mode}_${dist_mode}" \ "train_all_future.py" \ diff --git a/train_extra_info_assessment_all_multiseed_linux.sh b/train_extra_info_assessment_all_multiseed_linux.sh index 8f755fa..1f30fec 100755 --- a/train_extra_info_assessment_all_multiseed_linux.sh +++ b/train_extra_info_assessment_all_multiseed_linux.sh @@ -236,7 +236,6 @@ run_job() { --time_mode relative --dist_mode weibull --disease_history_mode timed - --continuous_value_scaling robust --extra_info_types_file "$EXTRA_INFO_TYPES_FILE" ) diff --git a/train_key_models_multiseed_linux.sh b/train_key_models_multiseed_linux.sh index d7dd459..efa719b 100644 --- a/train_key_models_multiseed_linux.sh +++ b/train_key_models_multiseed_linux.sh @@ -59,6 +59,7 @@ Options: 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 diff --git a/train_next_step.py b/train_next_step.py index e6c8b76..77b232d 100644 --- a/train_next_step.py +++ b/train_next_step.py @@ -23,10 +23,12 @@ from model_architectures import ( SUPPORTED_MODEL_ARCHITECTURES, ) from models import DeepHealth, DeepHealthOutput -from targets import CHECKUP_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, @@ -120,7 +122,17 @@ def parse_args() -> argparse.Namespace: return args -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, @@ -131,6 +143,8 @@ 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="absolute", @@ -143,7 +157,7 @@ def build_model(args: argparse.Namespace, dataset: HealthDataset) -> DeepHealth: def build_next_step_loss(args: argparse.Namespace): return build_loss( "delphi2m", - ignored_tokens={PAD_IDX, CHECKUP_IDX}, + ignored_tokens={PAD_IDX, RESERVED_IDX}, t_min=args.t_min, max_exp_input=args.max_exp_input, ce_weight=args.ce_weight, @@ -244,7 +258,6 @@ def build_augmented_next_step_targets( def compute_next_step_loss( - args: argparse.Namespace, model: DeepHealth, criterion, batch: Dict[str, torch.Tensor], @@ -309,7 +322,7 @@ def run_epoch( for batch_idx, batch in enumerate(progress): try: loss, parts = compute_next_step_loss( - args, model, criterion, batch, device + model, criterion, batch, device ) if is_train: if optimizer is None: @@ -352,6 +365,7 @@ def build_metadata( train_subset, val_subset, test_subset, + scaler_stats: ContinuousRobustScalerStats, ) -> Dict[str, Any]: return { "run_name": run_name, @@ -361,6 +375,8 @@ def build_metadata( "model_architecture": args.model_architecture, "model_target_mode": "next_token", "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": ( @@ -369,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), @@ -376,6 +394,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)), @@ -406,6 +426,7 @@ 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("Continuous value scaling: RobustScale (required)") logger.info("time_mode=absolute, readout=token, target_mode=delphi2m") dataset = HealthDataset( @@ -441,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, @@ -472,7 +507,7 @@ 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: " @@ -489,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( diff --git a/train_util.py b/train_util.py index 33f0614..9d13ecf 100644 --- a/train_util.py +++ b/train_util.py @@ -43,7 +43,7 @@ class ContinuousRobustScalerStats: def fit_continuous_robust_scaler( - dataset: AllFutureHealthDataset, + dataset: HealthDataset | AllFutureHealthDataset, subset: Subset, *, quantile_range: tuple[float, float] = (25.0, 75.0), @@ -90,13 +90,22 @@ def fit_continuous_robust_scaler( ) 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 = dataset.patients[int(patient_index)] + 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)