"""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_row_index.npy int64 window row per real disease event, -1 if missing 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, capacity (N, 1826, 4); first M rows are sequential matched windows channels: tmean, tmax, tmin, rhmean exposure_monthly.npy float32 memmap, capacity (N, 241, 2); first M rows are sequential matched windows channels: tmean, rhmean exposure_quality.npy float32 memmap, capacity (N, 4); first M rows are matched-window quality stats 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 from concurrent.futures import FIRST_COMPLETED, ProcessPoolExecutor, wait import json import os 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_matching_parquet_rows( path: Path, columns: list[str], wanted: pd.DataFrame, ) -> pd.DataFrame: 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) available = [col for col in columns if col in set(parquet_file.schema.names)] key_cols = ["eid", "onset_date", "token"] missing_keys = [col for col in key_cols if col not in available] if missing_keys: raise ValueError(f"{path} is missing key columns: {missing_keys}") wanted_keys = wanted[["eid", "onset_date", "token", "position"]].copy() wanted_keys["eid"] = wanted_keys["eid"].astype(np.int64) wanted_keys["token"] = wanted_keys["token"].astype(np.int64) wanted_keys["onset_date"] = pd.to_datetime( wanted_keys["onset_date"], errors="coerce", ).dt.normalize() chunks: list[pd.DataFrame] = [] for row_group_idx in range(parquet_file.num_row_groups): key_frame = parquet_file.read_row_group( row_group_idx, columns=key_cols, ).to_pandas() if key_frame.empty: continue key_frame = key_frame.copy() key_frame["_row_in_group"] = np.arange(len(key_frame), dtype=np.int64) key_frame["eid"] = key_frame["eid"].astype(np.int64) key_frame["token"] = key_frame["token"].astype(np.int64) key_frame["onset_date"] = pd.to_datetime( key_frame["onset_date"], errors="coerce", ).dt.normalize() matches = key_frame.merge( wanted_keys, on=key_cols, how="inner", sort=False, ) if matches.empty: continue row_values = parquet_file.read_row_group( row_group_idx, columns=available, ).to_pandas() selected = row_values.iloc[ matches["_row_in_group"].to_numpy(dtype=np.int64) ].copy() selected["position"] = matches["position"].to_numpy(dtype=np.int64) chunks.append(selected) if not chunks: return pd.DataFrame(columns=[*columns, "position"]) return pd.concat(chunks, ignore_index=True) 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 _quality_column(df: pd.DataFrame, name: str, n_rows: int) -> np.ndarray: if name not in df: return np.full(n_rows, np.nan, dtype=np.float32) return df[name].to_numpy(dtype=np.float32, copy=True) def _process_exposure_task(task: dict) -> dict: daily_file = Path(task["daily_path"]) monthly_file = Path(task["monthly_path"]) wanted = task["wanted"] daily_cols = task["daily_cols"] monthly_cols = task["monthly_cols"] 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_matching_parquet_rows(daily_file, daily_read_cols, wanted) monthly_df = _read_matching_parquet_rows(monthly_file, monthly_read_cols, wanted) if daily_df.empty: return {"positions": np.empty(0, dtype=np.int64)} common_positions = np.intersect1d( daily_df["position"].to_numpy(dtype=np.int64), monthly_df["position"].to_numpy(dtype=np.int64), ) if len(common_positions) == 0: return {"positions": np.empty(0, dtype=np.int64)} daily_df = daily_df.set_index("position").loc[common_positions].reset_index() monthly_df = monthly_df.set_index("position").loc[common_positions].reset_index() n_match = len(common_positions) quality = np.stack( [ _quality_column(daily_df, "n_days_nonmissing", n_match), _quality_column(daily_df, "n_rh_days_nonmissing", n_match), _quality_column(monthly_df, "n_months_nonmissing", n_match), _quality_column(monthly_df, "n_rh_months_nonmissing", n_match), ], axis=1, ) return { "positions": common_positions.astype(np.int64), "daily": _reshape_window( daily_df, daily_cols, DAILY_LENGTH, len(DAILY_CHANNELS), ), "monthly": _reshape_window( monthly_df, monthly_cols, MONTHLY_LENGTH, len(MONTHLY_CHANNELS), ), "quality": quality, } 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)) 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}") return summary def _load_label_token_map(labels_file: str | Path) -> dict[str, int]: out: dict[str, int] = {} with Path(labels_file).open(encoding="utf-8") as handle: for idx, line in enumerate(handle): parts = line.strip().split(" ") if parts and parts[0]: out[parts[0]] = idx + 2 return out 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", labels_file: str | Path = "labels.csv", summary_file: str = "summary.csv", workers: int = 1, max_in_flight: int = 0, 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_row_index.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" ) 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") label_token_map = _load_label_token_map(labels_file) summary = _load_summary( exposure_dir, summary_file, show_progress=show_progress, ) summary["raw_token"] = summary["label_code"].map(label_token_map) needed_tokens = set(sequence_rows["token"].astype(np.int64).unique().tolist()) summary = summary[ summary["raw_token"].notna() & summary["raw_token"].astype(np.int64).isin(needed_tokens) ].copy() summary["raw_token"] = summary["raw_token"].astype(np.int64) if summary.empty: raise ValueError( "No exposure summary rows match disease tokens in " f"{data_prefix}_event_data.npy. Check --summary-file and --labels-file." ) 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" row_index_path = output_dir / "exposure_row_index.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) row_index_mm = np.lib.format.open_memmap( row_index_path, mode="w+", dtype=np.int64, shape=(n_rows,), ) row_index_mm[:] = -1 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() wanted_by_token = { int(token): frame.reset_index(drop=True) for token, frame in sequence_rows.groupby("token", sort=False) } write_offset = 0 tasks: list[dict] = [] for row in summary.itertuples(index=False): token = int(row.raw_token) wanted = wanted_by_token.get(token) if wanted is None or wanted.empty: continue tasks.append( { "daily_path": str(row.daily_path), "monthly_path": str(row.monthly_path), "wanted": wanted, "daily_cols": daily_cols, "monthly_cols": monthly_cols, } ) workers = max(1, int(workers)) max_in_flight = int(max_in_flight) def write_result(result: dict) -> None: nonlocal write_offset positions = result["positions"] if len(positions) == 0: return n_match = len(positions) end_offset = write_offset + n_match daily_mm[write_offset:end_offset] = result["daily"] monthly_mm[write_offset:end_offset] = result["monthly"] quality_mm[write_offset:end_offset] = result["quality"] row_index_mm[positions] = np.arange(write_offset, end_offset, dtype=np.int64) write_offset = end_offset if workers == 1: iterator = tqdm( map(_process_exposure_task, tasks), total=len(tasks), desc="Reading exposure parquet and writing cache", unit="file", disable=not show_progress, ) for result in iterator: write_result(result) else: with ProcessPoolExecutor(max_workers=workers) as executor: task_iter = iter(tasks) iterator = tqdm( total=len(tasks), desc=f"Reading exposure parquet ({workers} workers)", unit="file", disable=not show_progress, ) if max_in_flight <= 0: in_flight = { executor.submit(_process_exposure_task, task) for task in task_iter } while in_flight: done, in_flight = wait(in_flight, return_when=FIRST_COMPLETED) for future in done: write_result(future.result()) iterator.update(1) else: in_flight = { executor.submit(_process_exposure_task, task) for task in [next(task_iter, None) for _ in range(max_in_flight)] if task is not None } while in_flight: done, in_flight = wait(in_flight, return_when=FIRST_COMPLETED) for future in done: write_result(future.result()) iterator.update(1) task = next(task_iter, None) if task is not None: in_flight.add(executor.submit(_process_exposure_task, task)) iterator.close() row_index_mm.flush() daily_mm.flush() monthly_mm.flush() quality_mm.flush() manifest = { "storage": "eid_sequence_npy", "source_dir": str(exposure_dir.resolve()), "data_prefix": data_prefix, "labels_file": str(Path(labels_file).resolve()), "n_rows": int(n_rows), "window_capacity_rows": int(n_rows), "matched_rows": int(write_offset), "missing_rows": int(n_rows - write_offset), "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)], "active_daily_shape": [int(write_offset), DAILY_LENGTH, len(DAILY_CHANNELS)], "daily_channels": list(DAILY_CHANNELS), "monthly_shape": [int(n_rows), MONTHLY_LENGTH, len(MONTHLY_CHANNELS)], "active_monthly_shape": [int(write_offset), 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("--labels-file", default="labels.csv") parser.add_argument("--summary-file", default="summary.csv") parser.add_argument( "--workers", type=int, default=max(1, os.cpu_count() or 1), help=( "Worker processes for parquet reading and window extraction. " "The main process remains the only writer to the output memmaps." ), ) parser.add_argument( "--max-in-flight", type=int, default=0, help=( "Maximum submitted parquet tasks waiting/running at once. " "Use 0 to submit all tasks, which is the default for high-memory servers." ), ) 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, labels_file=args.labels_file, summary_file=args.summary_file, workers=args.workers, max_in_flight=args.max_in_flight, 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()