From fb17f5e1c648c73fbb30aad632e66249999cb8b8 Mon Sep 17 00:00:00 2001
From: Logiquo <fanyongda2012@gmail.com>
Date: Tue, 27 Oct 2026 02:42:12 -0400
Subject: [PATCH 1/3] Add timeseries_binary and timeseries_regression modes

Per-step (time series) labels become model modes of their own, next to
binary, multiclass, multilabel and regression, so modules that assume one
prediction per sample (calibration, prediction sets) do not mistake them
for binary or regression labels.

- BaseModel: the two modes in VALID_MODES; a per-step BCE/MSE loss
  averaged over the steps whose mask is 1, for labels batched as
  (value, mask); per-step probabilities.
- RNN: predicts at every input step for these modes (features without a
  time axis are repeated over time) and rejects bidirectional=True, which
  would use future steps.
- Trainer: inference returns one array per sample over its labeled steps;
  get_metrics_fn maps the modes to the new timeseries_binary_metrics_fn
  and timeseries_regression_metrics_fn, which pool the labeled steps.

The (value, mask) label handling is a temporary workaround that keeps the
existing single-tensor label API; unified masked tensor handling is a
follow-up. Tests use a minimal mock label processor; the label processors
themselves follow in a separate change.
---
 pyhealth/metrics/__init__.py        |   6 +-
 pyhealth/metrics/binary.py          |  42 ++++++++-
 pyhealth/metrics/regression.py      |  38 +++++++-
 pyhealth/models/base_model.py       |  61 +++++++++++-
 pyhealth/models/rnn.py              |  35 +++++--
 pyhealth/trainer.py                 |  32 +++++--
 tests/core/test_timeseries_modes.py | 138 ++++++++++++++++++++++++++++
 7 files changed, 334 insertions(+), 18 deletions(-)
 create mode 100644 tests/core/test_timeseries_modes.py

diff --git a/pyhealth/metrics/__init__.py b/pyhealth/metrics/__init__.py
index 444bc8212..2bb03e594 100644
--- a/pyhealth/metrics/__init__.py
+++ b/pyhealth/metrics/__init__.py
@@ -1,4 +1,4 @@
-from .binary import binary_metrics_fn
+from .binary import binary_metrics_fn, timeseries_binary_metrics_fn
 from .bootstrap import bootstrap_ci, paired_bootstrap_diff
 from .drug_recommendation import ddi_rate_score
 from .fairness import fairness_metrics_fn
@@ -20,7 +20,7 @@
 from .multiclass import multiclass_metrics_fn
 from .multilabel import multilabel_metrics_fn
 from .ranking import ranking_metrics_fn
-from .regression import regression_metrics_fn
+from .regression import regression_metrics_fn, timeseries_regression_metrics_fn
 
 __all__ = [
     "ComprehensivenessMetric",
@@ -43,4 +43,6 @@
     "paired_bootstrap_diff",
     "ranking_metrics_fn",
     "regression_metrics_fn",
+    "timeseries_binary_metrics_fn",
+    "timeseries_regression_metrics_fn",
 ]
diff --git a/pyhealth/metrics/binary.py b/pyhealth/metrics/binary.py
index ee340dc4c..9670ebc91 100644
--- a/pyhealth/metrics/binary.py
+++ b/pyhealth/metrics/binary.py
@@ -1,4 +1,4 @@
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Sequence
 
 import numpy as np
 import sklearn.metrics as sklearn_metrics
