Add exposure cache and keep absolute time only
This commit is contained in:
300
prepare_exposure_cache.py
Normal file
300
prepare_exposure_cache.py
Normal file
@@ -0,0 +1,300 @@
|
||||
"""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
|
||||
<NO_EVENT> 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()
|
||||
Reference in New Issue
Block a user