[Logiquo]

Summary

Adds the five MIMIC-IV tasks of YAIB (Yet Another ICU Benchmark, van de Water et al., ICLR 2024) to pyhealth.tasks. The three hourly tasks use the timeseries_binary / timeseries_regression label modes and processors.

Task	Class	Label
ICU mortality after 24 hours	YAIBMortalityMIMIC4	binary
Kidney function (day 2 creatinine)	YAIBKidneyFunctionMIMIC4	regression
Remaining length of stay	YAIBLengthOfStayMIMIC4	timeseries_regression (per hour)
Acute kidney injury (KDIGO)	YAIBAKIMIMIC4	timeseries_binary (per hour)
Sepsis (Sepsis-3)	YAIBSepsisMIMIC4	timeseries_binary (per hour)

YAIB builds its cohorts with the YAIB-cohorts R scripts on top of the ricu R package. The tasks port that logic: the YAIB base cohort, the 48 hourly dynamic and 4 static features (ricu's item IDs, unit conversions and plausibility bounds, hourly binning from ICU admission, ricu's mapping of admission-level labs onto ICU stays), the task-specific exclusions and labels, and for AKI and sepsis the KDIGO staging, the SOFA score and the suspected-infection rule of the ricu version YAIB pins (prockenschaub/ricu-package at 0969a52, 0.6.0). The tasks declare YAIB's preprocessing in their input_schema; fit it on the training patients with set_task(task, split=...).

The MIMIC-IV EHR config gains the transfers, microbiologyevents, chartevents, outputevents, inputevents and procedureevents tables, which are only loaded when listed in tables, and the deathtime column of admissions, which is always loaded (see Test fixture below).

from pyhealth.datasets import MIMIC4EHRDataset, PatientSplit
from pyhealth.tasks import YAIBSepsisMIMIC4

dataset = MIMIC4EHRDataset(
    root="/path/to/mimiciv/3.1",
    tables=["labevents", "chartevents", "outputevents", "transfers",
            "inputevents", "procedureevents", "microbiologyevents"],
    num_workers=4,
)
train, val, test = dataset.set_task(
    YAIBSepsisMIMIC4(), split=PatientSplit(ratios=(0.7, 0.1, 0.2), seed=0), num_workers=4
)

Verification against YAIB

YAIB-cohorts was run unchanged in its Docker image on the MIMIC-IV demo (2.2), and each task's cohort, labels and features were compared with its output. With ricu 0.6.0's definitions (see the deviations below), every task matches exactly:

Task	Stays	Labels	Hourly feature values	Static values
Mortality	96 / 96	0 differing	0 of 115,200 differing	0 differing
Kidney function	73 / 73	0 differing	0 of 87,600 differing	0 differing
Length of stay	137 / 137	0 of 10,161 differing	0 of 487,728 differing	0 differing
AKI	117 / 117	0 of 5,432 differing	0 of 260,736 differing	0 differing
Sepsis	117 / 117	0 of 6,581 differing	0 of 315,888 differing	0 differing

For sepsis, the intermediate values were also compared with ricu's own output for every demo ICU stay: all 28,714 hourly SOFA values, the continuous antibiotic starts, the sample hours, the suspected infection times and the onsets are identical.

On full MIMIC-IV 3.1 (YAIB used 2.2, which has fewer ICU stays), the cohorts scale consistently with the paper:

Task	YAIB paper (2.2)	This PR (3.1)
Mortality	52,045 stays, 7.3% died	64,675 stays, 7.6% died
Kidney function	35,657 stays	44,605 stays
Length of stay	75,652 stays	92,972 stays
AKI	66,032 stays, 41.7% with AKI	80,431 stays, 44.5% with AKI
Sepsis	67,056 stays, 5.6% with sepsis	83,050 stays, 4.7% with sepsis

The verification scripts and the YAIB-cohorts setup are not part of this PR.

Deviations from YAIB

Two ricu 0.6.0 bugs that YAIB inherits are fixed and documented in the task docstrings. With the buggy behaviour swapped back in, the results above match YAIB exactly.

• o2sat: ricu 0.6.0 lists the blood gas item 50817 under chartevents (it is a labevents item), so arterial oxygen saturation never matches, and the charted arterial saturation (220227) is missing. Both items exist in MIMIC-IV 2.2, and ricu's o2sat for eICU and AUMCdb already combined pulse oximetry with arterial saturation; ricu fixed this in July 2024. On the demo, this changes 57 of 115,200 mortality feature values, all o2sat.
• SOFA grid: ricu 0.6.0's fill_gaps takes the duration of a stay's time range as its end, so the hourly SOFA grid runs to last - first instead of the last hour of data. The SOFA score is now computed up to the last hour; on the demo this changes no sepsis onset.

