JiaruiLi/DeepHealth
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Simplify event-free survival data loading

Browse Source
This commit is contained in:
Jiarui Li
2026-06-27 12:02:48 +08:00
parent 3cd1109249
commit 307704f74d
1 changed files with 42 additions and 175 deletions
Download Patch File Download Diff File Expand all files Collapse all files

217
evaluate_event_free_survival.py
View File

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