Add training-fitted normalization to time-series processors

docs/api/processors.rst

@@ -299,6 +299,64 @@ Common string keys for automatic processor selection:
 - ``"raw"``: For raw/unprocessed data
 - ``"graph"``: For knowledge graph subgraphs
 
+Training-only time-series normalization
+--------------------------------------
+
+``TimeseriesProcessor`` and ``TemporalTimeseriesProcessor`` accept
+``normalize_strategy="standard"`` for per-feature z-score normalization.
+The default, ``None``, keeps existing behavior. Min–max scaling is not supported.
+
+``fit()`` computes the training mean and population standard deviation
+(``ddof=0``) **after resampling and imputation**. Every resampled time step has
+equal weight, including filled gaps; longer sequences therefore contribute
+more observations. Constant features use a scale of 1. ``process()`` applies
+these fixed statistics without updating them. The temporal processor normalizes
+only ``"value"``; its ``"time"`` tensor remains in elapsed hours.
+
+.. important::
+
+   Split raw samples by patient before fitting normalization. Calling
+   ``set_task()`` with normalization enabled on the full dataset and then
+   splitting the resulting samples leaks validation/test statistics into training.
+
+Create the training dataset first, then explicitly transfer its fitted processors:
+
+.. code-block:: python
+
+    from pyhealth.datasets import create_sample_dataset
+
+    schema = {"vitals": ("timeseries", {"normalize_strategy": "standard"})}
+    train = create_sample_dataset(
+        samples=train_samples, input_schema=schema, output_schema={}
+    )
+    validation = create_sample_dataset(
+        samples=validation_samples,
+        input_schema=schema,
+        output_schema={},
+        input_processors=train.input_processors,
+    )
+
+Use the same transfer for the test split. For temporal data, replace
+``"timeseries"`` with ``"temporal_timeseries"``. Passing a pre-fitted processor
+instance in ``input_schema`` alone does **not** prevent refitting; use the
+``input_processors`` argument instead. The processor cannot infer which samples
+are training data, so the caller must supply the intended training partition.
+
+Normalization requires a successful ``fit()``; otherwise processing raises
+``RuntimeError``. Refitting replaces all previous statistics, and failed fitting
+leaves normalization unfitted. Missing or ``None`` training fields are skipped;
+no usable training samples, malformed present samples, and infinite values raise
+``ValueError``. NaNs are handled by the existing imputation step, including its
+zero initialization for leading or entirely missing features. Resampled and
+imputed values must be finite. Existing input shapes and imputation behavior
+are preserved: the ordinary processor expects ``(T, F)`` values, while the
+temporal processor also accepts ``(T,)``.
+
+Fitted statistics are saved with the existing processor metadata and reused by
+cached datasets; no separate statistics file is necessary. See
+``examples/timeseries_normalization.py`` for a complete synthetic example that
+can be run with ``pixi run -e test python examples/timeseries_normalization.py``.
+
 Writing Custom FeatureProcessors
 ---------------------------------
 
@@ -496,4 +554,4 @@ API Reference
     processors/pyhealth.processors.MultiHotProcessor
     processors/pyhealth.processors.StageNetProcessor
     processors/pyhealth.processors.StageNetTensorProcessor
-    processors/pyhealth.processors.GraphProcessor
+    processors/pyhealth.processors.GraphProcessor

examples/timeseries_normalization.py

