Data

ECGDataset

ECGDataset dataclass

ECGDataset(data: ndarray, lead_names: List[str], fs: float, source_path: Path, label: str, file_names: List[str] = list())

Container for a collection of ECG recordings loaded from a folder.

Attributes:

Name	Type	Description
data	ndarray	Shape (n_samples, n_leads, n_timepoints).
lead_names	List[str]	Ordered list of lead names.
fs	float	Sampling frequency in Hz.
source_path	Path	Folder the data was loaded from.
label	str	Human-readable label used in plots/reports (defaults to folder name).
file_names	List[str]	Original filenames in load order.

time_axis property

time_axis: ndarray

Time axis in seconds.

select_leads

select_leads(leads: List[str]) -> 'ECGDataset'

Return a new ECGDataset restricted to the requested leads.

Source code in ecgen_eval/data/dataset.py

def select_leads(self, leads: List[str]) -> "ECGDataset":
    """Return a new ECGDataset restricted to the requested leads."""
    indices = [self.lead_names.index(l) for l in leads]
    return ECGDataset(
        data=self.data[:, indices, :],
        lead_names=leads,
        fs=self.fs,
        source_path=self.source_path,
        label=self.label,
        file_names=self.file_names,
    )

sample

sample(n: int, seed: Optional[int] = 42) -> 'ECGDataset'

Return a new ECGDataset with n randomly selected recordings.

Source code in ecgen_eval/data/dataset.py

def sample(self, n: int, seed: Optional[int] = 42) -> "ECGDataset":
    """Return a new ECGDataset with n randomly selected recordings."""
    rng = np.random.default_rng(seed)
    idx = rng.choice(self.n_samples, size=min(n, self.n_samples), replace=False)
    return ECGDataset(
        data=self.data[idx],
        lead_names=self.lead_names,
        fs=self.fs,
        source_path=self.source_path,
        label=self.label,
        file_names=[self.file_names[i] for i in idx] if self.file_names else [],
    )

LEADS_12 module-attribute

LEADS_12 = ['I', 'II', 'III', 'aVR', 'aVL', 'aVF', 'V1', 'V2', 'V3', 'V4', 'V5', 'V6']

LEADS_8 module-attribute

LEADS_8 = ['I', 'II', 'V1', 'V2', 'V3', 'V4', 'V5', 'V6']

Loaders

load_ecg_npy

load_ecg_npy(path: Union[str, Path], fs: float = 500.0, label: Optional[str] = None, lead_names: Optional[List[str]] = None, max_samples: Optional[int] = None) -> ECGDataset

Load a NumPy array of ECGs saved with np.save / np.savez.

The array must have shape (N, L, T) — samples × leads × timepoints. This is the native output format of ECGEN-FM, ECGEN-VAE, and Pulse2Pulse.

Parameters:

Name	Type	Description	Default
path	str or Path	Path to a .npy file (single array) or .npz file (archive). For .npz the first array found is used unless label matches a key.	required
fs	float	Sampling frequency in Hz. Default: 500.	500.0
label	str	Display label. Defaults to stem of the file.	None
lead_names	list of str	Lead name list. Auto-assigned from L if 8 or 12 leads.	None
max_samples	int	Cap number of recordings loaded.	None

Returns:

Type	Description
ECGDataset

Source code in ecgen_eval/data/loader.py

def load_ecg_npy(
    path: Union[str, Path],
    fs: float = 500.0,
    label: Optional[str] = None,
    lead_names: Optional[List[str]] = None,
    max_samples: Optional[int] = None,
) -> ECGDataset:
    """
    Load a NumPy array of ECGs saved with ``np.save`` / ``np.savez``.

    The array must have shape ``(N, L, T)`` — samples × leads × timepoints.
    This is the native output format of ECGEN-FM, ECGEN-VAE, and Pulse2Pulse.

    Parameters
    ----------
    path : str or Path
        Path to a ``.npy`` file (single array) or ``.npz`` file (archive).
        For ``.npz`` the first array found is used unless *label* matches a key.
    fs : float
        Sampling frequency in Hz. Default: 500.
    label : str, optional
        Display label. Defaults to stem of the file.
    lead_names : list of str, optional
        Lead name list. Auto-assigned from L if 8 or 12 leads.
    max_samples : int, optional
        Cap number of recordings loaded.

    Returns
    -------
    ECGDataset
    """
    path = Path(path)

    if path.suffix.lower() == ".npz":
        archive = np.load(path)
        key = label if label in archive else list(archive.keys())[0]
        arr = archive[key]
    else:
        arr = np.load(path, allow_pickle=False)

    if arr.ndim == 2:
        # Single recording (L, T) → (1, L, T)
        arr = arr[np.newaxis]
    if arr.ndim != 3:
        raise ValueError(f"Expected array of shape (N, L, T), got {arr.shape}")

    if max_samples is not None:
        arr = arr[:max_samples]

    arr = arr.astype(np.float32)
    n_leads = arr.shape[1]

    if lead_names is None:
        lead_names = _default_lead_names(n_leads)
    if label is None:
        label = path.stem

    return ECGDataset(
        data=arr,
        lead_names=lead_names,
        fs=fs,
        source_path=path.parent,
        label=label,
        file_names=[path.name] * arr.shape[0],
    )

load_ecg_wfdb

load_ecg_wfdb(record_dir: Union[str, Path], fs: float = 500.0, label: Optional[str] = None, lead_names: Optional[List[str]] = None, max_samples: Optional[int] = None, resample_to: Optional[float] = None) -> ECGDataset

Load WFDB ECG records from a directory (e.g. PTB-XL).

Requires the wfdb package (pip install wfdb).

Parameters:

Name	Type	Description	Default
record_dir	str or Path	Directory containing .hea / .dat WFDB record pairs, or a single record path (without extension).	required
fs	float	Expected sampling frequency. Used for metadata only unless resample_to is set.	500.0
label	str	Display label. Defaults to directory name.	None
lead_names	list of str	Override lead names (e.g. to match a specific lead ordering).	None
max_samples	int	Cap on number of records loaded.	None
resample_to	float	If provided, resample each record to this sampling rate using scipy.	None

Returns:

Type	Description
ECGDataset

Source code in ecgen_eval/data/loader.py

def load_ecg_wfdb(
    record_dir: Union[str, Path],
    fs: float = 500.0,
    label: Optional[str] = None,
    lead_names: Optional[List[str]] = None,
    max_samples: Optional[int] = None,
    resample_to: Optional[float] = None,
) -> ECGDataset:
    """
    Load WFDB ECG records from a directory (e.g. PTB-XL).

    Requires the ``wfdb`` package (``pip install wfdb``).

    Parameters
    ----------
    record_dir : str or Path
        Directory containing ``.hea`` / ``.dat`` WFDB record pairs,
        or a single record path (without extension).
    fs : float
        Expected sampling frequency. Used for metadata only unless *resample_to* is set.
    label : str, optional
        Display label. Defaults to directory name.
    lead_names : list of str, optional
        Override lead names (e.g. to match a specific lead ordering).
    max_samples : int, optional
        Cap on number of records loaded.
    resample_to : float, optional
        If provided, resample each record to this sampling rate using scipy.

    Returns
    -------
    ECGDataset
    """
    try:
        import wfdb
    except ImportError as e:
        raise ImportError("wfdb package required: pip install wfdb") from e

    record_dir = Path(record_dir)

    # Support single record path (no extension)
    if not record_dir.is_dir():
        parent = record_dir.parent
        stem = record_dir.stem
        record_paths = [parent / stem]
    else:
        hea_files = sorted(record_dir.glob("*.hea"))
        if not hea_files:
            raise FileNotFoundError(f"No WFDB .hea files found in {record_dir}")
        record_paths = [f.parent / f.stem for f in hea_files]

    if label is None:
        label = record_dir.name if record_dir.is_dir() else record_dir.stem

    if max_samples is not None:
        record_paths = record_paths[:max_samples]

    recordings = []
    detected_lead_names: Optional[List[str]] = lead_names
    actual_fs = fs

    for rpath in tqdm(record_paths, desc=f"Loading {label} (WFDB)", unit="rec"):
        try:
            rec = wfdb.rdrecord(str(rpath))
            sig = rec.p_signal  # (T, L) float64
            if sig is None:
                continue
            sig = sig.T.astype(np.float32)  # (L, T)

            if resample_to is not None and rec.fs != resample_to:
                sig = _resample_leads(sig, rec.fs, resample_to)
                actual_fs = resample_to
            else:
                actual_fs = rec.fs or fs

            if detected_lead_names is None:
                detected_lead_names = rec.sig_name

            recordings.append(sig)
        except Exception:
            continue

    if not recordings:
        raise ValueError(f"No valid WFDB records loaded from {record_dir}")

    data = np.stack(recordings, axis=0)  # (N, L, T)
    n_leads = data.shape[1]
    if detected_lead_names is None:
        detected_lead_names = _default_lead_names(n_leads)

    return ECGDataset(
        data=data,
        lead_names=detected_lead_names,
        fs=float(actual_fs),
        source_path=record_dir if record_dir.is_dir() else record_dir.parent,
        label=label,
        file_names=[p.name for p in record_paths],
    )

