From 46d530fad059dfb5f7636198201f3d5125cbd056 Mon Sep 17 00:00:00 2001 From: Jiarui Li Date: Wed, 8 Jul 2026 12:17:30 +0800 Subject: [PATCH] Use eid-aligned exposure cache --- dataset.py | 304 +++++++------------- prepare_exposure_cache.py | 582 ++++++++++++-------------------------- train_next_step.py | 19 +- 3 files changed, 275 insertions(+), 630 deletions(-) diff --git a/dataset.py b/dataset.py index 84cc6ef..952e679 100644 --- a/dataset.py +++ b/dataset.py @@ -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 diff --git a/prepare_exposure_cache.py b/prepare_exposure_cache.py index 453cc76..345969a 100644 --- a/prepare_exposure_cache.py +++ b/prepare_exposure_cache.py @@ -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 - 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 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__": diff --git a/train_next_step.py b/train_next_step.py index b5c9b5a..e2c0e44 100644 --- a/train_next_step.py +++ b/train_next_step.py @@ -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(