diff --git a/docs/api/datasets.rst b/docs/api/datasets.rst index c9a88b7ff..933fa34a5 100644 --- a/docs/api/datasets.rst +++ b/docs/api/datasets.rst @@ -243,6 +243,7 @@ Available Datasets datasets/pyhealth.datasets.ChestXray14Dataset datasets/pyhealth.datasets.PhysioNetDeIDDataset datasets/pyhealth.datasets.EEGBCIDataset + PTB-XL Dataset datasets/pyhealth.datasets.TUABDataset datasets/pyhealth.datasets.TUEVDataset datasets/pyhealth.datasets.ClinVarDataset diff --git a/docs/api/datasets/pyhealth.datasets.PTBXLDataset.rst b/docs/api/datasets/pyhealth.datasets.PTBXLDataset.rst new file mode 100644 index 000000000..fca949e18 --- /dev/null +++ b/docs/api/datasets/pyhealth.datasets.PTBXLDataset.rst @@ -0,0 +1,34 @@ +pyhealth.datasets.PTBXLDataset +============================== + +Overview +-------- + +PTB-XL is a large publicly available 12-lead ECG dataset from PhysioNet +(version 1.0.3). It contains **21,799** clinical ECG records of 10 seconds +from **18,869** patients, with multi-label SCP-ECG statements, official +stratified folds (``strat_fold``), and demographic / site / device metadata +suited to shift-aware evaluation. Only v1.0.3 is supported (v1.0.1 figures +of 21,837 records / 18,885 patients come from Wagner et al., Scientific +Data 2020, and must not be mixed with the PhysioNet v1.0.3 counts). + +For more information see `PhysioNet PTB-XL v1.0.3 +`_ and Wagner et al., +`Scientific Data 2020 `_. + +Optional dependency: install waveform I/O with ``pip install 'pyhealth[ptbxl]'``. + +``PTBXLDataset.root`` is the version directory passed to the constructor. +Derived metadata CSVs are written under ``metadata_cache_dir`` (default +``~/.cache/pyhealth/datasets/ptbxl/``), not into ``root``, so read-only +data mounts stay untouched. ``BaseDataset.load_table`` concatenates +``root / file_path`` and cannot take an absolute path; PTB-XL overrides +``load_table`` to read the derived CSV from the metadata cache. + +API Reference +------------- + +.. autoclass:: pyhealth.datasets.PTBXLDataset + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/api/tasks.rst b/docs/api/tasks.rst index bdaa9599a..bd7a4a627 100644 --- a/docs/api/tasks.rst +++ b/docs/api/tasks.rst @@ -225,6 +225,7 @@ Available Tasks Sleep Staging (SleepEDF) Temple University EEG Tasks EEGBCI Tasks + PTB-XL Tasks Sleep Staging v2 Benchmark EHRShot ChestX-ray14 Binary Classification diff --git a/docs/api/tasks/pyhealth.tasks.ptbxl.rst b/docs/api/tasks/pyhealth.tasks.ptbxl.rst new file mode 100644 index 000000000..8e8fba74b --- /dev/null +++ b/docs/api/tasks/pyhealth.tasks.ptbxl.rst @@ -0,0 +1,24 @@ +pyhealth.tasks.ptbxl +==================== + +Overview +-------- + +Task helpers for PTB-XL multi-label ECG diagnosis. This module currently +implements the official **5-diagnostic-superclass** classification task +(``NORM``, ``MI``, ``STTC``, ``CD``, ``HYP``) by aggregating diagnostic SCP +statements via ``scp_statements.csv``. + +Empty superclass label sets (≈400 records on v1.0.1, Wagner et al. +Table 9; mainly pacemaker ECGs) are dropped by default. On the full +v1.0.3 corpus, ``MultiLabelProcessor`` emits labels in alphabetical +order: ``CD, HYP, MI, NORM, STTC``. Official fold splitting is provided +by :func:`pyhealth.datasets.split_by_strat_fold`. + +API Reference +------------- + +.. automodule:: pyhealth.tasks.ptbxl + :members: + :undoc-members: + :show-inheritance: diff --git a/examples/ecg/ptbxl/ptbxl_superclass_quickstart.py b/examples/ecg/ptbxl/ptbxl_superclass_quickstart.py new file mode 100644 index 000000000..21922c238 --- /dev/null +++ b/examples/ecg/ptbxl/ptbxl_superclass_quickstart.py @@ -0,0 +1,77 @@ +"""PTB-XL 5-superclass multi-label quickstart. + +Requires:: + + pip install 'pyhealth[ptbxl]' + +Download PTB-XL v1.0.3 from https://physionet.org/content/ptb-xl/1.0.3/ +and point ``--root`` at the extracted version directory (contains +``ptbxl_database.csv``, ``scp_statements.csv``, ``records100/``). + +Example:: + + python examples/ecg/ptbxl/ptbxl_superclass_quickstart.py \\ + --root /data/ptb-xl/1.0.3 --dev +""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +from pyhealth.datasets import PTBXLDataset, split_by_strat_fold +from pyhealth.tasks import PTBXLSuperclassClassification + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--root", + type=Path, + required=True, + help="PTB-XL version root (ptbxl_database.csv + records*/)", + ) + parser.add_argument( + "--sampling-rate", + type=int, + default=100, + choices=(100, 500), + help="Waveform sampling rate (default: 100 Hz / filename_lr)", + ) + parser.add_argument( + "--dev", + action="store_true", + help="Limit patients via BaseDataset.dev mode", + ) + args = parser.parse_args() + + dataset = PTBXLDataset( + root=str(args.root), + sampling_rate=args.sampling_rate, + dev=args.dev, + ) + task = PTBXLSuperclassClassification( + scp_statements_path=str(args.root / "scp_statements.csv"), + ) + samples = dataset.set_task(task) + train, val, test = split_by_strat_fold(samples) + print( + f"samples={len(samples)} " + f"train={len(train)} val={len(val)} test={len(test)}" + ) + if len(samples): + first = samples[0] + print( + "first sample:", + { + "patient_id": first["patient_id"], + "record_id": first["record_id"], + "labels": first["labels"], + "strat_fold": first["strat_fold"], + "signal_shape": tuple(first["signal"].shape), + }, + ) + + +if __name__ == "__main__": + main() diff --git a/pyhealth/datasets/__init__.py b/pyhealth/datasets/__init__.py index 57a9956c2..d489a8beb 100644 --- a/pyhealth/datasets/__init__.py +++ b/pyhealth/datasets/__init__.py @@ -80,10 +80,12 @@ def __init__(self, *args, **kwargs): split_by_sample_conformal, split_by_sample_conformal_tuh, split_by_sample_tuh, + split_by_strat_fold as split_by_strat_fold, split_by_visit, split_by_visit_conformal, ) from .eegbci import EEGBCIDataset as EEGBCIDataset # noqa: E402 +from .ptbxl import PTBXLDataset as PTBXLDataset from .tuab import TUABDataset from .tuev import TUEVDataset from .utils import ( diff --git a/pyhealth/datasets/configs/ptbxl.yaml b/pyhealth/datasets/configs/ptbxl.yaml new file mode 100644 index 000000000..322114542 --- /dev/null +++ b/pyhealth/datasets/configs/ptbxl.yaml @@ -0,0 +1,22 @@ +version: "1.0.3" # PhysioNet PTB-XL v1.0.3 only; earlier releases are not supported. +tables: + records: + # Template only: PTBXLDataset writes a resolved config whose file_path + # is the derived cache CSV (rate + data-root hash) before BaseDataset init. + file_path: "ptbxl-pyhealth-100hz.csv" + patient_id: "patient_id" + # recording_date from ptbxl_database.csv (e.g. 1984-11-09 09:17:34) + timestamp: "recording_date" + timestamp_format: "%Y-%m-%d %H:%M:%S" + attributes: + - "record_id" + - "signal_file" + - "sampling_rate" + - "strat_fold" + - "age" + - "age_is_censored" + - "age_is_missing" + - "sex" # 0 = female, 1 = male + - "site" + - "device" + - "scp_codes" diff --git a/pyhealth/datasets/ptbxl.py b/pyhealth/datasets/ptbxl.py new file mode 100644 index 000000000..1f49f6375 --- /dev/null +++ b/pyhealth/datasets/ptbxl.py @@ -0,0 +1,641 @@ +""" +PyHealth dataset for PTB-XL (12-lead ECG, PhysioNet). + +Dataset link: + https://physionet.org/content/ptb-xl/1.0.3/ + +Dataset paper: (please cite if you use this dataset) + Patrick Wagner, Nils Strodthoff, Ralf-Dieter Bousseljot, Dieter Kreiseler, + Fatima I. Lunze, Wojciech Samek, and Tobias Schaeffter. "PTB-XL, a large + publicly available electrocardiography dataset." Scientific Data 7, 154 + (2020). + +Dataset paper link: + https://www.nature.com/articles/s41597-020-0495-6 + +Author: + AxelNoun (GitHub: @AxelNoun) — external contributor, no NetID + +Description: + Implements ``PTBXLDataset`` (BaseDataset + YAML) for PTB-XL v1.0.3, + including metadata preparation, HIPAA age-censor handling, and lazy + ``wfdb`` signal loading. Task and split helpers live in separate modules. +""" + +from __future__ import annotations + +import ast +import hashlib +import logging +import os +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import numpy as np +import pandas as pd +import yaml +from filelock import FileLock + +from .base_dataset import BaseDataset +from .configs.config import load_yaml_config +from .utils import MODULE_CACHE_PATH + +if TYPE_CHECKING: + from pyhealth.tasks.ptbxl import PTBXLSuperclassClassification + +logger = logging.getLogger(__name__) + +# HIPAA: ages ≥ 90 are encoded as this sentinel in ptbxl_database.csv. +AGE_CENSOR_SENTINEL = 300 + +PTBXL_DATABASE_CSV = "ptbxl_database.csv" +PTBXL_SCP_STATEMENTS_CSV = "scp_statements.csv" + +_DEFAULT_METADATA_CACHE = Path(MODULE_CACHE_PATH) / "ptbxl" + +def _atomic_replace(tmp_path: Path, dest: Path) -> None: + """Replace ``dest`` with ``tmp_path`` (same-filesystem atomic on POSIX/NT).""" + os.replace(tmp_path, dest) + + +def _write_csv_atomic(df: pd.DataFrame, dest: Path) -> None: + tmp_path = dest.with_name(dest.name + ".tmp") + try: + df.to_csv(tmp_path, index=False) + _atomic_replace(tmp_path, dest) + except Exception: + if tmp_path.exists(): + tmp_path.unlink(missing_ok=True) + raise + + +def _write_yaml_atomic(payload: dict[str, Any], dest: Path) -> None: + tmp_path = dest.with_name(dest.name + ".tmp") + try: + with open(tmp_path, "w", encoding="utf-8") as handle: + yaml.safe_dump(payload, handle, sort_keys=False) + _atomic_replace(tmp_path, dest) + except Exception: + if tmp_path.exists(): + tmp_path.unlink(missing_ok=True) + raise + + +def format_patient_id(value: Any, *, ecg_id: Any | None = None) -> str: + """Cast PTB-XL ``patient_id`` (stored as float, e.g. ``15709.0``) to ``str``. + + Without an explicit int cast, stringification yields ``\"15709.0\"`` and + silently breaks patient-level splits. Called from + :meth:`PTBXLDataset.prepare_metadata`. + + Args: + value (Any): Raw ``patient_id`` cell (float, int, or numeric string). + ecg_id (Any | None): Optional ``ecg_id`` included in error messages. + + Returns: + str: Integer patient id as a string (e.g. ``\"15709\"``). + + Examples: + >>> format_patient_id(15709.0) + '15709' + """ + if value is None or (isinstance(value, float) and np.isnan(value)): + suffix = f" (ecg_id={ecg_id})" if ecg_id is not None else "" + raise ValueError(f"patient_id is missing{suffix}") + try: + return str(int(float(value))) + except (TypeError, ValueError) as exc: + suffix = f" (ecg_id={ecg_id})" if ecg_id is not None else "" + raise ValueError( + f"patient_id is not numeric: {value!r}{suffix}" + ) from exc + + +def parse_scp_codes(scp_codes: Any) -> dict[str, float]: + """Parse the stringified ``scp_codes`` dict from ``ptbxl_database.csv``. + + PhysioNet stores entries as ``statement: likelihood`` where likelihood is + in ``[0, 100]`` and **0 means unknown confidence, not absence**. All keys + present in the dict are therefore treated as positive statements; only an + empty dict means no statements. Used by task aggregation helpers. + + Args: + scp_codes (Any): Raw cell value (dict or stringified dict). + + Returns: + dict[str, float]: Mapping of SCP statement acronym → likelihood. + + Examples: + >>> parse_scp_codes("{'IMI': 80.0, 'SR': 0.0}")["SR"] + 0.0 + """ + if scp_codes is None or (isinstance(scp_codes, float) and np.isnan(scp_codes)): + return {} + if isinstance(scp_codes, dict): + return {str(k): float(v) for k, v in scp_codes.items()} + text = str(scp_codes).strip() + if not text or text.lower() in {"nan", "none", "{}"}: + return {} + parsed = ast.literal_eval(text) + if not isinstance(parsed, dict): + raise TypeError( + f"Expected scp_codes dict, got {type(parsed)!r}: {scp_codes!r}" + ) + return {str(k): float(v) for k, v in parsed.items()} + + +def is_age_censored(age: Any, sentinel: int = AGE_CENSOR_SENTINEL) -> bool: + """Return True when age is the HIPAA ≥90 sentinel (default 300). + + Distinct from missing age (NaN), which is not censored. Used when writing + ``age_is_censored`` in metadata. + + Args: + age (Any): Raw age cell from ``ptbxl_database.csv``. + sentinel (int): Censor value (default ``300``). + + Returns: + bool: True if ``age`` equals the HIPAA sentinel. + + Examples: + >>> is_age_censored(300), is_age_censored(float("nan")) + (True, False) + """ + if age is None or (isinstance(age, float) and np.isnan(age)): + return False + if isinstance(age, str) and not age.strip(): + return False + try: + return int(float(age)) == int(sentinel) + except (TypeError, ValueError): + return False + + +def is_age_missing(age: Any) -> bool: + """Return True when age is genuinely missing (NaN / empty), not censored. + + Args: + age (Any): Raw age cell from ``ptbxl_database.csv``. + + Returns: + bool: True if age is missing (not the 300 sentinel). + + Examples: + >>> is_age_missing(float("nan")), is_age_missing(300) + (True, False) + """ + if age is None: + return True + if isinstance(age, float) and np.isnan(age): + return True + return isinstance(age, str) and not str(age).strip() + + +def root_cache_key(data_root: str | Path) -> str: + """Return a short stable hash of the resolved data root path. + + Args: + data_root (str | Path): PTB-XL version root. + + Returns: + str: 10-character hex digest used in metadata filenames. + + Examples: + >>> len(root_cache_key("/data/ptb-xl/1.0.3")) + 10 + """ + resolved = str(Path(data_root).resolve()) + return hashlib.sha1(resolved.encode("utf-8")).hexdigest()[:10] + + +def metadata_filename( + sampling_rate: int, + data_root: str | Path, + source_key: str | None = None, +) -> str: + """Return the rate-, root-, and source-specific derived metadata CSV name. + + ``source_key`` is a content hash of ``ptbxl_database.csv``. Including it + in the filename forces ``prepare_metadata`` to regenerate when the + official CSV is replaced in-place (e.g. extracting v1.0.3 over v1.0.1). + + Args: + sampling_rate (int): ``100`` or ``500``. + data_root (str | Path): Absolute/relative PTB-XL data root. + source_key (str | None): Optional SHA1[:10] of the source + ``ptbxl_database.csv`` bytes. + + Returns: + str: Filename such as ``ptbxl-pyhealth-100hz-[-].csv``. + + Examples: + >>> name = metadata_filename(100, "/data/ptb-xl/1.0.3") + >>> name.startswith("ptbxl-pyhealth-100hz-") + True + >>> keyed = metadata_filename(100, "/data/ptb-xl/1.0.3", "abc123def0") + >>> "abc123def0" in keyed + True + """ + name = ( + f"ptbxl-pyhealth-{int(sampling_rate)}hz-" + f"{root_cache_key(data_root)}" + ) + if source_key: + name = f"{name}-{source_key}" + return f"{name}.csv" + + +def load_ptbxl_record(record_path: str | Path) -> np.ndarray: + """Load a PTB-XL WFDB record as ``(n_leads, n_samples)``. + + ``wfdb.rdsamp`` returns ``(n_samples, n_channels)`` (e.g. ``(1000, 12)`` at + 100 Hz). This helper **transposes** to ``(n_leads, n_samples)`` so the layout + matches PyHealth signal tasks such as EEGBCI / SleepEDF, which use + ``mne.io.BaseRaw.get_data()`` → ``(n_channels, n_times)``. The + ``\"tensor\"`` processor preserves that shape. + + The path must be a WFDB *record base* without ``.hea`` / ``.dat`` (as in + ``filename_lr`` / ``filename_hr``). Extensions are stripped only if the + caller accidentally includes them; they are never appended. + + Args: + record_path (str | Path): WFDB record base path (no extension). + + Returns: + np.ndarray: Signal of shape ``(n_leads, n_samples)``, float32, + physical units. + + Raises: + ImportError: If the optional ``wfdb`` extra is not installed. + FileNotFoundError: If the ``.hea`` header cannot be found. + + Examples: + >>> # doctest: +SKIP + >>> signal = load_ptbxl_record("/data/ptb-xl/1.0.3/records100/00000/00001_lr") + >>> signal.shape[0] + 12 + """ + try: + import wfdb + except ImportError as exc: + raise ImportError( + "Reading PTB-XL waveforms requires the optional 'wfdb' dependency. " + "Install it with: pip install 'pyhealth[ptbxl]'" + ) from exc + + path = Path(record_path) + # Strip accidental extensions; never add .hea/.dat for rdsamp. + if path.suffix.lower() in {".hea", ".dat"}: + path = path.with_suffix("") + header = Path(str(path) + ".hea") + if not header.is_file(): + raise FileNotFoundError( + f"PTB-XL WFDB header not found for record base: {record_path}" + ) + + signals, _ = wfdb.rdsamp(str(path)) + # (n_samples, n_leads) → (n_leads, n_samples) to match mne.get_data(). + return np.asarray(signals, dtype=np.float32).T + + +class PTBXLDataset(BaseDataset): + """PhysioNet PTB-XL ECG dataset (v1.0.3). + + Dataset: https://physionet.org/content/ptb-xl/1.0.3/ + + Expects ``root`` (``data_root``) to point at the extracted **v1.0.3** + directory containing ``ptbxl_database.csv``, ``scp_statements.csv``, + ``records100/``, and ``records500/``. Raw data must live outside the git + repo. Earlier PhysioNet releases (v1.0.1 / v1.0.2) are not supported: + required columns and record counts differ. + + Derived metadata CSVs are written under PyHealth's dataset cache + (``~/.cache/pyhealth/datasets/ptbxl/`` by default), **not** into ``root``, + so read-only / shared data mounts stay untouched. Override with + ``metadata_cache_dir``. Filenames include the sampling rate, a short + hash of the resolved data root, and a content hash of + ``ptbxl_database.csv`` so different roots and source versions never share + a derived CSV. ``BaseDataset`` cache identity (``global_event_df``) further + includes those hashes via ``dataset_name`` (same pattern as EEGBCI). + + Args: + root (str): Version root of the PTB-XL download (signal + official CSVs). + dataset_name (str | None): Optional name prefix; defaults to + ``ptbxl_{sampling_rate}hz``. Root and metadata content hashes are + appended so two data roots cannot share a ``BaseDataset`` cache. + config_path (str | Path | None): Optional YAML config; defaults to + ``configs/ptbxl.yaml``. + sampling_rate (int): ``100`` (default) or ``500``. + metadata_cache_dir (str | Path | None): Directory for derived + ``ptbxl-pyhealth-*.csv``. Defaults to + ``MODULE_CACHE_PATH / \"ptbxl\"``. + **kwargs: Forwarded to :class:`BaseDataset` (``cache_dir``, ``dev``, …). + + Attributes: + root (str): Same as the constructor ``root`` (the PTB-XL version + directory). Derived metadata is *not* stored here. + data_root (Path): User-provided PTB-XL version root (waveforms + CSVs). + sampling_rate (int): Selected waveform rate (100 or 500). + metadata_cache_dir (Path): Directory holding derived metadata CSVs. + + Note: + Age missing (NaN) vs HIPAA-censored (``age == 300`` for ≥90) are + exposed as ``age_is_missing`` / ``age_is_censored``. Sex is encoded + as ``0`` = female, ``1`` = male. ``scp_codes`` keeps the original + stringified dict; likelihood ``0`` means unknown confidence. Label + aggregation lives in task modules. ``recording_date`` is used as the + event timestamp. + + Examples: + >>> dataset = PTBXLDataset(root="/data/ptb-xl/1.0.3") # doctest: +SKIP + >>> dataset.stats() # doctest: +SKIP + >>> patient = dataset.get_patient(dataset.unique_patient_ids[0]) # doctest: +SKIP + >>> event = patient.get_events(event_type="records")[0] # doctest: +SKIP + >>> signal = load_ptbxl_record(event.signal_file) # doctest: +SKIP + """ + + def __init__( + self, + root: str, + dataset_name: str | None = None, + config_path: str | Path | None = None, + sampling_rate: int = 100, + metadata_cache_dir: str | Path | None = None, + **kwargs, + ) -> None: + if sampling_rate not in {100, 500}: + raise ValueError( + f"sampling_rate must be 100 or 500, got {sampling_rate}" + ) + package_config = ( + Path(config_path) + if config_path is not None + else Path(__file__).parent / "configs" / "ptbxl.yaml" + ) + + self.data_root = Path(root) + self.sampling_rate = int(sampling_rate) + self.metadata_cache_dir = Path( + metadata_cache_dir + if metadata_cache_dir is not None + else _DEFAULT_METADATA_CACHE + ) + db_path = self.data_root / PTBXL_DATABASE_CSV + if not db_path.is_file(): + raise FileNotFoundError( + f"Expected {PTBXL_DATABASE_CSV} under root={self.data_root}. " + "Download PTB-XL from https://physionet.org/content/ptb-xl/1.0.3/" + ) + self._source_key = hashlib.sha1(db_path.read_bytes()).hexdigest()[:10] + self.metadata_file_name = metadata_filename( + self.sampling_rate, self.data_root, self._source_key + ) + self.prepare_metadata() + + # BaseDataset only accepts config_path (not an in-memory config). Write + # a resolved YAML whose file_path matches the derived CSV *before* + # super().__init__ — loading is lazy via global_event_df, but the + # config must already be correct when that first happens. + resolved_config_path = self._write_resolved_config(package_config) + + # BaseDataset._init_cache_dir keys on {root, tables, dataset_name, dev}. + # root is the user data path (see load_table). dataset_name still + # includes a content hash of the derived CSV so replacing + # ptbxl_database.csv in-place does not reuse a stale global_event_df. + base_name = dataset_name or f"ptbxl_{self.sampling_rate}hz" + dataset_name = ( + f"{base_name}_{root_cache_key(self.data_root)}_" + f"{self._metadata_cache_key()}" + ) + + # Keep self.root as the user-provided PTB-XL path so stats()/logs + # match what the caller passed. Derived CSVs live in + # metadata_cache_dir (read-only data mounts). BaseDataset.load_table + # builds csv_path as f"{self.root}/{file_path}" and does not accept + # an absolute file_path, so load_table (below) temporarily points + # root at the metadata cache for the scan only. + super().__init__( + root=str(self.data_root), + tables=["records"], + dataset_name=dataset_name, + config_path=str(resolved_config_path), + **kwargs, + ) + + def load_table(self, table_name: str): + """Load tables from ``metadata_cache_dir``, not the user data root. + + See the comment on ``super().__init__`` in :meth:`__init__`. + + Args: + table_name (str): Table name from the YAML config. + + Returns: + The Dask dataframe produced by :meth:`BaseDataset.load_table`. + """ + original_root = self.root + self.root = str(self.metadata_cache_dir) + try: + return super().load_table(table_name) + finally: + self.root = original_root + + def _write_resolved_config(self, package_config: Path) -> Path: + """Write a cache-local YAML with the derived metadata CSV filename. + + Args: + package_config (Path): Packaged ``configs/ptbxl.yaml`` template. + + Returns: + Path: Written config path passed to :class:`BaseDataset`. + """ + config = load_yaml_config(str(package_config)) + config.tables["records"].file_path = self.metadata_file_name + out_path = self.metadata_cache_dir / ( + f"ptbxl-config-{self.sampling_rate}hz-" + f"{root_cache_key(self.data_root)}-{self._source_key}.yaml" + ) + self.metadata_cache_dir.mkdir(parents=True, exist_ok=True) + lock_path = out_path.with_name(out_path.name + ".lock") + with FileLock(str(lock_path)): + _write_yaml_atomic(config.model_dump(), out_path) + return out_path + + @property + def scp_statements_path(self) -> Path: + """Path to official ``scp_statements.csv`` under the data root. + + Returns: + Path: Absolute path to ``scp_statements.csv``. + """ + return self.data_root / PTBXL_SCP_STATEMENTS_CSV + + def _metadata_cache_key(self) -> str: + """SHA1[:10] of the derived metadata CSV bytes (EEGBCI pattern). + + Must be called after :meth:`prepare_metadata` so the file exists. + Injected into ``dataset_name`` so ``BaseDataset._init_cache_dir`` + does not collapse two data roots onto one ``global_event_df``. + """ + csv_path = self.metadata_cache_dir / self.metadata_file_name + return hashlib.sha1(csv_path.read_bytes()).hexdigest()[:10] + + def prepare_metadata(self) -> None: + """Build rate-/root-specific metadata CSV under ``metadata_cache_dir``. + + Returns: + None + + Examples: + >>> # doctest: +SKIP + >>> ds = PTBXLDataset(root="/data/ptb-xl/1.0.3") + >>> ds.prepare_metadata() + """ + csv_path = self.metadata_cache_dir / self.metadata_file_name + if csv_path.exists() and self._metadata_matches_request(csv_path): + return + + self.metadata_cache_dir.mkdir(parents=True, exist_ok=True) + lock_path = csv_path.with_name(csv_path.name + ".lock") + with FileLock(str(lock_path)): + # Another process may have finished while we waited. + if csv_path.exists() and self._metadata_matches_request(csv_path): + return + self._write_derived_metadata(csv_path) + + def _write_derived_metadata(self, csv_path: Path) -> None: + """Generate the derived metadata CSV (caller holds the file lock).""" + db_path = self.data_root / PTBXL_DATABASE_CSV + if not db_path.is_file(): + raise FileNotFoundError( + f"Expected {PTBXL_DATABASE_CSV} under root={self.data_root}. " + "Download PTB-XL from https://physionet.org/content/ptb-xl/1.0.3/" + ) + + db = pd.read_csv(db_path) + required = { + "ecg_id", + "patient_id", + "age", + "sex", + "site", + "device", + "scp_codes", + "strat_fold", + "filename_lr", + "filename_hr", + "recording_date", + } + missing = required - set(db.columns) + if missing: + raise ValueError( + f"ptbxl_database.csv missing columns: {sorted(missing)}. " + "PTBXLDataset supports PhysioNet PTB-XL v1.0.3 only " + "(https://physionet.org/content/ptb-xl/1.0.3/)." + ) + + filename_col = "filename_lr" if self.sampling_rate == 100 else "filename_hr" + rows: list[dict[str, Any]] = [] + for row in db.to_dict(orient="records"): + ecg_id = row["ecg_id"] + rel = str(row[filename_col]).strip() + # Record base path only — no .hea/.dat (matches wfdb.rdsamp). + if rel.endswith((".hea", ".dat")): + rel = rel.rsplit(".", 1)[0] + signal_file = str((self.data_root / rel).resolve()) + + age = row["age"] + missing_age = is_age_missing(age) + censored = is_age_censored(age) + if missing_age: + age_out: Any = pd.NA + else: + # Keep numeric age (including sentinel 300). + age_out = int(float(age)) + + rows.append( + { + "patient_id": format_patient_id( + row["patient_id"], ecg_id=ecg_id + ), + "record_id": str(int(float(ecg_id))), + "signal_file": signal_file, + "sampling_rate": self.sampling_rate, + "strat_fold": int(row["strat_fold"]), + "age": age_out, + "age_is_censored": int(censored), + "age_is_missing": int(missing_age), + "sex": int(row["sex"]), + "site": row["site"], + "device": row["device"], + "scp_codes": row["scp_codes"], + "recording_date": row["recording_date"], + } + ) + + out = pd.DataFrame(rows) + out["patient_id"] = out["patient_id"].astype(str) + out["record_id"] = out["record_id"].astype(str) + out.sort_values( + ["patient_id", "record_id"], + key=lambda col: col.astype(int), + inplace=True, + ) + out.reset_index(drop=True, inplace=True) + _write_csv_atomic(out, csv_path) + logger.info( + "Wrote PTB-XL metadata (%d records, %d Hz) to %s", + len(out), + self.sampling_rate, + csv_path, + ) + + def _metadata_matches_request(self, csv_path: Path) -> bool: + """Reuse cached metadata when the derived CSV has the expected schema. + + Filename already encodes sampling rate, data-root path, and source CSV + content, so this only guards against a truncated or partial write. + + Args: + csv_path (Path): Candidate derived metadata CSV. + + Returns: + bool: True if the cache is safe to reuse. + """ + try: + df = pd.read_csv(csv_path, nrows=0) + except (OSError, ValueError, pd.errors.ParserError): + return False + needed = { + "patient_id", + "record_id", + "signal_file", + "sampling_rate", + "strat_fold", + "age", + "age_is_censored", + "age_is_missing", + "sex", + "site", + "device", + "scp_codes", + "recording_date", + } + return needed.issubset(df.columns) + + @property + def default_task(self) -> PTBXLSuperclassClassification: + """Return the 5-superclass multi-label task wired to this data root. + + Matches other signal datasets (EEGBCI / TUAB / SleepEDF): BaseDataset + exposes ``default_task`` as a read-only property and never assigns to + it in ``__init__``. + + Returns: + PTBXLSuperclassClassification: Default task instance. + """ + from pyhealth.tasks.ptbxl import PTBXLSuperclassClassification + + return PTBXLSuperclassClassification( + scp_statements_path=str(self.scp_statements_path), + ) diff --git a/pyhealth/datasets/splitter.py b/pyhealth/datasets/splitter.py index 2dbc94186..3e3d33235 100644 --- a/pyhealth/datasets/splitter.py +++ b/pyhealth/datasets/splitter.py @@ -180,6 +180,132 @@ def split_by_patient( return train_dataset, val_dataset, test_dataset +def split_by_strat_fold( + dataset: SampleDataset, + train_folds: tuple[int, ...] | list[int] = tuple(range(1, 9)), + val_folds: tuple[int, ...] | list[int] = (9,), + test_folds: tuple[int, ...] | list[int] = (10,), + fold_field: str = "strat_fold", + folds: tuple[int, ...] | list[int] | None = None, + check_patient_disjoint: bool = False, +): + """Split a sample dataset using integer fold assignments. + + The recommended PTB-XL protocol uses folds 1–8 for training, fold 9 for + validation, and fold 10 for testing. Official PTB-XL folds are + patient-disjoint; this helper does not assume that of arbitrary data. + Pass ``check_patient_disjoint=True`` to verify no ``patient_id`` appears + in more than one of train/val/test (uses ``dataset.patient_to_index``). + + Pass ``folds`` (one integer per sample, aligned with dataset indices) to + avoid decoding every sample via ``dataset[i]``. When omitted, fold values + are read from ``sample[fold_field]``. + + Args: + dataset (SampleDataset): A :class:`~pyhealth.datasets.SampleDataset` + (or compatible object exposing ``subset`` and ``__getitem__``). + train_folds (Tuple[int, ...] | List[int]): Folds assigned to train + (default ``1..8``). + val_folds (Tuple[int, ...] | List[int]): Folds assigned to validation + (default ``9``). + test_folds (Tuple[int, ...] | List[int]): Folds assigned to test + (default ``10``). + fold_field (str): Sample key holding the fold integer when ``folds`` + is not provided. Defaults to ``\"strat_fold\"``. + folds (Tuple[int, ...] | List[int] | None): Optional precomputed fold + for each index. Length must equal ``len(dataset)``. + check_patient_disjoint (bool): If True, raise ``ValueError`` when a + ``patient_id`` has samples in more than one split. Requires + ``dataset.patient_to_index``. Defaults to False. + + Returns: + tuple: ``(train_dataset, val_dataset, test_dataset)`` subsets. + + Raises: + KeyError: If ``folds`` is omitted and a sample is missing ``fold_field``. + TypeError: If ``check_patient_disjoint`` is True but the dataset has + no ``patient_to_index``. + ValueError: If fold sets overlap, ``folds`` has the wrong length, a + sample fold is unassigned, or a patient leaks across splits. + + Examples: + >>> # doctest: +SKIP + >>> from pyhealth.datasets import split_by_strat_fold + >>> train, val, test = split_by_strat_fold(sample_dataset) + """ + train_set = {int(f) for f in train_folds} + val_set = {int(f) for f in val_folds} + test_set = {int(f) for f in test_folds} + if train_set & val_set or train_set & test_set or val_set & test_set: + raise ValueError("train_folds, val_folds, and test_folds must be disjoint") + + n_samples = len(dataset) + if folds is not None: + if len(folds) != n_samples: + raise ValueError( + f"folds has length {len(folds)} but dataset has {n_samples} samples" + ) + fold_values = [int(f) for f in folds] + else: + fold_values = [] + for i in range(n_samples): + sample = dataset[i] + if fold_field not in sample: + raise KeyError( + f"sample is missing {fold_field!r}; pass folds= to avoid " + "decoding samples, or emit this field from the task" + ) + fold_values.append(int(sample[fold_field])) + + train_index: list[int] = [] + val_index: list[int] = [] + test_index: list[int] = [] + for i, fold in enumerate(fold_values): + if fold in train_set: + train_index.append(i) + elif fold in val_set: + val_index.append(i) + elif fold in test_set: + test_index.append(i) + else: + raise ValueError( + f"{fold_field}={fold} is not in train/val/test fold sets" + ) + + if check_patient_disjoint: + patient_to_index = getattr(dataset, "patient_to_index", None) + if patient_to_index is None: + raise TypeError( + "check_patient_disjoint=True requires dataset.patient_to_index" + ) + split_of_index = {} + for i in train_index: + split_of_index[i] = "train" + for i in val_index: + split_of_index[i] = "val" + for i in test_index: + split_of_index[i] = "test" + leaks = [] + for pid, indices in patient_to_index.items(): + splits = {split_of_index[i] for i in indices if i in split_of_index} + if len(splits) > 1: + leaks.append((str(pid), splits)) + if leaks: + details = "; ".join( + f"{pid} in {sorted(splits)}" for pid, splits in leaks + ) + raise ValueError( + "check_patient_disjoint: " + f"{len(leaks)} patient(s) appear in multiple splits: {details}" + ) + + return ( + dataset.subset(train_index), # type: ignore + dataset.subset(val_index), # type: ignore + dataset.subset(test_index), # type: ignore + ) + + def split_by_sample( dataset: SampleDataset, ratios: Union[Tuple[float, float, float], List[float]], diff --git a/pyhealth/tasks/__init__.py b/pyhealth/tasks/__init__.py index df8411db0..8f30912cb 100644 --- a/pyhealth/tasks/__init__.py +++ b/pyhealth/tasks/__init__.py @@ -88,4 +88,8 @@ def __getattr__(name: str): from .mpf_clinical_prediction import MPFClinicalPredictionTask return MPFClinicalPredictionTask + if name == "PTBXLSuperclassClassification": + from .ptbxl import PTBXLSuperclassClassification + + return PTBXLSuperclassClassification raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/pyhealth/tasks/ptbxl.py b/pyhealth/tasks/ptbxl.py new file mode 100644 index 000000000..daa3041ed --- /dev/null +++ b/pyhealth/tasks/ptbxl.py @@ -0,0 +1,273 @@ +""" +PyHealth tasks for PTB-XL multi-label ECG diagnosis. + +Dataset link: + https://physionet.org/content/ptb-xl/1.0.3/ + +Dataset paper: (please cite if you use this dataset) + Patrick Wagner, Nils Strodthoff, Ralf-Dieter Bousseljot, Dieter Kreiseler, + Fatima I. Lunze, Wojciech Samek, and Tobias Schaeffter. "PTB-XL, a large + publicly available electrocardiography dataset." Scientific Data 7, 154 + (2020). + +Dataset paper link: + https://www.nature.com/articles/s41597-020-0495-6 + +Author: + AxelNoun (GitHub: @AxelNoun) — external contributor, no NetID + +Description: + Implements the official 5-diagnostic-superclass multi-label task and + aggregation helpers. Separable for a future ``pyhealth.benchmarks`` + package. Official fold splitting lives in + ``pyhealth.datasets.split_by_strat_fold``. +""" + +from __future__ import annotations + +import functools +from collections.abc import Mapping +from pathlib import Path +from typing import Any, ClassVar + +import pandas as pd + +from pyhealth.data import Patient +from pyhealth.tasks.base_task import BaseTask + +# Official diagnostic superclasses (duplicated here to avoid importing the +# dataset package at module import time — that would circular-import via +# ``datasets.__init__`` → ``BaseDataset`` → ``tasks``). +PTBXL_DIAGNOSTIC_SUPERCLASSES = ("NORM", "MI", "STTC", "CD", "HYP") + +# Missing age is emitted as this integer so litdata / default_collate can +# batch the field. ``age_is_missing`` disambiguates it from a real age. +AGE_MISSING_SENTINEL = -1 +# HIPAA-censored ages (raw 300 in ptbxl_database.csv) are clipped to 90 in +# the sample. ``age_is_censored`` remains True; the raw 300 stays in the +# derived metadata CSV. Matches the common PTB-XL literature convention. +HIPAA_AGE_CLIP = 90 + + +@functools.cache +def _load_diagnostic_class_map_cached(path: str) -> tuple[tuple[str, str], ...]: + """Load diagnostic_class map as a hashable tuple; cached on resolved path.""" + csv_path = Path(path) + if not csv_path.is_file(): + raise FileNotFoundError(f"scp_statements.csv not found: {csv_path}") + df = pd.read_csv(csv_path, index_col=0) + if "diagnostic" not in df.columns or "diagnostic_class" not in df.columns: + raise ValueError( + f"{csv_path} must contain 'diagnostic' and 'diagnostic_class' columns" + ) + diag = df[df["diagnostic"] == 1] + mapping: list[tuple[str, str]] = [] + for code, row in diag.iterrows(): + cls = row["diagnostic_class"] + if pd.isna(cls): + continue + mapping.append((str(code), str(cls))) + return tuple(mapping) + + +def load_diagnostic_class_map( + scp_statements_path: str | Path, +) -> dict[str, str]: + """Map SCP acronym → ``diagnostic_class`` for diagnostic statements only. + + Called when constructing :class:`PTBXLSuperclassClassification`. + + Args: + scp_statements_path (str | Path): Path to official ``scp_statements.csv``. + + Returns: + dict[str, str]: Mapping used for 5-superclass aggregation. + + Examples: + >>> # doctest: +SKIP + >>> mapping = load_diagnostic_class_map("/data/ptb-xl/1.0.3/scp_statements.csv") + >>> mapping["NORM"] + 'NORM' + """ + resolved = str(Path(scp_statements_path).resolve()) + return dict(_load_diagnostic_class_map_cached(resolved)) + + +def aggregate_diagnostic_superclasses( + scp_codes: Any, + diagnostic_class_map: Mapping[str, str], +) -> list[str]: + """Aggregate stringified ``scp_codes`` to unique diagnostic superclasses. + + Includes every SCP key present in the dict (likelihood ``0`` = unknown + confidence still counts). Only codes with ``diagnostic == 1`` in + ``scp_statements.csv`` contribute a superclass. + + Args: + scp_codes (Any): Raw / stringified SCP dict from PTB-XL metadata. + diagnostic_class_map (Mapping[str, str]): From + :func:`load_diagnostic_class_map`. + + Returns: + list[str]: Sorted unique superclass labels (subset of + NORM/MI/STTC/CD/HYP). + + Examples: + >>> aggregate_diagnostic_superclasses( + ... "{'IMI': 80.0, 'SR': 0.0}", {"IMI": "MI"} + ... ) + ['MI'] + """ + from pyhealth.datasets.ptbxl import parse_scp_codes + + codes = parse_scp_codes(scp_codes) + labels = { + diagnostic_class_map[code] + for code in codes + if code in diagnostic_class_map + } + return sorted(labels) + + +def _to_bool(value: Any) -> bool: + """Coerce PTB-XL flag cells (0/1, bool, ``\"1\"``/``\"True\"``) to ``bool``. + + Args: + value (Any): Raw event attribute. + + Returns: + bool: Parsed flag. + + Examples: + >>> _to_bool(1), _to_bool("true"), _to_bool(0) + (True, True, False) + """ + if isinstance(value, bool): + return value + if isinstance(value, (int, float)): + if pd.isna(value): + return False + return int(value) != 0 + text = str(value).strip() + return text in {"1", "True", "true"} + + +class PTBXLSuperclassClassification(BaseTask): + """5-superclass multi-label classification on PTB-XL. + + Labels are the official diagnostic superclasses + ``NORM``, ``MI``, ``STTC``, ``CD``, ``HYP`` obtained by mapping diagnostic + SCP statements via ``scp_statements.csv``. + + Output label order + ------------------ + ``MultiLabelProcessor.fit`` sorts **observed** labels alphabetically, so + on the full v1.0.3 corpus the multi-hot vector is + ``CD, HYP, MI, NORM, STTC``. The constructor does not pin this order. + + Empty label sets + ---------------- + After aggregation, about 400 records have no diagnostic superclass + (Wagner et al., Scientific Data Table 9, counted on v1.0.1's 21,837 + records; mainly pacemaker ECGs). By default these are **dropped** + (``drop_empty_labels=True``) because an all-zero multi-hot target breaks + typical multi-label nonconformity scores used in CPBench, and matches + common PTB-XL literature practice. + + Each sample also carries ``strat_fold``, ``site``, ``device``, ``age``, + ``sex``, ``age_is_censored``, and ``age_is_missing`` for official splits + and downstream shift evaluations. + + Age convention + -------------- + ``age`` is always an integer so litdata and ``default_collate`` can batch + it. Missing age is ``-1`` (see ``age_is_missing``). HIPAA-censored ages + (≥90, stored as 300 in ``ptbxl_database.csv``) are clipped to ``90`` in + the sample while ``age_is_censored`` stays ``True``. The raw 300 remains + in the derived metadata CSV. + + Note: + 71-SCP multi-label classification and age regression are intentionally + not implemented here; add them in this module in a follow-up PR. + + Args: + scp_statements_path (str | Path | None): Path to ``scp_statements.csv``. + Required unless set later via the dataset's ``default_task``. + drop_empty_labels (bool): Drop records with no superclass after + aggregation. Defaults to ``True``. + + Examples: + >>> # doctest: +SKIP + >>> from pyhealth.datasets import PTBXLDataset + >>> from pyhealth.tasks import PTBXLSuperclassClassification + >>> ds = PTBXLDataset(root="/data/ptb-xl/1.0.3") + >>> samples = ds.set_task(PTBXLSuperclassClassification( + ... scp_statements_path="/data/ptb-xl/1.0.3/scp_statements.csv" + ... )) + """ + + task_name: str = "PTBXLSuperclassClassification" + input_schema: ClassVar[dict[str, str]] = {"signal": "tensor"} + output_schema: ClassVar[dict[str, str]] = {"labels": "multilabel"} + # Official five; MultiLabelProcessor still sorts observed labels. + superclass_names: ClassVar[tuple[str, ...]] = PTBXL_DIAGNOSTIC_SUPERCLASSES + + def __init__( + self, + scp_statements_path: str | Path | None = None, + drop_empty_labels: bool = True, + ) -> None: + self.scp_statements_path = ( + Path(scp_statements_path) if scp_statements_path is not None else None + ) + self.drop_empty_labels = drop_empty_labels + super().__init__() + + def _class_map(self) -> dict[str, str]: + if self.scp_statements_path is None: + raise ValueError( + "scp_statements_path is required. Pass it to " + "PTBXLSuperclassClassification(...) or use " + "PTBXLDataset.default_task." + ) + return load_diagnostic_class_map(self.scp_statements_path) + + def __call__(self, patient: Patient) -> list[dict[str, Any]]: + """Build one sample per ECG record for the patient.""" + from pyhealth.datasets.ptbxl import load_ptbxl_record + + class_map = self._class_map() + samples: list[dict[str, Any]] = [] + for event in patient.get_events(event_type="records"): + labels = aggregate_diagnostic_superclasses(event.scp_codes, class_map) + if not labels and self.drop_empty_labels: + continue + + signal = load_ptbxl_record(event.signal_file) + age_missing = _to_bool(event.age_is_missing) + age_censored = _to_bool(event.age_is_censored) + age_raw = event.age + if age_missing or age_raw is None or str(age_raw).strip() == "": + age: int = AGE_MISSING_SENTINEL + elif age_censored: + age = HIPAA_AGE_CLIP + else: + age = int(float(age_raw)) + + samples.append( + { + "patient_id": patient.patient_id, + "record_id": str(event.record_id), + "signal": signal, + "labels": labels, + "strat_fold": int(float(event.strat_fold)), + "site": event.site, + "device": event.device, + "age": age, + "age_is_censored": age_censored, + "age_is_missing": age_missing, + # sex: 0 = female, 1 = male + "sex": int(float(event.sex)), + } + ) + return samples diff --git a/pyproject.toml b/pyproject.toml index b4626e649..1976c7a66 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -69,6 +69,9 @@ nlp = [ "rouge_score~=0.1.2", "nltk~=3.9.1", ] +ptbxl = [ + "wfdb>=4.1.0", +] lint = [ "ruff~=0.15", ] @@ -134,7 +137,7 @@ python = "~=3.13" pyhealth = { path = ".", editable = true } [tool.pixi.feature.test.pypi-dependencies] -pyhealth = { path = ".", editable = true } +pyhealth = { path = ".", editable = true, extras = ["ptbxl"] } [tool.pixi.feature.nlp.pypi-dependencies] pyhealth = { path = ".", editable = true } diff --git a/test-resources/core/ptbxl/README.md b/test-resources/core/ptbxl/README.md new file mode 100644 index 000000000..7b4762849 --- /dev/null +++ b/test-resources/core/ptbxl/README.md @@ -0,0 +1,15 @@ +# Synthetic PTB-XL-shaped fixture for unit tests (NOT real PhysioNet data). +# +# Contents: +# - ptbxl_database.csv / scp_statements.csv: tiny CSV stubs (committed) +# - records100 / records500: synthesized at test time by +# tests/core/test_ptbxl.py::_materialize_fixture (12 leads x 50 samples +# at 100 Hz; 12 x 250 at 500 Hz — unequal dims so a missing transpose +# fails the shape assertion). Waveform binaries are not committed. +# +# Edge-case coverage: +# - ecg_id=2: age=300 (HIPAA censored ≥90); likelihood 0 (SR) +# - ecg_id=3: missing age + non-empty non-diagnostic dict (PACE-only) +# - ecg_id=4: true multi-label (IMI+LVH → MI and HYP) +# - ecg_id=5: empty scp_codes dict {} +# - ecg_id=1 and 2 share patient_id (used by check_patient_disjoint tests) diff --git a/test-resources/core/ptbxl/ptbxl_database.csv b/test-resources/core/ptbxl/ptbxl_database.csv new file mode 100644 index 000000000..b72999108 --- /dev/null +++ b/test-resources/core/ptbxl/ptbxl_database.csv @@ -0,0 +1,6 @@ +ecg_id,patient_id,age,sex,site,device,recording_date,scp_codes,strat_fold,filename_lr,filename_hr +1,15709.0,65.0,1,0,CS-12,1990-01-01 10:00:00,"{'NORM': 100.0}",1,records100/00000/00001_lr,records500/00000/00001_hr +2,15709.0,300.0,0,1,CS-12,1991-06-15 11:30:00,"{'IMI': 80.0, 'SR': 0.0}",9,records100/00000/00002_lr,records500/00000/00002_hr +3,99.0,,1,2,CS-12,1992-03-20 09:15:00,"{'PACE': 100.0}",10,records100/00000/00003_lr,records500/00000/00003_hr +4,42.0,55.0,0,0,CS-12,1993-01-01 08:00:00,"{'IMI': 100.0, 'LVH': 80.0}",2,records100/00000/00004_lr,records500/00000/00004_hr +5,7.0,40.0,1,1,CS-12,1994-01-01 08:00:00,"{}",3,records100/00000/00005_lr,records500/00000/00005_hr diff --git a/test-resources/core/ptbxl/scp_statements.csv b/test-resources/core/ptbxl/scp_statements.csv new file mode 100644 index 000000000..2e08b8d55 --- /dev/null +++ b/test-resources/core/ptbxl/scp_statements.csv @@ -0,0 +1,9 @@ +,description,diagnostic,form,rhythm,diagnostic_class,diagnostic_subclass,Statement Category,SCP-ECG Statement Description,AHA code,aECG REFID,CDISC Code,DICOM Code +NORM,normal ECG,1.0,,,NORM,NORM,Normal/abnormal,normal ECG,1,,,F-000B7 +IMI,inferior myocardial infarction,1.0,,,MI,IMI,Myocardial Infarction,inferior myocardial infarction,161,,, +LVH,left ventricular hypertrophy,1.0,,,HYP,LVH,Ventricular Hypertrophy,left ventricular hypertrophy,142,,C71076, +SR,sinus rhythm,,,1.0,,,Statements related to impulse formation (abnormalities),sinus rhythm,20,MDC_ECG_RHY_SINUS_RHY,, +PACE,normal functioning artificial pacemaker,,,1.0,,,Pacemaker types and pacemaker function,normal functioning artificial pacemaker,,,, +STTC,non-specific ST-T changes,1.0,,,STTC,STTC,other ST-T descriptive statements,non-specific ST-T changes,,,, +CD,conduction disturbance,1.0,,,CD,CD,Intraventricular and intra-atrial Conduction disturbances,conduction disturbance,,,, +HYP,hypertrophy,1.0,,,HYP,HYP,Ventricular Hypertrophy,hypertrophy,,,, diff --git a/tests/core/test_ptbxl.py b/tests/core/test_ptbxl.py new file mode 100644 index 000000000..238806628 --- /dev/null +++ b/tests/core/test_ptbxl.py @@ -0,0 +1,691 @@ +""" +Unit tests for PTBXLDataset, PTBXLSuperclassClassification, and +split_by_strat_fold. + +Uses small synthetic WFDB fixtures under test-resources/core/ptbxl/ (not real +PhysioNet records). Waveforms are generated at test time. Covers censored +age (300), missing age, and empty diagnostic-superclass labels. + +Author: + AxelNoun (GitHub: @AxelNoun) — external contributor, no NetID +""" + +from __future__ import annotations + +import shutil +import tempfile +import unittest +from pathlib import Path +from unittest.mock import patch + +import numpy as np +import pandas as pd +import torch + +from pyhealth.datasets.ptbxl import ( + AGE_CENSOR_SENTINEL, + PTBXLDataset, + format_patient_id, + is_age_censored, + is_age_missing, + load_ptbxl_record, + metadata_filename, + parse_scp_codes, +) +from pyhealth.datasets.splitter import split_by_strat_fold +from pyhealth.tasks.ptbxl import ( + PTBXLSuperclassClassification, + _to_bool, + aggregate_diagnostic_superclasses, + load_diagnostic_class_map, +) + +FIXTURE_ROOT = ( + Path(__file__).resolve().parents[2] / "test-resources" / "core" / "ptbxl" +) + + +def _write_dummy_wfdb( + record_base: Path, + n_leads: int = 12, + n_samples: int = 50, + fs: int = 100, +) -> None: + """Write a minimal WFDB record (header + int16 dat) for tests.""" + record_base.parent.mkdir(parents=True, exist_ok=True) + data = np.zeros((n_samples, n_leads), dtype=np.int16) + for lead in range(n_leads): + data[:, lead] = np.arange(n_samples, dtype=np.int16) + lead + Path(str(record_base) + ".dat").write_bytes(data.tobytes()) + name = record_base.name + lines = [f"{name} {n_leads} {fs} {n_samples}"] + for i in range(n_leads): + lines.append(f"{name}.dat 16 1000.0(0)/uV 16 0 0 0 0 {i}") + Path(str(record_base) + ".hea").write_text("\n".join(lines) + "\n", encoding="utf-8") + + +def _materialize_fixture(dest: Path) -> Path: + """Copy committed CSVs and synthesize WFDB records under ``dest``.""" + dest.mkdir(parents=True, exist_ok=True) + shutil.copy(FIXTURE_ROOT / "ptbxl_database.csv", dest / "ptbxl_database.csv") + shutil.copy(FIXTURE_ROOT / "scp_statements.csv", dest / "scp_statements.csv") + db = pd.read_csv(dest / "ptbxl_database.csv") + for col, fs, n_samples in ( + ("filename_lr", 100, 50), + ("filename_hr", 500, 250), + ): + for rel in db[col]: + _write_dummy_wfdb(dest / str(rel), n_samples=n_samples, fs=fs) + return dest + + +class TestPTBXLHelpers(unittest.TestCase): + def test_format_patient_id_strips_float(self): + self.assertEqual(format_patient_id(15709.0), "15709") + self.assertEqual(format_patient_id("15709.0"), "15709") + + def test_format_patient_id_includes_ecg_id_in_errors(self): + with self.assertRaisesRegex(ValueError, r"ecg_id=99"): + format_patient_id(None, ecg_id=99) + with self.assertRaisesRegex(ValueError, r"ecg_id=7"): + format_patient_id("not-a-number", ecg_id=7) + + def test_parse_scp_codes_keeps_likelihood_zero(self): + codes = parse_scp_codes("{'IMI': 80.0, 'SR': 0.0}") + self.assertEqual(codes["SR"], 0.0) + self.assertIn("IMI", codes) + + def test_parse_scp_codes_empty(self): + self.assertEqual(parse_scp_codes("{}"), {}) + self.assertEqual(parse_scp_codes(None), {}) + + def test_age_missing_vs_censored(self): + self.assertTrue(is_age_censored(AGE_CENSOR_SENTINEL)) + self.assertFalse(is_age_missing(AGE_CENSOR_SENTINEL)) + self.assertTrue(is_age_missing(float("nan"))) + self.assertTrue(is_age_missing("")) + self.assertFalse(is_age_censored(float("nan"))) + self.assertFalse(is_age_censored(65)) + + def test_to_bool_parses_flag_cells(self): + self.assertTrue(_to_bool(1)) + self.assertTrue(_to_bool("True")) + self.assertTrue(_to_bool("true")) + self.assertTrue(_to_bool(True)) + self.assertFalse(_to_bool(0)) + self.assertFalse(_to_bool("0")) + self.assertFalse(_to_bool(False)) + + def test_metadata_filename_includes_rate_and_root(self): + name_a = metadata_filename(100, "/data/ptb-xl/a") + name_b = metadata_filename(100, "/data/ptb-xl/b") + self.assertTrue(name_a.startswith("ptbxl-pyhealth-100hz-")) + self.assertTrue(name_a.endswith(".csv")) + self.assertNotEqual(name_a, name_b) + self.assertNotEqual( + metadata_filename(100, "/tmp/x"), metadata_filename(500, "/tmp/x") + ) + keyed_a = metadata_filename(100, "/tmp/x", "aaaaaaaaaa") + keyed_b = metadata_filename(100, "/tmp/x", "bbbbbbbbbb") + self.assertIn("aaaaaaaaaa", keyed_a) + self.assertNotEqual(keyed_a, keyed_b) + + +class TestPTBXLAggregation(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.class_map = load_diagnostic_class_map(FIXTURE_ROOT / "scp_statements.csv") + + def test_aggregate_superclasses(self): + labels = aggregate_diagnostic_superclasses( + "{'NORM': 100.0}", self.class_map + ) + self.assertEqual(labels, ["NORM"]) + + labels = aggregate_diagnostic_superclasses( + "{'IMI': 80.0, 'SR': 0.0}", self.class_map + ) + self.assertEqual(labels, ["MI"]) + + # Non-empty dict with no diagnostic statements (PACE is rhythm-only). + labels = aggregate_diagnostic_superclasses( + "{'PACE': 100.0}", self.class_map + ) + self.assertEqual(labels, []) + + # Empty dict — distinct path from PACE-only. + labels = aggregate_diagnostic_superclasses("{}", self.class_map) + self.assertEqual(labels, []) + + # True multi-label: two distinct diagnostic superclasses. + labels = aggregate_diagnostic_superclasses( + "{'IMI': 100.0, 'LVH': 80.0}", self.class_map + ) + self.assertEqual(labels, ["HYP", "MI"]) + + +class TestPTBXLDatasetMetadata(unittest.TestCase): + def setUp(self): + self.tmp = Path(tempfile.mkdtemp()) + self.data_root = _materialize_fixture(self.tmp / "data") + self.cache_dir = self.tmp / "meta_cache" + + def tearDown(self): + shutil.rmtree(self.tmp, ignore_errors=True) + + def test_prepare_metadata_writes_cache_not_root(self): + ds = PTBXLDataset( + root=str(self.data_root), + metadata_cache_dir=self.cache_dir, + sampling_rate=100, + cache_dir=self.tmp / "pyhealth_cache", + ) + csv_path = self.cache_dir / ds.metadata_file_name + self.assertTrue(csv_path.is_file()) + self.assertFalse((self.data_root / ds.metadata_file_name).exists()) + meta = pd.read_csv(csv_path) + self.assertEqual( + sorted(meta["patient_id"].astype(str).unique()), + ["15709", "42", "7", "99"], + ) + self.assertNotIn("15709.0", set(meta["patient_id"].astype(str))) + self.assertEqual(len(meta), 5) + # Age flags: record 2 censored, record 3 missing + row2 = meta[meta["record_id"].astype(str) == "2"].iloc[0] + row3 = meta[meta["record_id"].astype(str) == "3"].iloc[0] + self.assertEqual(int(row2["age_is_censored"]), 1) + self.assertEqual(int(row2["age_is_missing"]), 0) + self.assertEqual(int(row2["age"]), 300) + self.assertEqual(int(row3["age_is_missing"]), 1) + self.assertEqual(int(row3["age_is_censored"]), 0) + # Signal paths are extension-free record bases + for path in meta["signal_file"]: + self.assertFalse(str(path).endswith(".hea")) + self.assertFalse(str(path).endswith(".dat")) + self.assertEqual(ds.sampling_rate, 100) + + def test_100_and_500_hz_coexist(self): + ds = PTBXLDataset( + root=str(self.data_root), + metadata_cache_dir=self.cache_dir, + sampling_rate=100, + cache_dir=self.tmp / "c100", + ) + self.assertTrue((self.cache_dir / ds.metadata_file_name).is_file()) + # Second rate gets a distinct filename (rate + root hash). + ds500 = PTBXLDataset( + root=str(self.data_root), + metadata_cache_dir=self.cache_dir, + sampling_rate=500, + cache_dir=self.tmp / "c500", + ) + self.assertTrue((self.cache_dir / ds500.metadata_file_name).is_file()) + self.assertNotEqual(ds.metadata_file_name, ds500.metadata_file_name) + + def test_end_to_end_reads_event_from_fixture(self): + """Instantiate PTBXLDataset and read a real event (not just helpers).""" + ds = PTBXLDataset( + root=str(self.data_root), + metadata_cache_dir=self.cache_dir, + sampling_rate=100, + cache_dir=self.tmp / "pyhealth_cache_e2e", + ) + # Config must already point at the derived CSV before any load. + self.assertEqual( + ds.config.tables["records"].file_path, ds.metadata_file_name + ) + self.assertEqual(Path(ds.root).resolve(), self.data_root.resolve()) + self.assertTrue( + (ds.metadata_cache_dir / ds.config.tables["records"].file_path).is_file() + ) + + patient_ids = ds.unique_patient_ids + self.assertGreaterEqual(len(patient_ids), 1) + patient = ds.get_patient(patient_ids[0]) + events = patient.get_events(event_type="records") + self.assertGreaterEqual(len(events), 1) + event = events[0] + self.assertTrue(hasattr(event, "signal_file")) + self.assertTrue(hasattr(event, "strat_fold")) + self.assertTrue(hasattr(event, "scp_codes")) + self.assertTrue(str(event.signal_file)) + # Absolute waveform path must live under the fixture data root. + Path(str(event.signal_file)).resolve().relative_to(self.data_root.resolve()) + # load_table must restore self.root after scanning the cache CSV. + self.assertEqual(Path(ds.root).resolve(), self.data_root.resolve()) + + +@unittest.skipUnless( + __import__("importlib").util.find_spec("wfdb") is not None, + "wfdb optional extra not installed", +) +class TestPTBXLSignalIO(unittest.TestCase): + def setUp(self): + self.tmp = Path(tempfile.mkdtemp()) + self.data_root = _materialize_fixture(self.tmp / "data") + + def tearDown(self): + shutil.rmtree(self.tmp, ignore_errors=True) + + def test_load_fixture_shape_channels_time(self): + record = self.data_root / "records100" / "00000" / "00001_lr" + self.assertTrue(Path(str(record) + ".hea").is_file()) + self.assertTrue(Path(str(record) + ".dat").is_file()) + signal = load_ptbxl_record(record) + # 12 leads != 50 samples: catches a missing transpose. + self.assertEqual(signal.shape, (12, 50)) + self.assertNotEqual(signal.shape[0], signal.shape[1]) + + def test_load_strips_extension_never_appends(self): + record = self.data_root / "records100" / "00000" / "00001_lr" + signal = load_ptbxl_record(Path(str(record) + ".hea")) + self.assertEqual(signal.shape, (12, 50)) + self.assertNotEqual(signal.shape[0], signal.shape[1]) + + +class TestPTBXLTaskAndSplit(unittest.TestCase): + def setUp(self): + self.tmp = Path(tempfile.mkdtemp()) + self.data_root = _materialize_fixture(self.tmp / "data") + self.cache_dir = self.tmp / "meta_cache" + self.dataset = PTBXLDataset( + root=str(self.data_root), + metadata_cache_dir=self.cache_dir, + sampling_rate=100, + cache_dir=self.tmp / "pyhealth_cache", + ) + + def tearDown(self): + shutil.rmtree(self.tmp, ignore_errors=True) + + def test_task_drops_empty_labels_by_default(self): + fake_signal = np.zeros((12, 50), dtype=np.float32) + + def _fake_load(_path): + return fake_signal + + task = PTBXLSuperclassClassification( + scp_statements_path=self.data_root / "scp_statements.csv", + drop_empty_labels=True, + ) + with patch( + "pyhealth.datasets.ptbxl.load_ptbxl_record", side_effect=_fake_load + ): + samples = [] + for patient in self.dataset.iter_patients(): + samples.extend(task(patient)) + # Keep: 1 NORM, 2 MI, 4 HYP+MI. Drop: 3 PACE-only and 5 empty {}. + self.assertEqual(len(samples), 3) + labels = {tuple(s["labels"]) for s in samples} + self.assertEqual(labels, {("NORM",), ("MI",), ("HYP", "MI")}) + multi = next(s for s in samples if len(s["labels"]) > 1) + self.assertEqual(multi["labels"], ["HYP", "MI"]) + for s in samples: + self.assertIn("strat_fold", s) + self.assertIn("site", s) + self.assertIn("device", s) + self.assertIn("sex", s) + + def test_task_keeps_empty_when_disabled(self): + fake_signal = np.zeros((12, 50), dtype=np.float32) + task = PTBXLSuperclassClassification( + scp_statements_path=self.data_root / "scp_statements.csv", + drop_empty_labels=False, + ) + with patch( + "pyhealth.datasets.ptbxl.load_ptbxl_record", return_value=fake_signal + ): + samples = [] + for patient in self.dataset.iter_patients(): + samples.extend(task(patient)) + self.assertEqual(len(samples), 5) + empty = [s for s in samples if s["labels"] == []] + # Both PACE-only and empty-dict {} yield empty superclass sets. + self.assertEqual(len(empty), 2) + + def test_split_by_strat_fold(self): + # Lightweight fake SampleDataset-like object + class _FakeDS: + def __init__(self, samples): + self._samples = samples + + def __len__(self): + return len(self._samples) + + def __getitem__(self, i): + return self._samples[i] + + def subset(self, indices): + return _FakeDS([self._samples[i] for i in indices]) + + samples = [ + {"strat_fold": 1, "id": "a"}, + {"strat_fold": 9, "id": "b"}, + {"strat_fold": 10, "id": "c"}, + {"strat_fold": 3, "id": "d"}, + ] + train, val, test = split_by_strat_fold(_FakeDS(samples)) + self.assertEqual([s["id"] for s in train._samples], ["a", "d"]) + self.assertEqual([s["id"] for s in val._samples], ["b"]) + self.assertEqual([s["id"] for s in test._samples], ["c"]) + + def test_split_by_strat_fold_check_patient_disjoint(self): + class _FakeDS: + def __init__(self, samples): + self._samples = samples + self.patient_to_index = {} + for i, sample in enumerate(samples): + self.patient_to_index.setdefault(sample["patient_id"], []).append( + i + ) + + def __len__(self): + return len(self._samples) + + def __getitem__(self, i): + return self._samples[i] + + def subset(self, indices): + return _FakeDS([self._samples[i] for i in indices]) + + leaked = [ + {"strat_fold": 1, "patient_id": "15709", "id": "a"}, + {"strat_fold": 9, "patient_id": "15709", "id": "b"}, + {"strat_fold": 10, "patient_id": "99", "id": "c"}, + {"strat_fold": 3, "patient_id": "42", "id": "d"}, + ] + train, val, test = split_by_strat_fold(_FakeDS(leaked)) + self.assertEqual([s["id"] for s in train._samples], ["a", "d"]) + self.assertEqual([s["id"] for s in val._samples], ["b"]) + self.assertEqual([s["id"] for s in test._samples], ["c"]) + with self.assertRaisesRegex(ValueError, r"15709"): + split_by_strat_fold(_FakeDS(leaked), check_patient_disjoint=True) + + class _NoMap: + def __len__(self): + return 1 + + def __getitem__(self, i): + return {"strat_fold": 1} + + def subset(self, indices): + return indices + + with self.assertRaises(TypeError): + split_by_strat_fold( + _NoMap(), + train_folds=(1,), + val_folds=(9,), + test_folds=(10,), + check_patient_disjoint=True, + ) + + def test_split_by_strat_fold_uses_precomputed_folds(self): + class _BoomDS: + def __init__(self, n): + self._n = n + + def __len__(self): + return self._n + + def __getitem__(self, i): + raise AssertionError( + "dataset[i] must not be called when folds= is provided" + ) + + def subset(self, indices): + return list(indices) + + train, val, test = split_by_strat_fold( + _BoomDS(4), + train_folds=(1, 3), + val_folds=(9,), + test_folds=(10,), + folds=[1, 9, 10, 3], + ) + self.assertEqual(train, [0, 3]) + self.assertEqual(val, [1]) + self.assertEqual(test, [2]) + + +@unittest.skipUnless( + __import__("importlib").util.find_spec("wfdb") is not None, + "wfdb optional extra not installed", +) +class TestPTBXLSetTaskE2E(unittest.TestCase): + """Exercise the real set_task / litdata / SampleDataset path (no mocks).""" + + def setUp(self): + self.tmp = Path(tempfile.mkdtemp()) + self.data_root = _materialize_fixture(self.tmp / "data") + self.cache_dir = self.tmp / "meta_cache" + self.dataset = PTBXLDataset( + root=str(self.data_root), + metadata_cache_dir=self.cache_dir, + sampling_rate=100, + cache_dir=self.tmp / "pyhealth_cache", + ) + self.task = PTBXLSuperclassClassification( + scp_statements_path=self.data_root / "scp_statements.csv", + ) + + def tearDown(self): + shutil.rmtree(self.tmp, ignore_errors=True) + + def test_set_task_builds_samples_without_mocks(self): + samples = self.dataset.set_task(self.task, num_workers=1) + # Keep: ecg 1 NORM, 2 MI, 4 HYP+MI. Drop: 3 PACE-only and 5 empty {}. + self.assertEqual(len(samples), 3) + + sample = samples[0] + self.assertEqual(tuple(sample["signal"].shape), (12, 50)) + self.assertNotEqual(sample["signal"].shape[0], sample["signal"].shape[1]) + self.assertIn("strat_fold", sample) + self.assertIn("labels", sample) + # Multi-hot from MultiLabelProcessor, not a list of strings. + self.assertEqual(tuple(sample["labels"].shape), (3,)) + self.assertTrue(torch.is_tensor(sample["labels"])) + + def test_multilabel_vocab_order_is_alphabetical_observed(self): + samples = self.dataset.set_task(self.task, num_workers=1) + vocab = samples.output_processors["labels"].label_vocab + # Fixture only observes HYP, MI, NORM — not CD/STTC. Processor sorts. + self.assertEqual(list(vocab.keys()), ["HYP", "MI", "NORM"]) + self.assertEqual(vocab, {"HYP": 0, "MI": 1, "NORM": 2}) + + by_record = {str(s["record_id"]): s for s in samples} + # ecg 1 = NORM → [0, 0, 1]; ecg 2 = MI → [0, 1, 0]; ecg 4 = HYP+MI → [1, 1, 0] + self.assertEqual(by_record["1"]["labels"].tolist(), [0.0, 0.0, 1.0]) + self.assertEqual(by_record["2"]["labels"].tolist(), [0.0, 1.0, 0.0]) + self.assertEqual(by_record["4"]["labels"].tolist(), [1.0, 1.0, 0.0]) + + def test_split_by_strat_fold_on_sample_dataset(self): + samples = self.dataset.set_task(self.task, num_workers=1) + # Kept folds are 1 (ecg 1), 2 (ecg 4), 9 (ecg 2). Fold 10 was dropped. + train, val, test = split_by_strat_fold( + samples, + train_folds=(1,), + val_folds=(2,), + test_folds=(9,), + ) + self.assertEqual(len(train), 1) + self.assertEqual(len(val), 1) + self.assertEqual(len(test), 1) + self.assertEqual(int(train[0]["strat_fold"]), 1) + self.assertEqual(int(val[0]["strat_fold"]), 2) + self.assertEqual(int(test[0]["strat_fold"]), 9) + self.assertEqual(tuple(train[0]["signal"].shape), (12, 50)) + + def test_split_by_strat_fold_detects_fixture_patient_leak(self): + samples = self.dataset.set_task(self.task, num_workers=1) + # Fixture: ecg 1 (fold 1) and ecg 2 (fold 9) share patient_id 15709. + with self.assertRaisesRegex(ValueError, r"15709"): + split_by_strat_fold( + samples, + train_folds=(1,), + val_folds=(2,), + test_folds=(9,), + check_patient_disjoint=True, + ) + + def test_age_sentinel_through_set_task(self): + """Missing age must be an int sentinel; censored age is clipped to 90.""" + task = PTBXLSuperclassClassification( + scp_statements_path=self.data_root / "scp_statements.csv", + drop_empty_labels=False, + ) + samples = self.dataset.set_task(task, num_workers=1) + by_record = {str(s["record_id"]): s for s in samples} + + missing = by_record["3"] + self.assertTrue(missing["age_is_missing"]) + self.assertIsInstance(missing["age"], (int, np.integer)) + self.assertEqual(int(missing["age"]), -1) + self.assertIsNotNone(missing["age"]) + + censored = by_record["2"] + self.assertTrue(censored["age_is_censored"]) + self.assertFalse(censored["age_is_missing"]) + self.assertEqual(int(censored["age"]), 90) + + +class TestPTBXLCacheIsolation(unittest.TestCase): + def setUp(self): + self.tmp = Path(tempfile.mkdtemp()) + + def tearDown(self): + shutil.rmtree(self.tmp, ignore_errors=True) + + def test_distinct_data_roots_get_distinct_cache_dirs(self): + root_a = _materialize_fixture(self.tmp / "a") + root_b = _materialize_fixture(self.tmp / "b") + db_b = pd.read_csv(root_b / "ptbxl_database.csv") + db_b.loc[0, "scp_codes"] = "{'LVH': 100.0}" + db_b.to_csv(root_b / "ptbxl_database.csv", index=False) + + shared_meta = self.tmp / "meta" + shared_cache = self.tmp / "pyhealth_cache" + ds_a = PTBXLDataset( + root=str(root_a), + metadata_cache_dir=shared_meta, + sampling_rate=100, + cache_dir=shared_cache, + ) + ds_b = PTBXLDataset( + root=str(root_b), + metadata_cache_dir=shared_meta, + sampling_rate=100, + cache_dir=shared_cache, + ) + self.assertNotEqual(ds_a.cache_dir, ds_b.cache_dir) + self.assertNotEqual(ds_a.dataset_name, ds_b.dataset_name) + + def test_source_csv_change_invalidates_derived_metadata(self): + data_root = _materialize_fixture(self.tmp / "data") + meta = self.tmp / "meta" + ds1 = PTBXLDataset( + root=str(data_root), + metadata_cache_dir=meta, + sampling_rate=100, + cache_dir=self.tmp / "c1", + ) + name1 = ds1.metadata_file_name + derived1 = pd.read_csv(meta / name1) + + db = pd.read_csv(data_root / "ptbxl_database.csv") + db.loc[0, "scp_codes"] = "{'LVH': 100.0}" + db.to_csv(data_root / "ptbxl_database.csv", index=False) + + ds2 = PTBXLDataset( + root=str(data_root), + metadata_cache_dir=meta, + sampling_rate=100, + cache_dir=self.tmp / "c2", + ) + self.assertNotEqual(ds2.metadata_file_name, name1) + derived2 = pd.read_csv(meta / ds2.metadata_file_name) + codes1 = str( + derived1.loc[derived1["record_id"].astype(str) == "1", "scp_codes"].iloc[0] + ) + codes2 = str( + derived2.loc[derived2["record_id"].astype(str) == "1", "scp_codes"].iloc[0] + ) + self.assertNotEqual(codes1, codes2) + + +class TestPTBXLTaskStability(unittest.TestCase): + def setUp(self): + self.tmp = Path(tempfile.mkdtemp()) + self.data_root = _materialize_fixture(self.tmp / "data") + self.dataset = PTBXLDataset( + root=str(self.data_root), + metadata_cache_dir=self.tmp / "meta_cache", + sampling_rate=100, + cache_dir=self.tmp / "pyhealth_cache", + ) + + def tearDown(self): + shutil.rmtree(self.tmp, ignore_errors=True) + + def test_vars_task_stable_after_call(self): + task = PTBXLSuperclassClassification( + scp_statements_path=self.data_root / "scp_statements.csv", + ) + before = dict(vars(task)) + fake_signal = np.zeros((12, 50), dtype=np.float32) + with patch( + "pyhealth.datasets.ptbxl.load_ptbxl_record", return_value=fake_signal + ): + patient = self.dataset.get_patient(self.dataset.unique_patient_ids[0]) + task(patient) + after = dict(vars(task)) + self.assertEqual(before, after) + + def test_age_sentinel_from_task_call(self): + task = PTBXLSuperclassClassification( + scp_statements_path=self.data_root / "scp_statements.csv", + drop_empty_labels=False, + ) + fake_signal = np.zeros((12, 50), dtype=np.float32) + with patch( + "pyhealth.datasets.ptbxl.load_ptbxl_record", return_value=fake_signal + ): + samples = [] + for patient in self.dataset.iter_patients(): + samples.extend(task(patient)) + by_record = {str(s["record_id"]): s for s in samples} + + missing = by_record["3"] + self.assertTrue(missing["age_is_missing"]) + self.assertIsInstance(missing["age"], int) + self.assertEqual(missing["age"], -1) + + censored = by_record["2"] + self.assertTrue(censored["age_is_censored"]) + self.assertEqual(censored["age"], 90) + + present = by_record["1"] + self.assertFalse(present["age_is_missing"]) + self.assertFalse(present["age_is_censored"]) + self.assertEqual(present["age"], 65) + + +class TestWFDBImportError(unittest.TestCase): + def test_import_error_points_at_extra(self): + import builtins + + real_import = builtins.__import__ + + def _guard(name, *args, **kwargs): + if name == "wfdb" or name.startswith("wfdb."): + raise ImportError("blocked") + return real_import(name, *args, **kwargs) + + with patch("builtins.__import__", side_effect=_guard): + with self.assertRaises(ImportError) as ctx: + load_ptbxl_record("/tmp/does-not-matter") + self.assertIn("pyhealth[ptbxl]", str(ctx.exception)) + + +if __name__ == "__main__": + unittest.main()