Other YAIB details are kept as they are, for example that AKI urine output is not divided by weight for stays without any weight record (dividing by 75 kg instead changes 3 demo stays relative to YAIB).

Tests

tests/core/test_yaib_mimic4.py covers each task on synthetic patients: the base cohort exclusions, labels, unit conversions, lab attribution to stays, ICU death attribution, the kidney function window, the AKI creatinine and urine stages, and the sepsis suspected infection, PaO2/FiO2 pairing, GCS, 24-hour urine, range merging and onset.

Test fixture

The mortality task needs the time of death, so the config now also loads admissions.deathtime, and the dataset loader requires every configured column to exist. The real MIMIC-IV hosp/admissions table has this column, right after dischtime, in both MIMIC-IV 2.2 and 3.1:

	Admissions	In-hospital deaths (hospital_expire_flag = 1)	With deathtime
MIMIC-IV demo 2.2	275	15	15
MIMIC-IV 3.1	546,028	11,801	11,790

deathtime is only set for in-hospital deaths. The hand-made fixture in test-resources/core/mimic4demo/hosp/admissions.csv omitted the column, so every test loading admissions from it failed with KeyError: "['admissions/deathtime'] not in index" (24 tests in test_mimic4_dka, test_mimic4_los, test_mimic4_readmission_prediction, test_stagenet_task_leakage_prevention and test_drug_recommendation_atc3). This PR adds the column to the fixture at the same position as in MIMIC-IV: the two deceased rows get their dischtime (in MIMIC-IV 3.1, deathtime equals dischtime for 93% of in-hospital deaths) and the others are empty. None of those tasks read deathtime, so their results are unchanged.

[github-actions]

Test Coverage

Total: 73.64% (20110 / 27310 lines)

Modified file	Coverage	Lines
pyhealth/tasks/__init__.py	89.47%	34 / 38
pyhealth/tasks/yaib_mimic4.py	91.43%	512 / 560

Workflow run · 6ac17e7

[Logiquo]

/review

[github-actions]

Previous review

_{Reviewed commit a504416}

[github-actions]

Previous review

AI Review

suggestion: need change
decision: human review

Thank you for this contribution. It adds the five YAIB MIMIC-IV benchmark tasks, with a very thorough parity check against YAIB-cohorts and clear documentation of the deviations. Coverage is 85.15% for yaib_mimic4.py (above the 80% bar) and 72.69% for the whole repository (above the 60% bar).

A maintainer will take a look as well. The change adds a required column (admissions.deathtime) and six new tables to the shared mimic4_ehr.yaml config, which every MIMIC4EHRDataset user loads, so the loader's requirements change for existing users. That touches a core area.

To get this merged:

1. YAIBMortalityMIMIC4 leaves _EXTRA_TABLES empty, but _icu_deaths needs the transfers table. If a user omits transfers, pre_filter passes and the task then fails with a raw polars column error for any patient with an in-hospital death. Set _EXTRA_TABLES = ("transfers",) for the mortality task so it raises the same clear ValueError as the others. Add a test for this.
2. All tests call task.__call__ on hand-built Patient objects. Please add at least one integration test that goes through the dataset load path for the new yaml tables and set_task, including processor fitting. For example, use the mimic4demo fixture, as the existing MIMIC-IV task tests do. This would check the new config entries and the declared input_schema.

Optional suggestions:

• State in the AKI and sepsis docstrings that onsets before ICU admission (negative hours) also lead to exclusion.
• Requiring deathtime means admissions files without that column will fail to load. Consider mentioning this in the docs, or making the column optional.

_{Reviewed commit a504416}

[Logiquo]

/review

[github-actions]

Previous review

_{Reviewed commit 79dfb54}

[github-actions]

Previous review

AI Review

suggestion: need change
decision: human review

Thank you for this contribution. Porting the five YAIB MIMIC-IV cohorts, with exact-match checks against YAIB-cohorts on the demo data, is a substantial addition. Coverage is above the bar: yaib_mimic4.py is at 86.4% (bar 80%), and the repository total is 73.3% (bar 60%).

