# dataset.py from __future__ import annotations import json from collections import OrderedDict from pathlib import Path from typing import Dict, Iterable, List, Literal, Optional, Tuple import numpy as np import pandas as pd import torch from torch.nn.utils.rnn import pad_sequence from torch.utils.data import Dataset from targets import ( CHECKUP_IDX, DAYS_PER_YEAR, NO_EVENT_IDX, PAD_IDX, build_all_targets, ) ONE_DAY_YEARS = 1.0 / DAYS_PER_YEAR DAILY_EXPOSURE_SHAPE = (1826, 4) MONTHLY_EXPOSURE_SHAPE = (241, 2) DAILY_EXPOSURE_CHANNELS = ("tmean", "tmax", "tmin", "rhmean") MONTHLY_EXPOSURE_CHANNELS = ("tmean", "rhmean") def _daily_exposure_columns() -> list[str]: cols: list[str] = [] for name in DAILY_EXPOSURE_CHANNELS: cols.extend(f"{name}_d{idx:04d}" for idx in range(DAILY_EXPOSURE_SHAPE[0])) return cols def _monthly_exposure_columns() -> list[str]: cols: list[str] = [] for name in MONTHLY_EXPOSURE_CHANNELS: cols.extend(f"{name}_m{idx:03d}" for idx in range(MONTHLY_EXPOSURE_SHAPE[0])) return cols class ExposureCache: """Random-access view over files produced by prepare_exposure_cache.py.""" def __init__(self, cache_dir: str | Path, row_group_cache_size: int = 4): cache_dir = Path(cache_dir) self.cache_dir = cache_dir manifest_path = cache_dir / "exposure_manifest.json" self.manifest = ( json.loads(manifest_path.read_text(encoding="utf-8")) if manifest_path.is_file() else {} ) self.storage = self.manifest.get("storage", "dense_npy") self._row_group_cache_size = int(row_group_cache_size) self._row_group_cache: OrderedDict[tuple[str, int, int], pd.DataFrame] = OrderedDict() self._parquet_files: dict[tuple[str, int], object] = {} self._parquet_columns: dict[tuple[str, int], list[str]] = {} eid_path = cache_dir / "exposure_eid.npy" token_path = cache_dir / "exposure_token.npy" onset_date_path = cache_dir / "exposure_onset_date.npy" if not (eid_path.is_file() and token_path.is_file() and onset_date_path.is_file()): raise FileNotFoundError( "Exposure cache must contain exposure_eid.npy, " "exposure_token.npy, and exposure_onset_date.npy. " "Regenerate it with the current prepare_exposure_cache.py." ) self.eids = np.load(eid_path, mmap_mode="r") self.raw_tokens = np.load(token_path, mmap_mode="r") self.onset_dates = np.load(onset_date_path, mmap_mode="r") self.daily = None self.monthly = None if self.storage == "dense_npy": self.daily = np.load(cache_dir / "exposure_daily.npy", mmap_mode="r") self.monthly = np.load(cache_dir / "exposure_monthly.npy", mmap_mode="r") elif self.storage == "parquet_index": self.daily_file_ids = np.load(cache_dir / "exposure_daily_file_id.npy", mmap_mode="r") self.daily_row_groups = np.load(cache_dir / "exposure_daily_row_group.npy", mmap_mode="r") self.daily_row_in_groups = np.load( cache_dir / "exposure_daily_row_in_group.npy", mmap_mode="r" ) self.monthly_file_ids = np.load( cache_dir / "exposure_monthly_file_id.npy", mmap_mode="r" ) self.monthly_row_groups = np.load( cache_dir / "exposure_monthly_row_group.npy", mmap_mode="r" ) self.monthly_row_in_groups = np.load( cache_dir / "exposure_monthly_row_in_group.npy", mmap_mode="r" ) self.daily_files = [Path(path) for path in self.manifest["daily_files"]] self.monthly_files = [Path(path) for path in self.manifest["monthly_files"]] else: raise ValueError(f"Unknown exposure cache storage mode: {self.storage!r}") quality_path = cache_dir / "exposure_quality.npy" self.quality = np.load(quality_path, mmap_mode="r") if quality_path.is_file() else None if self.storage == "dense_npy": if self.daily.ndim != 3 or self.daily.shape[1:] != DAILY_EXPOSURE_SHAPE: raise ValueError( f"exposure_daily.npy must have shape (N, {DAILY_EXPOSURE_SHAPE[0]}, " f"{DAILY_EXPOSURE_SHAPE[1]}), got {self.daily.shape}" ) if self.monthly.ndim != 3 or self.monthly.shape[1:] != MONTHLY_EXPOSURE_SHAPE: raise ValueError( f"exposure_monthly.npy must have shape (N, {MONTHLY_EXPOSURE_SHAPE[0]}, " f"{MONTHLY_EXPOSURE_SHAPE[1]}), got {self.monthly.shape}" ) n_rows = len(self.eids) if len(self.raw_tokens) != n_rows or len(self.onset_dates) != n_rows: raise ValueError("Exposure cache metadata/daily/monthly row counts do not match") if self.storage == "dense_npy": if self.daily.shape[0] != n_rows or self.monthly.shape[0] != n_rows: raise ValueError("Exposure cache metadata/daily/monthly row counts do not match") else: indexed_lengths = [ len(self.daily_file_ids), len(self.daily_row_groups), len(self.daily_row_in_groups), len(self.monthly_file_ids), len(self.monthly_row_groups), len(self.monthly_row_in_groups), ] if any(length != n_rows for length in indexed_lengths): raise ValueError("Exposure parquet index row counts do not match metadata") self._key_to_index: dict[tuple[int, int, int], int] | None = None def locality_key(self, indices: np.ndarray) -> tuple[int, int]: """Return a stable parquet locality key for sampler-side batching.""" indices = np.asarray(indices, dtype=np.int64) valid = indices[indices >= 0] if len(valid) == 0: return (2**31 - 1, 2**31 - 1) if self.storage != "parquet_index": return (0, int(valid[0] // 1024)) file_ids = np.asarray(self.daily_file_ids[valid], dtype=np.int64) row_groups = np.asarray(self.daily_row_groups[valid], dtype=np.int64) groups = np.stack([file_ids, row_groups], axis=1) unique_groups, counts = np.unique(groups, axis=0, return_counts=True) best = unique_groups[int(np.argmax(counts))] return (int(best[0]), int(best[1])) def build_age_index(self, birth_date_by_eid: dict[int, np.datetime64]) -> None: keys: dict[tuple[int, int, int], int] = {} eids = np.asarray(self.eids, dtype=np.int64) tokens = np.asarray(self.raw_tokens, dtype=np.int64) onset_dates = np.asarray(self.onset_dates, dtype="datetime64[D]") for idx, (eid, token, onset_date) in enumerate(zip(eids, tokens, onset_dates)): birth_date = birth_date_by_eid.get(int(eid)) if birth_date is None or np.isnat(onset_date) or np.isnat(birth_date): continue age_days = int((onset_date - birth_date).astype("timedelta64[D]").astype(np.int64)) if age_days < 0: continue keys[(int(eid), int(token), age_days)] = idx self._key_to_index = keys def lookup_indices(self, eid: int, raw_tokens: np.ndarray, age_days: np.ndarray) -> np.ndarray: if self._key_to_index is None: raise RuntimeError("ExposureCache.build_age_index must be called before lookup") out = np.full(len(raw_tokens), -1, dtype=np.int64) real = raw_tokens > 1 if not np.any(real): return out real_pos = np.nonzero(real)[0] out[real_pos] = [ self._key_to_index.get((int(eid), int(raw_tokens[pos]), int(round(float(age_days[pos])))), -1) for pos in real_pos ] return out def daily_window(self, index: int) -> np.ndarray: return self.daily_windows(np.asarray([index], dtype=np.int64))[0] def monthly_window(self, index: int) -> np.ndarray: return self.monthly_windows(np.asarray([index], dtype=np.int64))[0] def daily_windows(self, indices: np.ndarray) -> np.ndarray: return self._windows("daily", indices) def monthly_windows(self, indices: np.ndarray) -> np.ndarray: return self._windows("monthly", indices) def _windows( self, kind: Literal["daily", "monthly"], indices: np.ndarray, ) -> np.ndarray: indices = np.asarray(indices, dtype=np.int64) shape = DAILY_EXPOSURE_SHAPE if kind == "daily" else MONTHLY_EXPOSURE_SHAPE out = np.full((len(indices), shape[0], shape[1]), np.nan, dtype=np.float32) valid_pos = np.nonzero(indices >= 0)[0] if len(valid_pos) == 0: return out valid_indices = indices[valid_pos] if self.storage == "dense_npy": source = self.daily if kind == "daily" else self.monthly out[valid_pos] = np.asarray(source[valid_indices], dtype=np.float32) return out if kind == "daily": file_ids = np.asarray(self.daily_file_ids[valid_indices], dtype=np.int64) row_groups = np.asarray(self.daily_row_groups[valid_indices], dtype=np.int64) row_in_groups = np.asarray(self.daily_row_in_groups[valid_indices], dtype=np.int64) columns = _daily_exposure_columns() else: file_ids = np.asarray(self.monthly_file_ids[valid_indices], dtype=np.int64) row_groups = np.asarray(self.monthly_row_groups[valid_indices], dtype=np.int64) row_in_groups = np.asarray( self.monthly_row_in_groups[valid_indices], dtype=np.int64, ) columns = _monthly_exposure_columns() group_keys = np.stack([file_ids, row_groups], axis=1) unique_groups, inverse = np.unique(group_keys, axis=0, return_inverse=True) for group_idx, (file_id, row_group) in enumerate(unique_groups): group_pos = np.nonzero(inverse == group_idx)[0] frame = self._read_parquet_row_group( kind, int(file_id), int(row_group), columns, ) row_values = frame.iloc[row_in_groups[group_pos]].reindex(columns=columns) values = ( row_values.to_numpy(dtype=np.float32, copy=True) .reshape(len(group_pos), shape[1], shape[0]) .transpose(0, 2, 1) ) out[valid_pos[group_pos]] = values return out def _parquet_window(self, kind: Literal["daily", "monthly"], index: int) -> np.ndarray: if kind == "daily": file_id = int(self.daily_file_ids[index]) row_group = int(self.daily_row_groups[index]) row_in_group = int(self.daily_row_in_groups[index]) shape = DAILY_EXPOSURE_SHAPE columns = _daily_exposure_columns() else: file_id = int(self.monthly_file_ids[index]) row_group = int(self.monthly_row_groups[index]) row_in_group = int(self.monthly_row_in_groups[index]) shape = MONTHLY_EXPOSURE_SHAPE columns = _monthly_exposure_columns() frame = self._read_parquet_row_group(kind, file_id, row_group, columns) row = frame.iloc[row_in_group].reindex(columns) n_channels = shape[1] return ( row.to_numpy(dtype=np.float32, copy=True) .reshape(n_channels, shape[0]) .transpose(1, 0) ) def _read_parquet_row_group( self, kind: Literal["daily", "monthly"], file_id: int, row_group: int, columns: list[str], ) -> pd.DataFrame: cache_key = (kind, file_id, row_group) cached = self._row_group_cache.get(cache_key) if cached is not None: self._row_group_cache.move_to_end(cache_key) return cached try: import pyarrow.parquet as pq except ImportError as exc: raise ImportError( "Parquet exposure index loading requires pyarrow. Install requirements " "or use a dense numpy exposure cache." ) from exc parquet_key = (kind, file_id) parquet_file = self._parquet_files.get(parquet_key) if parquet_file is None: path = self.daily_files[file_id] if kind == "daily" else self.monthly_files[file_id] parquet_file = pq.ParquetFile(path) self._parquet_files[parquet_key] = parquet_file available_columns = self._parquet_columns.get(parquet_key) if available_columns is None: available = set(parquet_file.schema.names) available_columns = [col for col in columns if col in available] self._parquet_columns[parquet_key] = available_columns table = parquet_file.read_row_group(row_group, columns=available_columns) frame = table.to_pandas() if available_columns != columns: frame = frame.reindex(columns=columns) self._row_group_cache[cache_key] = frame self._row_group_cache.move_to_end(cache_key) while len(self._row_group_cache) > self._row_group_cache_size: self._row_group_cache.popitem(last=False) return frame def load_label_vocab( labels_file: str, include_no_event: bool = True, ) -> Tuple[Dict[str, int], Dict[int, str]]: label_id_to_code: Dict[int, str] = { PAD_IDX: "", CHECKUP_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 label_code_to_id: Dict[str, int] = {} with open(labels_file, encoding="utf-8") as f: for i, line in enumerate(f): parts = line.strip().split() if not parts: continue idx = offset + i code = parts[0] label_code_to_id[code] = idx label_id_to_code[idx] = code return label_code_to_id, label_id_to_code def _insert_gap_no_event_tokens( times_days: np.ndarray, labels: np.ndarray, interval_years: float = 5.0, ) -> Tuple[np.ndarray, np.ndarray]: if len(times_days) < 2: return times_days, labels step_days = interval_years * DAYS_PER_YEAR unique_times = np.unique(times_days.astype(np.float64)) extra_times: List[float] = [] for i in range(len(unique_times) - 1): t_left = float(unique_times[i]) t_right = float(unique_times[i + 1]) if t_right - t_left <= step_days: continue first = np.ceil((t_left + 1e-6) / step_days) * step_days t = first while t < t_right - 1e-6: extra_times.append(t) t += step_days if not extra_times: return times_days, labels extra_arr = np.array(extra_times, dtype=np.float32) no_event_labels = np.full(len(extra_arr), NO_EVENT_IDX, dtype=np.int64) all_times = np.concatenate([times_days.astype(np.float32), extra_arr]) all_labels = np.concatenate([labels.astype(np.int64), no_event_labels]) order = np.lexsort((all_labels, all_times)) return all_times[order], all_labels[order] class _ExpoBaseDataset(Dataset): def __init__( self, data_prefix: str = "ukb", labels_file: str = "labels.csv", no_event_interval_years: float = 5.0, include_no_event_in_uts_target: bool = False, exposure_cache_dir: str | Path | None = None, mask_onset_exposure: bool = False, exposure_row_group_cache_size: int = 4, ) -> None: self.data_prefix = data_prefix self.labels_file = labels_file self.no_event_interval_years = float(no_event_interval_years) self.include_no_event_in_uts_target = bool(include_no_event_in_uts_target) self.exposure_cache = ( ExposureCache( exposure_cache_dir, row_group_cache_size=exposure_row_group_cache_size, ) if exposure_cache_dir is not None else None ) self.mask_onset_exposure = bool(mask_onset_exposure) self.label_code_to_id, self.label_id_to_code = load_label_vocab( labels_file, include_no_event=True, ) event_data = np.load(f"{data_prefix}_event_data.npy") if event_data.ndim != 2 or event_data.shape[1] < 3: raise ValueError(f"event_data must have shape (N, 3+), got {event_data.shape}") event_data = event_data[:, :3].copy() order = np.lexsort((event_data[:, 2], event_data[:, 1], event_data[:, 0])) self.event_data = event_data[order] basic_table = pd.read_csv(f"{data_prefix}_basic_info.csv", index_col=0) basic_table.index = basic_table.index.astype(np.int64) unique_eids = np.unique(self.event_data[:, 0].astype(np.int64)) basic_table = basic_table.loc[unique_eids] self._prepare_sex(basic_table, unique_eids) self._prepare_birth_dates(basic_table, unique_eids) if self.exposure_cache is not None: self.exposure_cache.build_age_index(self.birth_date_mapping) max_id_in_vocab = max(self.label_id_to_code.keys()) max_id_in_data = int(self.event_data[:, 2].max()) if len(self.event_data) > 0 else 0 max_id_in_data += 1 self.vocab_size = max(max_id_in_vocab, max_id_in_data) + 1 if not self.include_no_event_in_uts_target: self.ignored_uts_target_ids = {PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX} else: self.ignored_uts_target_ids = {PAD_IDX, CHECKUP_IDX} def _prepare_sex(self, basic_table: pd.DataFrame, unique_eids: np.ndarray) -> None: sex_values = pd.to_numeric(basic_table["sex"], errors="coerce").to_numpy() if np.isnan(sex_values).any(): raise ValueError("sex column contains missing or non-numeric values") sex_values = sex_values.astype(np.int64) sex_unique = np.unique(sex_values) if np.all(np.isin(sex_unique, [0, 1])): sex01 = sex_values elif np.all(np.isin(sex_unique, [1, 2])): sex01 = sex_values - 1 else: raise ValueError( f"Unexpected sex values: {sex_unique.tolist()}. Expected {{0,1}} or {{1,2}}." ) self.sex_mapping = {int(eid): int(s) for eid, s in zip(unique_eids, sex01)} def _prepare_birth_dates(self, basic_table: pd.DataFrame, unique_eids: np.ndarray) -> None: if "date_of_birth" not in basic_table.columns: if self.exposure_cache is None: self.birth_date_mapping = {} return raise ValueError( "Exposure alignment requires ukb_basic_info.csv to contain " "'date_of_birth'. Regenerate it with the current prepare_data.py." ) birth = pd.to_datetime(basic_table["date_of_birth"], errors="coerce") if birth.isna().any() and self.exposure_cache is not None: raise ValueError("date_of_birth contains missing or invalid values") birth_np = birth.to_numpy(dtype="datetime64[D]") self.birth_date_mapping = { int(eid): np.datetime64(date, "D") for eid, date in zip(unique_eids, birth_np) if not np.isnat(date) } def _iter_patient_events( self, *, impute_no_event_gaps: bool, ) -> Iterable[tuple[int, np.ndarray, np.ndarray]]: unique_eids, starts = np.unique(self.event_data[:, 0], return_index=True) ends = np.concatenate([starts[1:], [len(self.event_data)]]) for eid_raw, start, end in zip(unique_eids, starts, ends): eid = int(eid_raw) rows = self.event_data[start:end] times_days_raw = rows[:, 1].astype(np.float32) labels_raw = rows[:, 2].astype(np.int64) if len(labels_raw) == 0: yield eid, times_days_raw, labels_raw continue labels_raw = np.where(labels_raw >= NO_EVENT_IDX, labels_raw + 1, labels_raw) if not impute_no_event_gaps: yield eid, times_days_raw, labels_raw continue times_days, labels = _insert_gap_no_event_tokens( times_days_raw, labels_raw, interval_years=self.no_event_interval_years, ) yield eid, times_days, labels def _split_features(self, eid: int) -> Optional[Dict]: if eid not in self.sex_mapping: return None return { "sex": self.sex_mapping[eid], } def _raw_tokens_from_model_tokens(self, model_tokens: np.ndarray) -> np.ndarray: raw_tokens = np.full(len(model_tokens), -1, dtype=np.int64) real = model_tokens > NO_EVENT_IDX raw_tokens[real] = model_tokens[real].astype(np.int64) - 1 return raw_tokens def _exposure_indices_for_inputs( self, eid: int, input_events: np.ndarray, input_times_days: np.ndarray, ) -> np.ndarray | None: if self.exposure_cache is None: return None raw_tokens = self._raw_tokens_from_model_tokens(input_events) return self.exposure_cache.lookup_indices( eid=eid, raw_tokens=raw_tokens, age_days=input_times_days, ) def _load_exposure_windows(self, exposure_index: np.ndarray) -> tuple[torch.Tensor, torch.Tensor]: if self.exposure_cache is None: raise RuntimeError("Exposure cache is not enabled") daily = self.exposure_cache.daily_windows(exposure_index).astype( np.float32, copy=False, ) monthly = self.exposure_cache.monthly_windows(exposure_index).astype( np.float32, copy=False, ) if self.mask_onset_exposure: daily[:, 0, :] = np.nan monthly[:, 0, :] = np.nan return torch.from_numpy(daily).float(), torch.from_numpy(monthly).float() class NextStepHealthDataset(_ExpoBaseDataset): """ Dataset for next-token and next-time-point losses with unified other-info tokens. Returned targets cover both: - Delphi2MLoss: target_event_seq, target_time_seq - UniqueTimeSetExponentialLoss: readout_mask, target_dt_unique, target_multi_hot """ CACHE_VERSION = 3 def __init__( self, data_prefix: str = "ukb", labels_file: str = "labels.csv", no_event_interval_years: float = 5.0, include_no_event_in_uts_target: bool = False, exposure_cache_dir: str | Path | None = None, mask_onset_exposure: bool = False, exposure_row_group_cache_size: int = 4, ) -> None: super().__init__( data_prefix=data_prefix, labels_file=labels_file, no_event_interval_years=no_event_interval_years, include_no_event_in_uts_target=include_no_event_in_uts_target, exposure_cache_dir=exposure_cache_dir, mask_onset_exposure=mask_onset_exposure, exposure_row_group_cache_size=exposure_row_group_cache_size, ) self.samples: List[Dict] = [] for eid, times_days, labels in self._iter_patient_events( impute_no_event_gaps=True, ): if len(labels) < 2: continue features = self._split_features(eid) if features is None: continue target_pack = build_all_targets( labels=labels, times_days=times_days, vocab_size=self.vocab_size, ignored_uts_target_ids=self.ignored_uts_target_ids, require_sorted=True, ) sample = { "eid": eid, "event_seq": target_pack.next_token.input_events, "time_seq": target_pack.next_token.input_times_years, "target_event_seq": target_pack.next_token.target_events, "target_time_seq": target_pack.next_token.target_times_years, "readout_mask": target_pack.unique_time_set.readout_mask, "target_dt_unique": target_pack.unique_time_set.target_dt_unique, "target_multi_hot": target_pack.unique_time_set.target_multi_hot, **features, } exposure_index = self._exposure_indices_for_inputs( eid=eid, input_events=target_pack.next_token.input_events, input_times_days=times_days[:-1], ) if exposure_index is not None: sample["exposure_index"] = exposure_index self.samples.append(sample) def __len__(self) -> int: return len(self.samples) def __getitem__(self, idx: int) -> Dict: s = self.samples[idx] out = { "event_seq": torch.from_numpy(s["event_seq"]).long(), "time_seq": torch.from_numpy(s["time_seq"]).float(), "sex": torch.tensor(s["sex"], dtype=torch.long), "target_event_seq": torch.from_numpy(s["target_event_seq"]).long(), "target_time_seq": torch.from_numpy(s["target_time_seq"]).float(), "readout_mask": torch.from_numpy(s["readout_mask"]).bool(), "target_dt_unique": torch.from_numpy(s["target_dt_unique"]).float(), "target_multi_hot": torch.from_numpy(s["target_multi_hot"]).bool(), } if "exposure_index" in s: daily, monthly = self._load_exposure_windows(s["exposure_index"]) out["exposure_daily"] = daily out["exposure_monthly"] = monthly return out class AllFutureHealthDataset(_ExpoBaseDataset): """ Dataset with unified other-info tokens and DeepHealthV2-style all-future targets. Train samples one query time per patient at each __getitem__ call. Valid/test use random-but-fixed query points. For each patient with N real disease events, N - 2 query points are sampled from the eligible observed time range, with at least one future event after every query. """ CACHE_VERSION = 5 def __init__( self, data_prefix: str = "ukb", labels_file: str = "labels.csv", split: Literal["train", "valid", "test"] = "train", no_event_interval_years: float = 5.0, include_no_event_in_uts_target: bool = False, min_history_events: int = 1, min_future_events: int = 1, validation_query_seed: int = 42, exposure_cache_dir: str | Path | None = None, mask_onset_exposure: bool = False, exposure_row_group_cache_size: int = 4, ) -> None: if split not in {"train", "valid", "test"}: raise ValueError(f"split must be train/valid/test, got {split!r}") super().__init__( data_prefix=data_prefix, labels_file=labels_file, no_event_interval_years=no_event_interval_years, include_no_event_in_uts_target=include_no_event_in_uts_target, exposure_cache_dir=exposure_cache_dir, mask_onset_exposure=mask_onset_exposure, exposure_row_group_cache_size=exposure_row_group_cache_size, ) self.split = split self.min_history_events = int(min_history_events) self.min_future_events = int(min_future_events) self.validation_query_seed = int(validation_query_seed) self.patients: List[Dict] = [] self.valid_queries: List[Tuple[int, float]] = [] validation_rng = None if split in {"valid", "test"}: split_offset = 0 if split == "valid" else 1_000_003 validation_rng = np.random.RandomState(self.validation_query_seed + split_offset) for eid, times_days, labels in self._iter_patient_events( impute_no_event_gaps=False, ): times_years = (times_days / DAYS_PER_YEAR).astype(np.float32) unique_times = np.unique(times_years) if len(labels) < 2 or len(unique_times) < 2: continue features = self._split_features(eid) if features is None: continue patient = { "eid": eid, "times": times_years, "times_days": times_days.astype(np.float32), "labels": labels.astype(np.int64), "t_obs": float(times_years.max()), **features, } pidx = len(self.patients) self.patients.append(patient) if split in {"valid", "test"}: if validation_rng is None: raise RuntimeError("validation_rng was not initialized") self.valid_queries.extend( (pidx, t_query) for t_query in self._sample_fixed_validation_queries( patient, validation_rng, ) ) if split in {"valid", "test"} and not self.valid_queries: raise ValueError("No random-but-fixed validation query points were built.") def _is_valid_query(self, patient: Dict, t_query: float) -> bool: times = patient["times"] labels = patient["labels"] real_event_mask = ~np.isin( labels, np.array([PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX], dtype=np.int64), ) n_hist = int((times <= t_query).sum()) n_future = int(((times > t_query) & real_event_mask).sum()) return ( n_hist >= self.min_history_events and n_future >= self.min_future_events and patient["t_obs"] > t_query ) def _sample_fixed_validation_queries( self, patient: Dict, rng: np.random.RandomState, ) -> List[float]: times = np.asarray(patient["times"], dtype=np.float32) labels = np.asarray(patient["labels"], dtype=np.int64) real_event_mask = ~np.isin( labels, np.array([PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX], dtype=np.int64), ) real_times = np.sort(times[real_event_mask].astype(np.float32, copy=False)) n_real_events = int(real_times.size) n_queries = max(0, n_real_events - 2) if n_queries == 0: return [] min_hist = int(self.min_history_events) min_future = int(self.min_future_events) if n_real_events < min_hist + min_future: return [] left = float(real_times[min_hist - 1]) right_event_time = float(real_times[n_real_events - min_future]) right = np.nextafter(np.float32(right_event_time), np.float32(-np.inf)) if not np.isfinite(left) or not np.isfinite(right) or float(right) <= left: return [] queries: List[float] = [] max_attempts = max(100, n_queries * 50) for _ in range(max_attempts): if len(queries) >= n_queries: break t_query = float(rng.uniform(left, float(right))) if self._is_valid_query(patient, t_query): queries.append(t_query) return queries def _sample_train_query(self, patient: Dict) -> float: unique_times = np.unique(patient["times"]) if len(unique_times) < 2: raise RuntimeError("Training patient has fewer than two unique times.") j = np.random.randint(1, len(unique_times)) left = float(unique_times[j - 1]) right = float(unique_times[j]) if right - left <= ONE_DAY_YEARS: t_query = right - ONE_DAY_YEARS else: t_query = np.random.uniform(left, right - ONE_DAY_YEARS) if not self._is_valid_query(patient, t_query): t_query = right - 1e-6 return float(t_query) def _build_item(self, patient: Dict, t_query: float) -> Dict: times = patient["times"] times_days = patient["times_days"] labels = patient["labels"] hist = times <= t_query fut = times > t_query out = { "event_seq": torch.from_numpy(labels[hist]).long(), "time_seq": torch.from_numpy(times[hist]).float(), "t_query": torch.tensor(t_query, dtype=torch.float32), "future_targets": torch.from_numpy(labels[fut]).long(), "future_dt": torch.from_numpy(times[fut] - np.float32(t_query)).float(), "exposure": torch.tensor(np.float32(patient["t_obs"] - t_query), dtype=torch.float32), "sex": torch.tensor(patient["sex"], dtype=torch.long), } if self.exposure_cache is not None: exposure_index = self._exposure_indices_for_inputs( eid=int(patient["eid"]), input_events=labels[hist], input_times_days=times_days[hist], ) if exposure_index is not None: daily, monthly = self._load_exposure_windows(exposure_index) out["exposure_daily"] = daily out["exposure_monthly"] = monthly return out def __len__(self) -> int: if self.split == "train": return len(self.patients) return len(self.valid_queries) def __getitem__(self, idx: int) -> Dict: if self.split == "train": patient = self.patients[idx] t_query = self._sample_train_query(patient) else: pidx, t_query = self.valid_queries[idx] patient = self.patients[pidx] return self._build_item(patient, t_query) def _collate_common_static(batch: List[Dict]) -> Dict: return { "sex": torch.stack([s["sex"] for s in batch]), } def _pad_exposure(batch: List[Dict], key: str, shape: tuple[int, int]) -> torch.Tensor: max_len = max(int(s["event_seq"].numel()) for s in batch) out = torch.full( (len(batch), max_len, shape[0], shape[1]), float("nan"), dtype=torch.float32, ) for idx, sample in enumerate(batch): value = sample.get(key) if value is None: continue seq_len = int(value.size(0)) out[idx, :seq_len] = value return out def next_step_collate_fn(batch: List[Dict]) -> Dict: event_seq = pad_sequence( [s["event_seq"] for s in batch], batch_first=True, padding_value=PAD_IDX, ) time_seq = pad_sequence( [s["time_seq"] for s in batch], batch_first=True, padding_value=0.0, ) target_event_seq = pad_sequence( [s["target_event_seq"] for s in batch], batch_first=True, padding_value=PAD_IDX, ) target_time_seq = pad_sequence( [s["target_time_seq"] for s in batch], batch_first=True, padding_value=0.0, ) readout_mask = pad_sequence( [s["readout_mask"] for s in batch], batch_first=True, padding_value=False, ) target_dt_unique = pad_sequence( [s["target_dt_unique"] for s in batch], batch_first=True, padding_value=0.0, ) target_multi_hot = pad_sequence( [s["target_multi_hot"] for s in batch], batch_first=True, padding_value=False, ) out = { "event_seq": event_seq, "time_seq": time_seq, "padding_mask": event_seq > PAD_IDX, "target_event_seq": target_event_seq, "target_time_seq": target_time_seq, "readout_mask": readout_mask, "target_dt_unique": target_dt_unique, "target_multi_hot": target_multi_hot, } out.update(_collate_common_static(batch)) if any("exposure_daily" in s for s in batch): out["exposure_daily"] = _pad_exposure(batch, "exposure_daily", DAILY_EXPOSURE_SHAPE) out["exposure_monthly"] = _pad_exposure(batch, "exposure_monthly", MONTHLY_EXPOSURE_SHAPE) return out def all_future_collate_fn(batch: List[Dict]) -> Dict: event_seq = pad_sequence( [s["event_seq"] for s in batch], batch_first=True, padding_value=PAD_IDX, ) time_seq = pad_sequence( [s["time_seq"] for s in batch], batch_first=True, padding_value=0.0, ) future_targets = pad_sequence( [s["future_targets"] for s in batch], batch_first=True, padding_value=PAD_IDX, ) future_dt = pad_sequence( [s["future_dt"] for s in batch], batch_first=True, padding_value=0.0, ) out = { "event_seq": event_seq, "time_seq": time_seq, "padding_mask": event_seq > PAD_IDX, "t_query": torch.stack([s["t_query"] for s in batch]), "future_targets": future_targets, "future_dt": future_dt, "exposure": torch.stack([s["exposure"] for s in batch]), } out.update(_collate_common_static(batch)) if any("exposure_daily" in s for s in batch): out["exposure_daily"] = _pad_exposure(batch, "exposure_daily", DAILY_EXPOSURE_SHAPE) out["exposure_monthly"] = _pad_exposure(batch, "exposure_monthly", MONTHLY_EXPOSURE_SHAPE) return out HealthDataset = NextStepHealthDataset collate_fn = next_step_collate_fn