Add WatchSleepNet full pipeline: IBISleepDataset + SleepStagingIBI + WatchSleepNet

[gregschrock]

Contributor

• Manasa Hegde (manasah3@illinois.edu)
• Janani Jayanth (jananij2@illinois.edu)
• Greg Schrock (gregs2@illinois.edu)

Contribution Type

Option 4: Full Pipeline — Dataset + Task + Model

Paper

Wang et al. (2025). WatchSleepNet: A Scalable Deep Learning Model for Wearable Sleep Staging. CHIL 2025, PMLR 287:1–20.
https://proceedings.mlr.press/v287/wang25a.html

Description

Implements the complete PyHealth pipeline for IBI-based sleep staging from consumer wearables, following the WatchSleepNet paper methodology. The paper introduces a deep learning architecture for sleep stage classification from inter-beat interval (IBI) signals derived from wrist-worn PPG/ECG wearables (e.g., Apple Watch, Empatica E4), demonstrating competitive accuracy against PSG-supervised baselines while remaining deployable on low-power consumer devices.

This PR adds three new components following the established SleepEDFDataset / SleepStagingSleepEDF pattern:

IBISleepDataset — a single, cohort-agnostic BaseDataset subclass that unifies DREAMT, SHHS, and MESA under one API via a source parameter, in contrast to the existing per-cohort classes which each load raw, cohort-specific file formats. Operates on a common intermediate NPZ representation (25 Hz IBI time series + per-sample stage labels) produced by the included preprocessing scripts.

SleepStagingIBI — BaseTask subclass following the SleepStagingSleepEDF API convention. Converts patient records into per-epoch sample dicts: 750-sample (30 s × 25 Hz) IBI windows with mapped sleep stage labels. Supports both 3-class (Wake / NREM / REM) and 5-class (W / N1 / N2 / N3 / REM) via num_classes.

WatchSleepNet — BaseModel subclass implementing the paper's architecture: 4-block 1D ResNet → dilated TCN → BiLSTM → Multi-head Attention → global average pooling → linear classifier. Accepts signal tensor of shape (B, 750) and returns {"loss", "y_prob", "y_true"}. Supports configurable num_classes, hidden_dim, and attn_heads.

File Guide

Core implementation

• pyhealth/datasets/ibi_sleep.py — IBISleepDataset
• pyhealth/datasets/configs/ibi_sleep.yaml — dataset YAML schema
• pyhealth/tasks/sleep_staging_ibi.py — SleepStagingIBI
• pyhealth/models/watchsleepnet.py — WatchSleepNet, ResidualBlock

Example / ablation

• examples/ibi_sleep_staging_ibi_watchsleepnet.py — end-to-end two-phase transfer learning pipeline; run with --synthetic for no real data required

Preprocessing scripts

• examples/preprocess_dreamt_to_ibi.py — Empatica E4 BVP CSV → NPZ (neurokit2)
• examples/preprocess_shhs_to_ibi.py — SHHS EDF + profusion XML → NPZ (biosppy)
• examples/preprocess_mesa_to_ibi.py — MESA EDF → NPZ (same ECG pipeline as SHHS)

Tests (synthetic data only, complete in milliseconds)

• tests/core/test_ibi_sleep_dataset.py — 12 tests: all three sources, patient ID
  extraction, AHI NaN passthrough, corrupt NPZ skipping, missing AHI key, dev mode, empty directory, set_task integration (3-class and 5-class)
• tests/core/test_sleep_staging_ibi.py — 9 tests: label mapping, invalid label skipping, max-epoch cap, short-record empty return, signal shape, wrong fs error
• tests/core/test_watchsleepnet.py — 8 tests: instantiation, forward shape, loss with/without labels, wrong input length, gradient flow, output dict keys

Docs

• docs/api/datasets/pyhealth.datasets.IBISleepDataset.rst
• docs/api/tasks/pyhealth.tasks.SleepStagingIBI.rst
• docs/api/models/pyhealth.models.WatchSleepNet.rst
• Index entries added to docs/api/datasets.rst, docs/api/tasks.rst, docs/api/models.rst