@@ -0,0 +1,59 @@
+"""Fit z-score normalization on training patients and reuse it on held-out data.
+
+Run with: pixi run -e test python examples/timeseries_normalization.py
+All samples are synthetic; no downloads or model training are required.
+"""
+
+from datetime import datetime, timedelta
+
+import numpy as np
+
+from pyhealth.datasets import create_sample_dataset
+
+
+def main() -> None:
+    """Demonstrate ordinary and temporal normalization without split leakage."""
+    start = datetime(2026, 1, 1)
+    raw_samples = [
+        {
+            "patient_id": "train-patient",
+            "vitals": ([start, start + timedelta(hours=1)], np.array([[10.0], [30.0]])),
+        },
+        {"patient_id": "validation-patient", "vitals": ([start], np.array([[40.0]]))},
+        {"patient_id": "test-patient", "vitals": ([start], np.array([[1000.0]]))},
+    ]
+    # Split raw samples by patient BEFORE fitting any normalization statistics.
+    split_patients = {
+        "train": {"train-patient"},
+        "validation": {"validation-patient"},
+        "test": {"test-patient"},
+    }
+    splits = {
+        name: [s for s in raw_samples if s["patient_id"] in patients]
+        for name, patients in split_patients.items()
+    }
+
+    for alias in ("timeseries", "temporal_timeseries"):
+        schema = {"vitals": (alias, {"normalize_strategy": "standard"})}
+        train = create_sample_dataset(
+            samples=splits["train"], input_schema=schema, output_schema={}
+        )
+        for name, expected in (("validation", 2.0), ("test", 98.0)):
+            held_out = create_sample_dataset(
+                samples=splits[name],
+                input_schema=schema,
+                output_schema={},
+                # Supplying processors here prevents fitting them on held-out data.
+                input_processors=train.input_processors,
+            )
+            value = held_out[0]["vitals"]
+            if isinstance(value, dict):
+                np.testing.assert_array_equal(value["time"], [0.0])
+                value = value["value"]
+            # Training mean=20, scale=10. The test outlier does not change them.
+            np.testing.assert_array_equal(value, [[expected]])
+            print(f"{alias}: {name} normalized values = {value.tolist()}")
+
+
+if __name__ == "__main__":
+    main()

pyhealth/processors/temporal_timeseries_processor.py

@@ -8,14 +8,20 @@
 ``UnifiedMultimodalEmbeddingModel`` can sort and align events across modalities
 on a shared timeline.
 """
+
 from datetime import datetime, timedelta
-from typing import Any, Iterable, Dict, List, Tuple
+from typing import Any, Iterable, Dict, List, Literal, Tuple
 
 import numpy as np
 import torch
 
 from . import register_processor
 from .base_processor import ModalityType, TemporalFeatureProcessor
+from .timeseries_processor import (
+    _fit_standardization,
+    _standardize,
+    _validate_normalization_input,
+)
 
 
 @register_processor("temporal_timeseries")
@@ -38,6 +44,10 @@ class TemporalTimeseriesProcessor(TemporalFeatureProcessor):
     Args:
         sampling_rate: Uniform re-sampling interval.  Defaults to 1 hour.
         impute_strategy: Currently only ``"forward_fill"`` is supported.
+        normalize_strategy: ``None`` (default) or ``"standard"``. Z-scores use
+            training-only statistics after resampling and imputation, weighting
+            each time step equally. Constant features use scale 1. Only the
+            ``"value"`` tensor is normalized; timestamps are unchanged.
 
     Example::
 
@@ -48,21 +58,57 @@ class TemporalTimeseriesProcessor(TemporalFeatureProcessor):
         out = proc.process_temporal((ts, val))
         # out["value"].shape  → (5, 2)   ← 5 two-hour steps over 8 h
         # out["time"].shape   → (5,)     ← [0., 2., 4., 6., 8.] hours
+
+    Examples:
+        >>> times = [datetime(2026, 1, 1), datetime(2026, 1, 1, 1)]
+        >>> values = np.array([[10.0], [30.0]])
+        >>> processor = TemporalTimeseriesProcessor(normalize_strategy="standard")
+        >>> processor.fit([{"vitals": (times, values)}], "vitals")
+        >>> output = processor.process((times, values))
+        >>> output["value"].tolist()
+        [[-1.0], [1.0]]
+        >>> output["time"].tolist()
+        [0.0, 1.0]
     """
 
     def __init__(
         self,
         sampling_rate: timedelta = timedelta(hours=1),
         impute_strategy: str = "forward_fill",
+        normalize_strategy: Literal["standard"] | None = None,
     ):
+        if normalize_strategy not in (None, "standard"):
+            raise ValueError("normalize_strategy must be None or 'standard'.")
+        if normalize_strategy is not None and sampling_rate <= timedelta(0):
+            raise ValueError("sampling_rate must be positive for normalization.")
         self.sampling_rate = sampling_rate
         self.impute_strategy = impute_strategy
         self.n_features: int | None = None
