Add calendar-dated event index utilities

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", "<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()