[Logiquo]

Summary

Adds support for labels given per time step to the models, the Trainer and the metrics, for tasks that predict at every step of a time series (for example at every hour of an ICU stay) rather than once per sample.

It introduces two model modes for such labels, timeseries_binary and timeseries_regression, next to binary, multiclass, multilabel and regression. They are separate modes on purpose: modules that assume one prediction per sample (calibration, prediction sets, interpretability) check for "binary" or "regression" and therefore do not silently accept per-step models.

A per-step label is batched as (value, mask), each of shape (batch, steps), with mask 0 for padding and for steps without a label. Values at masked steps are ignored and may be NaN. The mode is taken from the label's output_schema entry as usual.

Usage

The new modes need an output processor whose process() returns (value, mask) per sample, each of shape (steps,), with mask 0 for steps without a label. This PR does not include such a processor: until a built-in one is available, use a custom processor and pass it as a pre-fitted output processor, as the tests do:

dataset = create_sample_dataset(
    samples=samples,  # sample["label"] = (values, mask)
    input_schema={"vitals": "timeseries"},
    output_schema={"label": "timeseries_binary"},
    output_processors={"label": MyStepLabelProcessor()},
)

Changes

• BaseModel: the two modes in VALID_MODES; for them, get_loss_function returns BCE or MSE per step averaged over the steps whose mask is 1 (masked values, including NaN, reach neither the loss nor its gradient), and prepare_y_prob returns per-step probabilities of shape (batch, steps, 1).
• RNN: for these modes, predicts at every input step instead of from the last hidden state (features without a time axis, such as static tensors, are repeated over time), and rejects bidirectional=True, which would let a prediction see future steps.
• Trainer: inference returns, for these modes, one array per sample over its labeled steps (samples differ in length); get_metrics_fn maps them to the new metric functions.
• pyhealth.metrics: timeseries_binary_metrics_fn and timeseries_regression_metrics_fn pool the labeled steps of all samples and score them with binary_metrics_fn / regression_metrics_fn.

Existing modes and models are unchanged.

The tuple-based label handling is a temporary workaround that keeps the existing single-tensor label API; a more unified handling of padded and masked tensors is a follow-up (noted inline in the code).

Tests

tests/core/test_timeseries_modes.py uses a minimal mock label processor that returns (value, mask). It covers mode resolution, that the loss ignores unlabeled and padded steps, that a NaN label on an unlabeled step gives a finite loss and finite gradients, that predictions do not see future steps, the errors for a label longer than the input and for time series inputs on different grids, the bidirectional guard, Trainer inference and metrics for both modes, and the metric functions. All new executable lines are covered. Existing model, Trainer, calibration and embedding tests pass.

[github-actions]

Test Coverage

Total: 71.82% (18951 / 26386 lines)

Modified file	Coverage	Lines
pyhealth/metrics/__init__.py	100.00%	11 / 11
pyhealth/metrics/binary.py	75.00%	63 / 84
pyhealth/metrics/regression.py	55.56%	20 / 36
pyhealth/models/base_model.py	93.01%	133 / 143
pyhealth/models/rnn.py	89.18%	173 / 194
pyhealth/trainer.py	72.69%	197 / 271

Workflow run · d0d0f5c

[Logiquo]

/review

[github-actions]

Previous review

_{Reviewed commit afeaed4}

[github-actions]

Previous review

AI Review

suggestion: need change
decision: human review

Thank you for this contribution. Per-time-step label modes (timeseries_binary and timeseries_regression), with masked losses, a causal per-step RNN, Trainer inference and pooled metrics, are a useful addition. Your tests are thoughtful, and the existing modes keep their code paths. I found no prompt-injection text in the description or diff.

A maintainer will also take a look. The change modifies BaseModel (VALID_MODES, get_loss_function, prepare_y_prob) and Trainer.inference, which are shared core code. Existing modes look unchanged to me.

To get this merged:

