"""Build a random-access exposure index/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. 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 default index directory contains: 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. """ 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 _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) 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_index( *, 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) 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 = 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) 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,), ) daily_files: list[str] = [] monthly_files: list[str] = [] offset = 0 for file_id, row in enumerate(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_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): raise ValueError( f"Daily/monthly row count mismatch for {row.label_code}: " f"{len(daily_df)} vs {len(monthly_df)}" ) 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 {row.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 {row.label_code}" ) monthly_rg = monthly_rg_all[monthly_pos] monthly_row = monthly_row_all[monthly_pos] end = offset + n if end > n_rows: raise RuntimeError("Exposure index row count exceeded preallocated size") 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_file_id_mm[offset:end] = file_id daily_row_group_mm[offset:end] = daily_rg daily_row_in_group_mm[offset:end] = daily_row monthly_file_id_mm[offset:end] = file_id monthly_row_group_mm[offset:end] = monthly_rg monthly_row_in_group_mm[offset:end] = monthly_row daily_files.append(str(daily_file.resolve())) monthly_files.append(str(monthly_file.resolve())) offset = end if offset != n_rows: raise RuntimeError( f"Expected {n_rows} rows from summary but indexed {offset}. " "Regenerate summary.csv or remove n_cases before building." ) 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": daily_files, "monthly_files": monthly_files, "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", ) return int(n_rows) 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 = { "storage": "dense_npy", "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( "--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("--overwrite", action="store_true") args = parser.parse_args() 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, ) 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, ) print(f"Wrote {n_rows:,} dense exposure rows to {args.output_dir}") if __name__ == "__main__": main()