# dataset.py from __future__ import annotations 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) class ExposureCache: """Random-access view over files produced by prepare_exposure_cache.py.""" def __init__(self, cache_dir: str | Path): cache_dir = Path(cache_dir) self.cache_dir = cache_dir 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 = np.load(cache_dir / "exposure_daily.npy", mmap_mode="r") self.monthly = np.load(cache_dir / "exposure_monthly.npy", mmap_mode="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.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 or 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") self._key_to_index: dict[tuple[int, int, int], int] | None = None 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: if index < 0: return np.full(DAILY_EXPOSURE_SHAPE, np.nan, dtype=np.float32) return np.asarray(self.daily[index], dtype=np.float32) def monthly_window(self, index: int) -> np.ndarray: if index < 0: return np.full(MONTHLY_EXPOSURE_SHAPE, np.nan, dtype=np.float32) return np.asarray(self.monthly[index], dtype=np.float32) 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, ) -> 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) 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 = np.stack( [self.exposure_cache.daily_window(int(idx)) for idx in exposure_index], axis=0, ).astype(np.float32, copy=True) monthly = np.stack( [self.exposure_cache.monthly_window(int(idx)) for idx in exposure_index], axis=0, ).astype(np.float32, copy=True) 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, ) -> 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, ) 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, ) -> 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, ) 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