ecgen 0.4.0__tar.gz

ecgen-0.4.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Vajira Thambawita
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

ecgen-0.4.0/PKG-INFO

@@ -0,0 +1,37 @@
+Metadata-Version: 2.4
+Name: ecgen
+Version: 0.4.0
+Summary: ECG generation and modeling experiments
+License: MIT License
+        
+        Copyright (c) 2026 Vajira Thambawita
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy
+        of this software and associated documentation files (the "Software"), to deal
+        in the Software without restriction, including without limitation the rights
+        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+        copies of the Software, and to permit persons to whom the Software is
+        furnished to do so, subject to the following conditions:
+        
+        The above copyright notice and this permission notice shall be included in all
+        copies or substantial portions of the Software.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+        SOFTWARE.
+        
+Project-URL: Homepage, https://github.com/vlbthambawita/ECGEN
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch
+Requires-Dist: pyyaml
+Requires-Dist: pytorch-lightning
+Requires-Dist: pandas
+Requires-Dist: scikit-learn
+Requires-Dist: wfdb
+Dynamic: license-file

ecgen-0.4.0/pyproject.toml

@@ -0,0 +1,28 @@
+[build-system]
+requires = ["setuptools>=61", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "ecgen"
+dynamic = ["version"]
+description = "ECG generation and modeling experiments"
+readme = "README.md"
+license = { file = "LICENSE" }
+requires-python = ">=3.8"
+dependencies = [
+    "torch",
+    "pyyaml",
+    "pytorch-lightning",
+    "pandas",
+    "scikit-learn",
+    "wfdb",
+]
+
+[project.urls]
+Homepage = "https://github.com/vlbthambawita/ECGEN"
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[tool.setuptools.dynamic]
+version = { attr = "ecgen.__version__" }

ecgen-0.4.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

ecgen-0.4.0/src/ecgen/__init__.py

@@ -0,0 +1,7 @@
+"""
+Top-level package for ECG generation experiments.
+"""
+
+__version__ = "0.4.0"
+
+__all__ = ["data", "models", "training", "evaluation", "utils"]

ecgen-0.4.0/src/ecgen/data/datamodule.py

@@ -0,0 +1,32 @@
+from dataclasses import dataclass
+from typing import Optional
+
+from torch.utils.data import DataLoader, Dataset
+
+
+@dataclass
+class ECGDataModule:
+    """
+    Minimal placeholder datamodule describing train/val/test datasets.
+    """
+
+    train_dataset: Optional[Dataset] = None
+    val_dataset: Optional[Dataset] = None
+    test_dataset: Optional[Dataset] = None
+    batch_size: int = 32
+    num_workers: int = 4
+
+    def train_dataloader(self) -> DataLoader:
+        if self.train_dataset is None:
+            raise RuntimeError("train_dataset is not set.")
+        return DataLoader(self.train_dataset, batch_size=self.batch_size, num_workers=self.num_workers, shuffle=True)
+
+    def val_dataloader(self) -> DataLoader:
+        if self.val_dataset is None:
+            raise RuntimeError("val_dataset is not set.")
+        return DataLoader(self.val_dataset, batch_size=self.batch_size, num_workers=self.num_workers)
+
+    def test_dataloader(self) -> DataLoader:
+        if self.test_dataset is None:
+            raise RuntimeError("test_dataset is not set.")
+        return DataLoader(self.test_dataset, batch_size=self.batch_size, num_workers=self.num_workers)

ecgen-0.4.0/src/ecgen/data/mimic_dataset.py

@@ -0,0 +1,214 @@
+"""
+Native MIMIC-IV-ECG dataset implementation.
+
+Dataset structure:
+- ECG waveforms: files/p{XXXX}/p{subject_id}/s{study_id}/{study_id}.hea/.dat (WFDB format)
+- Machine measurements: machine_measurements.csv
+
+Expected columns in machine_measurements.csv:
+- subject_id, study_id
+- rr_interval, p_onset, p_end, qrs_onset, qrs_end, t_end
+- p_axis, qrs_axis, t_axis
+"""
+
+from __future__ import annotations
+
+import os
+from typing import Optional, Tuple
+
+import numpy as np
+import pandas as pd
+import torch
+import wfdb
+from sklearn.model_selection import train_test_split
+from torch.utils.data import Dataset
+
+
+class MIMICIVECGDataset(Dataset):
+    """
+    Dataset for MIMIC-IV-ECG signals with machine measurements conditioning.
+
+    Returns (ecg, features) where:
+    - ecg: (num_leads, seq_length) normalized ECG signal, float32
+    - features: (9,) normalized machine measurements, float32
+    """
+
+    FEATURE_NAMES = [
+        "rr_interval",
+        "p_onset",
+        "p_end",
+        "qrs_onset",
+        "qrs_end",
+        "t_end",
+        "p_axis",
+        "qrs_axis",
+        "t_axis",
+    ]
+
+    def __init__(
+        self,
+        mimic_path: str,
+        split: str = "train",
+        val_split: float = 0.1,
+        test_split: float = 0.1,
+        max_samples: Optional[int] = None,
+        seed: int = 42,
+        skip_missing_check: bool = False,
+        ecg_norm_eps: float = 1e-6,
+        ecg_norm_factor: Optional[float] = None,
+        num_leads: int = 12,
+        seq_length: int = 5000,
+    ) -> None:
+        self.mimic_path = mimic_path
+        self.ecg_norm_eps = ecg_norm_eps
+        self.ecg_norm_factor = ecg_norm_factor
+        self.split = split
+        self.seed = seed
+        self.skip_missing_check = skip_missing_check
+        self.num_leads = num_leads
+        self.seq_length = seq_length
+
+        self.load_measurements()
+        self.create_splits(val_split, test_split)
+        self.filter_by_split()
+
+        if not skip_missing_check:
+            self.filter_missing_files()
+        else:
+            import warnings
+            warnings.warn(
+                "Skipping missing file check. Some samples may fail during loading.",
+                UserWarning,
+                stacklevel=2,
+            )
+
+        if max_samples is not None:
+            self.measurements = self.measurements.head(max_samples).reset_index(drop=True)
+
+        self.compute_feature_stats()
+
+    def load_measurements(self) -> None:
+        path = os.path.join(self.mimic_path, "machine_measurements.csv")
+        if not os.path.isfile(path):
+            raise FileNotFoundError(
+                f"machine_measurements.csv not found at {path}. "
+                "Download from https://physionet.org/content/mimic-iv-ecg/1.0/"
+            )
+        self.measurements = pd.read_csv(path)
+
+        required = ["subject_id", "study_id"] + self.FEATURE_NAMES
+        missing = [c for c in required if c not in self.measurements.columns]
+        if missing:
+            raise ValueError(f"machine_measurements.csv missing columns: {missing}")
+
+        self.measurements = self.measurements.dropna(subset=self.FEATURE_NAMES).reset_index(drop=True)
+
+    def create_splits(self, val_split: float, test_split: float) -> None:
+        subjects = self.measurements["subject_id"].unique()
+
+        train_subjects, test_subjects = train_test_split(
+            subjects, test_size=test_split, random_state=self.seed
+        )
+        train_subjects, val_subjects = train_test_split(
+            train_subjects,
+            test_size=val_split / (1 - test_split),
+            random_state=self.seed,
+        )
+
+        def assign_split(row: pd.Series) -> str:
+            sid = row["subject_id"]
+            if sid in val_subjects:
+                return "val"
+            if sid in test_subjects:
+                return "test"
+            return "train"
+
+        self.measurements["split"] = self.measurements.apply(assign_split, axis=1)
+
+    def filter_by_split(self) -> None:
+        self.measurements = self.measurements[
+            self.measurements["split"] == self.split
+        ].reset_index(drop=True)
+
+    def filter_missing_files(self) -> None:
+        files_dir = os.path.join(self.mimic_path, "files")
+        if not os.path.isdir(files_dir):
+            return
+
+        valid = []
+        for idx in range(len(self.measurements)):
+            row = self.measurements.iloc[idx]
+            rec_path = self._ecg_record_path(row["subject_id"], row["study_id"])
+            if os.path.isfile(rec_path + ".hea"):
+                valid.append(idx)
+
+        self.measurements = self.measurements.iloc[valid].reset_index(drop=True)
+
+    def _ecg_record_path(self, subject_id: int, study_id: int) -> str:
+        sub_str = str(subject_id)
+        prefix = sub_str[:4]
+        return os.path.join(
+            self.mimic_path,
+            "files",
+            f"p{prefix}",
+            f"p{subject_id}",
+            f"s{study_id}",
+            str(study_id),
+        )
+
+    def compute_feature_stats(self) -> None:
+        self.feature_stats = {}
+        for name in self.FEATURE_NAMES:
+            vals = self.measurements[name].values
+            self.feature_stats[name] = {
+                "mean": float(np.mean(vals)),
+                "std": float(np.std(vals)) + 1e-6,
+            }
+
+    def load_ecg(self, idx: int) -> np.ndarray:
+        row = self.measurements.iloc[idx]
+        rec_path = self._ecg_record_path(row["subject_id"], row["study_id"])
+
+        record = wfdb.rdrecord(rec_path)
+        signal = record.p_signal  # (time, leads)
+        signal = signal.T.astype(np.float32)  # (leads, time)
+
+        if signal.shape[0] < self.num_leads:
+            pad = np.zeros((self.num_leads - signal.shape[0], signal.shape[1]), dtype=np.float32)
+            signal = np.vstack([signal, pad])
+        elif signal.shape[0] > self.num_leads:
+            signal = signal[: self.num_leads]
+
+        if signal.shape[1] < self.seq_length:
+            pad = np.zeros((signal.shape[0], self.seq_length - signal.shape[1]), dtype=np.float32)
+            signal = np.hstack([signal, pad])
+        elif signal.shape[1] > self.seq_length:
+            signal = signal[:, : self.seq_length]
+
+        return signal
+
+    def _get_features(self, idx: int) -> np.ndarray:
+        row = self.measurements.iloc[idx]
+        out = []
+        for name in self.FEATURE_NAMES:
+            val = row[name]
+            s = self.feature_stats[name]
+            out.append((float(val) - s["mean"]) / s["std"])
+        return np.array(out, dtype=np.float32)
+
+    def __len__(self) -> int:
+        return len(self.measurements)
+
+    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
+        ecg = self.load_ecg(idx)
+
+        ecg_mean = ecg.mean()
+        if self.ecg_norm_factor is not None:
+            scale = self.ecg_norm_factor
+        else:
+            scale = max(float(np.std(ecg)), self.ecg_norm_eps)
+        ecg = (ecg.astype(np.float32) - ecg_mean) / scale
+
+        features = self._get_features(idx)
+
+        return torch.from_numpy(ecg), torch.from_numpy(features)

ecgen-0.4.0/src/ecgen/data/pulse2pulse_mimic.py

@@ -0,0 +1,106 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+from torch.utils.data import DataLoader, Dataset
+
+import pytorch_lightning as pl
+
+
+class ECGDatasetAdapter(Dataset):
+    """
+    Wrap (ecg, cond) style dataset and return only ECG signals for GAN training.
+    """
+
+    def __init__(self, base_dataset: Dataset, num_leads: int = 8) -> None:
+        self.base = base_dataset
+        self.num_leads = num_leads
+
+    def __len__(self) -> int:
+        return len(self.base)  # type: ignore[arg-type]
+
+    def __getitem__(self, idx: int):
+        ecg, _ = self.base[idx]
+        if ecg.shape[0] > self.num_leads:
+            ecg = ecg[: self.num_leads]
+        ecg = ecg.float()
+        return {"ecg_signals": ecg}
+
+
+@dataclass
+class Pulse2PulseMIMICConfig:
+    data_dir: str
+    batch_size: int = 128
+    num_workers: int = 4
+    max_samples: Optional[int] = None
+    skip_missing_check: bool = True
+    num_channels: int = 8
+    seq_length: int = 5000
+
+
+class Pulse2PulseMIMICDataModule(pl.LightningDataModule):
+    """
+    LightningDataModule for MIMIC‑IV‑ECG ECG generation with Pulse2Pulse.
+
+    Uses native ecgen.data.mimic_dataset.MIMICIVECGDataset.
+    """
+
+    def __init__(self, config: Pulse2PulseMIMICConfig | dict) -> None:
+        super().__init__()
+        if isinstance(config, dict):
+            config = Pulse2PulseMIMICConfig(**config)
+        self.config = config
+
+        self.train_dataset: Optional[Dataset] = None
+        self.val_dataset: Optional[Dataset] = None
+
+    def setup(self, stage: Optional[str] = None) -> None:
+        if self.train_dataset is not None and self.val_dataset is not None:
+            return
+
+        from ecgen.data.mimic_dataset import MIMICIVECGDataset
+
+        train_base = MIMICIVECGDataset(
+            mimic_path=self.config.data_dir,
+            split="train",
+            max_samples=self.config.max_samples,
+            skip_missing_check=self.config.skip_missing_check,
+            num_leads=12,
+            seq_length=self.config.seq_length,
+        )
+        val_base = MIMICIVECGDataset(
+            mimic_path=self.config.data_dir,
+            split="val",
+            max_samples=min(self.config.max_samples or 1000, 1000),
+            skip_missing_check=self.config.skip_missing_check,
+            num_leads=12,
+            seq_length=self.config.seq_length,
+        )
+
+        self.train_dataset = ECGDatasetAdapter(train_base, num_leads=self.config.num_channels)
+        self.val_dataset = ECGDatasetAdapter(val_base, num_leads=self.config.num_channels)
+
+    def train_dataloader(self) -> DataLoader:
+        if self.train_dataset is None:
+            raise RuntimeError("train_dataset is not set. Did you forget to call setup()? ")
+        return DataLoader(
+            self.train_dataset,
+            batch_size=self.config.batch_size,
+            shuffle=True,
+            num_workers=self.config.num_workers,
+            pin_memory=torch.cuda.is_available(),
+        )
+
+    def val_dataloader(self) -> DataLoader:
+        if self.val_dataset is None:
+            raise RuntimeError("val_dataset is not set. Did you forget to call setup()? ")
+        return DataLoader(
+            self.val_dataset,
+            batch_size=self.config.batch_size,
+            shuffle=False,
+            num_workers=self.config.num_workers,
+            pin_memory=torch.cuda.is_available(),
+        )
+

ecgen-0.4.0/src/ecgen/models/__init__.py

@@ -0,0 +1,33 @@
+"""
+ECG generation models.
+"""
+
+from ecgen.models.pulse2pulse import (
+    Pulse2PulseConfig,
+    Pulse2PulseGAN,
+    WaveGANDiscriminator,
+    WaveGANGenerator,
+)
+from ecgen.models.vae import (
+    Decoder1D,
+    Encoder1D,
+    ResidualBlock1D,
+    VAE1D,
+    VAEConfig,
+    VAELightning,
+    vae_loss,
+)
+
+__all__ = [
+    "Pulse2PulseConfig",
+    "Pulse2PulseGAN",
+    "WaveGANDiscriminator",
+    "WaveGANGenerator",
+    "ResidualBlock1D",
+    "Encoder1D",
+    "Decoder1D",
+    "VAE1D",
+    "VAEConfig",
+    "VAELightning",
+    "vae_loss",
+]