Add calendar-dated event index utilities

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])

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