wav2sleep reproduction

[JollyJiant]

Bronze Frazer
bfrazer2
bfrazer2@illinois.edu

Dataset + Task + Model
"wav2sleep: A Unified Multi-Modal Approach to Sleep Stage Classification from Physiological Signals"
https://arxiv.org/pdf/2411.04644

Description
This PR integrates the wav2sleep model (arXiv:2411.04644) into PyHealth for automated sleep stage classification from polysomnography (PSG) biosignals. It adds a dataset, task, and model following the standard PyHealth pattern.

Wav2SleepDataset ingests EDF recordings and annotations from 7 sleepdata.org cohorts (SHHS, MESA, WSC, CHAT, CFS, CCSHS, MROS), auto-generating a metadata CSV on first load.

Wav2SleepStaging normalizes channel names across datasets, remaps annotations to 4 classes (Wake, Light, Deep, REM), and produces an availability mask for missing modalities.

Wav2Sleep implements the three-stage architecture: per-modality CNN Signal Encoders → a Transformer-based Epoch Mixer with stochastic modality masking → a dilated conv Sequence Mixer, ending in a 4-class linear classifier.

File Guide

• pyhealth/datasets/wav2sleep.py — Wav2SleepDataset
• pyhealth/datasets/configs/wav2sleep.yaml — Dataset YAML config
• pyhealth/tasks/wav2sleep_staging.py — Wav2SleepStaging
• pyhealth/models/wav2sleep.py — Wav2Sleep and all sub-modules
• examples/wav2sleep_sleep_staging_wav2sleep.py — End-to-end demo with mock data
• tests/test_wav2sleep.py — Unit tests
• docs/api/ — RST docs for dataset, task, and model
• pyhealth/{datasets,tasks,models}/init.py — Class registration