From 0138cbd95b1ba26d8cd83aaab590d2a428701546 Mon Sep 17 00:00:00 2001 From: Jiarui Li Date: Mon, 22 Jun 2026 17:36:04 +0800 Subject: [PATCH] Add calendar-dated event index utilities --- event_date_utils.py | 104 +++++++++++++++++++++ prepare_event_dates.py | 203 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 307 insertions(+) create mode 100644 event_date_utils.py create mode 100644 prepare_event_dates.py diff --git a/event_date_utils.py b/event_date_utils.py new file mode 100644 index 0000000..24e602e --- /dev/null +++ b/event_date_utils.py @@ -0,0 +1,104 @@ +"""Read and query calendar-dated disease-event arrays from prepare_event_dates.py.""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Iterable + +import numpy as np +import pandas as pd + + +REQUIRED_FIELDS = {"eid", "event_date", "token"} + + +def load_event_dates(path: str | Path) -> np.ndarray: + """Load and validate the structured ``.npy`` event array.""" + events = np.load(path) + if events.dtype.names is None or not REQUIRED_FIELDS.issubset(events.dtype.names): + raise ValueError( + "Expected a structured .npy with eid, event_date, token fields. " + "Create it with prepare_event_dates.py." + ) + return events + + +def load_token_labels(labels_file: str | Path) -> dict[int, str]: + """Load token -> human-readable code using the project label convention.""" + labels = {1: "CHECKUP"} + 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: + labels[index + 2] = code + return labels + + +@dataclass +class EventDateIndex: + """Small in-memory query wrapper for exposure-linkage and cohort scripts.""" + + events: np.ndarray + token_labels: dict[int, str] | None = None + + @classmethod + def from_files( + cls, + event_file: str | Path, + labels_file: str | Path | None = None, + ) -> "EventDateIndex": + labels = load_token_labels(labels_file) if labels_file is not None else None + return cls(load_event_dates(event_file), labels) + + def to_frame(self, events: np.ndarray | None = None) -> pd.DataFrame: + """Convert records to a convenient, calendar-dated DataFrame.""" + data = self.events if events is None else events + frame = pd.DataFrame( + { + "eid": data["eid"].astype("int64"), + "event_date": pd.to_datetime(data["event_date"]), + "token": data["token"].astype("int32"), + } + ) + if self.token_labels is not None: + frame["label_code"] = frame["token"].map(self.token_labels).fillna("UNKNOWN") + return frame.sort_values(["eid", "event_date", "token"], kind="stable").reset_index(drop=True) + + def for_eid(self, eid: int) -> pd.DataFrame: + """Return every stored disease/death event for one participant.""" + return self.to_frame(self.events[self.events["eid"] == int(eid)]) + + def between( + self, + start: str | pd.Timestamp, + end: str | pd.Timestamp, + *, + eids: Iterable[int] | None = None, + tokens: Iterable[int] | None = None, + ) -> pd.DataFrame: + """Query events in an inclusive calendar-date interval.""" + start_day = np.datetime64(pd.Timestamp(start).date(), "D") + end_day = np.datetime64(pd.Timestamp(end).date(), "D") + mask = (self.events["event_date"] >= start_day) & (self.events["event_date"] <= end_day) + if eids is not None: + mask &= np.isin(self.events["eid"], list(eids)) + if tokens is not None: + mask &= np.isin(self.events["token"], list(tokens)) + return self.to_frame(self.events[mask]) + + def anchors_before(self, eid: int, date: str | pd.Timestamp) -> pd.DataFrame: + """Return a participant's event history strictly before an exposure anchor.""" + day = np.datetime64(pd.Timestamp(date).date(), "D") + mask = (self.events["eid"] == int(eid)) & (self.events["event_date"] < day) + return self.to_frame(self.events[mask]) + + def first_event(self, token: int) -> pd.DataFrame: + """Return each participant's first date for a requested token.""" + selected = self.events[self.events["token"] == int(token)] + # Arrays produced by prepare_event_dates.py are already deduplicated; + # sorting makes this safe for externally produced compatible arrays too. + order = np.lexsort((selected["event_date"], selected["eid"])) + selected = selected[order] + _, first = np.unique(selected["eid"], return_index=True) + return self.to_frame(selected[first]) diff --git a/prepare_event_dates.py b/prepare_event_dates.py new file mode 100644 index 0000000..4d84ad3 --- /dev/null +++ b/prepare_event_dates.py @@ -0,0 +1,203 @@ +"""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", " 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()