204 lines
7.2 KiB
Python
204 lines
7.2 KiB
Python
|
"""Create a compact calendar-dated disease-event index from ``ukb_data.csv``.
|
||
|
|
|
||
|
|
Unlike ``prepare_data.py``, this ETL does not create model-ready relative-time
|
||
|
|
sequences or other-information tokens. It writes one structured ``.npy`` file
|
||
|
|
with exactly three fields:
|
||
|
|
|
||
|
|
eid int64
|
||
|
|
event_date datetime64[D]
|
||
|
|
token int32
|
||
|
|
|
||
|
|
``token`` follows the existing ``labels.csv`` convention used by
|
||
|
|
``prepare_data.py``: padding=0, checkup=1 (not emitted here), and the first
|
||
|
|
label in ``labels.csv`` receives token 2. Each ``(eid, token)`` is deduplicated
|
||
|
|
to the first known event date.
|
||
|
|
|
||
|
|
The output is intended for calendar-indexed temperature and air-pollution
|
||
|
|
queries. It contains no date of birth, sex, covariates, or checkup events.
|
||
|
|
|
||
|
|
Usage
|
||
|
|
-----
|
||
|
|
python prepare_event_dates.py --output ukb_disease_event_dates.npy
|
||
|
|
"""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
import pandas as pd
|
||
|
|
try:
|
||
|
|
from tqdm import tqdm
|
||
|
|
except ImportError: # Keep the ETL runnable in minimal Python environments.
|
||
|
|
tqdm = None
|
||
|
|
|
||
|
|
|
||
|
|
EVENT_DTYPE = np.dtype(
|
||
|
|
[
|
||
|
|
("eid", "<i8"),
|
||
|
|
("event_date", "datetime64[D]"),
|
||
|
|
("token", "<i4"),
|
||
|
|
]
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def load_label_tokens(labels_file: str | Path) -> dict[str, int]:
|
||
|
|
"""Return the label-code -> token mapping shared with ``prepare_data.py``."""
|
||
|
|
token_map: dict[str, int] = {}
|
||
|
|
with Path(labels_file).open(encoding="utf-8") as handle:
|
||
|
|
for index, line in enumerate(handle):
|
||
|
|
code = line.strip().split(" ", maxsplit=1)[0]
|
||
|
|
if code:
|
||
|
|
token_map[code] = index + 2
|
||
|
|
return token_map
|
||
|
|
|
||
|
|
|
||
|
|
def build_raw_column_map(
|
||
|
|
field_map_file: str | Path,
|
||
|
|
icd_map_file: str | Path,
|
||
|
|
) -> tuple[dict[str, str], list[str]]:
|
||
|
|
"""Build raw-column renames and the calendar-date event columns to inspect."""
|
||
|
|
field_df = pd.read_csv(field_map_file, low_memory=False)
|
||
|
|
required = {"field_instance", "var_name"}
|
||
|
|
missing = required - set(field_df.columns)
|
||
|
|
if missing:
|
||
|
|
raise ValueError(f"{field_map_file} is missing columns: {sorted(missing)}")
|
||
|
|
|
||
|
|
raw_to_name = dict(
|
||
|
|
zip(field_df["field_instance"].astype(str), field_df["var_name"].astype(str))
|
||
|
|
)
|
||
|
|
icd_date_columns: list[str] = []
|
||
|
|
with Path(icd_map_file).open(encoding="utf-8") as handle:
|
||
|
|
for line in handle:
|
||
|
|
parts = line.strip().split()
|
||
|
|
if len(parts) >= 6:
|
||
|
|
raw_to_name[parts[0]] = parts[5]
|
||
|
|
icd_date_columns.append(parts[5])
|
||
|
|
|
||
|
|
for slot in range(17):
|
||
|
|
raw_to_name[f"40005-{slot}.0"] = f"cancer_date_{slot}"
|
||
|
|
raw_to_name[f"40006-{slot}.0"] = f"cancer_type_{slot}"
|
||
|
|
return raw_to_name, icd_date_columns
|
||
|
|
|
||
|
|
|
||
|
|
def _records_from_icd_columns(
|
||
|
|
chunk: pd.DataFrame,
|
||
|
|
event_columns: list[str],
|
||
|
|
token_map: dict[str, int],
|
||
|
|
) -> list[pd.DataFrame]:
|
||
|
|
frames: list[pd.DataFrame] = []
|
||
|
|
for column in event_columns:
|
||
|
|
token = token_map.get(column)
|
||
|
|
if token is None or column not in chunk.columns:
|
||
|
|
continue
|
||
|
|
event_date = pd.to_datetime(chunk[column], format="%Y-%m-%d", errors="coerce")
|
||
|
|
valid = event_date.notna()
|
||
|
|
if valid.any():
|
||
|
|
frames.append(
|
||
|
|
pd.DataFrame(
|
||
|
|
{
|
||
|
|
"eid": chunk.index[valid].astype("int64"),
|
||
|
|
"event_date": event_date.loc[valid].to_numpy(),
|
||
|
|
"token": np.full(valid.sum(), token, dtype=np.int32),
|
||
|
|
}
|
||
|
|
)
|
||
|
|
)
|
||
|
|
return frames
|
||
|
|
|
||
|
|
|
||
|
|
def _records_from_cancer_columns(chunk: pd.DataFrame, token_map: dict[str, int]) -> list[pd.DataFrame]:
|
||
|
|
frames: list[pd.DataFrame] = []
|
||
|
|
for slot in range(17):
|
||
|
|
date_column = f"cancer_date_{slot}"
|
||
|
|
type_column = f"cancer_type_{slot}"
|
||
|
|
if date_column not in chunk.columns or type_column not in chunk.columns:
|
||
|
|
continue
|
||
|
|
event_date = pd.to_datetime(chunk[date_column], format="%Y-%m-%d", errors="coerce")
|
||
|
|
code = chunk[type_column].astype("string").str.slice(0, 3)
|
||
|
|
token = code.map(token_map)
|
||
|
|
valid = event_date.notna() & token.notna()
|
||
|
|
if valid.any():
|
||
|
|
frames.append(
|
||
|
|
pd.DataFrame(
|
||
|
|
{
|
||
|
|
"eid": chunk.index[valid].astype("int64"),
|
||
|
|
"event_date": event_date.loc[valid].to_numpy(),
|
||
|
|
"token": token.loc[valid].astype("int32").to_numpy(),
|
||
|
|
}
|
||
|
|
)
|
||
|
|
)
|
||
|
|
return frames
|
||
|
|
|
||
|
|
|
||
|
|
def prepare_event_dates(
|
||
|
|
*,
|
||
|
|
ukb_data_file: str | Path,
|
||
|
|
field_map_file: str | Path,
|
||
|
|
icd_map_file: str | Path,
|
||
|
|
labels_file: str | Path,
|
||
|
|
output_file: str | Path,
|
||
|
|
chunksize: int = 10_000,
|
||
|
|
) -> int:
|
||
|
|
"""Stream the raw UKB export, then write a sorted structured event array."""
|
||
|
|
token_map = load_label_tokens(labels_file)
|
||
|
|
raw_to_name, icd_date_columns = build_raw_column_map(field_map_file, icd_map_file)
|
||
|
|
event_columns = [*icd_date_columns, "Death"]
|
||
|
|
frames: list[pd.DataFrame] = []
|
||
|
|
|
||
|
|
reader = pd.read_csv(
|
||
|
|
ukb_data_file,
|
||
|
|
chunksize=chunksize,
|
||
|
|
index_col=0, # UKB participant ID / eid
|
||
|
|
low_memory=False,
|
||
|
|
)
|
||
|
|
chunk_iterator = tqdm(reader, desc="Extracting calendar-dated disease events") if tqdm else reader
|
||
|
|
for raw_chunk in chunk_iterator:
|
||
|
|
chunk = raw_chunk.rename(columns=raw_to_name)
|
||
|
|
frames.extend(_records_from_icd_columns(chunk, event_columns, token_map))
|
||
|
|
frames.extend(_records_from_cancer_columns(chunk, token_map))
|
||
|
|
|
||
|
|
if not frames:
|
||
|
|
result = np.empty(0, dtype=EVENT_DTYPE)
|
||
|
|
else:
|
||
|
|
events = pd.concat(frames, ignore_index=True)
|
||
|
|
events = events.dropna(subset=["eid", "event_date", "token"])
|
||
|
|
events = events.sort_values(["eid", "token", "event_date"], kind="stable")
|
||
|
|
# Match prepare_data.py: first occurrence of each disease/death token.
|
||
|
|
events = events.drop_duplicates(["eid", "token"], keep="first")
|
||
|
|
events = events.sort_values(["eid", "event_date", "token"], kind="stable")
|
||
|
|
|
||
|
|
result = np.empty(len(events), dtype=EVENT_DTYPE)
|
||
|
|
result["eid"] = events["eid"].to_numpy(dtype=np.int64)
|
||
|
|
result["event_date"] = events["event_date"].to_numpy(dtype="datetime64[D]")
|
||
|
|
result["token"] = events["token"].to_numpy(dtype=np.int32)
|
||
|
|
|
||
|
|
np.save(output_file, result)
|
||
|
|
return len(result)
|
||
|
|
|
||
|
|
|
||
|
|
def main() -> None:
|
||
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
||
|
|
parser.add_argument("--ukb-data", default="ukb_data.csv")
|
||
|
|
parser.add_argument("--field-map", default="field_ids_enriched.csv")
|
||
|
|
parser.add_argument("--icd-map", default="icd10_codes_mod.tsv")
|
||
|
|
parser.add_argument("--labels", default="labels.csv")
|
||
|
|
parser.add_argument("--output", default="ukb_disease_event_dates.npy")
|
||
|
|
parser.add_argument("--chunksize", type=int, default=10_000)
|
||
|
|
args = parser.parse_args()
|
||
|
|
if args.chunksize <= 0:
|
||
|
|
raise ValueError("chunksize must be positive")
|
||
|
|
count = prepare_event_dates(
|
||
|
|
ukb_data_file=args.ukb_data,
|
||
|
|
field_map_file=args.field_map,
|
||
|
|
icd_map_file=args.icd_map,
|
||
|
|
labels_file=args.labels,
|
||
|
|
output_file=args.output,
|
||
|
|
chunksize=args.chunksize,
|
||
|
|
)
|
||
|
|
print(f"Wrote {count:,} first disease/death events to {args.output}")
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|