A maintainer will also take a look. The change edits the shared mimic4_ehr.yaml, so every MIMIC4EHRDataset load now requires an admissions.deathtime column. Real MIMIC-IV has that column, but this changes the loading behavior for existing users. The shared test fixture changed too.

To get this merged:

1. Null care unit in the ICU death check. In _icu_deaths, ~pl.col("careunit").is_in(icu_wards).fill_null(False) treats a null careunit on the last ward before discharge as non-ICU. That drops the death from the label. The docstring says a death is non-ICU only when the ward is known and not an ICU ward. Either make null count as unknown, so the death stays an ICU death, or change the docstring to describe the current behavior. If the current behavior is what ricu does, please say so in a comment. Please also add a test for the null case.
2. Missing test for YAIBLengthOfStayMIMIC4. Please add a small test. It should check that the label is min(168, end - hour) on the same grid as dynamic, and that stays are capped at 7 days.

Optional suggestions:

• Add direct tests for the vasopressor, ventilation and intubation paths feeding SOFA.
• Add a docs entry or a short example for the new tasks.
• Add a test that loads the updated YAML config and the demo fixture, so the deathtime requirement is exercised.

_{Reviewed commit 79dfb54}

[Logiquo]

/review

[github-actions]

_{Reviewed commit 8baeb1e}

[github-actions]

[medium] New MIMIC-IV tables are not exercised by any dataset-level test

The six tables added here (transfers, microbiologyevents, chartevents, outputevents, inputevents, procedureevents) are only read through hand-built Patient objects in tests/core/test_yaib_mimic4.py. No test loads them through MIMIC4EHRDataset, runs YAIB's pre_filter through set_task, or checks the config column names against a real table layout. The demo fixture in test-resources/core/mimic4demo has none of these tables, while existing MIMIC-IV task tests (tests/core/test_mimic4_los.py, tests/core/test_mimic4_dka.py) run set_task on that fixture. Fix: add minimal CSV fixtures for these tables (or extend the demo fixture) and one set_task test per task family, including the missing-table ValueError path through the real loader.

[github-actions]

[low] o2sat includes the SpO2 desaturation alarm limit (226253)

Item 226253 is 'SpO2 Desat Limit', which the docstring itself says is a monitor alarm threshold rather than a measurement. Including it in the o2sat feature mixes a setting into the measured series. The PR says it is kept for consistency with ricu, so this is a documented choice, not a bug. Consider dropping it, or at least marking the choice in the feature table, once the YAIB-match claim is settled.

[github-actions]

[low] New public tasks have no API docs entry

Five public task classes are exported here, but no docs/api/tasks/*.rst file or docs/api/tasks.rst toctree entry is added, unlike other task modules such as InHospitalMortalityMIMIC4. Add an autoclass page per task (or one for the module) so the new tasks appear in the API reference.

[github-actions]

AI Review

suggestion: approve
decision: human review

Thank you for this contribution. It adds the five YAIB MIMIC-IV benchmark tasks (mortality, kidney function, remaining length of stay, AKI and sepsis), ported carefully from ricu and YAIB-cohorts. Your write-up of the verification against YAIB output and of the two deliberate deviations is very clear.

The code looks correct and complete as far as I can tell, and it does what the description says. It doesn't appear to duplicate existing PyHealth functionality. Coverage meets the bar: pyhealth/tasks/yaib_mimic4.py is at 91.43% (bar 80%), pyhealth/tasks/__init__.py is at 89.47%, and the whole repository is at 73.64% (bar 60%).

A maintainer will take a look before this is merged. The change adds admissions.deathtime to the shared MIMIC-IV EHR config as a column that is always loaded, and the dataset loader requires every configured column to exist. That changes what every MIMIC-IV EHR dataset needs from its input data, and it is why the test fixture had to change. This is a shared contract, so a maintainer should confirm it is acceptable. I can't promise the outcome.

Optional suggestions:

• Add a test that loads the new config tables through MIMIC4EHRDataset, then runs set_task with pre_filter and processor fitting. The current tests only call task(patient) on hand-built patients. This would exercise the dataset path end to end.
• Add a docs/api entry and an example script for the new tasks.
• o2sat includes item 226253, the SpO2 desat limit, which your docstring notes is an alarm threshold rather than a measurement. You keep it for consistency with ricu, which is reasonable. Please consider whether to keep it.

_{Reviewed commit 8baeb1e}

[jhnwu3]

lgtm