Add disease history ablation modes

dataset.py

@@ -20,6 +20,167 @@ from targets import (
 
 ONE_DAY_YEARS = 1.0 / DAYS_PER_YEAR
 
+DISEASE_HISTORY_MODE_TIMED = "timed"
+DISEASE_HISTORY_MODE_ORDERED = "ordered"
+DISEASE_HISTORY_MODE_SET = "set"
+DISEASE_HISTORY_MODES = (
+    DISEASE_HISTORY_MODE_TIMED,
+    DISEASE_HISTORY_MODE_ORDERED,
+    DISEASE_HISTORY_MODE_SET,
+)
+
+
+def normalize_disease_history_mode(mode: str | None) -> str:
+    value = DISEASE_HISTORY_MODE_TIMED if mode is None else str(mode).lower()
+    if value not in DISEASE_HISTORY_MODES:
+        raise ValueError(
+            "disease_history_mode must be one of "
+            f"{list(DISEASE_HISTORY_MODES)}, got {mode!r}"
+        )
+    return value
+
+
+def transform_disease_history(
+    event_seq: np.ndarray,
+    actual_time_seq: np.ndarray,
+    actual_t_query: float,
+    disease_history_mode: str,
+) -> Tuple[np.ndarray, np.ndarray, np.float32]:
+    """
+    Convert an already-truncated disease history into its model representation.
+
+    ``timed`` keeps the real event/query times. ``ordered`` preserves the
+    chronological event order but replaces calendar time with ordinal event-time
+    groups. Diseases first recorded on the same day share one ordinal position.
+    ``set`` removes both time and order by sorting the unique disease codes and
+    assigning every disease and the query the same model time.
+    """
+    mode = normalize_disease_history_mode(disease_history_mode)
+    events = np.asarray(event_seq, dtype=np.int64)
+    times = np.asarray(actual_time_seq, dtype=np.float32)
+    if events.ndim != 1 or times.ndim != 1 or events.shape != times.shape:
+        raise ValueError(
+            "event_seq and actual_time_seq must be aligned 1D arrays, got "
+            f"{events.shape} and {times.shape}"
+        )
+
+    if mode == DISEASE_HISTORY_MODE_TIMED:
+        return events, times, np.float32(actual_t_query)
+
+    special = events <= NO_EVENT_IDX
+    if np.any(special):
+        raise ValueError(
+            f"{mode} disease history must contain only disease events; "
+            f"found special token ids {np.unique(events[special]).tolist()}"
+        )
+
+    if mode == DISEASE_HISTORY_MODE_ORDERED:
+        _, ordinal_groups = np.unique(times, return_inverse=True)
+        model_times = ordinal_groups.astype(np.float32, copy=False)
+        n_groups = int(model_times.max()) + 1 if model_times.size else 0
+        return events, model_times, np.float32(n_groups)
+
+    set_events = np.unique(events)
+    model_times = np.zeros(set_events.size, dtype=np.float32)
+    return set_events, model_times, np.float32(0.0)
+
+
+def transform_disease_history_batch_at_position(
+    event_seq: torch.Tensor,
+    actual_time_seq: torch.Tensor,
+    padding_mask: torch.Tensor,
+    query_position: int,
+    disease_history_mode: str,
+    vocab_size: int,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    Build a model-visible prefix for token-position all-future evaluation.
+
+    Actual event times remain outside this return value for AUC bookkeeping.
+    For ordered/set modes, events after ``query_position`` are explicitly
+    masked so collapsing time cannot expose future diseases.
+    """
+    mode = normalize_disease_history_mode(disease_history_mode)
+    if event_seq.ndim != 2 or actual_time_seq.shape != event_seq.shape:
+        raise ValueError(
+            "event_seq and actual_time_seq must be aligned 2D tensors, got "
+            f"{tuple(event_seq.shape)} and {tuple(actual_time_seq.shape)}"
+        )
+    if padding_mask.shape != event_seq.shape:
+        raise ValueError(
+            "padding_mask must match event_seq, got "
+            f"{tuple(padding_mask.shape)} and {tuple(event_seq.shape)}"
+        )
+    if query_position < 0 or query_position >= event_seq.size(1):
+        raise ValueError(
+            f"query_position={query_position} is outside sequence length "
+            f"{event_seq.size(1)}"
+        )
+
+    padding_mask = padding_mask.to(device=event_seq.device, dtype=torch.bool)
+    if not torch.all(padding_mask[:, query_position]):
+        raise ValueError("query_position must be valid for every batch row")
+
+    if mode == DISEASE_HISTORY_MODE_TIMED:
+        return (
+            event_seq,
+            actual_time_seq,
+            padding_mask,
+            actual_time_seq[:, query_position],
+        )
+
+    positions = torch.arange(
+        event_seq.size(1),
+        device=event_seq.device,
+    )[None, :]
+    history_mask = padding_mask & (positions <= query_position)
+    visible_events = event_seq.masked_select(history_mask)
+    if torch.any(visible_events <= NO_EVENT_IDX):
+        special_ids = torch.unique(
+            visible_events[visible_events <= NO_EVENT_IDX]
+        ).detach().cpu().tolist()
+        raise ValueError(
+            f"{mode} disease history must contain only disease events; "
+            f"found special token ids {special_ids}"
+        )
+
+    if mode == DISEASE_HISTORY_MODE_ORDERED:
+        model_times = torch.zeros_like(actual_time_seq)
+        model_t_query = torch.zeros(
+            event_seq.size(0),
+            device=actual_time_seq.device,
+            dtype=actual_time_seq.dtype,
+        )
+        for row_idx in range(event_seq.size(0)):
+            row_mask = history_mask[row_idx]
+            _, ordinal_groups = torch.unique(
+                actual_time_seq[row_idx, row_mask],
+                sorted=True,
+                return_inverse=True,
+            )
+            model_times[row_idx, row_mask] = ordinal_groups.to(
+                dtype=actual_time_seq.dtype
+            )
+            model_t_query[row_idx] = float(
+                int(ordinal_groups.max().item()) + 1
+                if ordinal_groups.numel()
+                else 0
+            )
+        return event_seq, model_times, history_mask, model_t_query
+
+    sentinel = torch.full_like(event_seq, int(vocab_size))
+    sortable = torch.where(history_mask, event_seq, sentinel)
+    set_events = torch.sort(sortable, dim=1).values
+    set_mask = set_events != int(vocab_size)
+    set_events = set_events.masked_fill(~set_mask, PAD_IDX)
+    model_times = torch.zeros_like(actual_time_seq)
+    model_t_query = torch.zeros(
+        event_seq.size(0),
+        device=actual_time_seq.device,
+        dtype=actual_time_seq.dtype,
+    )
+    return set_events, model_times, set_mask, model_t_query
+
 
 def load_label_vocab(
     labels_file: str,
@@ -255,6 +416,15 @@ class _ExpoBaseDataset(Dataset):
             times_days_raw = rows[:, 1].astype(np.float32)
             labels_raw = rows[:, 2].astype(np.int64)
 
+            # CHECKUP is the assessment landmark for selected extra-info tokens.
+            # An explicitly empty selection represents a disease-only history,
+            # so retaining CHECKUP in that case would introduce an empty
+            # landmark token that is not part of the disease sequence.
+            if not self.extra_info_types:
+                keep = labels_raw != CHECKUP_IDX
+                times_days_raw = times_days_raw[keep]
+                labels_raw = labels_raw[keep]
+
             if len(labels_raw) == 0:
                 yield eid, times_days_raw, labels_raw
                 continue
@@ -368,6 +538,7 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
         min_future_events: int = 1,
         validation_query_seed: int = 42,
         extra_info_types: Iterable[int] | None = None,
+        disease_history_mode: str = DISEASE_HISTORY_MODE_TIMED,
     ) -> None:
         if split not in {"train", "valid", "test"}:
             raise ValueError(f"split must be train/valid/test, got {split!r}")
@@ -379,6 +550,18 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
             extra_info_types=extra_info_types,
         )
 
+        self.disease_history_mode = normalize_disease_history_mode(
+            disease_history_mode
+        )
+        if (
+            self.disease_history_mode != DISEASE_HISTORY_MODE_TIMED
+            and self.extra_info_types
+        ):
+            raise ValueError(
+                f"disease_history_mode={self.disease_history_mode!r} is only "
+                "supported with an explicitly empty extra-info selection"
+            )
+
         self.split = split
         self.min_history_events = int(min_history_events)
         self.min_future_events = int(min_future_events)
@@ -504,11 +687,17 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
         labels = patient["labels"]
         hist = times <= t_query
         fut = times > t_query
+        event_seq, model_time_seq, model_t_query = transform_disease_history(
+            event_seq=labels[hist],
+            actual_time_seq=times[hist],
+            actual_t_query=t_query,
+            disease_history_mode=self.disease_history_mode,
+        )
 
         return {
-            "event_seq": torch.from_numpy(labels[hist]).long(),
-            "time_seq": torch.from_numpy(times[hist]).float(),
-            "t_query": torch.tensor(t_query, dtype=torch.float32),
+            "event_seq": torch.from_numpy(event_seq).long(),
+            "time_seq": torch.from_numpy(model_time_seq).float(),
+            "t_query": torch.tensor(model_t_query, dtype=torch.float32),
             "future_targets": torch.from_numpy(labels[fut]).long(),
             "future_dt": torch.from_numpy(times[fut] - np.float32(t_query)).float(),
             "exposure": torch.tensor(np.float32(patient["t_obs"] - t_query), dtype=torch.float32),