Add exposure cache and keep absolute time only

eval_data.py

@@ -1,12 +1,18 @@
 from __future__ import annotations
 
+from pathlib import Path
 from typing import Any, Dict, List
 
 import numpy as np
 import torch
 from torch.nn.utils.rnn import pad_sequence
 
-from dataset import AllFutureHealthDataset, HealthDataset
+from dataset import (
+    DAILY_EXPOSURE_SHAPE,
+    MONTHLY_EXPOSURE_SHAPE,
+    AllFutureHealthDataset,
+    HealthDataset,
+)
 from targets import PAD_IDX
 
 
@@ -25,6 +31,8 @@ class AllFutureSequenceEvalDataset:
         labels_file: str,
         min_history_events: int = 1,
         min_future_events: int = 1,
+        exposure_cache_dir: str | Path | None = None,
+        mask_onset_exposure: bool = False,
     ) -> None:
         base = AllFutureHealthDataset(
             data_prefix=data_prefix,
@@ -32,6 +40,8 @@ class AllFutureSequenceEvalDataset:
             split="train",
             min_history_events=min_history_events,
             min_future_events=min_future_events,
+            exposure_cache_dir=exposure_cache_dir,
+            mask_onset_exposure=mask_onset_exposure,
         )
 
         self.base = base
@@ -43,11 +53,12 @@ class AllFutureSequenceEvalDataset:
         for patient in base.patients:
             labels = np.asarray(patient["labels"], dtype=np.int64)
             times = np.asarray(patient["times"], dtype=np.float32)
+            times_days = np.asarray(patient["times_days"], dtype=np.float32)
             if labels.size < 2:
                 continue
             input_len = int(labels.size - 1)
             self.samples.append(
-                {
+                sample := {
                     "eid": int(patient["eid"]),
                     "event_seq": labels[:-1],
                     "time_seq": times[:-1],
@@ -57,13 +68,21 @@ class AllFutureSequenceEvalDataset:
                     "sex": int(patient["sex"]),
                 }
             )
+            if base.exposure_cache is not None:
+                exposure_index = base._exposure_indices_for_inputs(
+                    eid=int(patient["eid"]),
+                    input_events=labels[:-1],
+                    input_times_days=times_days[:-1],
+                )
+                if exposure_index is not None:
+                    sample["exposure_index"] = exposure_index
 
     def __len__(self) -> int:
         return len(self.samples)
 
     def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
         s = self.samples[idx]
-        return {
+        out = {
             "event_seq": torch.from_numpy(s["event_seq"]).long(),
             "time_seq": torch.from_numpy(s["time_seq"]).float(),
             "target_event_seq": torch.from_numpy(s["target_event_seq"]).long(),
@@ -71,6 +90,11 @@ class AllFutureSequenceEvalDataset:
             "readout_mask": torch.from_numpy(s["readout_mask"]).bool(),
             "sex": torch.tensor(s["sex"], dtype=torch.long),
         }
+        if "exposure_index" in s:
+            daily, monthly = self.base._load_exposure_windows(s["exposure_index"])
+            out["exposure_daily"] = daily
+            out["exposure_monthly"] = monthly
+        return out
 
 
 def load_sequence_eval_dataset(
@@ -82,6 +106,8 @@ def load_sequence_eval_dataset(
     include_no_event_in_uts_target: bool,
     min_history_events: int,
     min_future_events: int,
+    exposure_cache_dir: str | Path | None = None,
+    mask_onset_exposure: bool = False,
 ):
     mode = str(model_target_mode).lower()
     if mode == "next_token":
@@ -90,6 +116,8 @@ def load_sequence_eval_dataset(
             labels_file=labels_file,
             no_event_interval_years=no_event_interval_years,
             include_no_event_in_uts_target=include_no_event_in_uts_target,
+            exposure_cache_dir=exposure_cache_dir,
+            mask_onset_exposure=mask_onset_exposure,
         )
     if mode == "all_future":
         return AllFutureSequenceEvalDataset(
@@ -97,6 +125,8 @@ def load_sequence_eval_dataset(
             labels_file=labels_file,
             min_history_events=min_history_events,
             min_future_events=min_future_events,
+            exposure_cache_dir=exposure_cache_dir,
+            mask_onset_exposure=mask_onset_exposure,
         )
     raise ValueError(f"Unknown model_target_mode: {model_target_mode!r}")
 
@@ -117,7 +147,7 @@ def sequence_eval_collate_fn(batch: List[Dict[str, torch.Tensor]]) -> Dict[str,
     readout_mask = pad_sequence(
         [s["readout_mask"] for s in batch], batch_first=True, padding_value=False
     )
-    return {
+    out = {
         "event_seq": event_seq,
         "time_seq": time_seq,
         "padding_mask": event_seq > PAD_IDX,
@@ -126,3 +156,34 @@ def sequence_eval_collate_fn(batch: List[Dict[str, torch.Tensor]]) -> Dict[str,
         "readout_mask": readout_mask,
         "sex": torch.stack([s["sex"] for s in batch]),
     }
+    if any("exposure_daily" in s for s in batch):
+        out["exposure_daily"] = _pad_eval_exposure(
+            batch,
+            "exposure_daily",
+            DAILY_EXPOSURE_SHAPE,
+        )
+        out["exposure_monthly"] = _pad_eval_exposure(
+            batch,
+            "exposure_monthly",
+            MONTHLY_EXPOSURE_SHAPE,
+        )
+    return out
+
+
+def _pad_eval_exposure(
+    batch: List[Dict[str, torch.Tensor]],
+    key: str,
+    shape: tuple[int, int],
+) -> torch.Tensor:
+    max_len = max(int(s["event_seq"].numel()) for s in batch)
+    out = torch.full(
+        (len(batch), max_len, shape[0], shape[1]),
+        float("nan"),
+        dtype=torch.float32,
+    )
+    for idx, sample in enumerate(batch):
+        value = sample.get(key)
+        if value is None:
+            continue
+        out[idx, : int(value.size(0))] = value
+    return out