Improve all-future first-onset training

train_util.py

@@ -18,6 +18,7 @@ from torch.utils.data import Subset
 
 from dataset import AllFutureHealthDataset, HealthDataset
 from models import DeepHealth
+from targets import NO_EVENT_IDX, PAD_IDX, RESERVED_IDX
 
 
 @dataclass(frozen=True)
@@ -42,6 +43,134 @@ class ContinuousRobustScalerStats:
         }
 
 
+@dataclass(frozen=True)
+class AllFutureBaselineStats:
+    """Training-only marginal first-onset rates for output initialization."""
+
+    event_count: np.ndarray
+    at_risk_exposure: np.ndarray
+    rate: np.ndarray
+    bias: np.ndarray
+    query_count: int
+    rate_floor: float
+    rate_ceiling: float
+
+    def as_metadata(self) -> Dict[str, Any]:
+        positive = self.rate[self.at_risk_exposure > 0]
+        if positive.size:
+            rate_summary = {
+                "min": float(positive.min()),
+                "median": float(np.median(positive)),
+                "max": float(positive.max()),
+            }
+        else:
+            rate_summary = {"min": 0.0, "median": 0.0, "max": 0.0}
+        return {
+            "method": "training_marginal_first_onset_rate",
+            "fitted_on": "one_seeded_query_per_training_patient",
+            "query_distribution": "patient_interval_time_uniform",
+            "query_count": int(self.query_count),
+            "observed_first_onsets": int(self.event_count.sum()),
+            "rate_floor": float(self.rate_floor),
+            "rate_ceiling": float(self.rate_ceiling),
+            "rate_per_year": rate_summary,
+            "checkpoint_parameter": "risk_head.bias",
+        }
+
+
+def fit_all_future_baseline(
+    dataset: AllFutureHealthDataset,
+    subset: Subset,
+    *,
+    seed: int,
+    ignored_idx: Iterable[int] = (PAD_IDX, RESERVED_IDX, NO_EVENT_IDX),
+    rate_floor: float = 1e-6,
+    rate_ceiling: float = 10.0,
+) -> AllFutureBaselineStats:
+    """Fit marginal per-outcome rates under the training query distribution."""
+    if subset.dataset is not dataset:
+        raise ValueError("subset must reference the all-future training dataset")
+    if rate_floor <= 0 or rate_ceiling <= rate_floor:
+        raise ValueError("Require 0 < rate_floor < rate_ceiling")
+
+    subset_indices = np.asarray(subset.indices, dtype=np.int64)
+    if subset_indices.ndim != 1 or subset_indices.size == 0:
+        raise ValueError("training subset must contain at least one patient")
+
+    vocab_size = int(dataset.vocab_size)
+    ignored = {
+        int(idx)
+        for idx in ignored_idx
+        if 0 <= int(idx) < vocab_size
+    }
+    event_count = np.zeros(vocab_size, dtype=np.int64)
+    exposure_adjustment = np.zeros(vocab_size, dtype=np.float64)
+    common_exposure = 0.0
+    rng = np.random.RandomState(int(seed))
+
+    for patient_index in subset_indices.tolist():
+        patient = dataset.patients[int(patient_index)]
+        t_query = dataset.sample_query(patient, rng)
+        times = np.asarray(patient["times"], dtype=np.float64)
+        labels = np.asarray(patient["labels"], dtype=np.int64)
+        censor_time = max(float(patient["t_obs"]) - t_query, 0.0)
+        common_exposure += censor_time
+
+        prevalent = {
+            int(label)
+            for label in labels[times <= t_query].tolist()
+            if 0 <= int(label) < vocab_size and int(label) not in ignored
+        }
+        for label in prevalent:
+            exposure_adjustment[label] -= censor_time
+
+        # Keep the earliest future occurrence defensively; prepared disease
+        # events are already deduplicated to their first occurrence.
+        first_future_dt: Dict[int, float] = {}
+        for label, event_time in zip(labels[times > t_query], times[times > t_query]):
+            label = int(label)
+            if label in ignored or label in prevalent or not (0 <= label < vocab_size):
+                continue
+            event_dt = max(float(event_time) - t_query, 0.0)
+            previous = first_future_dt.get(label)
+            if previous is None or event_dt < previous:
+                first_future_dt[label] = event_dt
+        for label, event_dt in first_future_dt.items():
+            event_count[label] += 1
+            exposure_adjustment[label] += event_dt - censor_time
+
+    at_risk_exposure = common_exposure + exposure_adjustment
+    at_risk_exposure = np.maximum(at_risk_exposure, 0.0)
+    rate = np.full(vocab_size, float(rate_floor), dtype=np.float64)
+    has_exposure = at_risk_exposure > 0
+    rate[has_exposure] = (
+        event_count[has_exposure].astype(np.float64)
+        / at_risk_exposure[has_exposure]
+    )
+    rate = np.clip(rate, float(rate_floor), float(rate_ceiling))
+
+    for idx in ignored:
+        at_risk_exposure[idx] = 0.0
+        rate[idx] = 0.0
+    bias = np.zeros(vocab_size, dtype=np.float64)
+    valid = np.ones(vocab_size, dtype=bool)
+    if ignored:
+        valid[np.asarray(sorted(ignored), dtype=np.int64)] = False
+    bias[valid] = np.log(np.expm1(rate[valid]))
+
+    if not np.isfinite(bias).all():
+        raise RuntimeError("All-future baseline initialization is non-finite")
+    return AllFutureBaselineStats(
+        event_count=event_count,
+        at_risk_exposure=at_risk_exposure.astype(np.float32),
+        rate=rate.astype(np.float32),
+        bias=bias.astype(np.float32),
+        query_count=int(subset_indices.size),
+        rate_floor=float(rate_floor),
+        rate_ceiling=float(rate_ceiling),
+    )
+
+
 def fit_continuous_robust_scaler(
     dataset: HealthDataset | AllFutureHealthDataset,
     subset: Subset,