Add timeseries_binary and timeseries_regression model modes

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",
 ]

pyhealth/metrics/binary.py

@@ -1,3 +1,4 @@
+from collections.abc import Sequence
 from typing import Dict, List, Optional
 
 import numpy as np
@@ -147,6 +148,46 @@ def binary_metrics_fn(
     return output
 
 
+
+def timeseries_binary_metrics_fn(
+    y_true: Sequence[np.ndarray],
+    y_prob: Sequence[np.ndarray],
+    metrics: list[str] | None = 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",

pyhealth/metrics/regression.py

@@ -1,3 +1,4 @@
+from collections.abc import Sequence
 from typing import Dict, List, Optional
 
 import numpy as np
@@ -70,6 +71,42 @@ def regression_metrics_fn(
     return output
 
 
+
+def timeseries_regression_metrics_fn(
+    y_true: Sequence[np.ndarray],
+    y_pred: Sequence[np.ndarray],
+    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
+    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)

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,49 @@
 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; they need an output processor
+# whose process() returns (value, mask), with mask 0 for steps without a label.
+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. 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
+    # 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]
+    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")
+    loss = torch.where(labeled, loss, torch.zeros_like(loss))
+    return loss.sum() / labeled.sum().clamp(min=1)
 
 
 class BaseModel(ABC, nn.Module):
@@ -38,7 +81,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 +313,16 @@ 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, 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``.
@@ -298,6 +352,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 +373,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 +397,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

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}

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
 
@@ -51,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":
@@ -59,6 +77,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 +357,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 +369,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)