From 307704f74df7df05444107a86e5184954c47456d Mon Sep 17 00:00:00 2001 From: Jiarui Li Date: Sat, 27 Jun 2026 12:02:48 +0800 Subject: [PATCH] Simplify event-free survival data loading --- evaluate_event_free_survival.py | 217 +++++++------------------------- 1 file changed, 42 insertions(+), 175 deletions(-) diff --git a/evaluate_event_free_survival.py b/evaluate_event_free_survival.py index 1f8f8c1..292d09b 100644 --- a/evaluate_event_free_survival.py +++ b/evaluate_event_free_survival.py @@ -15,7 +15,7 @@ from __future__ import annotations import argparse import json from pathlib import Path -from typing import Any, Dict, Iterable, List, Optional, Sequence +from typing import Any, Dict, List, Optional, Sequence import numpy as np import pandas as pd @@ -24,7 +24,8 @@ from torch.nn.utils.rnn import pad_sequence from torch.utils.data import DataLoader, Dataset from tqdm.auto import tqdm -from dataset import AllFutureHealthDataset, HealthDataset +from dataset import HealthDataset +from eval_data import load_sequence_eval_dataset from evaluate_auc_v2 import ( LandmarkDataset, build_model_from_dataset, @@ -32,6 +33,7 @@ from evaluate_auc_v2 import ( load_checkpoint_state_dict, load_json_config, load_model_state, + make_eval_indices, resolve_dist_mode_for_checkpoint, resolve_eval_device, validate_dataset_metadata, @@ -43,69 +45,12 @@ from future_event_free_survival import ( from models import DeepHealth from readouts import build_readout from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX -from train_util import ( - load_extra_info_types_file, - split_all_future_datasets, - split_all_future_datasets_by_eid_files, - split_dataset, - split_dataset_by_eid_files, -) +from train_util import load_eid_file, load_extra_info_types_file SPECIAL_TOKENS = {PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX} -class AllFutureSelectedSequenceDataset: - """Sequence-view dataset built from selected AllFutureHealthDataset patients.""" - - def __init__( - self, - base: AllFutureHealthDataset, - patient_indices: Iterable[int], - ) -> None: - self.base = base - self.label_code_to_id = base.label_code_to_id - self.label_id_to_code = base.label_id_to_code - self.vocab_size = base.vocab_size - self.n_types = base.n_types - self.n_cont_types = base.n_cont_types - self.n_categories = base.n_categories - self.cont_type_ids = base.cont_type_ids - self.extra_info_types = base.extra_info_types - - seen: set[int] = set() - self.samples: List[Dict[str, Any]] = [] - for pidx in patient_indices: - pidx = int(pidx) - if pidx in seen: - continue - seen.add(pidx) - patient = base.patients[pidx] - labels = np.asarray(patient["labels"], dtype=np.int64) - 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"]), - "event_seq": labels[:-1], - "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), - "other_value_kind": np.asarray(patient["other_value_kind"], dtype=np.int64), - "other_time": np.asarray(patient["other_time"], dtype=np.float32), - } - ) - - def __len__(self) -> int: - return len(self.samples) - - def parse_int_list(value: Any) -> Optional[List[int]]: if value is None: return None @@ -183,27 +128,7 @@ def normalize_eval_split(args: argparse.Namespace, cfg: Dict[str, Any]) -> str: return eval_split -def _subset_indices(subset: Any) -> np.ndarray: - if not hasattr(subset, "indices"): - raise TypeError(f"Expected a torch Subset-like object, got {type(subset).__name__}") - return np.asarray(subset.indices, dtype=np.int64) - - -def _patient_indices_from_all_future_subset( - dataset: AllFutureHealthDataset, - subset: Any, -) -> np.ndarray: - indices = _subset_indices(subset) - if dataset.split == "train": - return indices - patient_indices = [ - int(dataset.valid_queries[int(query_idx)][0]) - for query_idx in indices.tolist() - ] - return np.asarray(sorted(set(patient_indices)), dtype=np.int64) - - -def load_training_style_sequence_dataset( +def load_eval_sequence_dataset( args: argparse.Namespace, cfg: Dict[str, Any], ) -> tuple[Any, np.ndarray, str, str]: @@ -217,6 +142,18 @@ def load_training_style_sequence_dataset( 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") @@ -225,103 +162,33 @@ def load_training_style_sequence_dataset( for path in (train_eid_file, val_eid_file, test_eid_file) ) - if model_target_mode == "all_future": - print("Loading AllFutureHealthDataset objects using the training path...") - train_dataset = AllFutureHealthDataset( - data_prefix=data_prefix, - labels_file=labels_file, - split="train", - min_history_events=int(cfg.get("all_future_min_history_events", 1)), - min_future_events=int(cfg.get("all_future_min_future_events", 1)), - validation_query_seed=int(cfg.get("all_future_validation_query_seed", cfg.get("seed", 42))), - extra_info_types=extra_info_types, - ) - val_dataset = AllFutureHealthDataset( - data_prefix=data_prefix, - labels_file=labels_file, - split="valid", - min_history_events=int(cfg.get("all_future_min_history_events", 1)), - min_future_events=int(cfg.get("all_future_min_future_events", 1)), - validation_query_seed=int(cfg.get("all_future_validation_query_seed", cfg.get("seed", 42))), - extra_info_types=extra_info_types, - ) - test_dataset = AllFutureHealthDataset( - data_prefix=data_prefix, - labels_file=labels_file, - split="test", - min_history_events=int(cfg.get("all_future_min_history_events", 1)), - min_future_events=int(cfg.get("all_future_min_future_events", 1)), - validation_query_seed=int(cfg.get("all_future_validation_query_seed", cfg.get("seed", 42))), - extra_info_types=extra_info_types, - ) - if split_files_exist: - train_subset, val_subset, test_subset = split_all_future_datasets_by_eid_files( - train_dataset=train_dataset, - val_dataset=val_dataset, - test_dataset=test_dataset, - train_eid_file=train_eid_file, - val_eid_file=val_eid_file, - test_eid_file=test_eid_file, - ) - split_source = "eid_files" - else: - train_subset, val_subset, test_subset = split_all_future_datasets( - train_dataset=train_dataset, - val_dataset=val_dataset, - test_dataset=test_dataset, - 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)), - seed=int(cfg_get(args, cfg, "seed", 42)), - ) - split_source = "ratio_split" - - split_map = { - "train": (train_dataset, train_subset), - "val": (val_dataset, val_subset), - "test": (test_dataset, test_subset), + if eval_split != "all" and split_files_exist: + split_files = { + "train": train_eid_file, + "val": val_eid_file, + "test": test_eid_file, } - if eval_split == "all": - patient_indices = np.arange(len(train_dataset.patients), dtype=np.int64) - dataset = AllFutureSelectedSequenceDataset(train_dataset, patient_indices) - else: - source_dataset, subset = split_map[eval_split] - patient_indices = _patient_indices_from_all_future_subset(source_dataset, subset) - dataset = AllFutureSelectedSequenceDataset(source_dataset, patient_indices) - out = np.arange(len(dataset.samples), dtype=np.int64) + 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: - print("Loading HealthDataset using the training path...") - dataset = HealthDataset( - 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 split_files_exist: - train_subset, val_subset, test_subset = split_dataset_by_eid_files( - dataset=dataset, - train_eid_file=train_eid_file, - val_eid_file=val_eid_file, - test_eid_file=test_eid_file, - ) - split_source = "eid_files" + if eval_split == "all": + out = np.arange(len(dataset.samples), dtype=np.int64) + split_source = "all" else: - train_subset, val_subset, test_subset = split_dataset( - dataset=dataset, - 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)), - seed=int(cfg_get(args, cfg, "seed", 42)), - ) + out = make_eval_indices(dataset, args, cfg) split_source = "ratio_split" - split_map = { - "train": _subset_indices(train_subset), - "val": _subset_indices(val_subset), - "test": _subset_indices(test_subset), - "all": np.arange(len(dataset.samples), dtype=np.int64), - } - out = 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: @@ -555,7 +422,7 @@ def main() -> None: 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_training_style_sequence_dataset( + dataset, subset_indices, eval_split, split_source = load_eval_sequence_dataset( args, cfg, )