105 lines
4.1 KiB
Python
105 lines
4.1 KiB
Python
|
"""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])
|