+        self.normalize_strategy = normalize_strategy
+        self._normalization_mean: list[float] | None = None
+        self._normalization_scale: list[float] | None = None
 
     # ── FeatureProcessor interface ─────────────────────────────────────────
 
     def fit(self, samples: Iterable[Dict[str, Any]], field: str) -> None:
-        """Infer feature dimension from the first valid sample."""
+        """Infer feature count and optionally fit training-only z-score statistics.
+
+        Statistics use population standard deviation (ddof=0) after filling.
+        Each fit replaces previous statistics, and a failed fit leaves
+        normalization unfitted. Missing/None fields are skipped; normalization
+        requires at least one usable training sample.
+        """
+        self._normalization_mean = None
+        self._normalization_scale = None
+        if getattr(self, "normalize_strategy", None) == "standard":
+            self.n_features = None
+            (
+                self.n_features,
+                self._normalization_mean,
+                self._normalization_scale,
+            ) = _fit_standardization(samples, field, self._resample_and_impute)
+            return
+
         for sample in samples:
             if field in sample and sample[field] is not None:
                 _, values = sample[field]
@@ -83,7 +129,30 @@ def process(self, value: Tuple[List[datetime], np.ndarray]) -> dict:
 
         Returns:
             ``{"value": FloatTensor (S, F), "time": FloatTensor (S,)}``
+
+        Raises:
+            RuntimeError: If normalization is enabled but fitting has not succeeded.
+            ValueError: If the series is invalid for normalization.
         """
+        sampled_values = self._resample_and_impute(value)
+        if getattr(self, "normalize_strategy", None) == "standard":
+            sampled_values = _standardize(
+                sampled_values, self._normalization_mean, self._normalization_scale
+            )
+
+        hours_per_step = self.sampling_rate.total_seconds() / 3600.0
+        time_hours = np.array(
+            [i * hours_per_step for i in range(len(sampled_values))], dtype=np.float32
+        )
+        return {
+            "value": torch.tensor(sampled_values, dtype=torch.float32),
+            "time": torch.tensor(time_hours, dtype=torch.float32),
+        }
+
+    def _resample_and_impute(
+        self, value: Tuple[List[datetime], np.ndarray]
+    ) -> np.ndarray:
+        """Produce the same unnormalized grid for fitting and processing."""
         timestamps, values = value
 
         if len(timestamps) == 0:
@@ -93,6 +162,11 @@ def process(self, value: Tuple[List[datetime], np.ndarray]) -> dict:
         if values.ndim == 1:
             values = values[:, None]  # (T,) → (T, 1)
 
+        if getattr(self, "normalize_strategy", None) == "standard":
+            _validate_normalization_input(
+                timestamps, values, self.sampling_rate, self.n_features
+            )
+
         num_features = values.shape[1]
         start_time = timestamps[0]
         end_time = timestamps[-1]
@@ -114,16 +188,7 @@ def process(self, value: Tuple[List[datetime], np.ndarray]) -> dict:
                 else:
                     sampled_values[i, f] = last
 
-        # Build time tensor (hours from first observation)
-        hours_per_step = self.sampling_rate.total_seconds() / 3600.0
-        time_hours = np.array(
-            [i * hours_per_step for i in range(total_steps)], dtype=np.float32
-        )
-
-        return {
-            "value": torch.tensor(sampled_values, dtype=torch.float32),
-            "time":  torch.tensor(time_hours, dtype=torch.float32),
-        }
+        return sampled_values
 
     # process_temporal delegates to process (already returns dict)
     def process_temporal(self, value) -> dict:
@@ -157,8 +222,10 @@ def size(self) -> int | None:
         return self.n_features
 
     def __repr__(self) -> str:
+        strategy = getattr(self, "normalize_strategy", None)
+        normalization = f", normalize_strategy={strategy!r}" if strategy else ""
         return (
             f"TemporalTimeseriesProcessor("
             f"sampling_rate={self.sampling_rate}, "
-            f"n_features={self.n_features})"
+            f"n_features={self.n_features}{normalization})"
         )