468 lines
17 KiB
Python
468 lines
17 KiB
Python
"""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 <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.
|
|
"""
|
|
|
|
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_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 _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",
|
|
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"
|
|
)
|
|
|
|
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"
|
|
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)
|
|
token = int(row.raw_token)
|
|
wanted = wanted_by_token.get(token)
|
|
if wanted is None or wanted.empty:
|
|
continue
|
|
|
|
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:
|
|
continue
|
|
|
|
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:
|
|
continue
|
|
|
|
daily_df = daily_df.set_index("position").loc[common_positions].reset_index()
|
|
monthly_df = monthly_df.set_index("position").loc[common_positions].reset_index()
|
|
positions = common_positions.astype(np.int64)
|
|
daily_mm[positions] = _reshape_window(
|
|
daily_df,
|
|
daily_cols,
|
|
DAILY_LENGTH,
|
|
len(DAILY_CHANNELS),
|
|
)
|
|
monthly_mm[positions] = _reshape_window(
|
|
monthly_df,
|
|
monthly_cols,
|
|
MONTHLY_LENGTH,
|
|
len(MONTHLY_CHANNELS),
|
|
)
|
|
quality_mm[positions, 0] = daily_df.get("n_days_nonmissing", np.nan)
|
|
quality_mm[positions, 1] = daily_df.get("n_rh_days_nonmissing", np.nan)
|
|
quality_mm[positions, 2] = monthly_df.get("n_months_nonmissing", np.nan)
|
|
quality_mm[positions, 3] = monthly_df.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,
|
|
"labels_file": str(Path(labels_file).resolve()),
|
|
"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("--labels-file", default="labels.csv")
|
|
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,
|
|
labels_file=args.labels_file,
|
|
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()
|