sample_rate.py

#!/usr/bin/python3

from pathlib import Path
import numpy as np
from scipy.misc import electrocardiogram
from enum import IntEnum
from pyedflib import highlevel


class SleepStage(IntEnum):
    """
    Mapping of AASM sleep stages to integers.

    To facilitate hypnogram plotting, values start with zero and increase with wakefulness.
    """

    UNDEFINED = 0
    N3 = 1
    N2 = 2
    N1 = 3
    REM = 4
    WAKE = 5


def _dummy_nsrr_edf(filename: str, hours: float, ecg_channel: str):
    ECG_FS = 360
    ecg_5_min = electrocardiogram()
    seconds = int(hours * 60 * 60)
    ecg = np.tile(ecg_5_min, int(np.ceil(seconds / 300)))[np.newaxis, : seconds * ECG_FS]
    signal_headers = highlevel.make_signal_headers([ecg_channel], sample_frequency=ECG_FS)
    highlevel.write_edf(filename, ecg, signal_headers)


def _dummy_nsrr_xml(filename: str, hours: float, random_state: int):
    EPOCH_LENGTH = 30
    STAGES = [
        "Wake|0",
        "Stage 1 sleep|1",
        "Stage 2 sleep|2",
        "Stage 3 sleep|3",
        "Stage 4 sleep|4",
        "REM sleep|5",
        "Unscored|9",
    ]

    rng = np.random.default_rng(random_state)

    with open(filename, "w") as xml_file:
        xml_file.write(
            '<?xml version="1.0" encoding="UTF-8" standalone="no"?>\n'
            "<PSGAnnotation>\n"
            f"<EpochLength>{EPOCH_LENGTH}</EpochLength>\n"
            "<ScoredEvents>\n"
            "<ScoredEvent>\n"
            "<EventType/>\n"
            "<EventConcept>Recording Start Time</EventConcept>\n"
            "<ClockTime>01.01.85 20.29.59</ClockTime>\n"
            "</ScoredEvent>\n",
        )
        record_duration = hours * 60 * 60
        start = 0
        while True:
            if start > record_duration:
                break
            epoch_duration = rng.choice(np.arange(4, 21)) * EPOCH_LENGTH
            stage = rng.choice(STAGES)
            xml_file.write(
                "<ScoredEvent>\n"
                "<EventType>Stages|Stages</EventType>\n"
                f"<EventConcept>{stage}</EventConcept>\n"
                f"<Start>{start:.1f}</Start>\n"
                f"<Duration>{epoch_duration:.1f}</Duration>\n"
                "</ScoredEvent>\n",
            )
            start += epoch_duration

        xml_file.write(
            "</ScoredEvents>\n" "</PSGAnnotation>\n",
        )


def _create_dummy_shhs(data_dir: str, durations: list[float], random_state: int = 42):
    DB_SLUG = "shhs"
    ANNOTATION_DIRNAME = "polysomnography/annotations-events-nsrr"
    EDF_DIRNAME = "polysomnography/edfs"
    CSV_DIRNAME = "datasets"

    db_dir = Path(data_dir).expanduser() / DB_SLUG
    annotations_dir = db_dir / ANNOTATION_DIRNAME
    edf_dir = db_dir / EDF_DIRNAME
    csv_dir = db_dir / CSV_DIRNAME

    for directory in (annotations_dir, edf_dir):
        for visit in ("shhs1", "shhs2"):
            (directory / visit).mkdir(parents=True, exist_ok=True)
    csv_dir.mkdir(parents=True, exist_ok=True)

    for i, hours in enumerate(durations):
        for visit in ("shhs1", "shhs2"):
            record_id = f"{visit}/{visit}-20{i:04}"
            _dummy_nsrr_edf(f"{edf_dir}/{record_id}.edf", hours, ecg_channel="ECG")
            _dummy_nsrr_xml(f"{annotations_dir}/{record_id}-nsrr.xml", hours, random_state)

    with open(csv_dir / "shhs1-dataset-0.0.0.csv", "w") as csv:
        csv.write("nsrrid,age_s1,gender,weight\n")
        for i in range(len(durations)):
            csv.write(f"2{i:05},55,1,77\n")
    with open(csv_dir / "shhs2-dataset-0.0.0.csv", "w") as csv:
        csv.write("nsrrid,age_s2,gender,weight\n")
        for i in range(len(durations)):
            csv.write(f"2{i:05},61,2,\n")


def main():
    durations = [0.1, 0.2]  # hours
    valid_stages = {int(s) for s in SleepStage}
    
    _create_dummy_shhs(data_dir="/tmp", durations=durations)


if __name__ == '__main__':
    main()