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