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..9f73bc576 100644 --- a/pyhealth/metrics/binary.py +++ b/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", diff --git a/pyhealth/metrics/regression.py b/pyhealth/metrics/regression.py index 91f87028f..2d0d02ed5 100644 --- a/pyhealth/metrics/regression.py +++ b/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) diff --git a/pyhealth/models/base_model.py b/pyhealth/models/base_model.py index 8b9683763..40ca7a007 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,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 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..3726e2792 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 @@ -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) diff --git a/tests/core/test_timeseries_modes.py b/tests/core/test_timeseries_modes.py new file mode 100644 index 000000000..bb82858ca --- /dev/null +++ b/tests/core/test_timeseries_modes.py @@ -0,0 +1,190 @@ +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=(), 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}", + "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_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) + + 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()