"""Build a random-access exposure cache from disease-level parquet files. 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. This script converts those parquet files into a compact directory: 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. """ from __future__ import annotations import argparse import json from pathlib import Path from typing import Iterable import numpy as np import pandas as pd DAILY_LENGTH = 1826 MONTHLY_LENGTH = 241 DAILY_CHANNELS = ("tmean", "tmax", "tmin", "rhmean") MONTHLY_CHANNELS = ("tmean", "rhmean") QUALITY_COLUMNS = ( "n_days_nonmissing", "n_rh_days_nonmissing", "n_months_nonmissing", "n_rh_months_nonmissing", ) 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: cols.extend(f"{name}_d{idx:04d}" for idx in range(DAILY_LENGTH)) return cols def _monthly_columns() -> list[str]: cols: list[str] = [] for name in MONTHLY_CHANNELS: cols.extend(f"{name}_m{idx:03d}" for idx in range(MONTHLY_LENGTH)) return cols 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: raise ImportError( "prepare_exposure_cache.py requires pyarrow. Install requirements " "or run `pip install pyarrow`." ) from exc schema_names = set(pq.ParquetFile(path).schema.names) return [col for col in columns if col in schema_names] def _read_parquet_columns(path: Path, columns: list[str]) -> pd.DataFrame: return pd.read_parquet(path, columns=columns) 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) def _count_rows(summary: pd.DataFrame) -> int: if "n_cases" in summary.columns: return int(summary["n_cases"].sum()) return int(sum(pd.read_parquet(path, columns=["eid"]).shape[0] for path in summary["daily_path"])) def build_exposure_cache( *, exposure_dir: str | Path, output_dir: str | Path, summary_file: str = "summary.csv", overwrite: bool = False, ) -> int: exposure_dir = Path(exposure_dir) output_dir = Path(output_dir) summary_path = exposure_dir / summary_file if not summary_path.is_file(): raise FileNotFoundError(f"summary.csv not found: {summary_path}") 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, ] if any(path.exists() for path in outputs) and not overwrite: raise FileExistsError( f"{output_dir} already contains exposure cache files; pass --overwrite" ) summary = pd.read_csv(summary_path) required = {"label_code", "daily_file", "monthly_file"} missing = required - set(summary.columns) if missing: raise ValueError(f"{summary_path} is missing columns: {sorted(missing)}") summary = summary.copy() summary["daily_path"] = summary["daily_file"].map(lambda name: exposure_dir / str(name)) summary["monthly_path"] = summary["monthly_file"].map(lambda name: exposure_dir / str(name)) n_rows = _count_rows(summary) 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( onset_date_path, mode="w+", dtype="datetime64[D]", shape=(n_rows,), ) daily_mm = np.lib.format.open_memmap( daily_path, mode="w+", dtype=np.float32, shape=(n_rows, DAILY_LENGTH, len(DAILY_CHANNELS)), ) monthly_mm = np.lib.format.open_memmap( monthly_path, mode="w+", dtype=np.float32, shape=(n_rows, MONTHLY_LENGTH, len(MONTHLY_CHANNELS)), ) quality_mm = np.lib.format.open_memmap( quality_path, mode="w+", dtype=np.float32, shape=(n_rows, len(QUALITY_COLUMNS)), ) daily_cols = _daily_columns() monthly_cols = _monthly_columns() offset = 0 for row in summary.itertuples(index=False): 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", "onset_date", "token", *_safe_columns(daily_file, daily_cols), *_safe_columns(daily_file, ["n_days_nonmissing", "n_rh_days_nonmissing"]), ] monthly_read_cols = [ "eid", "onset_date", "token", *_safe_columns(monthly_file, monthly_cols), *_safe_columns(monthly_file, ["n_months_nonmissing", "n_rh_months_nonmissing"]), ] 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)}" ) 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") keys[offset:end] = encode_exposure_key( daily_df["eid"].to_numpy(dtype=np.int64), daily_df["token"].to_numpy(dtype=np.int64), ) 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, daily_cols, DAILY_LENGTH, len(DAILY_CHANNELS), ) monthly_mm[offset:end] = _reshape_window( monthly_df, 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 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}. " "Regenerate summary.csv or remove n_cases before building." ) manifest = { "source_dir": str(exposure_dir), "n_rows": int(n_rows), "legacy_key": "(eid << 16) | raw_token", "alignment_key": "(eid, raw_token, onset_date - date_of_birth)", "requires_basic_info_column": "date_of_birth", "daily_shape": [int(n_rows), DAILY_LENGTH, len(DAILY_CHANNELS)], "daily_channels": list(DAILY_CHANNELS), "monthly_shape": [int(n_rows), MONTHLY_LENGTH, len(MONTHLY_CHANNELS)], "monthly_channels": list(MONTHLY_CHANNELS), "quality_columns": list(QUALITY_COLUMNS), "raw_token_convention": "padding=0, checkup=1, labels.csv first row token=2", } manifest_path.write_text(json.dumps(manifest, indent=2), encoding="utf-8") return int(n_rows) 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("--summary-file", default="summary.csv") parser.add_argument("--overwrite", action="store_true") args = parser.parse_args() n_rows = build_exposure_cache( exposure_dir=args.exposure_dir, output_dir=args.output_dir, summary_file=args.summary_file, overwrite=args.overwrite, ) print(f"Wrote {n_rows:,} exposure rows to {args.output_dir}") if __name__ == "__main__": main()