load_ecg_folder

load_ecg_folder(folder: Union[str, Path], fs: float = 500.0, label: Optional[str] = None, n_leads: Optional[int] = None, n_timepoints: Optional[int] = None, lead_names: Optional[List[str]] = None, max_samples: Optional[int] = None) -> ECGDataset

Load ECGs from folder or a single file, auto-detecting the format.

Supported: .npy / .npz (NumPy), .hea / WFDB records, and .csv.

Parameters:

Name	Type	Description	Default
folder	str or Path	Directory or a single file path.	required
fs	float	Sampling frequency in Hz (default 500 Hz).	500.0
label	str	Display label for plots. Defaults to the folder/file name.	None
n_leads	int	Expected number of leads (8 or 12). Auto-detected if None. (CSV only)	None
n_timepoints	int	Expected number of timepoints per lead. Auto-detected if None. (CSV only)	None
lead_names	list of str	Explicit lead name list. Auto-detected or default-assigned if None.	None
max_samples	int	Cap the number of loaded recordings (useful for quick tests).	None

Returns:

Type	Description
ECGDataset

Source code in ecgen_eval/data/loader.py

def load_ecg_folder(
    folder: Union[str, Path],
    fs: float = 500.0,
    label: Optional[str] = None,
    n_leads: Optional[int] = None,
    n_timepoints: Optional[int] = None,
    lead_names: Optional[List[str]] = None,
    max_samples: Optional[int] = None,
) -> ECGDataset:
    """
    Load ECGs from *folder* or a single file, auto-detecting the format.

    Supported: ``.npy`` / ``.npz`` (NumPy), ``.hea`` / WFDB records, and ``.csv``.

    Parameters
    ----------
    folder : str or Path
        Directory or a single file path.
    fs : float
        Sampling frequency in Hz (default 500 Hz).
    label : str, optional
        Display label for plots. Defaults to the folder/file name.
    n_leads : int, optional
        Expected number of leads (8 or 12). Auto-detected if None. (CSV only)
    n_timepoints : int, optional
        Expected number of timepoints per lead. Auto-detected if None. (CSV only)
    lead_names : list of str, optional
        Explicit lead name list. Auto-detected or default-assigned if None.
    max_samples : int, optional
        Cap the number of loaded recordings (useful for quick tests).

    Returns
    -------
    ECGDataset
    """
    path = Path(folder)

    # ---- NumPy single file -------------------------------------------------
    if path.is_file() and path.suffix.lower() in (".npy", ".npz"):
        return load_ecg_npy(path, fs=fs, label=label, lead_names=lead_names,
                            max_samples=max_samples)

    # ---- NumPy folder (multiple .npy files) --------------------------------
    if path.is_dir():
        npy_files = sorted(path.glob("*.npy"))
        if npy_files:
            return _load_npy_folder(npy_files, fs=fs, label=label or path.name,
                                    lead_names=lead_names, max_samples=max_samples)

        npz_files = sorted(path.glob("*.npz"))
        if npz_files:
            return _load_npy_folder(npz_files, fs=fs, label=label or path.name,
                                    lead_names=lead_names, max_samples=max_samples)

    # ---- WFDB folder -------------------------------------------------------
    if path.is_dir() and any(path.glob("*.hea")):
        return load_ecg_wfdb(path, fs=fs, label=label, lead_names=lead_names,
                             max_samples=max_samples)

    # ---- CSV (legacy / interop) --------------------------------------------
    return _load_ecg_csv(path, fs=fs, label=label, n_leads=n_leads,
                         n_timepoints=n_timepoints, lead_names=lead_names,
                         max_samples=max_samples)