Add disease history ablation modes

eval_data.py

@@ -9,7 +9,12 @@ import numpy as np
 import torch
 from torch.nn.utils.rnn import pad_sequence
 
-from dataset import AllFutureHealthDataset, HealthDataset
+from dataset import (
+    DISEASE_HISTORY_MODE_TIMED,
+    AllFutureHealthDataset,
+    HealthDataset,
+    normalize_disease_history_mode,
+)
 from model_architectures import resolve_model_architecture
 from models import DeepHealth
 from targets import PAD_IDX
@@ -61,6 +66,40 @@ def validate_training_mode_config(cfg: Dict[str, Any]) -> None:
             "model_target_mode must be next_token or all_future, got "
             f"{model_target_mode!r}"
         )
+
+    disease_history_mode = normalize_disease_history_mode(
+        cfg.get("disease_history_mode", DISEASE_HISTORY_MODE_TIMED)
+    )
+    if disease_history_mode != DISEASE_HISTORY_MODE_TIMED:
+        expected = {
+            "model_target_mode": "all_future",
+            "time_mode": "relative",
+            "dist_mode": "weibull",
+            "model_architecture": "traj_mixer_v5",
+        }
+        actual = {
+            "model_target_mode": model_target_mode,
+            "time_mode": str(cfg.get("time_mode", "")).lower(),
+            "dist_mode": str(cfg.get("dist_mode", "")).lower(),
+            "model_architecture": str(
+                cfg.get("model_architecture", "")
+            ).lower(),
+        }
+        mismatches = [
+            f"{name}={actual[name]!r} (expected {value!r})"
+            for name, value in expected.items()
+            if actual[name] != value
+        ]
+        extra_info_types = cfg.get("extra_info_types", None)
+        if extra_info_types != []:
+            mismatches.append("extra_info_types must be []")
+        if mismatches:
+            raise ValueError(
+                f"disease_history_mode={disease_history_mode!r} is only valid "
+                "for the no-extra TrajMixer + all_future + relative + Weibull "
+                "ablation; " + "; ".join(mismatches)
+            )
+
     if model_target_mode != "next_token":
         return
 
@@ -208,9 +247,10 @@ class AllFutureSequenceEvalDataset:
     """
     Eval-only sequence view for all-future checkpoints.
 
-    All-future training uses the observed history, including CHECKUP state
-    tokens, without reusing the next-step view that contains imputed
-    <NO_EVENT> gap tokens.
+    All-future training uses the observed history without reusing the
+    next-step view that contains imputed <NO_EVENT> gap tokens. CHECKUP is
+    retained only when the experiment selects at least one extra-info type;
+    an explicitly empty selection is a disease-only history.
     """
 
     def __init__(
@@ -220,6 +260,7 @@ class AllFutureSequenceEvalDataset:
         min_history_events: int = 1,
         min_future_events: int = 1,
         extra_info_types: Iterable[int] | None = None,
+        disease_history_mode: str = DISEASE_HISTORY_MODE_TIMED,
     ) -> None:
         base = AllFutureHealthDataset(
             data_prefix=data_prefix,
@@ -228,6 +269,7 @@ class AllFutureSequenceEvalDataset:
             min_history_events=min_history_events,
             min_future_events=min_future_events,
             extra_info_types=extra_info_types,
+            disease_history_mode=disease_history_mode,
         )
 
         self.base = base
@@ -239,6 +281,7 @@ class AllFutureSequenceEvalDataset:
         self.n_categories = base.n_categories
         self.cont_type_ids = base.cont_type_ids
         self.extra_info_types = base.extra_info_types
+        self.disease_history_mode = base.disease_history_mode
 
         self.samples: List[Dict[str, Any]] = []
         for patient in base.patients:
@@ -288,6 +331,7 @@ def load_sequence_eval_dataset(
     min_history_events: int,
     min_future_events: int,
     extra_info_types: Iterable[int] | None,
+    disease_history_mode: str = DISEASE_HISTORY_MODE_TIMED,
 ):
     mode = str(model_target_mode).lower()
     if mode == "next_token":
@@ -304,6 +348,7 @@ def load_sequence_eval_dataset(
             min_history_events=min_history_events,
             min_future_events=min_future_events,
             extra_info_types=extra_info_types,
+            disease_history_mode=disease_history_mode,
         )
     raise ValueError(f"Unknown model_target_mode: {model_target_mode!r}")