858 lines
32 KiB
Python
858 lines
32 KiB
Python
# dataset.py
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
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
|
|
|
|
|
|
def _load_readonly_npy(path: Path) -> np.ndarray:
|
|
arr = np.load(path, mmap_mode="r")
|
|
arr.setflags(write=False)
|
|
return arr
|
|
|
|
|
|
class ExposureCache:
|
|
"""Eid-sequence-aligned exposure windows from prepare_exposure_cache.py."""
|
|
|
|
def __init__(self, cache_dir: str | Path):
|
|
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")
|
|
if self.storage != "eid_sequence_npy":
|
|
raise ValueError(
|
|
"Exposure cache must use storage='eid_sequence_npy'. "
|
|
"Regenerate it with the current prepare_exposure_cache.py."
|
|
)
|
|
|
|
eid_path = cache_dir / "exposure_eid.npy"
|
|
token_path = cache_dir / "exposure_token.npy"
|
|
age_path = cache_dir / "exposure_age_days.npy"
|
|
onset_date_path = cache_dir / "exposure_onset_date.npy"
|
|
row_index_path = cache_dir / "exposure_row_index.npy"
|
|
eid_index_path = cache_dir / "exposure_eid_index.npy"
|
|
eid_start_path = cache_dir / "exposure_eid_start.npy"
|
|
daily_path = cache_dir / "exposure_daily.npy"
|
|
monthly_path = cache_dir / "exposure_monthly.npy"
|
|
required_paths = [
|
|
eid_path,
|
|
token_path,
|
|
age_path,
|
|
onset_date_path,
|
|
row_index_path,
|
|
eid_index_path,
|
|
eid_start_path,
|
|
daily_path,
|
|
monthly_path,
|
|
]
|
|
if any(not path.is_file() for path in required_paths):
|
|
raise FileNotFoundError(
|
|
"Exposure cache is missing one or more eid-sequence files. "
|
|
"Regenerate it with the current prepare_exposure_cache.py."
|
|
)
|
|
|
|
self.eids = _load_readonly_npy(eid_path)
|
|
self.raw_tokens = _load_readonly_npy(token_path)
|
|
self.age_days = _load_readonly_npy(age_path)
|
|
self.onset_dates = _load_readonly_npy(onset_date_path)
|
|
self.row_index = _load_readonly_npy(row_index_path)
|
|
self.eid_index = _load_readonly_npy(eid_index_path)
|
|
self.eid_start = _load_readonly_npy(eid_start_path)
|
|
self.daily = _load_readonly_npy(daily_path)
|
|
self.monthly = _load_readonly_npy(monthly_path)
|
|
quality_path = cache_dir / "exposure_quality.npy"
|
|
self.quality = _load_readonly_npy(quality_path) 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.age_days) != n_rows
|
|
or len(self.onset_dates) != n_rows
|
|
or len(self.row_index) != n_rows
|
|
):
|
|
raise ValueError("Exposure cache sequence metadata row counts do not match")
|
|
max_window_index = int(np.max(self.row_index)) if n_rows else -1
|
|
if (
|
|
max_window_index >= self.daily.shape[0]
|
|
or max_window_index >= self.monthly.shape[0]
|
|
):
|
|
raise ValueError("Exposure row index points past daily/monthly window arrays")
|
|
if len(self.eid_start) != len(self.eid_index) + 1:
|
|
raise ValueError("exposure_eid_start.npy must have len(eid_index) + 1")
|
|
if len(self.eid_start) and int(self.eid_start[-1]) != n_rows:
|
|
raise ValueError("Last exposure eid offset must equal exposure row count")
|
|
|
|
self._eid_to_pos = {
|
|
int(eid): idx
|
|
for idx, eid in enumerate(np.asarray(self.eid_index, dtype=np.int64))
|
|
}
|
|
|
|
def locality_key(self, indices: np.ndarray) -> tuple[int, int]:
|
|
"""Return a stable 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)
|
|
first = int(valid[0])
|
|
return (first // 1024, first % 1024)
|
|
|
|
def build_age_index(self, birth_date_by_eid: dict[int, np.datetime64]) -> None:
|
|
"""Kept for the dataset constructor; the new cache already stores age days."""
|
|
return None
|
|
|
|
def lookup_indices(self, eid: int, raw_tokens: np.ndarray, age_days: np.ndarray) -> np.ndarray:
|
|
out = np.full(len(raw_tokens), -1, dtype=np.int64)
|
|
real = raw_tokens > 1
|
|
if not np.any(real):
|
|
return out
|
|
eid_pos = self._eid_to_pos.get(int(eid))
|
|
if eid_pos is None:
|
|
return out
|
|
|
|
start = int(self.eid_start[eid_pos])
|
|
end = int(self.eid_start[eid_pos + 1])
|
|
if start == end:
|
|
return out
|
|
|
|
real_pos = np.nonzero(real)[0]
|
|
n_take = min(len(real_pos), end - start)
|
|
if n_take == 0:
|
|
return out
|
|
|
|
out[real_pos[:n_take]] = np.asarray(
|
|
self.row_index[start:start + n_take],
|
|
dtype=np.int64,
|
|
)
|
|
|
|
expected_tokens = np.asarray(self.raw_tokens[start:start + n_take], dtype=np.int64)
|
|
expected_age_days = np.asarray(self.age_days[start:start + n_take], dtype=np.int64)
|
|
actual_tokens = np.asarray(raw_tokens[real_pos[:n_take]], dtype=np.int64)
|
|
actual_age_days = np.rint(
|
|
np.asarray(age_days[real_pos[:n_take]], dtype=np.float64)
|
|
).astype(np.int64)
|
|
if (
|
|
not np.array_equal(expected_tokens, actual_tokens)
|
|
or not np.array_equal(expected_age_days, actual_age_days)
|
|
):
|
|
raise ValueError(
|
|
"Exposure cache is not aligned to the disease sequence for "
|
|
f"eid={eid}. Regenerate it with the same data_prefix and labels."
|
|
)
|
|
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]
|
|
source = self.daily if kind == "daily" else self.monthly
|
|
out[valid_pos] = np.asarray(source[valid_indices], dtype=np.float32)
|
|
return out
|
|
|
|
|
|
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: "<PAD>",
|
|
CHECKUP_IDX: "<CHECKUP>",
|
|
}
|
|
if include_no_event:
|
|
label_id_to_code[NO_EVENT_IDX] = "<NO_EVENT>"
|
|
|
|
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 = 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,
|
|
) -> 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
|