TaskAug ECG replication: PTBXLDataset, ECGBinaryClassification, TaskAugResNet

docs/api/datasets.rst

@@ -226,6 +226,7 @@ Available Datasets
     datasets/pyhealth.datasets.MIMIC4Dataset
     datasets/pyhealth.datasets.MedicalTranscriptionsDataset
     datasets/pyhealth.datasets.CardiologyDataset
+    datasets/pyhealth.datasets.PTBXLDataset
     datasets/pyhealth.datasets.eICUDataset
     datasets/pyhealth.datasets.ISRUCDataset
     datasets/pyhealth.datasets.MIMICExtractDataset

docs/api/datasets/pyhealth.datasets.PTBXLDataset.rst

@@ -0,0 +1,21 @@
+pyhealth.datasets.PTBXLDataset
+==============================
+
+PTB-XL is a large publicly available 12-lead ECG dataset containing 21 799
+10-second recordings from 18 869 patients, annotated with SCP-ECG diagnostic
+statements and mapped to four binary superclass labels: MI (myocardial
+infarction), HYP (hypertrophy), STTC (ST/T-change), and CD (conduction
+disturbance).
+
+Paper: Wagner et al. (2020). PTB-XL, a large publicly available
+electrocardiography dataset. *Scientific Data*, 7(154).
+https://physionet.org/content/ptb-xl/1.0.3/
+
+TaskAug reference: Raghu et al. (2022). Data Augmentation for
+Electrocardiograms. *CHIL*, PMLR 174.
+https://proceedings.mlr.press/v174/raghu22a.html
+
+.. autoclass:: pyhealth.datasets.PTBXLDataset
+   :members:
+   :undoc-members:
+   :show-inheritance:

docs/api/datasets/pyhealth.datasets.ptbxl.rst

@@ -0,0 +1,21 @@
+pyhealth.datasets.ptbxl
+=======================
+
+Overview
+--------
+PTB-XL is a large publicly available 12-lead ECG dataset containing 21 799
+10-second records from 18 869 patients (Wagner et al., 2020). This PyHealth
+wrapper maps each record to four binary diagnostic superclass labels —
+myocardial infarction (MI), hypertrophy (HYP), ST/T-change (STTC), and
+conduction disturbance (CD) — enabling direct use with
+:class:`~pyhealth.tasks.ecg_classification.ECGBinaryClassification`.
+
+Source: https://physionet.org/content/ptb-xl/1.0.3/
+
+API Reference
+-------------
+
+.. autoclass:: pyhealth.datasets.ptbxl.PTBXLDataset
+   :members:
+   :undoc-members:
+   :show-inheritance:

docs/api/models.rst

@@ -189,6 +189,7 @@ API Reference
     models/pyhealth.models.JambaEHR
     models/pyhealth.models.ContraWR
     models/pyhealth.models.SparcNet
+    models/pyhealth.models.TaskAugResNet
     models/pyhealth.models.StageNet
     models/pyhealth.models.StageAttentionNet
     models/pyhealth.models.AdaCare

docs/api/models/pyhealth.models.TaskAugResNet.rst

@@ -0,0 +1,28 @@
+pyhealth.models.TaskAugResNet
+==============================
+
+TaskAug is a differentiable, task-adaptive data augmentation framework for
+ECG binary classification. A K-stage Gumbel-Softmax augmentation policy
+selects from eight time-series operations with class-specific learnable
+magnitudes. The policy is trained jointly with a 1-D ResNet-18 backbone via
+bi-level optimisation (inner loop: backbone on augmented training data; outer
+loop: policy on clean validation loss).
+
+Paper: Raghu et al. (2022). Data Augmentation for Electrocardiograms.
+*Conference on Health, Inference, and Learning (CHIL)*, PMLR 174.
+https://proceedings.mlr.press/v174/raghu22a.html
+
+.. autoclass:: pyhealth.models.taskaug_resnet.TaskAugPolicy
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+.. autoclass:: pyhealth.models.taskaug_resnet._ResNet1D
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+.. autoclass:: pyhealth.models.TaskAugResNet
+   :members:
+   :undoc-members:
+   :show-inheritance:

docs/api/models/pyhealth.models.taskaug_resnet.rst

@@ -0,0 +1,28 @@
+pyhealth.models.taskaug_resnet
+==============================
+
+Overview
+--------
+Replication of the TaskAug framework from Raghu et al. (2022) *Data
+Augmentation for Electrocardiograms* (CHIL, PMLR 174). A K-stage
+differentiable augmentation policy (:class:`TaskAugPolicy`) selects among
+seven ECG-specific operations via Gumbel-Softmax with class-specific learnable
+magnitudes. The policy is trained jointly with a 1-D ResNet-18 backbone
+(:class:`_ResNet1D`) using a bi-level optimisation scheme: the inner loop
+updates backbone weights on augmented training data while the outer loop
+updates policy weights on clean validation loss.
+
+Paper: https://proceedings.mlr.press/v174/raghu22a.html
+
+API Reference
+-------------
+
+.. autoclass:: pyhealth.models.taskaug_resnet.TaskAugResNet
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+.. autoclass:: pyhealth.models.taskaug_resnet.TaskAugPolicy
+   :members:
+   :undoc-members:
+   :show-inheritance:

docs/api/tasks.rst

@@ -209,6 +209,7 @@ Available Tasks
     In-Hospital Mortality (MIMIC-IV) <tasks/pyhealth.tasks.InHospitalMortalityMIMIC4>
     MIMIC-III ICD-9 Coding <tasks/pyhealth.tasks.MIMIC3ICD9Coding>
     Cardiology Detection <tasks/pyhealth.tasks.cardiology_detect>
+    ECG Binary Classification (PTB-XL) <tasks/pyhealth.tasks.ECGBinaryClassification>
     COVID-19 CXR Classification <tasks/pyhealth.tasks.COVID19CXRClassification>
     DKA Prediction (MIMIC-IV) <tasks/pyhealth.tasks.dka>
     Drug Recommendation <tasks/pyhealth.tasks.drug_recommendation>

docs/api/tasks/pyhealth.tasks.ecg_classification.rst

@@ -0,0 +1,20 @@
+pyhealth.tasks.ecg_classification
+==================================
+
+Overview
+--------
+Binary ECG classification task for the PTB-XL dataset, implementing the
+task interface from Raghu et al. (2022) *Data Augmentation for
+Electrocardiograms* (CHIL, PMLR 174). Each ECG record is loaded from disk
+via WFDB, per-lead z-score normalised, and padded or truncated to a fixed
+time length. Supports four diagnostic superclasses: MI, HYP, STTC, and CD.
+
+Paper: https://proceedings.mlr.press/v174/raghu22a.html
+
+API Reference
+-------------
+
+.. autoclass:: pyhealth.tasks.ecg_classification.ECGBinaryClassification
+   :members:
+   :undoc-members:
+   :show-inheritance: