Improve all-future first-onset training

dataset.py

@@ -17,9 +17,6 @@ from targets import (
     build_next_token_targets,
 )
 
-
-ONE_DAY_YEARS = 1.0 / DAYS_PER_YEAR
-
 DISEASE_HISTORY_MODE_TIMED = "timed"
 DISEASE_HISTORY_MODE_ORDERED = "ordered"
 DISEASE_HISTORY_MODE_SET = "set"
@@ -518,10 +515,10 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
     Dataset with unified other-info tokens and DeepHealthV2-style all-future
     targets.
 
-    Train samples one query time per patient at each __getitem__ call.
-    Valid/test use random-but-fixed query points. For each patient with N real
-    disease events, N - 2 query points are sampled from the eligible observed
-    time range, with at least one future event after every query.
+    Every split uses the same patient-equal query distribution: choose one
+    eligible inter-event interval uniformly, then choose a time uniformly in
+    that interval. Train resamples on every ``__getitem__`` call; valid/test
+    keep one deterministic draw per patient.
     """
 
     CACHE_VERSION = 5
@@ -591,6 +588,11 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
                 **features,
             }
 
+            query_intervals = self._eligible_query_intervals(patient)
+            if not query_intervals:
+                continue
+            patient["query_intervals"] = query_intervals
+
             pidx = len(self.patients)
             self.patients.append(patient)
 
@@ -615,7 +617,7 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
             labels,
             np.array([PAD_IDX, RESERVED_IDX, NO_EVENT_IDX], dtype=np.int64),
         )
-        n_hist = int((times <= t_query).sum())
+        n_hist = int(((times <= t_query) & real_event_mask).sum())
         n_future = int(((times > t_query) & real_event_mask).sum())
         return (
             n_hist >= self.min_history_events
@@ -623,62 +625,59 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
             and patient["t_obs"] > t_query
         )
 
-    def _sample_fixed_validation_queries(
-        self,
-        patient: Dict,
-        rng: np.random.RandomState,
-    ) -> List[float]:
+    def _eligible_query_intervals(self, patient: Dict) -> List[Tuple[float, float]]:
         times = np.asarray(patient["times"], dtype=np.float32)
         labels = np.asarray(patient["labels"], dtype=np.int64)
         real_event_mask = ~np.isin(
             labels,
             np.array([PAD_IDX, RESERVED_IDX, NO_EVENT_IDX], dtype=np.int64),
         )
-        real_times = np.sort(times[real_event_mask].astype(np.float32, copy=False))
-        n_real_events = int(real_times.size)
-        n_queries = max(0, n_real_events - 2)
-        if n_queries == 0:
-            return []
+        unique_times = np.unique(times[real_event_mask])
+        intervals: List[Tuple[float, float]] = []
+        for j in range(1, len(unique_times)):
+            left = float(unique_times[j - 1])
+            right = float(unique_times[j])
+            probe = float(
+                np.nextafter(np.float32(right), np.float32(-np.inf))
+            )
+            if np.isfinite(left) and np.isfinite(probe) and probe >= left:
+                if self._is_valid_query(patient, probe):
+                    intervals.append((left, right))
+        return intervals
 
-        min_hist = int(self.min_history_events)
-        min_future = int(self.min_future_events)
-        if n_real_events < min_hist + min_future:
-            return []
+    def sample_query(
+        self,
+        patient: Dict,
+        rng,
+    ) -> float:
+        intervals = patient.get("query_intervals")
+        if intervals is None:
+            intervals = self._eligible_query_intervals(patient)
+        if not intervals:
+            raise RuntimeError("Patient has no eligible all-future query interval.")
 
-        left = float(real_times[min_hist - 1])
-        right_event_time = float(real_times[n_real_events - min_future])
-        right = np.nextafter(np.float32(right_event_time), np.float32(-np.inf))
-        if not np.isfinite(left) or not np.isfinite(right) or float(right) <= left:
-            return []
+        interval_idx = int(rng.randint(0, len(intervals)))
+        left, right_event_time = intervals[interval_idx]
+        right = float(
+            np.nextafter(np.float32(right_event_time), np.float32(-np.inf))
+        )
+        if right <= left:
+            t_query = float(left)
+        else:
+            t_query = float(rng.uniform(left, right))
+        if not self._is_valid_query(patient, t_query):
+            raise RuntimeError("Sampled an invalid all-future query time.")
+        return t_query
 
-        queries: List[float] = []
-        max_attempts = max(100, n_queries * 50)
-        for _ in range(max_attempts):
-            if len(queries) >= n_queries:
-                break
-            t_query = float(rng.uniform(left, float(right)))
-            if self._is_valid_query(patient, t_query):
-                queries.append(t_query)
-
-        return queries
+    def _sample_fixed_validation_queries(
+        self,
+        patient: Dict,
+        rng: np.random.RandomState,
+    ) -> List[float]:
+        return [self.sample_query(patient, rng)]
 
     def _sample_train_query(self, patient: Dict) -> float:
-        unique_times = np.unique(patient["times"])
-        if len(unique_times) < 2:
-            raise RuntimeError("Training patient has fewer than two unique times.")
-
-        j = np.random.randint(1, len(unique_times))
-        left = float(unique_times[j - 1])
-        right = float(unique_times[j])
-
-        if right - left <= ONE_DAY_YEARS:
-            t_query = right - ONE_DAY_YEARS
-        else:
-            t_query = np.random.uniform(left, right - ONE_DAY_YEARS)
-
-        if not self._is_valid_query(patient, t_query):
-            t_query = right - 1e-6
-        return float(t_query)
+        return self.sample_query(patient, np.random)
 
     def _build_item(self, patient: Dict, t_query: float) -> Dict:
         times = patient["times"]