Simplify event-free survival data loading

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 (
+from train_util import load_eid_file, load_extra_info_types_file
-    load_extra_info_types_file,
-    split_all_future_datasets,
-    split_all_future_datasets_by_eid_files,
-    split_dataset,
-    split_dataset_by_eid_files,
-)
 
 
 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:
+def load_eval_sequence_dataset(
-    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(
     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":
+    if eval_split != "all" and split_files_exist:
-        print("Loading AllFutureHealthDataset objects using the training path...")
+        split_files = {
-        train_dataset = AllFutureHealthDataset(
+            "train": train_eid_file,
-            data_prefix=data_prefix,
+            "val": val_eid_file,
-            labels_file=labels_file,
+            "test": test_eid_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":
+        selected_eids = load_eid_file(split_files[eval_split])
-            patient_indices = np.arange(len(train_dataset.patients), dtype=np.int64)
+        out = np.asarray(
-            dataset = AllFutureSelectedSequenceDataset(train_dataset, patient_indices)
+            [
-        else:
+                idx
-            source_dataset, subset = split_map[eval_split]
+                for idx, sample in enumerate(dataset.samples)
-            patient_indices = _patient_indices_from_all_future_subset(source_dataset, subset)
+                if int(sample["eid"]) in selected_eids
-            dataset = AllFutureSelectedSequenceDataset(source_dataset, patient_indices)
+            ],
-        out = np.arange(len(dataset.samples), dtype=np.int64)
+            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...")
+        if eval_split == "all":
-        dataset = HealthDataset(
+            out = np.arange(len(dataset.samples), dtype=np.int64)
-            data_prefix=data_prefix,
+            split_source = "all"
-            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"
         else:
-            train_subset, val_subset, test_subset = split_dataset(
+            out = make_eval_indices(dataset, args, cfg)
-                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)),
-            )
             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,
     )