Add train-split robust scaling for continuous values

train_util.py

@@ -5,6 +5,7 @@ import logging
 import sys
 import time
 import csv
+from dataclasses import dataclass
 from datetime import datetime
 import math
 from pathlib import Path
@@ -19,6 +20,134 @@ from dataset import AllFutureHealthDataset, HealthDataset
 from models import DeepHealth
 
 
+@dataclass(frozen=True)
+class ContinuousRobustScalerStats:
+    """Train-split robust scaling statistics aligned to ``cont_type_ids``."""
+
+    cont_type_ids: tuple[int, ...]
+    center: np.ndarray
+    scale: np.ndarray
+    observation_count: np.ndarray
+    quantile_range: tuple[float, float] = (25.0, 75.0)
+
+    def as_metadata(self) -> Dict[str, Any]:
+        return {
+            "method": "robust",
+            "fitted_on": "train_subset",
+            "quantile_range": [float(x) for x in self.quantile_range],
+            "cont_type_ids": [int(x) for x in self.cont_type_ids],
+            "observation_count": [int(x) for x in self.observation_count.tolist()],
+            "center_buffer": "tokenizer.continuous_value_center",
+            "scale_buffer": "tokenizer.continuous_value_scale",
+        }
+
+
+def fit_continuous_robust_scaler(
+    dataset: AllFutureHealthDataset,
+    subset: Subset,
+    *,
+    quantile_range: tuple[float, float] = (25.0, 75.0),
+    scale_epsilon: float = 1e-6,
+) -> ContinuousRobustScalerStats:
+    """Fit median/IQR statistics using only patients in the training subset.
+
+    The prepared arrays remain unchanged. Scaling is performed later inside the
+    model tokenizer so the fitted center and scale can live in the checkpoint.
+    """
+
+    if subset.dataset is not dataset:
+        raise ValueError("subset must reference the dataset used to fit the scaler")
+    low, high = (float(quantile_range[0]), float(quantile_range[1]))
+    if not (0.0 <= low < high <= 100.0):
+        raise ValueError(
+            "quantile_range must satisfy 0 <= low < high <= 100, got "
+            f"{quantile_range!r}"
+        )
+    if scale_epsilon <= 0:
+        raise ValueError("scale_epsilon must be > 0")
+
+    cont_type_ids = tuple(int(x) for x in dataset.cont_type_ids)
+    n_cont_types = len(cont_type_ids)
+    if n_cont_types == 0:
+        empty = np.zeros(0, dtype=np.float32)
+        return ContinuousRobustScalerStats(
+            cont_type_ids=cont_type_ids,
+            center=empty.copy(),
+            scale=empty.copy(),
+            observation_count=np.zeros(0, dtype=np.int64),
+            quantile_range=(low, high),
+        )
+
+    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")
+
+    type_to_column = np.full(int(dataset.n_types), -1, dtype=np.int64)
+    for column, type_id in enumerate(cont_type_ids):
+        if type_id <= 0 or type_id >= len(type_to_column):
+            raise ValueError(
+                f"continuous type id {type_id} is outside [1, {len(type_to_column)})"
+            )
+        type_to_column[type_id] = column
+
+    values = np.full(
+        (int(subset_indices.size), n_cont_types),
+        np.nan,
+        dtype=np.float32,
+    )
+    for row, patient_index in enumerate(subset_indices.tolist()):
+        patient = dataset.patients[int(patient_index)]
+        other_type = np.asarray(patient["other_type"], dtype=np.int64)
+        other_value = np.asarray(patient["other_value"], dtype=np.float32)
+        other_kind = np.asarray(patient["other_value_kind"], dtype=np.int64)
+        continuous = other_kind == 1
+        if not np.any(continuous):
+            continue
+        selected_type = other_type[continuous]
+        selected_value = other_value[continuous]
+        valid_type = (selected_type > 0) & (selected_type < len(type_to_column))
+        columns = np.full(selected_type.shape, -1, dtype=np.int64)
+        columns[valid_type] = type_to_column[selected_type[valid_type]]
+        valid = (columns >= 0) & np.isfinite(selected_value)
+        values[row, columns[valid]] = selected_value[valid]
+
+    observation_count = np.isfinite(values).sum(axis=0).astype(np.int64)
+    missing_types = [
+        type_id
+        for type_id, count in zip(cont_type_ids, observation_count.tolist())
+        if count == 0
+    ]
+    if missing_types:
+        raise ValueError(
+            "Training subset has no finite observations for continuous type ids: "
+            f"{missing_types}"
+        )
+
+    low_value, center, high_value = np.nanpercentile(
+        values,
+        [low, 50.0, high],
+        axis=0,
+    )
+    scale = high_value - low_value
+    near_constant = (~np.isfinite(scale)) | (np.abs(scale) <= float(scale_epsilon))
+    scale[near_constant] = 1.0
+
+    center = np.asarray(center, dtype=np.float32)
+    scale = np.asarray(scale, dtype=np.float32)
+    if not np.isfinite(center).all():
+        raise RuntimeError("Robust scaler produced non-finite center values")
+    if not np.isfinite(scale).all() or np.any(scale <= 0):
+        raise RuntimeError("Robust scaler produced invalid scale values")
+
+    return ContinuousRobustScalerStats(
+        cont_type_ids=cont_type_ids,
+        center=center,
+        scale=scale,
+        observation_count=observation_count,
+        quantile_range=(low, high),
+    )
+
+
 def create_unique_run_dir(name_fn, runs_root: Path = Path("runs")) -> tuple[Path, str]:
     while True:
         timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")