"""Build an eid-sequence-aligned exposure cache for DeepHealth training. 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``. The output 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 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 import json from pathlib import Path from typing import Iterable import numpy as np import pandas as pd 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 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 _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]: 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 _parquet_row_count(path: Path) -> int: 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 return int(pq.ParquetFile(path).metadata.num_rows) 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 _load_summary( exposure_dir: Path, summary_file: str, *, show_progress: bool, ) -> pd.DataFrame: summary_path = exposure_dir / summary_file if not summary_path.is_file(): raise FileNotFoundError(f"summary.csv not found: {summary_path}") 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)) counts: list[int] = [] iterator = tqdm( summary.itertuples(index=False), total=len(summary), desc="Counting exposure rows", unit="file", disable=not show_progress, ) 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_count = _parquet_row_count(daily_file) monthly_count = _parquet_row_count(monthly_file) if daily_count != monthly_count: raise ValueError( f"Daily/monthly row count mismatch for {row.label_code}: " f"{daily_count} vs {monthly_count}" ) counts.append(daily_count) summary["n_rows"] = counts return summary 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] 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"{data_prefix}_basic_info.csv must contain date_of_birth for exposure alignment" ) 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) birth = pd.to_datetime( basic_table.loc[rows["eid"].to_numpy(), "date_of_birth"].to_numpy(), errors="coerce", ) 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, ) -> 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_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 output_paths) and not overwrite: raise FileExistsError( f"{output_dir} already contains exposure cache files; pass --overwrite" ) summary = _load_summary( exposure_dir, summary_file, show_progress=show_progress, ) 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, 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+", 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_mm[:] = np.nan monthly_mm[:] = np.nan quality_mm[:] = np.nan daily_cols = _daily_columns() monthly_cols = _monthly_columns() 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) iterator = tqdm( summary.itertuples(index=False), total=len(summary), desc="Writing eid-sequence exposure cache", unit="file", disable=not show_progress, ) for row in iterator: daily_file = Path(row.daily_path) monthly_file = Path(row.monthly_path) 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)}" ) 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() 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 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, ) 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[positions] = _reshape_window( monthly_df.iloc[source_rows], monthly_cols, MONTHLY_LENGTH, len(MONTHLY_CHANNELS), ) 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 daily_mm.flush() monthly_mm.flush() quality_mm.flush() manifest = { "storage": "eid_sequence_npy", "source_dir": str(exposure_dir.resolve()), "data_prefix": data_prefix, "n_rows": int(n_rows), "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), "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("--data-prefix", default="ukb") parser.add_argument("--summary-file", default="summary.csv") parser.add_argument( "--no-progress", action="store_true", help="Disable tqdm progress bars.", ) 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, 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__": main()