1. In _masked_step_loss (base_model.py), (loss * mask).sum() gives a NaN loss if an unlabeled step holds NaN in value. Real processors often store NaN for missing labels. Make the loss robust, for example by zeroing the loss where mask is 0 with torch.where(mask > 0, loss, 0), or by replacing NaN values before computing the loss. Add a test with a NaN value on an unlabeled step.
2. The docstring in get_loss_function and the comment on VALID_MODES refer to a TimeseriesLabelProcessor, but I can't see one in this diff. Either add it, or reword the references to say that a processor returning (value, mask) is required. As it stands, the new modes cannot be used without a custom processor, so please state this clearly in the docs or the description.
3. Coverage is below the 80% bar for some changed files. metrics/binary.py is at 75%, metrics/regression.py at 55.56% and trainer.py at 72.69%. Much of the gap is probably in existing __main__ or untested code, but please confirm that your new lines are covered. The whole-repo total of 71.81% meets the 60% bar.

Optional suggestions:

• Raise a clear error in the timeseries metric functions when no sample has a labeled step. At present np.concatenate fails with an unclear message.
• Add a regression-masking test for the loss.
• List the new metric functions in docs/api/metrics.

_{Reviewed commit afeaed4}

[Logiquo]

/review

[github-actions]

_{Reviewed commit d0d0f5c}

[github-actions]

[medium] Per-step modes are accepted by every model but only RNN implements them

VALID_MODES and get_loss_function now resolve timeseries_binary / timeseries_regression for any BaseModel subclass. Only RNN.forward produces per-step logits and tuple labels. Other models (e.g. MLP, Transformer, TCN, which call self.get_loss_function()(logits, y_true) with a tuple label) fail at loss time with an AttributeError (.to() on a tuple) or an unpacking/IndexError inside _masked_step_loss, instead of a clear message. Fix: reject per-step modes in models that do not implement them, for example in BaseModel.mode resolution with an opt-in class flag that only RNN sets, and add a test that another model raises a clear ValueError. Alternatively document RNN as the only supported model.

[github-actions]

[low] Single-step time series are silently broadcast like static features

The per-step branch repeats any feature whose batch time length is 1 across all steps. A genuine time series that has a single step in this batch is indistinguishable from a static tensor, so a grid mismatch with a longer series passes without the intended 'same time grid' ValueError. Fix: decide which features are static from the input path (the x_dim_orig == 2 branch) rather than from the output length, and raise on any 3D feature whose length differs from the others. Add a test with one feature of length 1 and one of length 5.

[github-actions]

AI Review

suggestion: need change
decision: human review

Thank you for this contribution. Per-step labels with masked losses, per-sample Trainer inference and pooled metric functions are a useful addition, and the tests are thorough for the RNN path. I found no text aimed at reviewers in the PR.

A maintainer will also take a look, because the change modifies BaseModel (VALID_MODES, get_loss_function, prepare_y_prob) and Trainer.inference. Every model and the calibrators depend on these, so they count as the core part of the codebase.

Coverage: the whole repository is at 71.82%, which meets the 60% bar. Per file, base_model.py (93%) and rnn.py (89%) meet the 80% bar. binary.py (75%), regression.py (56%) and trainer.py (73%) are below it. Their uncovered lines look like existing code, such as the __main__ blocks and unrelated Trainer paths, and not the new lines, which you say are all covered. That is not a blocker by itself.

To get this merged:

1. BaseModel now accepts timeseries_binary and timeseries_regression for every model, but only RNN implements the per-step forward. With any other model, these modes fail at loss time with an obscure AttributeError or unpacking error. Please either reject the modes with a clear error message in the models that don't support them, or add a guard in BaseModel. Add a test that covers this.

Optional suggestions:

• In rnn.py, a time series whose time length is 1 is treated as static and broadcast. A real grid mismatch can slip through when one feature has a single step. Distinguishing true static features from length-1 series would make the check stricter.
• TemperatureScaling doesn't check the mode and fails with an AttributeError on the tuple label. This contradicts the description's claim that calibration modules don't silently accept per-step models. A clear mode check would fix that.
• The interpretability removal metrics only warn and continue for these modes. Rejecting them explicitly would be safer.
• Consider a test for Trainer.train and for a single-step sample.

_{Reviewed commit d0d0f5c}

[jhnwu3]

I feel that this is a long time coming which is we should remove the need for model modes and just assume it's tied to the output processor here and that as long as our forward pass produces a loss, our trainers should be compatible iirc. What do you think?

[jhnwu3]

lgtm