@@ -147,6 +147,46 @@ def binary_metrics_fn(
     return output
 
 
+
+def timeseries_binary_metrics_fn(
+    y_true: Sequence[np.ndarray],
+    y_prob: Sequence[np.ndarray],
+    metrics: Optional[List[str]] = None,
+    threshold: float = 0.5,
+) -> Dict[str, float]:
+    """Computes metrics for binary labels given per time step.
+
+    Each sample contributes the labels and predicted probabilities of its
+    labeled time steps. The steps of all samples are pooled and scored with
+    :func:`binary_metrics_fn`, as in benchmarks that evaluate every hour of
+    every ICU stay (e.g. YAIB).
+
+    Args:
+        y_true: One array of 0/1 labels per sample, over its labeled steps.
+        y_prob: One array of predicted probabilities per sample, matching
+            ``y_true``.
+        metrics: Metrics to compute, see :func:`binary_metrics_fn`.
+        threshold: Decision threshold, see :func:`binary_metrics_fn`.
+
+    Returns:
+        Dictionary of metrics whose keys are the metric names and values are
+            the metric values.
+
+    Examples:
+        >>> import numpy as np
+        >>> from pyhealth.metrics import timeseries_binary_metrics_fn
+        >>> y_true = [np.array([0, 0, 1]), np.array([1, 0])]
+        >>> y_prob = [np.array([0.1, 0.4, 0.8]), np.array([0.7, 0.2])]
+        >>> timeseries_binary_metrics_fn(y_true, y_prob, metrics=["roc_auc"])
+        {'roc_auc': 1.0}
+    """
+    return binary_metrics_fn(
+        np.concatenate([np.ravel(y) for y in y_true]),
+        np.concatenate([np.ravel(p) for p in y_prob]),
+        metrics=metrics,
+        threshold=threshold,
+    )
+
 if __name__ == "__main__":
     all_metrics = [
         "pr_auc",
diff --git a/pyhealth/metrics/regression.py b/pyhealth/metrics/regression.py
index 91f87028f..1d05bb9c1 100644
--- a/pyhealth/metrics/regression.py
+++ b/pyhealth/metrics/regression.py
@@ -1,4 +1,4 @@
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Sequence
 
 import numpy as np
 import sklearn.metrics as sklearn_metrics
@@ -70,6 +70,42 @@ def regression_metrics_fn(
     return output
 
 
+
+def timeseries_regression_metrics_fn(
+    y_true: Sequence[np.ndarray],
+    y_pred: Sequence[np.ndarray],
+    metrics: Optional[List[str]] = None,
+) -> Dict[str, float]:
+    """Computes metrics for real-valued labels given per time step.
+
+    Each sample contributes the targets and predictions of its labeled time
+    steps. The steps of all samples are pooled and scored with
+    :func:`regression_metrics_fn`, as in benchmarks that evaluate every hour
+    of every ICU stay (e.g. YAIB).
+
+    Args:
+        y_true: One array of targets per sample, over its labeled steps.
+        y_pred: One array of predictions per sample, matching ``y_true``.
+        metrics: Metrics to compute, see :func:`regression_metrics_fn`.
+
+    Returns:
+        Dictionary of metrics whose keys are the metric names and values are
+            the metric values.
+
+    Examples:
+        >>> import numpy as np
+        >>> from pyhealth.metrics import timeseries_regression_metrics_fn
+        >>> y_true = [np.array([3.0, 2.0]), np.array([1.0])]
+        >>> y_pred = [np.array([2.0, 2.0]), np.array([2.0])]
+        >>> timeseries_regression_metrics_fn(y_true, y_pred, metrics=["mae"])
+        {'mae': 0.6666666666666666}
+    """
+    return regression_metrics_fn(
+        np.concatenate([np.ravel(y) for y in y_true]),
+        np.concatenate([np.ravel(p) for p in y_pred]),
+        metrics=metrics,
+    )
+
 if __name__ == "__main__":
     x = np.random.random(size=10000)
     x_rec = np.random.random(size=10000)
diff --git a/pyhealth/models/base_model.py b/pyhealth/models/base_model.py
index 8b9683763..9b83241f1 100644
--- a/pyhealth/models/base_model.py
+++ b/pyhealth/models/base_model.py
@@ -1,6 +1,7 @@
 from abc import ABC
 from collections.abc import Iterable, Sequence
 from typing import Callable, Any, Optional
+import functools
 import inspect
 import logging
 
@@ -14,7 +15,42 @@
 logger = logging.getLogger(__name__)
 
 # Label kinds understood by the loss helpers, Trainer and the calibration methods.
-VALID_MODES = ("binary", "multiclass", "multilabel", "regression")
+# The timeseries modes are labels per time step (see TimeseriesLabelProcessor).
+VALID_MODES = (
+    "binary",
+    "multiclass",
+    "multilabel",
+    "regression",
+    "timeseries_binary",
+    "timeseries_regression",
+)
+
+
+def _masked_step_loss(
+    loss_fn: Callable,
+    logits: torch.Tensor,
+    target: tuple[torch.Tensor, torch.Tensor],
+) -> torch.Tensor:
+    """Averages a per-step loss over the time steps that have a label.
+
+    Args:
+        loss_fn: Elementwise loss accepting ``reduction="none"``.
+        logits: Per-step logits of shape (batch, steps, 1).
+        target: Batched ``(value, mask)`` of shape (batch, label steps), where
+            padding and unlabeled steps have mask 0.
+    """
+    # Temporary workaround: per-step labels arrive as a (value, mask) tuple so
+    # the existing single-tensor label API stays unchanged. A more unified
+    # handling of padded/masked label tensors across models is a follow-up.
+    value, mask = target
+    if value.shape[1] > logits.shape[1]:
+        raise ValueError(
+            f"The label has {value.shape[1]} time steps but the input only "
+            f"{logits.shape[1]}; label and input must share the time grid."
+        )
+    logits = logits[:, : value.shape[1], 0]
+    loss = loss_fn(logits, value, reduction="none")
+    return (loss * mask).sum() / mask.sum().clamp(min=1)
 
 
 class BaseModel(ABC, nn.Module):
@@ -38,7 +74,8 @@ class BaseModel(ABC, nn.Module):
     Mode
     --------
         ``mode`` is always one of ``"binary"``, ``"multiclass"``, ``"multilabel"``,
-        ``"regression"`` or ``None``. It is resolved from the single label's
+        ``"regression"``, ``"timeseries_binary"``, ``"timeseries_regression"``
+        or ``None``. It is resolved from the single label's
         ``output_schema`` entry, which may be a string, a processor class, a
         processor instance or a ``(name, kwargs)`` tuple. Assigning ``self.mode``
         in a subclass resolves the value the same way.
@@ -269,6 +306,14 @@ def get_loss_function(self) -> Callable:
             - multiclass: `F.cross_entropy`
             - multilabel: `F.binary_cross_entropy_with_logits`
             - regression: `F.mse_loss`
+            - timeseries_binary: `F.binary_cross_entropy_with_logits` per
+              time step, averaged over the labeled steps
+            - timeseries_regression: `F.mse_loss` per time step, averaged
+              over the labeled steps
+
+        The timeseries losses take per-step logits of shape (batch, steps, 1)
+        and the batched ``(value, mask)`` label of a
+        :class:`~pyhealth.processors.TimeseriesLabelProcessor`.
 
         For binary and multilabel labels, a weight set with
         :meth:`set_pos_weight` is passed as ``pos_weight``.
@@ -298,6 +343,10 @@ def weighted_bce(input, target, **kwargs):
             return F.binary_cross_entropy_with_logits
         elif mode == "regression":
             return F.mse_loss
+        elif mode == "timeseries_binary":
+            return functools.partial(_masked_step_loss, F.binary_cross_entropy_with_logits)
+        elif mode == "timeseries_regression":
+            return functools.partial(_masked_step_loss, F.mse_loss)
         else:
             raise ValueError(f"Invalid mode: {mode}")
 
@@ -315,6 +364,10 @@ def prepare_y_prob(self, logits: torch.Tensor) -> torch.Tensor:
             - multilabel: a tensor of shape (batch_size, num_labels) with values
                 in [0, 1], which is obtained with `torch.sigmoid()`
             - regression: a tensor of shape (batch_size, 1) with raw logits
+            - timeseries_binary: a tensor of shape (batch_size, steps, 1)
+                with values in [0, 1] per time step, from `torch.sigmoid()`
+            - timeseries_regression: a tensor of shape (batch_size, steps, 1)
+                with raw logits per time step
 
         Args:
             logits (torch.Tensor): The predicted logit tensor.
@@ -335,6 +388,10 @@ def prepare_y_prob(self, logits: torch.Tensor) -> torch.Tensor:
             y_prob = torch.sigmoid(logits)
         elif mode in ["regression"]:
             y_prob = logits
+        elif mode in ["timeseries_binary"]:
+            y_prob = torch.sigmoid(logits)
+        elif mode in ["timeseries_regression"]:
+            y_prob = logits
         else:
             raise NotImplementedError
         return y_prob
diff --git a/pyhealth/models/rnn.py b/pyhealth/models/rnn.py
index ef0b6f6bc..8e728b622 100644
--- a/pyhealth/models/rnn.py
+++ b/pyhealth/models/rnn.py
@@ -225,6 +225,13 @@ def __init__(
             raise ValueError("hidden_size is determined by hidden_dim")
         assert len(self.label_keys) == 1, "Only one label key is supported if RNN is initialized"
         self.label_key = self.label_keys[0]
+        # Per-step (timeseries) labels predict at every input step.
+        self._per_step = self.mode in ("timeseries_binary", "timeseries_regression")
+        if self._per_step and kwargs.get("bidirectional", False):
+            raise ValueError(
+                "A per-step label needs a causal model; bidirectional=True would "
+                "use future time steps."
+            )
 
         self.embedding_model = EmbeddingModel(dataset, embedding_dim)
 
@@ -302,14 +309,30 @@ def forward(self, **kwargs) -> Dict[str, torch.Tensor]:
                 else:
                     mask = (torch.abs(x).sum(dim=-1) != 0).int()
 
-            _, x = self.rnn[feature_key](x, mask)
-            patient_emb.append(x)
-
-        patient_emb = torch.cat(patient_emb, dim=1)
-        # (patient, label_size)
+            outputs, x = self.rnn[feature_key](x, mask)
+            # Per-step labels use the output at every step, else the last one.
+            patient_emb.append(outputs if self._per_step else x)
+
+        if self._per_step:
+            # Features without a time axis (one step) are repeated over time.
+            steps = max(emb.shape[1] for emb in patient_emb)
+            for i, emb in enumerate(patient_emb):
+                if emb.shape[1] == 1:
+                    patient_emb[i] = emb.expand(-1, steps, -1)
+                elif emb.shape[1] != steps:
+                    raise ValueError(
+                        "Per-step labels need all time series features on the "
+                        "same time grid."
+                    )
+        patient_emb = torch.cat(patient_emb, dim=-1)
+        # (patient, label_size), or (patient, steps, label_size) per step
         logits = self.fc(patient_emb)
         # obtain y_true, loss, y_prob
-        y_true = kwargs[self.label_key].to(self.device)
+        if self._per_step:
+            # (value, mask) per step; the loss averages over labeled steps.
+            y_true = tuple(t.to(self.device) for t in kwargs[self.label_key])
+        else:
+            y_true = kwargs[self.label_key].to(self.device)
         loss = self.get_loss_function()(logits, y_true)
         y_prob = self.prepare_y_prob(logits)
         results = {"loss": loss, "y_prob": y_prob, "y_true": y_true, "logit": logits}
diff --git a/pyhealth/trainer.py b/pyhealth/trainer.py
index 91e7e8d04..efedd1b54 100644
--- a/pyhealth/trainer.py
+++ b/pyhealth/trainer.py
@@ -17,6 +17,8 @@
     multiclass_metrics_fn,
     multilabel_metrics_fn,
     regression_metrics_fn,
+    timeseries_binary_metrics_fn,
+    timeseries_regression_metrics_fn,
 )
 from pyhealth.utils import create_directory
 
@@ -59,6 +61,10 @@ def get_metrics_fn(mode: str) -> Callable:
         return multilabel_metrics_fn
     elif mode == "regression":
         return regression_metrics_fn
+    elif mode == "timeseries_binary":
+        return timeseries_binary_metrics_fn
+    elif mode == "timeseries_regression":
+        return timeseries_regression_metrics_fn
     else:
         raise ValueError(f"Mode {mode} is not supported")
 
@@ -335,6 +341,8 @@ def inference(self, dataloader, additional_outputs=None,
         Returns:
             y_true_all: List of true labels.
             y_prob_all: List of predicted probabilities.
+                For the per-step ``timeseries_*`` modes, both are lists with
+                one array per sample, over its labeled time steps.
             loss_mean: Mean loss over batches.
             additional_outputs (only if requested): Dict of additional results.
             patient_ids (only if requested): List of patient ids in the same order as y_true_all/y_prob_all.
@@ -345,24 +353,36 @@ def inference(self, dataloader, additional_outputs=None,
         patient_ids = []
         if additional_outputs is not None:
             additional_outputs = {k: [] for k in additional_outputs}
+        per_step = self.model.mode in ("timeseries_binary", "timeseries_regression")
         self.model.eval()
         for data in tqdm(dataloader, desc="Evaluation"):
             with torch.no_grad():
                 output = self.model(**data)
                 loss = output["loss"]
-                y_true = output["y_true"].cpu().numpy()
-                y_prob = output["y_prob"].cpu().numpy()
                 loss_all.append(loss.item())
-                y_true_all.append(y_true)
-                y_prob_all.append(y_prob)
+                if per_step:
+                    # Temporary workaround to avoid breaking the existing API:
+                    # samples differ in length, so keep each sample's labeled
+                    # steps as a separate array. A more unified handling of
+                    # padded/masked outputs is a follow-up.
+                    value, mask = (t.cpu().numpy() for t in output["y_true"])
+                    y_prob = output["y_prob"].cpu().numpy()[:, : value.shape[1], 0]
+                    for i in range(len(value)):
+                        labeled = mask[i] > 0
+                        y_true_all.append(value[i][labeled])
+                        y_prob_all.append(y_prob[i][labeled])
+                else:
+                    y_true_all.append(output["y_true"].cpu().numpy())
+                    y_prob_all.append(output["y_prob"].cpu().numpy())
                 if additional_outputs is not None:
                     for key in additional_outputs.keys():
                         additional_outputs[key].append(output[key].cpu().numpy())
             if return_patient_ids:
                 patient_ids.extend(data["patient_id"])
         loss_mean = sum(loss_all) / len(loss_all)
-        y_true_all = np.concatenate(y_true_all, axis=0)
-        y_prob_all = np.concatenate(y_prob_all, axis=0)
+        if not per_step:
+            y_true_all = np.concatenate(y_true_all, axis=0)
+            y_prob_all = np.concatenate(y_prob_all, axis=0)
         outputs = [y_true_all, y_prob_all, loss_mean]
         if additional_outputs is not None:
             additional_outputs = {key: np.concatenate(val)
diff --git a/tests/core/test_timeseries_modes.py b/tests/core/test_timeseries_modes.py
new file mode 100644
index 000000000..677884487
--- /dev/null
+++ b/tests/core/test_timeseries_modes.py
@@ -0,0 +1,138 @@
+import unittest
+from datetime import datetime, timedelta
+
+import numpy as np
+import torch
+
+from pyhealth.datasets import create_sample_dataset, get_dataloader
+from pyhealth.metrics import (
+    timeseries_binary_metrics_fn,
+    timeseries_regression_metrics_fn,
+)
+from pyhealth.models import RNN
+from pyhealth.processors.base_processor import FeatureProcessor
+from pyhealth.trainer import Trainer
+
+T0 = datetime(2020, 1, 1)
+
+
+class _StepLabelProcessor(FeatureProcessor):
+    """Minimal per-step label processor: ``(values, mask)`` lists to tensors."""
+
+    def process(self, value):
+        values, mask = value
+        return (
+            torch.tensor(values, dtype=torch.float),
+            torch.tensor(mask, dtype=torch.float),
+        )
+
+    def size(self):
+        return 1
+
+    def is_token(self):
+        return False
+
+    def schema(self):
+        return ("value", "mask")
+
+    def dim(self):
+        return (1, 1)
+
+    def spatial(self):
+        return (True,)
+
+
+def _dataset(mode="timeseries_binary", unlabeled=()):
+    rng = np.random.default_rng(0)
+    samples = []
+    for i, steps in enumerate([5, 3, 6, 4]):
+        values = [float((i + h) % 2) for h in range(steps)]
+        mask = [0.0 if (i, h) in unlabeled else 1.0 for h in range(steps)]
+        samples.append(
+            {
+                "patient_id": f"p{i}",
+                "dynamic": (
+                    [T0 + timedelta(hours=h) for h in range(steps)],
+                    rng.normal(size=(steps, 2)),
+                ),
+                "static": [float(i), 1.0],
+                "label": (values, mask),
+            }
+        )
+    return create_sample_dataset(
+        samples=samples,
+        input_schema={"dynamic": "timeseries", "static": "tensor"},
+        output_schema={"label": mode},
+        output_processors={"label": _StepLabelProcessor()},
+        dataset_name="timeseries_modes",
+    )
+
+
+class TestTimeseriesModes(unittest.TestCase):
+    def test_modes_are_separate_from_binary_and_regression(self):
+        for mode in ("timeseries_binary", "timeseries_regression"):
+            self.assertEqual(RNN(dataset=_dataset(mode)).mode, mode)
+
+    def test_loss_ignores_unlabeled_and_padded_steps(self):
+        dataset = _dataset(unlabeled={(1, 0)})
+        model = RNN(dataset=dataset, hidden_dim=8, embedding_dim=8).eval()
+        batch = next(iter(get_dataloader(dataset, batch_size=4)))
+        with torch.no_grad():
+            out = model(**batch)
+        value, mask = out["y_true"]
+        self.assertEqual(tuple(out["y_prob"].shape), (4, 6, 1))
+        self.assertEqual(int(mask.sum()), 17)  # 5 + 3 + 6 + 4 steps, minus one
+
+        logits = out["logit"].clone()
+        logits[mask == 0] = 1e3
+        loss = model.get_loss_function()(logits, (value, mask))
+        torch.testing.assert_close(loss, out["loss"])
+
+    def test_predictions_do_not_see_future_steps(self):
+        dataset = _dataset()
+        model = RNN(dataset=dataset, hidden_dim=8, embedding_dim=8).eval()
+        batch = next(iter(get_dataloader(dataset, batch_size=1)))
+        with torch.no_grad():
+            before = model(**batch)["y_prob"][0, :, 0]
+            batch["dynamic"] = batch["dynamic"].clone()
+            batch["dynamic"][:, -1] += 100.0
+            after = model(**batch)["y_prob"][0, :, 0]
+        torch.testing.assert_close(before[:-1], after[:-1])
+        self.assertFalse(torch.allclose(before[-1], after[-1]))
+
+    def test_bidirectional_is_rejected(self):
+        with self.assertRaises(ValueError):
+            RNN(dataset=_dataset(), bidirectional=True)
+
+    def test_trainer_scores_labeled_steps_per_sample(self):
+        dataset = _dataset()
+        model = RNN(dataset=dataset, hidden_dim=8, embedding_dim=8)
+        trainer = Trainer(model=model, metrics=["roc_auc"], enable_logging=False)
+        loader = get_dataloader(dataset, batch_size=3)
+        y_true, y_prob, _ = trainer.inference(loader)
+        self.assertEqual(sorted(len(y) for y in y_true), [3, 4, 5, 6])
+        self.assertEqual([len(p) for p in y_prob], [len(y) for y in y_true])
+        self.assertIn("roc_auc", trainer.evaluate(loader))
+
+        dataset = _dataset("timeseries_regression")
+        model = RNN(dataset=dataset, hidden_dim=8, embedding_dim=8)
+        trainer = Trainer(model=model, metrics=["mae"], enable_logging=False)
+        self.assertIn("mae", trainer.evaluate(get_dataloader(dataset, batch_size=3)))
+
+
+class TestTimeseriesMetrics(unittest.TestCase):
+    def test_binary_pools_steps_of_all_samples(self):
+        y_true = [np.array([0, 0, 1]), np.array([1, 0])]
+        y_prob = [np.array([0.1, 0.4, 0.8]), np.array([0.7, 0.2])]
+        scores = timeseries_binary_metrics_fn(y_true, y_prob, metrics=["roc_auc"])
+        self.assertEqual(scores, {"roc_auc": 1.0})
+
+    def test_regression_pools_steps_of_all_samples(self):
+        y_true = [np.array([3.0, 2.0]), np.array([1.0])]
+        y_pred = [np.array([2.0, 2.0]), np.array([2.0])]
+        scores = timeseries_regression_metrics_fn(y_true, y_pred, metrics=["mae"])
+        self.assertAlmostEqual(scores["mae"], 2 / 3)
+
+
+if __name__ == "__main__":
+    unittest.main()

From afeaed40909bcc0b25c2e7156ced5d4a593da1d7 Mon Sep 17 00:00:00 2001
From: Logiquo <fanyongda2012@gmail.com>
Date: Tue, 27 Oct 2026 03:03:09 -0400
Subject: [PATCH 2/3] Use built-in generics and add a usage example to
 get_metrics_fn

---
 pyhealth/metrics/binary.py     |  7 ++++---
 pyhealth/metrics/regression.py |  7 ++++---
 pyhealth/trainer.py            | 16 ++++++++++++++++
 3 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/pyhealth/metrics/binary.py b/pyhealth/metrics/binary.py
index 9670ebc91..9f73bc576 100644
--- a/pyhealth/metrics/binary.py
+++ b/pyhealth/metrics/binary.py
@@ -1,4 +1,5 @@
-from typing import Dict, List, Optional, Sequence
+from collections.abc import Sequence
+from typing import Dict, List, Optional
 
 import numpy as np
 import sklearn.metrics as sklearn_metrics
@@ -151,9 +152,9 @@ def binary_metrics_fn(
 def timeseries_binary_metrics_fn(
     y_true: Sequence[np.ndarray],
     y_prob: Sequence[np.ndarray],
-    metrics: Optional[List[str]] = None,
+    metrics: list[str] | None = None,
     threshold: float = 0.5,
-) -> Dict[str, float]:
+) -> dict[str, float]:
     """Computes metrics for binary labels given per time step.
 
     Each sample contributes the labels and predicted probabilities of its
diff --git a/pyhealth/metrics/regression.py b/pyhealth/metrics/regression.py
index 1d05bb9c1..2d0d02ed5 100644
--- a/pyhealth/metrics/regression.py
+++ b/pyhealth/metrics/regression.py
@@ -1,4 +1,5 @@
-from typing import Dict, List, Optional, Sequence
+from collections.abc import Sequence
+from typing import Dict, List, Optional
 
 import numpy as np
 import sklearn.metrics as sklearn_metrics
@@ -74,8 +75,8 @@ def regression_metrics_fn(
 def timeseries_regression_metrics_fn(
     y_true: Sequence[np.ndarray],
     y_pred: Sequence[np.ndarray],
-    metrics: Optional[List[str]] = None,
-) -> Dict[str, float]:
+    metrics: list[str] | None = None,
+) -> dict[str, float]:
     """Computes metrics for real-valued labels given per time step.
 
     Each sample contributes the targets and predictions of its labeled time
diff --git a/pyhealth/trainer.py b/pyhealth/trainer.py
index efedd1b54..3726e2792 100644
--- a/pyhealth/trainer.py
+++ b/pyhealth/trainer.py
@@ -53,6 +53,22 @@ def set_logger(log_path: str) -> None:
 
 
 def get_metrics_fn(mode: str) -> Callable:
+    """Returns the metrics function for a model mode.
+
+    Args:
+        mode: One of ``"binary"``, ``"multiclass"``, ``"multilabel"``,
+            ``"regression"``, ``"timeseries_binary"`` or
+            ``"timeseries_regression"``.
+
+    Returns:
+        The function computing the metrics of that mode. The ``timeseries_*``
+        modes take one array per sample over its labeled steps.
+
+    Examples:
+        >>> from pyhealth.trainer import get_metrics_fn
+        >>> get_metrics_fn("timeseries_binary").__name__
+        'timeseries_binary_metrics_fn'
+    """
     if mode == "binary":
         return binary_metrics_fn
     elif mode == "multiclass":

From d0d0f5cd0a8d478d318556a150cd184890d90c4c Mon Sep 17 00:00:00 2001
From: Logiquo <fanyongda2012@gmail.com>
Date: Tue, 27 Oct 2026 03:36:18 -0400
Subject: [PATCH 3/3] Ignore NaN labels on unlabeled steps and test the
 per-step error paths

---
 pyhealth/models/base_model.py       | 19 +++++++---
 tests/core/test_timeseries_modes.py | 54 ++++++++++++++++++++++++++++-
 2 files changed, 67 insertions(+), 6 deletions(-)

diff --git a/pyhealth/models/base_model.py b/pyhealth/models/base_model.py
index 9b83241f1..40ca7a007 100644
--- a/pyhealth/models/base_model.py
+++ b/pyhealth/models/base_model.py
@@ -15,7 +15,8 @@
 logger = logging.getLogger(__name__)
 
 # Label kinds understood by the loss helpers, Trainer and the calibration methods.
-# The timeseries modes are labels per time step (see TimeseriesLabelProcessor).
+# The timeseries modes are labels per time step; they need an output processor
+# whose process() returns (value, mask), with mask 0 for steps without a label.
 VALID_MODES = (
     "binary",
     "multiclass",
@@ -37,7 +38,8 @@ def _masked_step_loss(
         loss_fn: Elementwise loss accepting ``reduction="none"``.
         logits: Per-step logits of shape (batch, steps, 1).
         target: Batched ``(value, mask)`` of shape (batch, label steps), where
-            padding and unlabeled steps have mask 0.
+            padding and unlabeled steps have mask 0. Values at those steps are
+            ignored and may be NaN.
     """
     # Temporary workaround: per-step labels arrive as a (value, mask) tuple so
     # the existing single-tensor label API stays unchanged. A more unified
@@ -49,8 +51,13 @@ def _masked_step_loss(
             f"{logits.shape[1]}; label and input must share the time grid."
         )
     logits = logits[:, : value.shape[1], 0]
+    labeled = mask > 0
+    # Replace unlabeled values before the loss so a NaN there reaches neither
+    # the loss nor its gradient, then drop those steps from the sum.
+    value = torch.where(labeled, value, torch.zeros_like(value))
     loss = loss_fn(logits, value, reduction="none")
-    return (loss * mask).sum() / mask.sum().clamp(min=1)
+    loss = torch.where(labeled, loss, torch.zeros_like(loss))
+    return loss.sum() / labeled.sum().clamp(min=1)
 
 
 class BaseModel(ABC, nn.Module):
@@ -312,8 +319,10 @@ def get_loss_function(self) -> Callable:
               over the labeled steps
 
         The timeseries losses take per-step logits of shape (batch, steps, 1)
-        and the batched ``(value, mask)`` label of a
-        :class:`~pyhealth.processors.TimeseriesLabelProcessor`.
+        and the batched ``(value, mask)`` label, each of shape (batch, steps),
+        produced by an output processor whose ``process()`` returns
+        ``(value, mask)`` per sample with mask 0 for steps without a label.
+        Values at masked steps are ignored and may be NaN.
 
         For binary and multilabel labels, a weight set with
         :meth:`set_pos_weight` is passed as ``pos_weight``.
diff --git a/tests/core/test_timeseries_modes.py b/tests/core/test_timeseries_modes.py
index 677884487..bb82858ca 100644
--- a/tests/core/test_timeseries_modes.py
+++ b/tests/core/test_timeseries_modes.py
@@ -42,12 +42,14 @@ def spatial(self):
         return (True,)
 
 
-def _dataset(mode="timeseries_binary", unlabeled=()):
+def _dataset(mode="timeseries_binary", unlabeled=(), unlabeled_value=None):
     rng = np.random.default_rng(0)
     samples = []
     for i, steps in enumerate([5, 3, 6, 4]):
         values = [float((i + h) % 2) for h in range(steps)]
         mask = [0.0 if (i, h) in unlabeled else 1.0 for h in range(steps)]
+        if unlabeled_value is not None:
+            values = [unlabeled_value if m == 0 else v for v, m in zip(values, mask)]
         samples.append(
             {
                 "patient_id": f"p{i}",
@@ -100,6 +102,56 @@ def test_predictions_do_not_see_future_steps(self):
         torch.testing.assert_close(before[:-1], after[:-1])
         self.assertFalse(torch.allclose(before[-1], after[-1]))
 
+    def test_nan_on_unlabeled_step_is_ignored(self):
+        for mode in ("timeseries_binary", "timeseries_regression"):
+            nan = _dataset(mode, unlabeled={(1, 0)}, unlabeled_value=float("nan"))
+            zero = _dataset(mode, unlabeled={(1, 0)}, unlabeled_value=0.0)
+            torch.manual_seed(0)
+            model = RNN(dataset=nan, hidden_dim=8, embedding_dim=8)
+            batch = next(iter(get_dataloader(nan, batch_size=4)))
+            out = model(**batch)
+            self.assertTrue(torch.isfinite(out["loss"]))
+            out["loss"].backward()
+            for parameter in model.parameters():
+                if parameter.grad is not None:
+                    self.assertTrue(torch.isfinite(parameter.grad).all())
+
+            # Same loss as with any other value on the unlabeled step.
+            loss_fn = model.get_loss_function()
+            logits = out["logit"].detach()
+            zero_batch = next(iter(get_dataloader(zero, batch_size=4)))
+            torch.testing.assert_close(
+                loss_fn(logits, batch["label"]), loss_fn(logits, zero_batch["label"])
+            )
+
+    def test_label_longer_than_input_is_rejected(self):
+        model = RNN(dataset=_dataset(), hidden_dim=8, embedding_dim=8)
+        logits = torch.zeros(1, 3, 1)
+        label = (torch.zeros(1, 4), torch.ones(1, 4))
+        with self.assertRaises(ValueError):
+            model.get_loss_function()(logits, label)
+
+    def test_time_series_inputs_must_share_the_time_grid(self):
+        samples = [
+            {
+                "patient_id": "p0",
+                "a": ([T0 + timedelta(hours=h) for h in range(3)], np.zeros((3, 1))),
+                "b": ([T0 + timedelta(hours=h) for h in range(5)], np.zeros((5, 1))),
+                "label": ([0.0] * 3, [1.0] * 3),
+            }
+        ]
+        dataset = create_sample_dataset(
+            samples=samples,
+            input_schema={"a": "timeseries", "b": "timeseries"},
+            output_schema={"label": "timeseries_binary"},
+            output_processors={"label": _StepLabelProcessor()},
+            dataset_name="timeseries_grids",
+        )
+        model = RNN(dataset=dataset, hidden_dim=8, embedding_dim=8)
+        batch = next(iter(get_dataloader(dataset, batch_size=1)))
+        with self.assertRaises(ValueError):
+            model(**batch)
+
     def test_bidirectional_is_rejected(self):
         with self.assertRaises(ValueError):
             RNN(dataset=_dataset(), bidirectional=True)