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

582
prepare_exposure_cache.py
View File

@@ -1,65 +1,49 @@
"""Build a random-access exposure index/cache from disease-level parquet files.
"""Build an eid-sequence-aligned exposure cache for DeepHealth training.
The README-described exposure dataset is stored as one daily and one monthly
parquet file per disease. That layout is good for disease-specific analysis but
too expensive for mini-batch training, where we need exposure windows aligned
to arbitrary event sequences.
The source exposure dataset is stored as one daily and one monthly parquet file
per disease. That layout is inconvenient for mini-batch training because the
model consumes per-participant disease sequences. This script materializes one
large numpy cache ordered exactly like ``{data_prefix}_event_data.npy`` after
sorting by ``eid, age_days, token``.
By default this script builds a lightweight parquet index. It does not copy the
daily/monthly exposure windows; it only records which source parquet file,
row-group, and row each exposure event lives in. Dataset loading then reads the
original parquet row groups on demand.
The output directory contains:
The default index directory contains:
exposure_eid.npy int64 eid per real disease event
exposure_token.npy int32 raw disease token per real disease event
exposure_age_days.npy int32 age in days per real disease event
exposure_onset_date.npy datetime64[D] onset date per real disease event
exposure_eid_index.npy int64 unique eids in cache order
exposure_eid_start.npy int64 start offsets, length len(eid_index) + 1
exposure_daily.npy float32 memmap, shape (N, 1826, 4)
channels: tmean, tmax, tmin, rhmean
exposure_monthly.npy float32 memmap, shape (N, 241, 2)
channels: tmean, rhmean
exposure_quality.npy float32 memmap, shape (N, 4)
n_days, n_rh_days, n_months, n_rh_months
exposure_manifest.json metadata
exposure_eid.npy int64 eid per exposure row
exposure_token.npy int32 raw disease token per exposure row
exposure_onset_date.npy datetime64[D] onset date per exposure row
exposure_daily_file_id.npy int32 source daily file id per row
exposure_daily_row_group.npy int32 source daily row group per row
exposure_daily_row_in_group.npy int32 row offset inside daily row group
exposure_monthly_file_id.npy int32 source monthly file id per row
exposure_monthly_row_group.npy int32 source monthly row group per row
exposure_monthly_row_in_group.npy
int32 row offset inside monthly row group
exposure_manifest.json metadata and source parquet paths
For faster but much larger training storage, ``--mode dense`` materializes a
full dense numpy cache:
exposure_keys.npy uint64 legacy keys, key = (eid << 16) | raw_token
exposure_eid.npy int64 eid per exposure row
exposure_token.npy int32 raw disease token per exposure row
exposure_onset_date.npy datetime64[D] onset date per exposure row
exposure_daily.npy float32 memmap, shape (N, 1826, 4)
channels: tmean, tmax, tmin, rhmean
exposure_monthly.npy float32 memmap, shape (N, 241, 2)
channels: tmean, rhmean
exposure_quality.npy float32 memmap, shape (N, 4)
n_days, n_rh_days, n_months, n_rh_months
exposure_manifest.json metadata
The raw token convention follows the exposure README: padding=0, checkup=1,
and the first row of labels.csv is token=2. The model dataset inserts
<NO_EVENT> at token 2 and shifts real disease tokens by +1 internally; dataset
lookup converts back to these raw tokens before reading this cache. Dataset
alignment uses (eid, raw_token, onset_date - date_of_birth) so that raw
calendar dates in the exposure files match the age-day event times used by the
model.
Rows without matching exposure parquet records are kept as NaN windows. The
raw token convention follows the exposure README: padding=0, checkup=1, and
the first row of labels.csv is token=2. The model dataset inserts <NO_EVENT> at
token 2 and shifts real disease tokens by +1 internally; dataset lookup
converts back to these raw tokens before reading this cache.
"""
from __future__ import annotations
import argparse
from concurrent.futures import ProcessPoolExecutor, as_completed
import json
import os
from pathlib import Path
from typing import Iterable
import numpy as np
import pandas as pd
from tqdm.auto import tqdm
try:
from tqdm.auto import tqdm
except ImportError:
def tqdm(iterable=None, **kwargs):
return iterable if iterable is not None else range(kwargs.get("total", 0))
DAILY_LENGTH = 1826
@@ -74,14 +58,6 @@ QUALITY_COLUMNS = (
)
def encode_exposure_key(eid: np.ndarray, raw_token: np.ndarray) -> np.ndarray:
eid_u64 = np.asarray(eid, dtype=np.uint64)
token_u64 = np.asarray(raw_token, dtype=np.uint64)
if np.any(token_u64 >= (1 << 16)):
raise ValueError("raw_token must fit in 16 bits")
return (eid_u64 << np.uint64(16)) | token_u64
def _daily_columns() -> list[str]:
cols: list[str] = []
for name in DAILY_CHANNELS:
@@ -97,7 +73,6 @@ def _monthly_columns() -> list[str]:
def _safe_columns(path: Path, columns: Iterable[str]) -> list[str]:
"""Return the subset of requested columns present in a parquet file."""
try:
import pyarrow.parquet as pq
except ImportError as exc:
@@ -125,28 +100,6 @@ def _parquet_row_count(path: Path) -> int:
return int(pq.ParquetFile(path).metadata.num_rows)
def _row_group_positions(path: Path) -> tuple[np.ndarray, np.ndarray]:
"""Return row_group and row-in-group vectors for every parquet row."""
try:
import pyarrow.parquet as pq
except ImportError as exc:
raise ImportError(
"prepare_exposure_cache.py requires pyarrow. Install requirements "
"or run `pip install pyarrow`."
) from exc
parquet_file = pq.ParquetFile(path)
row_groups: list[np.ndarray] = []
row_offsets: list[np.ndarray] = []
for row_group_idx in range(parquet_file.num_row_groups):
n = parquet_file.metadata.row_group(row_group_idx).num_rows
row_groups.append(np.full(n, row_group_idx, dtype=np.int32))
row_offsets.append(np.arange(n, dtype=np.int32))
if not row_groups:
return np.empty(0, dtype=np.int32), np.empty(0, dtype=np.int32)
return np.concatenate(row_groups), np.concatenate(row_offsets)
def _reshape_window(df: pd.DataFrame, cols: list[str], length: int, n_channels: int) -> np.ndarray:
arr = df.reindex(columns=cols).to_numpy(dtype=np.float32, copy=True)
return arr.reshape(len(df), n_channels, length).transpose(0, 2, 1)
@@ -173,9 +126,8 @@ def _load_summary(
summary["monthly_path"] = summary["monthly_file"].map(lambda name: exposure_dir / str(name))
counts: list[int] = []
iterator = summary.itertuples(index=False)
iterator = tqdm(
iterator,
summary.itertuples(index=False),
total=len(summary),
desc="Counting exposure rows",
unit="file",
@@ -198,237 +150,63 @@ def _load_summary(
counts.append(daily_count)
summary["n_rows"] = counts
summary["offset"] = np.cumsum([0, *counts[:-1]], dtype=np.int64)
return summary
def _process_index_file_pair(task: tuple[int, str, str, str]) -> dict:
file_id, label_code, daily_path, monthly_path = task
daily_file = Path(daily_path)
monthly_file = Path(monthly_path)
def _load_sequence_rows(data_prefix: str) -> pd.DataFrame:
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]))
event_data = event_data[order]
daily_df = _read_parquet_columns(daily_file, ["eid", "onset_date", "token"])
monthly_df = _read_parquet_columns(monthly_file, ["eid", "onset_date", "token"])
if len(daily_df) != len(monthly_df):
basic_table = pd.read_csv(f"{data_prefix}_basic_info.csv", index_col=0)
basic_table.index = basic_table.index.astype(np.int64)
if "date_of_birth" not in basic_table.columns:
raise ValueError(
f"Daily/monthly row count mismatch for {label_code}: "
f"{len(daily_df)} vs {len(monthly_df)}"
f"{data_prefix}_basic_info.csv must contain date_of_birth for exposure alignment"
)
daily_rg, daily_row = _row_group_positions(daily_file)
monthly_rg_all, monthly_row_all = _row_group_positions(monthly_file)
n = len(daily_df)
if len(daily_rg) != n or len(monthly_rg_all) != n:
raise ValueError(f"Parquet row-group metadata row count mismatch for {label_code}")
daily_index = pd.MultiIndex.from_frame(daily_df[["eid", "onset_date", "token"]])
monthly_index = pd.MultiIndex.from_frame(monthly_df[["eid", "onset_date", "token"]])
monthly_pos = monthly_index.get_indexer(daily_index)
if np.any(monthly_pos < 0):
raise ValueError(f"Monthly parquet is missing daily exposure keys for {label_code}")
return {
"file_id": int(file_id),
"label_code": label_code,
"n_rows": int(n),
"eid": daily_df["eid"].to_numpy(dtype=np.int64),
"token": daily_df["token"].to_numpy(dtype=np.int32),
"onset_date": pd.to_datetime(
daily_df["onset_date"],
errors="coerce",
).to_numpy(dtype="datetime64[D]"),
"daily_row_group": daily_rg,
"daily_row_in_group": daily_row,
"monthly_row_group": monthly_rg_all[monthly_pos],
"monthly_row_in_group": monthly_row_all[monthly_pos],
}
def build_exposure_index(
*,
exposure_dir: str | Path,
output_dir: str | Path,
summary_file: str = "summary.csv",
overwrite: bool = False,
workers: int = 1,
show_progress: bool = True,
) -> int:
exposure_dir = Path(exposure_dir)
output_dir = Path(output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
output_paths = [
output_dir / "exposure_eid.npy",
output_dir / "exposure_token.npy",
output_dir / "exposure_onset_date.npy",
output_dir / "exposure_daily_file_id.npy",
output_dir / "exposure_daily_row_group.npy",
output_dir / "exposure_daily_row_in_group.npy",
output_dir / "exposure_monthly_file_id.npy",
output_dir / "exposure_monthly_row_group.npy",
output_dir / "exposure_monthly_row_in_group.npy",
output_dir / "exposure_manifest.json",
]
if any(path.exists() for path in output_paths) and not overwrite:
raise FileExistsError(
f"{output_dir} already contains exposure index files; pass --overwrite"
)
summary = _load_summary(
exposure_dir,
summary_file,
show_progress=show_progress,
)
n_rows = int(summary["n_rows"].sum())
eids_mm = np.lib.format.open_memmap(
output_dir / "exposure_eid.npy", mode="w+", dtype=np.int64, shape=(n_rows,)
)
tokens_mm = np.lib.format.open_memmap(
output_dir / "exposure_token.npy", mode="w+", dtype=np.int32, shape=(n_rows,)
)
onset_dates_mm = np.lib.format.open_memmap(
output_dir / "exposure_onset_date.npy",
mode="w+",
dtype="datetime64[D]",
shape=(n_rows,),
)
daily_file_id_mm = np.lib.format.open_memmap(
output_dir / "exposure_daily_file_id.npy",
mode="w+",
dtype=np.int32,
shape=(n_rows,),
)
daily_row_group_mm = np.lib.format.open_memmap(
output_dir / "exposure_daily_row_group.npy",
mode="w+",
dtype=np.int32,
shape=(n_rows,),
)
daily_row_in_group_mm = np.lib.format.open_memmap(
output_dir / "exposure_daily_row_in_group.npy",
mode="w+",
dtype=np.int32,
shape=(n_rows,),
)
monthly_file_id_mm = np.lib.format.open_memmap(
output_dir / "exposure_monthly_file_id.npy",
mode="w+",
dtype=np.int32,
shape=(n_rows,),
)
monthly_row_group_mm = np.lib.format.open_memmap(
output_dir / "exposure_monthly_row_group.npy",
mode="w+",
dtype=np.int32,
shape=(n_rows,),
)
monthly_row_in_group_mm = np.lib.format.open_memmap(
output_dir / "exposure_monthly_row_in_group.npy",
mode="w+",
dtype=np.int32,
shape=(n_rows,),
rows = pd.DataFrame(
{
"eid": event_data[:, 0].astype(np.int64),
"age_days": np.rint(event_data[:, 1].astype(np.float64)).astype(np.int32),
"token": event_data[:, 2].astype(np.int32),
}
)
rows = rows[rows["token"] > 1].reset_index(drop=True)
rows["position"] = np.arange(len(rows), dtype=np.int64)
tasks = [
(
int(file_id),
str(row.label_code),
str(Path(row.daily_path)),
str(Path(row.monthly_path)),
)
for file_id, row in enumerate(summary.itertuples(index=False))
]
workers = max(1, int(workers))
def write_result(result: dict) -> None:
file_id = int(result["file_id"])
row = summary.iloc[file_id]
offset = int(row.offset)
expected_n = int(row.n_rows)
n = int(result["n_rows"])
if n != expected_n:
raise RuntimeError(
f"Expected {expected_n} rows for {result['label_code']} "
f"from metadata but indexed {n}"
)
end = offset + n
if end > n_rows:
raise RuntimeError("Exposure index row count exceeded preallocated size")
eids_mm[offset:end] = result["eid"]
tokens_mm[offset:end] = result["token"]
onset_dates_mm[offset:end] = result["onset_date"]
daily_file_id_mm[offset:end] = file_id
daily_row_group_mm[offset:end] = result["daily_row_group"]
daily_row_in_group_mm[offset:end] = result["daily_row_in_group"]
monthly_file_id_mm[offset:end] = file_id
monthly_row_group_mm[offset:end] = result["monthly_row_group"]
monthly_row_in_group_mm[offset:end] = result["monthly_row_in_group"]
if workers == 1:
iterator = map(_process_index_file_pair, tasks)
for result in tqdm(
iterator,
total=len(tasks),
desc="Indexing exposure parquet",
unit="file",
disable=not show_progress,
):
write_result(result)
else:
with ProcessPoolExecutor(max_workers=workers) as executor:
futures = [executor.submit(_process_index_file_pair, task) for task in tasks]
for future in tqdm(
as_completed(futures),
total=len(futures),
desc=f"Indexing exposure parquet ({workers} workers)",
unit="file",
disable=not show_progress,
):
write_result(future.result())
for memmap in (
eids_mm,
tokens_mm,
onset_dates_mm,
daily_file_id_mm,
daily_row_group_mm,
daily_row_in_group_mm,
monthly_file_id_mm,
monthly_row_group_mm,
monthly_row_in_group_mm,
):
memmap.flush()
manifest = {
"storage": "parquet_index",
"source_dir": str(exposure_dir.resolve()),
"n_rows": int(n_rows),
"alignment_key": "(eid, raw_token, onset_date - date_of_birth)",
"requires_basic_info_column": "date_of_birth",
"daily_files": [
str(Path(path).resolve()) for path in summary["daily_path"].tolist()
],
"monthly_files": [
str(Path(path).resolve()) for path in summary["monthly_path"].tolist()
],
"daily_shape_per_row": [DAILY_LENGTH, len(DAILY_CHANNELS)],
"daily_channels": list(DAILY_CHANNELS),
"monthly_shape_per_row": [MONTHLY_LENGTH, len(MONTHLY_CHANNELS)],
"monthly_channels": list(MONTHLY_CHANNELS),
"raw_token_convention": "padding=0, checkup=1, labels.csv first row token=2",
}
(output_dir / "exposure_manifest.json").write_text(
json.dumps(manifest, indent=2),
encoding="utf-8",
birth = pd.to_datetime(
basic_table.loc[rows["eid"].to_numpy(), "date_of_birth"].to_numpy(),
errors="coerce",
)
return int(n_rows)
if birth.isna().any():
raise ValueError("date_of_birth contains missing or invalid values")
rows["onset_date"] = (
birth.to_numpy(dtype="datetime64[D]")
+ rows["age_days"].to_numpy(dtype="timedelta64[D]")
)
rows["onset_date"] = pd.to_datetime(rows["onset_date"]).dt.normalize()
return rows
def _write_eid_offsets(rows: pd.DataFrame, output_dir: Path) -> None:
eids = rows["eid"].to_numpy(dtype=np.int64)
unique_eids, starts = np.unique(eids, return_index=True)
starts = starts.astype(np.int64)
ends = np.concatenate([starts[1:], np.asarray([len(rows)], dtype=np.int64)])
eid_start = np.concatenate([starts, ends[-1:]]).astype(np.int64)
np.save(output_dir / "exposure_eid_index.npy", unique_eids.astype(np.int64))
np.save(output_dir / "exposure_eid_start.npy", eid_start)
def build_exposure_cache(
*,
exposure_dir: str | Path,
output_dir: str | Path,
data_prefix: str = "ukb",
summary_file: str = "summary.csv",
overwrite: bool = False,
show_progress: bool = True,
@@ -436,25 +214,20 @@ def build_exposure_cache(
exposure_dir = Path(exposure_dir)
output_dir = Path(output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
keys_path = output_dir / "exposure_keys.npy"
eid_path = output_dir / "exposure_eid.npy"
token_path = output_dir / "exposure_token.npy"
onset_date_path = output_dir / "exposure_onset_date.npy"
daily_path = output_dir / "exposure_daily.npy"
monthly_path = output_dir / "exposure_monthly.npy"
quality_path = output_dir / "exposure_quality.npy"
manifest_path = output_dir / "exposure_manifest.json"
outputs = [
keys_path,
eid_path,
token_path,
onset_date_path,
daily_path,
monthly_path,
quality_path,
manifest_path,
output_paths = [
output_dir / "exposure_eid.npy",
output_dir / "exposure_token.npy",
output_dir / "exposure_age_days.npy",
output_dir / "exposure_onset_date.npy",
output_dir / "exposure_eid_index.npy",
output_dir / "exposure_eid_start.npy",
output_dir / "exposure_daily.npy",
output_dir / "exposure_monthly.npy",
output_dir / "exposure_quality.npy",
output_dir / "exposure_manifest.json",
]
if any(path.exists() for path in outputs) and not overwrite:
if any(path.exists() for path in output_paths) and not overwrite:
raise FileExistsError(
f"{output_dir} already contains exposure cache files; pass --overwrite"
)
@@ -464,16 +237,29 @@ def build_exposure_cache(
summary_file,
show_progress=show_progress,
)
n_rows = int(summary["n_rows"].sum())
keys = np.lib.format.open_memmap(keys_path, mode="w+", dtype=np.uint64, shape=(n_rows,))
eids_mm = np.lib.format.open_memmap(eid_path, mode="w+", dtype=np.int64, shape=(n_rows,))
tokens_mm = np.lib.format.open_memmap(token_path, mode="w+", dtype=np.int32, shape=(n_rows,))
onset_dates_mm = np.lib.format.open_memmap(
sequence_rows = _load_sequence_rows(data_prefix)
n_rows = len(sequence_rows)
if n_rows == 0:
raise ValueError(f"{data_prefix}_event_data.npy contains no real disease events")
eid_path = output_dir / "exposure_eid.npy"
token_path = output_dir / "exposure_token.npy"
age_path = output_dir / "exposure_age_days.npy"
onset_date_path = output_dir / "exposure_onset_date.npy"
daily_path = output_dir / "exposure_daily.npy"
monthly_path = output_dir / "exposure_monthly.npy"
quality_path = output_dir / "exposure_quality.npy"
manifest_path = output_dir / "exposure_manifest.json"
np.save(eid_path, sequence_rows["eid"].to_numpy(dtype=np.int64))
np.save(token_path, sequence_rows["token"].to_numpy(dtype=np.int32))
np.save(age_path, sequence_rows["age_days"].to_numpy(dtype=np.int32))
np.save(
onset_date_path,
mode="w+",
dtype="datetime64[D]",
shape=(n_rows,),
sequence_rows["onset_date"].to_numpy(dtype="datetime64[D]"),
)
_write_eid_offsets(sequence_rows, output_dir)
daily_mm = np.lib.format.open_memmap(
daily_path,
mode="w+",
@@ -492,25 +278,28 @@ def build_exposure_cache(
dtype=np.float32,
shape=(n_rows, len(QUALITY_COLUMNS)),
)
daily_mm[:] = np.nan
monthly_mm[:] = np.nan
quality_mm[:] = np.nan
daily_cols = _daily_columns()
monthly_cols = _monthly_columns()
offset = 0
wanted_by_token = {
int(token): frame.reset_index(drop=True)
for token, frame in sequence_rows.groupby("token", sort=False)
}
matched = np.zeros(n_rows, dtype=bool)
rows = tqdm(
iterator = tqdm(
summary.itertuples(index=False),
total=len(summary),
desc="Materializing dense exposure cache",
desc="Writing eid-sequence exposure cache",
unit="file",
disable=not show_progress,
)
for row in rows:
for row in iterator:
daily_file = Path(row.daily_path)
monthly_file = Path(row.monthly_path)
if not daily_file.is_file():
raise FileNotFoundError(f"Missing daily parquet: {daily_file}")
if not monthly_file.is_file():
raise FileNotFoundError(f"Missing monthly parquet: {monthly_file}")
daily_read_cols = [
"eid",
@@ -528,70 +317,76 @@ def build_exposure_cache(
]
daily_df = _read_parquet_columns(daily_file, daily_read_cols)
monthly_df = _read_parquet_columns(monthly_file, monthly_read_cols)
if len(daily_df) != len(monthly_df):
raise ValueError(
f"Daily/monthly row count mismatch for {row.label_code}: "
f"{len(daily_df)} vs {len(monthly_df)}"
)
daily_df = daily_df.copy()
monthly_df = monthly_df.copy()
daily_df["_source_row"] = np.arange(len(daily_df), dtype=np.int64)
daily_df["onset_date"] = pd.to_datetime(
daily_df["onset_date"],
errors="coerce",
).dt.normalize()
monthly_df["onset_date"] = pd.to_datetime(
monthly_df["onset_date"],
errors="coerce",
).dt.normalize()
monthly_df = monthly_df.set_index(["eid", "onset_date", "token"]).reindex(
pd.MultiIndex.from_frame(daily_df[["eid", "onset_date", "token"]])
).reset_index()
n = len(daily_df)
end = offset + n
if end > n_rows:
raise RuntimeError("Exposure cache row count exceeded preallocated size")
tokens = daily_df["token"].dropna().astype(np.int64).unique()
wanted = pd.concat(
[wanted_by_token[int(token)] for token in tokens if int(token) in wanted_by_token],
ignore_index=True,
) if len(tokens) else pd.DataFrame()
if wanted.empty:
continue
keys[offset:end] = encode_exposure_key(
daily_df["eid"].to_numpy(dtype=np.int64),
daily_df["token"].to_numpy(dtype=np.int64),
matches = daily_df[["eid", "onset_date", "token", "_source_row"]].merge(
wanted[["eid", "onset_date", "token", "position"]],
on=["eid", "onset_date", "token"],
how="inner",
sort=False,
)
eids_mm[offset:end] = daily_df["eid"].to_numpy(dtype=np.int64)
tokens_mm[offset:end] = daily_df["token"].to_numpy(dtype=np.int32)
onset_dates_mm[offset:end] = pd.to_datetime(
daily_df["onset_date"],
errors="coerce",
).to_numpy(dtype="datetime64[D]")
daily_mm[offset:end] = _reshape_window(
daily_df,
if matches.empty:
continue
source_rows = matches["_source_row"].to_numpy(dtype=np.int64)
positions = matches["position"].to_numpy(dtype=np.int64)
daily_mm[positions] = _reshape_window(
daily_df.iloc[source_rows],
daily_cols,
DAILY_LENGTH,
len(DAILY_CHANNELS),
)
monthly_mm[offset:end] = _reshape_window(
monthly_df,
monthly_mm[positions] = _reshape_window(
monthly_df.iloc[source_rows],
monthly_cols,
MONTHLY_LENGTH,
len(MONTHLY_CHANNELS),
)
quality_mm[offset:end, 0] = daily_df.get("n_days_nonmissing", np.nan)
quality_mm[offset:end, 1] = daily_df.get("n_rh_days_nonmissing", np.nan)
quality_mm[offset:end, 2] = monthly_df.get("n_months_nonmissing", np.nan)
quality_mm[offset:end, 3] = monthly_df.get("n_rh_months_nonmissing", np.nan)
offset = end
quality_mm[positions, 0] = daily_df.iloc[source_rows].get("n_days_nonmissing", np.nan)
quality_mm[positions, 1] = daily_df.iloc[source_rows].get("n_rh_days_nonmissing", np.nan)
quality_mm[positions, 2] = monthly_df.iloc[source_rows].get("n_months_nonmissing", np.nan)
quality_mm[positions, 3] = monthly_df.iloc[source_rows].get("n_rh_months_nonmissing", np.nan)
matched[positions] = True
if offset != n_rows:
keys.flush()
eids_mm.flush()
tokens_mm.flush()
onset_dates_mm.flush()
daily_mm.flush()
monthly_mm.flush()
quality_mm.flush()
keys = np.lib.format.open_memmap(keys_path, mode="r+", dtype=np.uint64, shape=(offset,))
raise RuntimeError(
f"Expected {n_rows} rows from summary but wrote {offset}. "
"Check parquet metadata and regenerate summary.csv before building."
)
daily_mm.flush()
monthly_mm.flush()
quality_mm.flush()
manifest = {
"storage": "dense_npy",
"source_dir": str(exposure_dir),
"storage": "eid_sequence_npy",
"source_dir": str(exposure_dir.resolve()),
"data_prefix": data_prefix,
"n_rows": int(n_rows),
"legacy_key": "(eid << 16) | raw_token",
"alignment_key": "(eid, raw_token, onset_date - date_of_birth)",
"matched_rows": int(matched.sum()),
"missing_rows": int((~matched).sum()),
"alignment_key": "(eid, raw_token, date_of_birth + age_days)",
"requires_basic_info_column": "date_of_birth",
"daily_shape": [int(n_rows), DAILY_LENGTH, len(DAILY_CHANNELS)],
"daily_channels": list(DAILY_CHANNELS),
@@ -608,25 +403,8 @@ def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--exposure-dir", required=True)
parser.add_argument("--output-dir", default="ukb_exposure_cache")
parser.add_argument("--data-prefix", default="ukb")
parser.add_argument("--summary-file", default="summary.csv")
parser.add_argument(
"--mode",
choices=("index", "dense"),
default="index",
help=(
"index writes only lightweight parquet row pointers; dense copies "
"all exposure windows into numpy memmaps."
),
)
parser.add_argument(
"--workers",
type=int,
default=max(1, min(8, (os.cpu_count() or 1))),
help=(
"Number of worker processes for --mode index. Dense mode remains "
"single-writer to avoid concurrent writes to the same memmap."
),
)
parser.add_argument(
"--no-progress",
action="store_true",
@@ -634,26 +412,16 @@ def main() -> None:
)
parser.add_argument("--overwrite", action="store_true")
args = parser.parse_args()
show_progress = not args.no_progress
if args.mode == "index":
n_rows = build_exposure_index(
exposure_dir=args.exposure_dir,
output_dir=args.output_dir,
summary_file=args.summary_file,
overwrite=args.overwrite,
workers=args.workers,
show_progress=show_progress,
)
print(f"Wrote {n_rows:,} exposure row pointers to {args.output_dir}")
else:
n_rows = build_exposure_cache(
exposure_dir=args.exposure_dir,
output_dir=args.output_dir,
summary_file=args.summary_file,
overwrite=args.overwrite,
show_progress=show_progress,
)
print(f"Wrote {n_rows:,} dense exposure rows to {args.output_dir}")
n_rows = build_exposure_cache(
exposure_dir=args.exposure_dir,
output_dir=args.output_dir,
data_prefix=args.data_prefix,
summary_file=args.summary_file,
overwrite=args.overwrite,
show_progress=not args.no_progress,
)
print(f"Wrote {n_rows:,} eid-sequence-aligned exposure rows to {args.output_dir}")
if __name__ == "__main__":

19
train_next_step.py
View File

@@ -102,8 +102,8 @@ class ExposureLocalityBatchSampler(Sampler[List[int]]):
exposure_cache = getattr(dataset, "exposure_cache", None)
if exposure_index is None or exposure_cache is None:
return (2**31 - 1, 2**31 - 1, raw_idx)
file_id, row_group = exposure_cache.locality_key(exposure_index)
return (file_id, row_group, raw_idx)
block_id, block_offset = exposure_cache.locality_key(exposure_index)
return (block_id, block_offset, raw_idx)
def parse_args() -> argparse.Namespace:
@@ -136,16 +136,6 @@ def parse_args() -> argparse.Namespace:
parser.add_argument("--exposure_conv_kernel_size", type=int, default=7)
parser.add_argument("--exposure_mlp_ratio", type=float, default=4.0)
parser.add_argument("--no_exposure_gate", action="store_true")
parser.add_argument(
"--exposure_row_group_cache_size",
type=int,
default=4,
help=(
"Number of parquet exposure row groups cached per DataLoader worker "
"when using indexed exposure storage."
),
)
parser.add_argument("--target_mode", type=str, default="uts",
choices=["delphi2m", "uts"])
parser.add_argument("--readout_name", type=str, default=None,
@@ -195,8 +185,6 @@ def parse_args() -> argparse.Namespace:
raise ValueError("train_ratio + val_ratio + test_ratio must equal 1.0")
if args.num_workers > 0 and args.prefetch_factor <= 0:
raise ValueError("prefetch_factor must be positive when num_workers > 0")
if args.exposure_row_group_cache_size < 0:
raise ValueError("exposure_row_group_cache_size must be non-negative")
if args.exposure_locality_buffer_size < 0:
raise ValueError("exposure_locality_buffer_size must be non-negative")
if args.target_mode == "uts":
@@ -447,7 +435,6 @@ def build_metadata(
"exposure_conv_kernel_size": int(args.exposure_conv_kernel_size),
"exposure_mlp_ratio": float(args.exposure_mlp_ratio),
"exposure_use_gate": not bool(args.no_exposure_gate),
"exposure_row_group_cache_size": int(args.exposure_row_group_cache_size),
"num_workers": int(args.num_workers),
"prefetch_factor": int(args.prefetch_factor),
"exposure_locality_buffer_size": int(args.exposure_locality_buffer_size),
@@ -485,7 +472,6 @@ def main() -> None:
"DataLoader IO: "
f"num_workers={args.num_workers}, "
f"prefetch_factor={args.prefetch_factor if args.num_workers > 0 else None}, "
f"exposure_row_group_cache_size={args.exposure_row_group_cache_size}, "
f"exposure_locality_buffer_size={args.exposure_locality_buffer_size}"
)
@@ -496,7 +482,6 @@ def main() -> None:
include_no_event_in_uts_target=args.include_no_event_in_uts_target,
exposure_cache_dir=args.exposure_cache_dir,
mask_onset_exposure=args.mask_onset_exposure,
exposure_row_group_cache_size=args.exposure_row_group_cache_size,
)
if args.train_eid_file and args.val_eid_file and args.test_eid_file:
train_subset, val_subset, test_subset = split_dataset_by_eid_files(