Activity count feature addition

src/sleepecg/io/sleep_readers.py

@@ -16,6 +16,8 @@
 from xml.etree import ElementTree
 
 import numpy as np
+import pandas as pd
+from pandas import read_csv
 
 from sleepecg.config import get_config_value
 from sleepecg.heartbeats import detect_heartbeats
@@ -86,6 +88,8 @@ class SleepRecord:
         Times of heartbeats relative to recording start in seconds, by default `None`.
     subject_data : SubjectData, optional
         Dataclass containing subject data (such as gender or age), by default `None`.
+    activity_counts: np.ndarray, optional
+        Activity counts according to actiwacth actigraphy, by default `None`.
     """
 
     sleep_stages: np.ndarray | None = None
@@ -94,12 +98,14 @@ class SleepRecord:
     recording_start_time: datetime.time | None = None
     heartbeat_times: np.ndarray | None = None
     subject_data: SubjectData | None = None
+    activity_counts: np.ndarray | None = None
 
 
 class _ParseNsrrXmlResult(NamedTuple):
     sleep_stages: np.ndarray
     sleep_stage_duration: int
     recording_start_time: datetime.time
+    recording_duration: float
 
 
 def _parse_nsrr_xml(xml_filepath: Path) -> _ParseNsrrXmlResult:
@@ -120,6 +126,8 @@ def _parse_nsrr_xml(xml_filepath: Path) -> _ParseNsrrXmlResult:
         Duration of each sleep stage in seconds.
     recording_start_time : datetime.time
         Time at which the recording was started.
+    recording_duration: float
+        Duration of the recording in seconds.
 
     """
     STAGE_MAPPING = {
@@ -139,6 +147,8 @@ def _parse_nsrr_xml(xml_filepath: Path) -> _ParseNsrrXmlResult:
         raise RuntimeError(f"EpochLength not found in {xml_filepath}.")
     epoch_length = int(epoch_length)
 
+    recording_duration = float(root.find(".//ScoredEvent").findtext(".//Duration"))
+
     start_time = None
     annot_stages = []
 
@@ -157,9 +167,7 @@ def _parse_nsrr_xml(xml_filepath: Path) -> _ParseNsrrXmlResult:
         raise RuntimeError(f"'Recording Start Time' not found in {xml_filepath}.")
 
     return _ParseNsrrXmlResult(
-        np.array(annot_stages, dtype=np.int8),
-        epoch_length,
-        start_time,
+        np.array(annot_stages, dtype=np.int8), epoch_length, start_time, recording_duration
     )
 
 
@@ -169,6 +177,7 @@ def read_mesa(
     offline: bool = False,
     keep_edfs: bool = False,
     data_dir: str | Path | None = None,
+    activity_source: str = None,
 ) -> Iterator[SleepRecord]:
     """
     Lazily read records from [MESA](https://sleepdata.org/datasets/mesa).
@@ -197,6 +206,10 @@ def read_mesa(
     data_dir : str | pathlib.Path, optional
         Directory where all datasets are stored. If `None` (default), the value will be
         taken from the configuration.
+    activity_source :{'actigraphy', 'cached', None'},  str, optional
+        If None (default), actigraphy data will not be downloaded. If 'actigraphy',
+        download actigraphy data from mesa dataset. If 'cached', get the cached activity
+        counts.
 
     Yields
     ------
@@ -208,7 +221,10 @@ def read_mesa(
     DB_SLUG = "mesa"
     ANNOTATION_DIRNAME = "polysomnography/annotations-events-nsrr"
     EDF_DIRNAME = "polysomnography/edfs"
+    ACTIVITY_DIRNAME = "actigraphy/"
+    OVERLAP_DIRNAME = "overlap/"
     HEARTBEATS_DIRNAME = "preprocessed/heartbeats"
+    ACTIVITY_COUNTS_DIRNAME = "preprocessed/activity_counts"
     RPOINTS_DIRNAME = "polysomnography/annotations-rpoints"
 
     GENDER_MAPPING = {0: Gender.FEMALE, 1: Gender.MALE}
@@ -220,6 +236,13 @@ def read_mesa(
             f"possible options: {heartbeats_source_options}"
         )
 
+    activitiy_source_options = {None, "cached", "actigraphy"}
+    if activity_source not in activitiy_source_options:
+        raise ValueError(
+            f"Invalid value for parameter `activity_source`: {activity_source}, "
+            f"possible options: {activitiy_source_options}"
+        )
+
     if data_dir is None:
         data_dir = get_config_value("data_dir")
 
@@ -231,6 +254,13 @@ def read_mesa(
     for directory in (annotations_dir, edf_dir, heartbeats_dir):
         directory.mkdir(parents=True, exist_ok=True)
 
+    if activity_source is not None:
+        activity_dir = db_dir / ACTIVITY_DIRNAME
+        activity_counts_dir = db_dir / ACTIVITY_COUNTS_DIRNAME
+        overlap_dir = db_dir / OVERLAP_DIRNAME
+        for directory in (activity_dir, activity_counts_dir, overlap_dir):
+            directory.mkdir(parents=True, exist_ok=True)
+
     if not offline:
         download_url = _get_nsrr_url(DB_SLUG)
 
@@ -271,10 +301,34 @@ def read_mesa(
             shallow=True,
         )
         checksums.update(rpoints_files)
+
+        if activity_source is not None:
+            activity_files = _list_nsrr(
+                DB_SLUG, ACTIVITY_DIRNAME, pattern=f"mesa-sleep-{records_pattern}.csv"
+            )
+            checksums.update(activity_files)
+            overlap_filename, overlap_checksum = _list_nsrr(
+                db_slug="mesa", subfolder="overlap", shallow=True
+            )[0]
+            overlap_filepath = db_dir / overlap_filename
+            _download_nsrr_file(
+                download_url + overlap_filename,
+                target_filepath=overlap_filepath,
+                checksum=overlap_checksum,
+            )
+
     else:
         subject_data_filepath = next((db_dir / "datasets").glob("mesa-sleep-dataset-*.csv"))
         xml_paths = annotations_dir.glob(f"mesa-sleep-{records_pattern}-nsrr.xml")
         requested_records = sorted([file.stem[:-5] for file in xml_paths])
+        if activity_source == "actigraphy":
+            overlap_filename = "mesa-actigraphy-psg-overlap.csv"
+            overlap_filepath = overlap_dir / overlap_filename
+            if not overlap_filepath.is_file():
+                raise RuntimeError(
+                    "Overlap file not found, make sure it is in the correct directory "
+                    "overlap/mesa-actigraphy-psg-overlap.csv."
+                )
 
     subject_data_array = np.loadtxt(
         subject_data_filepath,
@@ -291,6 +345,9 @@ def read_mesa(
             age=age,
         )
 
+    if activity_source == "actigraphy":
+        overlap_data = read_csv(overlap_filepath)
+
     for record_id in requested_records:
         heartbeats_file = heartbeats_dir / f"{record_id}.npy"
         if heartbeats_source == "annotation":
@@ -350,14 +407,77 @@ def read_mesa(
 
         parsed_xml = _parse_nsrr_xml(xml_filepath)
 
-        yield SleepRecord(
-            sleep_stages=parsed_xml.sleep_stages,
-            sleep_stage_duration=parsed_xml.sleep_stage_duration,
-            id=record_id,
-            recording_start_time=parsed_xml.recording_start_time,
-            heartbeat_times=heartbeat_times,
-            subject_data=subject_data[record_id],
-        )
+        if activity_source is not None:
+            if activity_source == "cached":
+                activity_counts_file = activity_dir / f"{record_id}-activity-counts.npy"
+                if not activity_counts_file.is_file():
+                    print(f"Skipping {record_id} due to missing cached activity counts.")
+                    continue
+                activity_counts = np.load(activity_counts_file)
+            else:
+                activity_filename = ACTIVITY_DIRNAME + f"/{record_id}.csv"
+                activity_filepath = db_dir / activity_filename
+                if not offline:
+                    _download_nsrr_file(
+                        download_url + activity_filename,
+                        activity_filepath,
+                        checksums[activity_filename],
+                    )
+
+                activity_data = pd.read_csv(activity_filepath)
+
+                recording_start_time = parsed_xml.recording_start_time
+                recording_duration = parsed_xml.recording_duration
+                recording_end_time = datetime.datetime.combine(
+                    datetime.datetime.today(), recording_start_time
+                ) + datetime.timedelta(seconds=recording_duration)
+
+                recording_end_time_seconds = recording_end_time.second
+                rounding_seconds = (
+                    (30 - recording_end_time_seconds % 30)
+                    if recording_end_time_seconds % 30 >= 15
+                    else -(recording_end_time_seconds % 30)
+                )
+
+                recording_end_time = recording_end_time + datetime.timedelta(
+                    seconds=rounding_seconds
+                )
+                recording_end_time = recording_end_time.strftime("%H:%M:%S").lstrip("0")
+
+                mesa_id = activity_data["mesaid"].iloc[0]
+
+                start_line = overlap_data.loc[
+                    overlap_data["mesaid"] == mesa_id, "line"
+                ].iloc[0]
+                end_line = activity_data.loc[
+                    activity_data["linetime"] == recording_end_time, "line"
+                ].iloc[0]
+
+                activity_counts = activity_data[
+                    (activity_data["line"] >= start_line)
+                    & (activity_data["line"] <= end_line)
+                ]["activity"].to_list()
+
+                activity_counts = np.array(activity_counts)
+
+            yield SleepRecord(
+                sleep_stages=parsed_xml.sleep_stages,
+                sleep_stage_duration=parsed_xml.sleep_stage_duration,
+                id=record_id,
+                recording_start_time=parsed_xml.recording_start_time,
+                heartbeat_times=heartbeat_times,
+                subject_data=subject_data[record_id],
+                activity_counts=activity_counts,
+            )
+        else:
+            yield SleepRecord(
+                sleep_stages=parsed_xml.sleep_stages,
+                sleep_stage_duration=parsed_xml.sleep_stage_duration,
+                id=record_id,
+                recording_start_time=parsed_xml.recording_start_time,
+                heartbeat_times=heartbeat_times,
+                subject_data=subject_data[record_id],
+            )
 
 
 def read_slpdb(

tests/test_sleep_readers.py

@@ -10,12 +10,45 @@
 from pathlib import Path
 
 import numpy as np
+import pandas as pd
 from edfio import Edf, EdfSignal
 
 from sleepecg import SleepStage, get_toy_ecg, read_mesa, read_shhs, read_slpdb
 from sleepecg.io.sleep_readers import Gender
 
 
+def _dummy_nsrr_overlap(filename: str, mesa_ids: list[int]):
+    overlap_data = {
+        "mesaid": mesa_ids,
+        "line": [563, 301],
+        "linetime": ["20:30:00", "21:00:00"],
+        "starttime_psg": ["20:29:59", "20:59:59"],
+    }
+
+    overlap_df = pd.DataFrame(overlap_data)
+    overlap_df.to_csv(filename, index=False)
+
+
+def _dummy_nsrr_actigraphy(filename: str, mesa_id: str):
+    base_time = datetime.datetime(2024, 1, 1, 20, 30, 0)
+
+    # Generate linetime values as strings
+    linetimes = [
+        (base_time + datetime.timedelta(seconds=30 * i)).strftime("%H:%M:%S")
+        for i in range(10)
+    ]
+
+    actigraphy_data = {
+        "mesaid": [mesa_id] * 10,
+        "line": list(range(563, 573)),
+        "linetime": linetimes,
+        "activity": [10] * 10,
+    }
+
+    actigraphy_df = pd.DataFrame(actigraphy_data)
+    actigraphy_df.to_csv(filename)
+
+
 def _dummy_nsrr_edf(filename: str, hours: float, ecg_channel: str):
     ecg_5_min, fs = get_toy_ecg()
     seconds = int(hours * 60 * 60)
@@ -25,6 +58,7 @@ def _dummy_nsrr_edf(filename: str, hours: float, ecg_channel: str):
 
 def _dummy_nsrr_xml(filename: str, hours: float, random_state: int):
     EPOCH_LENGTH = 30
+    RECORDING_DURATION = 154.0
     STAGES = [
         "Wake|0",
         "Stage 1 sleep|1",
@@ -46,6 +80,7 @@ def _dummy_nsrr_xml(filename: str, hours: float, random_state: int):
             "<ScoredEvent>\n"
             "<EventType/>\n"
             "<EventConcept>Recording Start Time</EventConcept>\n"
+            f"<Duration>{RECORDING_DURATION}</Duration>\n"
             "<ClockTime>01.01.85 20.29.59</ClockTime>\n"
             "</ScoredEvent>\n",
         )
@@ -71,24 +106,41 @@ def _dummy_nsrr_xml(filename: str, hours: float, random_state: int):
         )
 
 
-def _create_dummy_mesa(data_dir: str, durations: list[float], random_state: int = 42):
+def _create_dummy_mesa(
+    data_dir: str, durations: list[float], random_state: int = 42, actigraphy: bool = False
+):
     DB_SLUG = "mesa"
     ANNOTATION_DIRNAME = "polysomnography/annotations-events-nsrr"
     EDF_DIRNAME = "polysomnography/edfs"
     CSV_DIRNAME = "datasets"
+    OVERLAP_DIRNAME = "overlap"
+    ACTIVITY_DIRNAME = "actigraphy"
 
     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
+    overlap_dir = db_dir / OVERLAP_DIRNAME
+    activity_dir = db_dir / ACTIVITY_DIRNAME
 
     for directory in (annotations_dir, edf_dir, csv_dir):
         directory.mkdir(parents=True, exist_ok=True)
 
+    if actigraphy:
+        for directory in (overlap_dir, activity_dir):
+            directory.mkdir(parents=True, exist_ok=True)
+            record_ids = []
+
     for i, hours in enumerate(durations):
         record_id = f"mesa-sleep-{i:04}"
         _dummy_nsrr_edf(f"{edf_dir}/{record_id}.edf", hours, ecg_channel="EKG")
         _dummy_nsrr_xml(f"{annotations_dir}/{record_id}-nsrr.xml", hours, random_state)
+        if actigraphy:
+            _dummy_nsrr_actigraphy(f"{activity_dir}/{record_id}.csv", mesa_id=record_id)
+            record_ids.append(record_id)
+
+    if actigraphy:
+        _dummy_nsrr_overlap(f"{overlap_dir}/mesa-actigraphy-psg-overlap.csv", record_ids)
 
     with open(csv_dir / "mesa-sleep-dataset-0.0.0.csv", "w") as csv:
         csv.write("mesaid,examnumber,race1c,gender1,cucmcn1c,sleepage5c\n")
@@ -143,6 +195,29 @@ def test_read_mesa(tmp_path):
         assert set(rec.sleep_stages) - valid_stages == set()
 
 
+def test_read_mesa_actigraphy(tmp_path):
+    """Basic sanity checks for records read via read_mesa including actigraphy."""
+    durations = [0.1, 0.2]  # hours
+    valid_stages = {int(s) for s in SleepStage}
+
+    _create_dummy_mesa(data_dir=tmp_path, durations=durations, actigraphy=True)
+    records = list(
+        read_mesa(
+            data_dir=tmp_path,
+            heartbeats_source="ecg",
+            offline=True,
+            activity_source="actigraphy",
+        )
+    )
+
+    assert len(records) == 2
+
+    for rec in records:
+        assert rec.sleep_stage_duration == 30
+        assert set(rec.sleep_stages) - valid_stages == set()
+        assert len(rec.activity_counts) == 6
+
+
 def test_read_shhs(tmp_path):
     """Basic sanity checks for records read via read_shhs."""
     durations = [0.1, 0.2]  # hours