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

Add exposure cache and keep absolute time only

Browse Source
This commit is contained in:
Jiarui Li
2026-07-07 17:21:52 +08:00
parent a0379daf29
commit 45a857d1a6
9 changed files with 690 additions and 198 deletions
Download Patch File Download Diff File Expand all files Collapse all files

223
dataset.py
View File

@@ -1,7 +1,8 @@
# dataset.py
from __future__ import annotations
from typing import Dict, List, Literal, Optional, Tuple
from pathlib import Path
from typing import Dict, Iterable, List, Literal, Optional, Tuple
import numpy as np
import pandas as pd
@@ -19,6 +20,92 @@ from targets import (
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(
@@ -87,11 +174,19 @@ class _ExpoBaseDataset(Dataset):
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,
@@ -112,6 +207,9 @@ class _ExpoBaseDataset(Dataset):
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
@@ -140,6 +238,25 @@ class _ExpoBaseDataset(Dataset):
)
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,
*,
@@ -176,6 +293,46 @@ class _ExpoBaseDataset(Dataset):
"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
@@ -194,12 +351,16 @@ class NextStepHealthDataset(_ExpoBaseDataset):
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] = []
@@ -221,7 +382,7 @@ class NextStepHealthDataset(_ExpoBaseDataset):
require_sorted=True,
)
self.samples.append({
sample = {
"eid": eid,
"event_seq": target_pack.next_token.input_events,
"time_seq": target_pack.next_token.input_times_years,
@@ -231,14 +392,22 @@ class NextStepHealthDataset(_ExpoBaseDataset):
"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]
return {
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),
@@ -248,6 +417,11 @@ class NextStepHealthDataset(_ExpoBaseDataset):
"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):
@@ -273,6 +447,8 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
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}")
@@ -282,6 +458,8 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
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
@@ -310,6 +488,7 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
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,
@@ -406,11 +585,12 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
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
return {
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),
@@ -419,6 +599,17 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
"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":
@@ -441,6 +632,22 @@ def _collate_common_static(batch: List[Dict]) -> Dict:
}
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],
@@ -489,6 +696,9 @@ def next_step_collate_fn(batch: List[Dict]) -> Dict:
"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
@@ -524,6 +734,9 @@ def all_future_collate_fn(batch: List[Dict]) -> Dict:
"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