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

Use eid-aligned exposure cache

Browse Source
This commit is contained in:
Jiarui Li
2026-07-08 12:17:30 +08:00
parent 642e8aad67
commit 46d530fad0
3 changed files with 275 additions and 630 deletions
Download Patch File Download Diff File Expand all files Collapse all files

304
dataset.py
View File

@@ -2,7 +2,6 @@
from __future__ import annotations
import json
from collections import OrderedDict
from pathlib import Path
from typing import Dict, Iterable, List, Literal, Optional, Tuple
@@ -43,9 +42,9 @@ def _monthly_exposure_columns() -> list[str]:
class ExposureCache:
"""Random-access view over files produced by prepare_exposure_cache.py."""
"""Eid-sequence-aligned exposure windows from prepare_exposure_cache.py."""
def __init__(self, cache_dir: str | Path, row_group_cache_size: int = 4):
def __init__(self, cache_dir: str | Path):
cache_dir = Path(cache_dir)
self.cache_dir = cache_dir
manifest_path = cache_dir / "exposure_manifest.json"
@@ -54,123 +53,125 @@ class ExposureCache:
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. "
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"
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,
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 = np.load(eid_path, mmap_mode="r")
self.raw_tokens = np.load(token_path, mmap_mode="r")
self.age_days = np.load(age_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}")
self.eid_index = np.load(eid_index_path, mmap_mode="r")
self.eid_start = np.load(eid_start_path, mmap_mode="r")
self.daily = np.load(daily_path, mmap_mode="r")
self.monthly = np.load(monthly_path, 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.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}"
)
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:
if (
len(self.raw_tokens) != n_rows
or len(self.age_days) != 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")
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")
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._key_to_index: dict[tuple[int, int, int], int] | None = None
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 parquet locality key for sampler-side batching."""
"""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)
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]))
first = int(valid[0])
return (first // 1024, first % 1024)
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
"""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:
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
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]
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
]
n_take = min(len(real_pos), end - start)
if n_take == 0:
return out
out[real_pos[:n_take]] = np.arange(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:
@@ -198,111 +199,10 @@ class ExposureCache:
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
source = self.daily if kind == "daily" else self.monthly
out[valid_pos] = np.asarray(source[valid_indices], dtype=np.float32)
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,
@@ -372,17 +272,13 @@ class _ExpoBaseDataset(Dataset):
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,
)
ExposureCache(exposure_cache_dir)
if exposure_cache_dir is not None
else None
)
@@ -553,7 +449,6 @@ class NextStepHealthDataset(_ExpoBaseDataset):
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,
@@ -562,7 +457,6 @@ class NextStepHealthDataset(_ExpoBaseDataset):
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] = []
@@ -651,7 +545,6 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
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}")
@@ -663,7 +556,6 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
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