Data API

Auto-generated from docstrings.

ecgen.types

Data contracts shared by datasets, conditioners, and generative families.

ECGBatch dataclass

A collated batch of ECGs and their condition fields.

Attributes:

Name	Type	Description
signal	Tensor	[B, C, T] ECG signals.
lead_names	List[str]	Names of the C leads, in channel order.
fs	int	Sampling rate in Hz.
cond	Dict[str, Tensor]	Raw condition fields by name (e.g. "measurements": [B, F]). Empty for unconditional data.
cond_missing	Dict[str, Tensor]	Per-field missing masks, True = missing, same keys as cond.
obs_mask	Optional[Tensor]	[B, C, T] mask, 1 = observed, for partial-ECG conditioning.
meta	Dict[str, Any]	Record identifiers etc.; never fed to a model.

Source code in src/ecgen/types.py

@dataclass
class ECGBatch:
    """
    A collated batch of ECGs and their condition fields.

    Attributes:
        signal: ``[B, C, T]`` ECG signals.
        lead_names: Names of the ``C`` leads, in channel order.
        fs: Sampling rate in Hz.
        cond: Raw condition fields by name (e.g. ``"measurements": [B, F]``).
            Empty for unconditional data.
        cond_missing: Per-field missing masks, ``True`` = missing, same keys as ``cond``.
        obs_mask: ``[B, C, T]`` mask, 1 = observed, for partial-ECG conditioning.
        meta: Record identifiers etc.; never fed to a model.
    """

    signal: Tensor
    lead_names: List[str]
    fs: int
    cond: Dict[str, Tensor] = field(default_factory=dict)
    cond_missing: Dict[str, Tensor] = field(default_factory=dict)
    obs_mask: Optional[Tensor] = None
    meta: Dict[str, Any] = field(default_factory=dict)

    def __len__(self) -> int:
        return self.signal.shape[0]

ConditionOutput dataclass

Encoded conditions, in the slots a backbone can consume.

Attributes:

Name	Type	Description
vector	Optional[Tensor]	[B, D] global embedding (labels, tabular, pooled text); injected with FiLM / AdaGN / AdaLN or added to the time embedding.
tokens	Optional[Tensor]	[B, N, D] token sequence for cross-attention.
tokens_mask	Optional[Tensor]	[B, N] bool, True = valid token.
dense	Optional[Tensor]	[B, C', T] time-aligned conditions (masked ECG + mask, PPG, ...); concatenated to the input.

Source code in src/ecgen/types.py

@dataclass
class ConditionOutput:
    """
    Encoded conditions, in the slots a backbone can consume.

    Attributes:
        vector: ``[B, D]`` global embedding (labels, tabular, pooled text);
            injected with FiLM / AdaGN / AdaLN or added to the time embedding.
        tokens: ``[B, N, D]`` token sequence for cross-attention.
        tokens_mask: ``[B, N]`` bool, ``True`` = valid token.
        dense: ``[B, C', T]`` time-aligned conditions (masked ECG + mask, PPG, ...);
            concatenated to the input.
    """

    vector: Optional[Tensor] = None
    tokens: Optional[Tensor] = None
    tokens_mask: Optional[Tensor] = None
    dense: Optional[Tensor] = None

get_signal

get_signal(batch: Any) -> Tensor

Return the [B, C, T] ECG tensor from any supported batch format.

Accepts an :class:ECGBatch, a dict with "signal" or the legacy "ecg_signals" key, or a legacy (ecg, features) tuple.

Source code in src/ecgen/types.py

def get_signal(batch: Any) -> Tensor:
    """
    Return the ``[B, C, T]`` ECG tensor from any supported batch format.

    Accepts an :class:`ECGBatch`, a dict with ``"signal"`` or the legacy
    ``"ecg_signals"`` key, or a legacy ``(ecg, features)`` tuple.
    """
    if isinstance(batch, ECGBatch):
        return batch.signal
    if isinstance(batch, dict):
        if "signal" in batch:
            return batch["signal"]
        if "ecg_signals" in batch:
            return batch["ecg_signals"]
        raise KeyError(
            f"Batch dict has no 'signal' or 'ecg_signals' key: {sorted(batch)}"
        )
    if isinstance(batch, (tuple, list)):
        return batch[0]
    if isinstance(batch, Tensor):
        return batch
    raise TypeError(f"Unsupported batch type: {type(batch).__name__}")

ecgen.data.datasets.base

Dataset contract.

BaseECGDataset

Bases: Dataset, ABC

Base class for ECG datasets.

__getitem__ returns a dict::

{
    "signal": FloatTensor [C, T],
    "cond": {field_name: Tensor},          # may be empty
    "cond_missing": {field_name: BoolTensor},
    "meta": {...},                         # ids; not fed to models
}

and the dataset exposes :attr:lead_names and :attr:fs, which :class:ecgen.data.collate.ECGCollate attaches to every batch.

Source code in src/ecgen/data/datasets/base.py

class BaseECGDataset(Dataset, ABC):
    """
    Base class for ECG datasets.

    ``__getitem__`` returns a dict::

        {
            "signal": FloatTensor [C, T],
            "cond": {field_name: Tensor},          # may be empty
            "cond_missing": {field_name: BoolTensor},
            "meta": {...},                         # ids; not fed to models
        }

    and the dataset exposes :attr:`lead_names` and :attr:`fs`, which
    :class:`ecgen.data.collate.ECGCollate` attaches to every batch.
    """

    lead_names: List[str]
    fs: int
    seq_length: int
    feature_normalizer: Optional[FeatureNormalizer] = None

    @abstractmethod
    def __len__(self) -> int: ...

    @abstractmethod
    def __getitem__(self, idx: int) -> Dict[str, Any]: ...

    @property
    def num_leads(self) -> int:
        return len(self.lead_names)

    @property
    def field_normalizers(self) -> Dict[str, FeatureNormalizer]:
        """Normalizers of numeric ``cond`` fields, by field name.

        Saved in checkpoints so a model can normalize raw values at sampling
        time. Default: :attr:`feature_normalizer` for the ``measurements`` field.
        """
        if self.feature_normalizer is None or not self.feature_normalizer.columns:
            return {}
        return {"measurements": self.feature_normalizer}

field_normalizers property

field_normalizers: Dict[str, FeatureNormalizer]

Normalizers of numeric cond fields, by field name.

Saved in checkpoints so a model can normalize raw values at sampling time. Default: :attr:feature_normalizer for the measurements field.

scale_signal

scale_signal(ecg: ndarray, scale: Scale, eps: float = 1e-06) -> np.ndarray

Apply the dataset-level amplitude scaling to a [C, T] record.

Source code in src/ecgen/data/datasets/base.py

def scale_signal(ecg: np.ndarray, scale: Scale, eps: float = 1e-6) -> np.ndarray:
    """Apply the dataset-level amplitude scaling to a ``[C, T]`` record."""
    if scale == "none":
        return ecg.astype(np.float32)
    mean = ecg.mean()
    if scale == "zscore":
        divisor = max(float(np.std(ecg)), eps)
    elif isinstance(scale, (int, float)) and not isinstance(scale, bool):
        divisor = float(scale)
    else:
        raise ValueError(f"Unknown scale {scale!r}; use 'zscore', 'none', or a number")
    return ((ecg.astype(np.float32) - mean) / divisor).astype(np.float32)

ecgen.data.datasets.mimic

MIMIC-IV-ECG dataset (dict samples).

Layout under root:

• files/p{XXXX}/p{subject_id}/s{study_id}/{study_id}.hea/.dat (WFDB, 12 leads, 500 Hz, 10 s)
• machine_measurements.csv (one row per study; interval/axis measurements and report lines)

Records are stored in the lead order I, II, III, aVR, aVF, aVL, V1..V6 — note aVF before aVL, unlike the standard clinical order.

MIMICIVECG

Bases: BaseECGDataset

MIMIC-IV-ECG records with machine measurements as a numeric condition field.

Samples are dicts (see :class:~ecgen.data.datasets.base.BaseECGDataset) with cond["measurements"] = the normalized measurements columns.

The train/val/test split is by subject and seeded, and matches the legacy MIMICIVECGDataset when missing="drop" (the default). The measurement normalizer is always fit on the train split, whichever split is loaded.

Parameters:

Name	Type	Description	Default
root	str	MIMIC-IV-ECG directory (contains files/ and machine_measurements.csv).	required
split	str	"train", "val", or "test".	'train'
val_split	float	Fraction of subjects in the validation split.	0.1
test_split	float	Fraction of subjects in the test split.	0.1
seed	int	Split seed.	42
max_samples	Optional[int]	Keep only the first N rows of the split (for quick runs).	None
skip_missing_check	bool	Skip the check that each record's .hea exists.	False
leads	LeadSpec	None (file order, padded/truncated to num_leads), an alias ("standard12", "independent8"), or a list of lead names.	None
num_leads	int	Channel count when leads is None.	12
seq_length	int	Samples per record (padded/truncated).	5000
scale	Scale	"zscore" (per record, legacy), "none" (mV), or a number (subtract record mean, divide by it).	'zscore'
measurements	Optional[Sequence[str]]	Columns for cond["measurements"]; [] for none.	None
missing	str	"drop" rows with any missing measurement (legacy) or "mask" to keep them and report the gaps in cond_missing. "mask" keeps more subjects, so its split differs from "drop".	'drop'
feature_normalizer	Optional[FeatureNormalizer | Dict[str, Any]]	Use these statistics instead of fitting on train (e.g. those stored in a checkpoint).	None
report	bool	Add cond["report"]: the machine report lines joined by ". ".	False

Source code in src/ecgen/data/datasets/mimic.py

@DATASETS.register("mimic_iv_ecg")
class MIMICIVECG(BaseECGDataset):
    """
    MIMIC-IV-ECG records with machine measurements as a numeric condition field.

    Samples are dicts (see :class:`~ecgen.data.datasets.base.BaseECGDataset`) with
    ``cond["measurements"]`` = the normalized ``measurements`` columns.

    The train/val/test split is by subject and seeded, and matches the legacy
    ``MIMICIVECGDataset`` when ``missing="drop"`` (the default). The measurement
    normalizer is always fit on the *train* split, whichever split is loaded.

    Args:
        root: MIMIC-IV-ECG directory (contains ``files/`` and ``machine_measurements.csv``).
        split: ``"train"``, ``"val"``, or ``"test"``.
        val_split: Fraction of subjects in the validation split.
        test_split: Fraction of subjects in the test split.
        seed: Split seed.
        max_samples: Keep only the first N rows of the split (for quick runs).
        skip_missing_check: Skip the check that each record's ``.hea`` exists.
        leads: ``None`` (file order, padded/truncated to ``num_leads``), an alias
            (``"standard12"``, ``"independent8"``), or a list of lead names.
        num_leads: Channel count when ``leads`` is ``None``.
        seq_length: Samples per record (padded/truncated).
        scale: ``"zscore"`` (per record, legacy), ``"none"`` (mV), or a number
            (subtract record mean, divide by it).
        measurements: Columns for ``cond["measurements"]``; ``[]`` for none.
        missing: ``"drop"`` rows with any missing measurement (legacy) or
            ``"mask"`` to keep them and report the gaps in ``cond_missing``.
            ``"mask"`` keeps more subjects, so its split differs from ``"drop"``.
        feature_normalizer: Use these statistics instead of fitting on train
            (e.g. those stored in a checkpoint).
        report: Add ``cond["report"]``: the machine report lines joined by ". ".
    """

    fs = 500

    def __init__(
        self,
        root: str,
        split: str = "train",
        val_split: float = 0.1,
        test_split: float = 0.1,
        seed: int = 42,
        max_samples: Optional[int] = None,
        skip_missing_check: bool = False,
        leads: LeadSpec = None,
        num_leads: int = 12,
        seq_length: int = 5000,
        scale: Scale = "zscore",
        ecg_norm_eps: float = 1e-6,
        measurements: Optional[Sequence[str]] = None,
        missing: str = "drop",
        feature_normalizer: Optional[FeatureNormalizer | Dict[str, Any]] = None,
        report: bool = False,
    ) -> None:
        if split not in ("train", "val", "test"):
            raise ValueError(f"split must be train/val/test, got {split!r}")
        if missing not in ("drop", "mask"):
            raise ValueError(f"missing must be 'drop' or 'mask', got {missing!r}")

        self.root = root
        self.split = split
        self.seed = seed
        self.seq_length = seq_length
        self.scale = scale
        self.ecg_norm_eps = ecg_norm_eps
        self.use_report = report
        self.measurement_columns: List[str] = list(
            MEASUREMENT_COLUMNS if measurements is None else measurements
        )

        self._leads_spec = leads
        if leads is None:
            self.lead_names = list(MIMIC_LEAD_ORDER[:num_leads]) + [
                f"pad{i}" for i in range(max(0, num_leads - len(MIMIC_LEAD_ORDER)))
            ]
        else:
            self.lead_names = resolve_leads(leads, MIMIC_LEAD_ORDER)

        table = self._load_measurements(missing)
        table["split"] = self._assign_splits(table, val_split, test_split)

        if feature_normalizer is None:
            self.feature_normalizer = FeatureNormalizer.fit(
                table[table["split"] == "train"], self.measurement_columns
            )
        elif isinstance(feature_normalizer, dict):
            self.feature_normalizer = FeatureNormalizer.from_dict(feature_normalizer)
        else:
            self.feature_normalizer = feature_normalizer

        self.records = table[table["split"] == split].reset_index(drop=True)

        if skip_missing_check:
            warnings.warn(
                "Skipping missing file check. Some samples may fail during loading.",
                UserWarning,
                stacklevel=2,
            )
        else:
            self._filter_missing_files()

        if max_samples is not None:
            self.records = self.records.head(max_samples).reset_index(drop=True)

    # ------------------------------------------------------------------ setup

    def _load_measurements(self, missing: str) -> pd.DataFrame:
        path = os.path.join(self.root, "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/"
            )
        table = pd.read_csv(path, low_memory=False)

        # The legacy loader always required (and dropped NaNs in) the 9 standard
        # columns, which also decides which subjects enter the split.
        required = (
            ["subject_id", "study_id"]
            + list(MEASUREMENT_COLUMNS)
            + self.measurement_columns
        )
        missing_cols = sorted({c for c in required if c not in table.columns})
        if missing_cols:
            raise ValueError(
                f"machine_measurements.csv missing columns: {missing_cols}"
            )

        if missing == "drop":
            subset = list(
                dict.fromkeys(list(MEASUREMENT_COLUMNS) + self.measurement_columns)
            )
            table = table.dropna(subset=subset)
        return table.reset_index(drop=True)

    def _assign_splits(
        self, table: pd.DataFrame, val_split: float, test_split: float
    ) -> pd.Series:
        subjects = table["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,
        )
        split = pd.Series("train", index=table.index)
        split[table["subject_id"].isin(set(val_subjects))] = "val"
        split[table["subject_id"].isin(set(test_subjects))] = "test"
        return split

    def record_path(self, subject_id: int, study_id: int) -> str:
        sub = str(subject_id)
        return os.path.join(
            self.root, "files", f"p{sub[:4]}", f"p{sub}", f"s{study_id}", str(study_id)
        )

    def _filter_missing_files(self) -> None:
        if not os.path.isdir(os.path.join(self.root, "files")):
            return
        exists = [
            os.path.isfile(self.record_path(s, t) + ".hea")
            for s, t in zip(self.records["subject_id"], self.records["study_id"])
        ]
        self.records = self.records[np.asarray(exists, dtype=bool)].reset_index(
            drop=True
        )

    # ---------------------------------------------------------------- samples

    def load_ecg(self, idx: int) -> np.ndarray:
        """Raw ``[C, T]`` signal in mV, with the configured leads and length."""
        row = self.records.iloc[idx]
        record = wfdb.rdrecord(self.record_path(row["subject_id"], row["study_id"]))
        signal = record.p_signal.T.astype(np.float32)  # (leads, time)

        if self._leads_spec is None:
            n = len(self.lead_names)
            if signal.shape[0] < n:
                signal = np.vstack(
                    [
                        signal,
                        np.zeros(
                            (n - signal.shape[0], signal.shape[1]), dtype=np.float32
                        ),
                    ]
                )
            signal = signal[:n]
        else:
            signal = signal[lead_indices(self.lead_names, record.sig_name)]

        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])
        return signal[:, : self.seq_length]

    def __len__(self) -> int:
        return len(self.records)

    def __getitem__(self, idx: int) -> Dict[str, Any]:
        ecg = scale_signal(self.load_ecg(idx), self.scale, self.ecg_norm_eps)
        row = self.records.iloc[idx]
        sample: Dict[str, Any] = {
            "signal": torch.from_numpy(ecg),
            "cond": {},
            "cond_missing": {},
            "meta": {
                "subject_id": int(row["subject_id"]),
                "study_id": int(row["study_id"]),
            },
        }
        if self.use_report:
            lines = [row.get(f"report_{i}") for i in range(18)]
            sample["cond"]["report"] = ". ".join(
                str(x).strip() for x in lines if isinstance(x, str) and x.strip()
            )
        if self.measurement_columns:
            raw = pd.to_numeric(
                row[self.measurement_columns], errors="coerce"
            ).to_numpy(dtype=np.float64)
            values, missing = self.feature_normalizer.transform(raw)
            sample["cond"]["measurements"] = torch.from_numpy(values)
            sample["cond_missing"]["measurements"] = torch.from_numpy(missing)
        return sample

load_ecg

load_ecg(idx: int) -> np.ndarray

Raw [C, T] signal in mV, with the configured leads and length.

Source code in src/ecgen/data/datasets/mimic.py

def load_ecg(self, idx: int) -> np.ndarray:
    """Raw ``[C, T]`` signal in mV, with the configured leads and length."""
    row = self.records.iloc[idx]
    record = wfdb.rdrecord(self.record_path(row["subject_id"], row["study_id"]))
    signal = record.p_signal.T.astype(np.float32)  # (leads, time)

    if self._leads_spec is None:
        n = len(self.lead_names)
        if signal.shape[0] < n:
            signal = np.vstack(
                [
                    signal,
                    np.zeros(
                        (n - signal.shape[0], signal.shape[1]), dtype=np.float32
                    ),
                ]
            )
        signal = signal[:n]
    else:
        signal = signal[lead_indices(self.lead_names, record.sig_name)]

    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])
    return signal[:, : self.seq_length]

ecgen.data.datasets.synthetic

Synthetic ECG-like dataset for tests and smoke runs (no data download needed).

SyntheticECG

Bases: BaseECGDataset

Deterministic periodic "beats" with per-record heart rate and amplitude.

Every record carries the ground truth that produced it, so conditional models can be tested against known targets:

• cond["measurements"]: [heart_rate (bpm), amplitude (mV)], normalized with :attr:feature_normalizer (fixed statistics of the sampling ranges);
• cond["rhythm"]: class 0 = brady (< 60 bpm), 1 = normal, 2 = tachy (> 100 bpm);
• cond["labels"]: multi-hot [brady, tachy, high_amplitude (> 1.25 mV)];
• cond["report"]: a templated text, e.g. "sinus tachycardia, high voltage".

Parameters:

Name	Type	Description	Default
split	str	"train", "val", or "test" (different records per split).	'train'
num_samples	int	Records in this split.	64
leads	LeadSpec	Lead spec (see :func:ecgen.data.leads.resolve_leads); default all 12.	None
seq_length	int	Samples per record.	5000
fs	int	Sampling rate in Hz.	500
scale	Scale	As in :func:ecgen.data.datasets.base.scale_signal; default "none" keeps the amplitude, which "zscore" would normalize away.	'none'
seed	int	Base seed; record i of a split is always the same signal.	0

Source code in src/ecgen/data/datasets/synthetic.py

@DATASETS.register("synthetic")
class SyntheticECG(BaseECGDataset):
    """
    Deterministic periodic "beats" with per-record heart rate and amplitude.

    Every record carries the ground truth that produced it, so conditional
    models can be tested against known targets:

    - ``cond["measurements"]``: ``[heart_rate (bpm), amplitude (mV)]``, normalized
      with :attr:`feature_normalizer` (fixed statistics of the sampling ranges);
    - ``cond["rhythm"]``: class 0 = brady (< 60 bpm), 1 = normal, 2 = tachy (> 100 bpm);
    - ``cond["labels"]``: multi-hot ``[brady, tachy, high_amplitude (> 1.25 mV)]``;
    - ``cond["report"]``: a templated text, e.g. ``"sinus tachycardia, high voltage"``.

    Args:
        split: ``"train"``, ``"val"``, or ``"test"`` (different records per split).
        num_samples: Records in this split.
        leads: Lead spec (see :func:`ecgen.data.leads.resolve_leads`); default all 12.
        seq_length: Samples per record.
        fs: Sampling rate in Hz.
        scale: As in :func:`ecgen.data.datasets.base.scale_signal`; default ``"none"``
            keeps the amplitude, which ``"zscore"`` would normalize away.
        seed: Base seed; record ``i`` of a split is always the same signal.
    """

    HR_RANGE = (45.0, 150.0)
    AMP_RANGE = (0.5, 2.0)
    MEASUREMENTS = ("heart_rate", "amplitude")
    RHYTHMS = ("brady", "normal", "tachy")
    LABELS = ("brady", "tachy", "high_amplitude")

    def __init__(
        self,
        split: str = "train",
        num_samples: int = 64,
        leads: LeadSpec = None,
        seq_length: int = 5000,
        fs: int = 500,
        scale: Scale = "none",
        seed: int = 0,
    ) -> None:
        if split not in _SPLIT_OFFSET:
            raise ValueError(f"split must be train/val/test, got {split!r}")
        self.split = split
        self.num_samples = num_samples
        self.lead_names = resolve_leads(leads, STANDARD_12)
        self.seq_length = seq_length
        self.fs = fs
        self.scale = scale
        self.seed = seed
        ranges = (self.HR_RANGE, self.AMP_RANGE)
        self.feature_normalizer = FeatureNormalizer(
            columns=list(self.MEASUREMENTS),
            mean={c: float(np.mean(r)) for c, r in zip(self.MEASUREMENTS, ranges)},
            std={
                c: float(np.ptp(r) / np.sqrt(12))
                for c, r in zip(self.MEASUREMENTS, ranges)
            },
        )

    def __len__(self) -> int:
        return self.num_samples

    @classmethod
    def render(
        cls, hr: float, amp: float, num_leads: int, seq_length: int, fs: int, rng
    ) -> np.ndarray:
        """
        The ``[C, T]`` signal for a heart rate (bpm) and R amplitude (mV).

        Each beat has P, Q, R, S and T waves with fixed widths in seconds (the
        QT interval shortens with the RR interval, Bazett-style); every lead is
        the same beat train times a random gain and sign, plus white noise.
        """
        rr = 60.0 / hr
        t = np.arange(seq_length) / fs - rng.uniform(0, rr)

        def wave(center: float, width: float) -> np.ndarray:
            dt = (t - center + rr / 2) % rr - rr / 2  # time to the nearest wave
            return np.exp(-(dt**2) / (2 * width**2))

        beat = (
            0.12 * wave(-0.16, 0.025)  # P
            - 0.15 * wave(-0.03, 0.008)  # Q
            + wave(0.0, 0.012)  # R
            - 0.25 * wave(0.03, 0.010)  # S
            + 0.30 * wave(0.30 * np.sqrt(rr), 0.045)  # T
        )
        gains = rng.uniform(0.5, 1.0, size=(num_leads, 1)) * rng.choice(
            [-1.0, 1.0], size=(num_leads, 1)
        )
        noise = 0.02 * rng.standard_normal((num_leads, seq_length))
        return (amp * gains * beat[None, :] + noise).astype(np.float32)

    def __getitem__(self, idx: int) -> Dict[str, Any]:
        rng = np.random.default_rng(self.seed + _SPLIT_OFFSET[self.split] + idx)
        hr = rng.uniform(*self.HR_RANGE)
        amp = rng.uniform(*self.AMP_RANGE)
        ecg = self.render(hr, amp, len(self.lead_names), self.seq_length, self.fs, rng)

        features, missing = self.feature_normalizer.transform([hr, amp])
        rhythm = 0 if hr < 60 else (2 if hr > 100 else 1)
        labels = np.array([hr < 60, hr > 100, amp > 1.25], dtype=np.float32)
        rhythm_text = ("sinus bradycardia", "sinus rhythm", "sinus tachycardia")[rhythm]
        report = rhythm_text + (", high voltage" if amp > 1.25 else "")
        return {
            "signal": torch.from_numpy(scale_signal(ecg, self.scale)),
            "cond": {
                "measurements": torch.from_numpy(features),
                "rhythm": torch.tensor(rhythm, dtype=torch.long),
                "labels": torch.from_numpy(labels),
                "report": report,
            },
            "cond_missing": {"measurements": torch.from_numpy(missing)},
            "meta": {"index": idx, "heart_rate": float(hr), "amplitude": float(amp)},
        }

render classmethod

render(hr: float, amp: float, num_leads: int, seq_length: int, fs: int, rng) -> np.ndarray

The [C, T] signal for a heart rate (bpm) and R amplitude (mV).

Each beat has P, Q, R, S and T waves with fixed widths in seconds (the QT interval shortens with the RR interval, Bazett-style); every lead is the same beat train times a random gain and sign, plus white noise.

Source code in src/ecgen/data/datasets/synthetic.py

@classmethod
def render(
    cls, hr: float, amp: float, num_leads: int, seq_length: int, fs: int, rng
) -> np.ndarray:
    """
    The ``[C, T]`` signal for a heart rate (bpm) and R amplitude (mV).

    Each beat has P, Q, R, S and T waves with fixed widths in seconds (the
    QT interval shortens with the RR interval, Bazett-style); every lead is
    the same beat train times a random gain and sign, plus white noise.
    """
    rr = 60.0 / hr
    t = np.arange(seq_length) / fs - rng.uniform(0, rr)

    def wave(center: float, width: float) -> np.ndarray:
        dt = (t - center + rr / 2) % rr - rr / 2  # time to the nearest wave
        return np.exp(-(dt**2) / (2 * width**2))

    beat = (
        0.12 * wave(-0.16, 0.025)  # P
        - 0.15 * wave(-0.03, 0.008)  # Q
        + wave(0.0, 0.012)  # R
        - 0.25 * wave(0.03, 0.010)  # S
        + 0.30 * wave(0.30 * np.sqrt(rr), 0.045)  # T
    )
    gains = rng.uniform(0.5, 1.0, size=(num_leads, 1)) * rng.choice(
        [-1.0, 1.0], size=(num_leads, 1)
    )
    noise = 0.02 * rng.standard_normal((num_leads, seq_length))
    return (amp * gains * beat[None, :] + noise).astype(np.float32)

ecgen.data.leads

Lead names and lead selection.

resolve_leads

resolve_leads(spec: LeadSpec, available: Sequence[str]) -> List[str]

Turn a lead spec into a list of lead names.

Parameters:

Name	Type	Description	Default
spec	LeadSpec	None keeps available in its stored order; a string is an alias from :data:LEAD_ALIASES; a sequence lists lead names.	required
available	Sequence[str]	Leads the dataset stores, in file order.	required

Returns:

Type	Description
List[str]	Lead names in output channel order.

Source code in src/ecgen/data/leads.py

def resolve_leads(spec: LeadSpec, available: Sequence[str]) -> List[str]:
    """
    Turn a lead spec into a list of lead names.

    Args:
        spec: ``None`` keeps ``available`` in its stored order; a string is an
            alias from :data:`LEAD_ALIASES`; a sequence lists lead names.
        available: Leads the dataset stores, in file order.

    Returns:
        Lead names in output channel order.
    """
    if spec is None:
        return list(available)
    if isinstance(spec, str):
        if spec not in LEAD_ALIASES:
            raise ValueError(
                f"Unknown lead alias {spec!r}; known: {sorted(LEAD_ALIASES)}"
            )
        names = list(LEAD_ALIASES[spec])
    else:
        names = list(spec)
    missing = [n for n in names if n not in available]
    if missing:
        raise ValueError(
            f"Leads {missing} are not available; dataset has {list(available)}"
        )
    return names

lead_indices

lead_indices(names: Sequence[str], available: Sequence[str]) -> List[int]

Indices into available for each name in names.

Source code in src/ecgen/data/leads.py

def lead_indices(names: Sequence[str], available: Sequence[str]) -> List[int]:
    """Indices into ``available`` for each name in ``names``."""
    position = {n: i for i, n in enumerate(available)}
    return [position[n] for n in names]

maybe_lead_names

maybe_lead_names(n: int, names: Optional[Sequence[str]]) -> List[str]

names if given, else ["lead0", "lead1", ...].

Source code in src/ecgen/data/leads.py

def maybe_lead_names(n: int, names: Optional[Sequence[str]]) -> List[str]:
    """``names`` if given, else ``["lead0", "lead1", ...]``."""
    return list(names) if names is not None else [f"lead{i}" for i in range(n)]

derive_standard12

derive_standard12(x: Any, lead_names: Sequence[str]) -> Any

Complete 8 independent leads (I, II, V1–V6, any order) to the standard 12 using Einthoven / Goldberger: III = II − I, aVR = −(I + II)/2, aVL = I − II/2, aVF = II − I/2.

Parameters:

Name	Type	Description	Default
x	Any	[..., C, T] tensor or array with leads lead_names.	required

Returns:

Type	Description
Any	[..., 12, T] in :data:STANDARD_12 order (same type as x).

Source code in src/ecgen/data/leads.py

def derive_standard12(x: Any, lead_names: Sequence[str]) -> Any:
    """
    Complete 8 independent leads (I, II, V1–V6, any order) to the standard 12
    using Einthoven / Goldberger: III = II − I, aVR = −(I + II)/2,
    aVL = I − II/2, aVF = II − I/2.

    Args:
        x: ``[..., C, T]`` tensor or array with leads ``lead_names``.

    Returns:
        ``[..., 12, T]`` in :data:`STANDARD_12` order (same type as ``x``).
    """
    import numpy as np

    idx = {n: i for i, n in enumerate(lead_names)}
    missing = [n for n in INDEPENDENT_8 if n not in idx]
    if missing:
        raise ValueError(
            f"derive_standard12 needs leads {list(INDEPENDENT_8)}; missing {missing}"
        )
    i, ii = x[..., idx["I"], :], x[..., idx["II"], :]
    derived = {
        "III": ii - i,
        "aVR": -(i + ii) / 2,
        "aVL": i - ii / 2,
        "aVF": ii - i / 2,
    }
    leads = [derived[n] if n in derived else x[..., idx[n], :] for n in STANDARD_12]
    try:
        import torch

        if isinstance(x, torch.Tensor):
            return torch.stack(leads, dim=-2)
    except ImportError:  # pragma: no cover
        pass
    return np.stack(leads, axis=-2)

ecgen.data.normalizers

Normalization of numeric condition fields.

FeatureNormalizer dataclass

Per-column standardization, (x - mean) / std, with missing values.

Fit it on the training split only and reuse it for validation, test, and sampling; :meth:to_dict / :meth:from_dict make it storable in a checkpoint. Missing values (NaN) are ignored when fitting and come out of :meth:transform as 0 with missing=True.

Source code in src/ecgen/data/normalizers.py

@dataclass
class FeatureNormalizer:
    """
    Per-column standardization, ``(x - mean) / std``, with missing values.

    Fit it on the training split only and reuse it for validation, test, and
    sampling; :meth:`to_dict` / :meth:`from_dict` make it storable in a checkpoint.
    Missing values (NaN) are ignored when fitting and come out of
    :meth:`transform` as 0 with ``missing=True``.
    """

    columns: List[str]
    mean: Dict[str, float] = field(default_factory=dict)
    std: Dict[str, float] = field(default_factory=dict)
    eps: float = 1e-6

    @classmethod
    def fit(
        cls, df: pd.DataFrame, columns: Sequence[str], eps: float = 1e-6
    ) -> "FeatureNormalizer":
        norm = cls(columns=list(columns), eps=eps)
        for name in norm.columns:
            vals = pd.to_numeric(df[name], errors="coerce").to_numpy(dtype=np.float64)
            vals = vals[~np.isnan(vals)]
            if vals.size == 0:
                norm.mean[name], norm.std[name] = 0.0, 1.0
            else:
                norm.mean[name] = float(np.mean(vals))
                norm.std[name] = float(np.std(vals)) + eps
        return norm

    def transform(self, values: Sequence[float]) -> Tuple[np.ndarray, np.ndarray]:
        """
        Normalize one row of raw values (ordered as :attr:`columns`).

        Returns:
            ``(normalized float32 [F], missing bool [F])``
        """
        raw = np.asarray(values, dtype=np.float64)
        missing = np.isnan(raw)
        mean = np.array([self.mean[c] for c in self.columns])
        std = np.array([self.std[c] for c in self.columns])
        out = np.where(missing, 0.0, (raw - mean) / std)
        return out.astype(np.float32), missing

    def inverse_transform(self, values: np.ndarray) -> np.ndarray:
        mean = np.array([self.mean[c] for c in self.columns])
        std = np.array([self.std[c] for c in self.columns])
        return np.asarray(values) * std + mean

    def to_dict(self) -> Dict[str, Any]:
        return {
            "columns": list(self.columns),
            "mean": dict(self.mean),
            "std": dict(self.std),
            "eps": self.eps,
        }

    @classmethod
    def from_dict(cls, d: Dict[str, Any]) -> "FeatureNormalizer":
        return cls(
            columns=list(d["columns"]),
            mean=dict(d["mean"]),
            std=dict(d["std"]),
            eps=d.get("eps", 1e-6),
        )

transform

transform(values: Sequence[float]) -> Tuple[np.ndarray, np.ndarray]

Normalize one row of raw values (ordered as :attr:columns).

Returns:

Type	Description
Tuple[ndarray, ndarray]	(normalized float32 [F], missing bool [F])

Source code in src/ecgen/data/normalizers.py

def transform(self, values: Sequence[float]) -> Tuple[np.ndarray, np.ndarray]:
    """
    Normalize one row of raw values (ordered as :attr:`columns`).

    Returns:
        ``(normalized float32 [F], missing bool [F])``
    """
    raw = np.asarray(values, dtype=np.float64)
    missing = np.isnan(raw)
    mean = np.array([self.mean[c] for c in self.columns])
    std = np.array([self.std[c] for c in self.columns])
    out = np.where(missing, 0.0, (raw - mean) / std)
    return out.astype(np.float32), missing

ecgen.data.collate

Collation of dict samples into :class:ecgen.types.ECGBatch.

ECGCollate

collate_fn that stacks dict samples into an :class:ECGBatch.

Tensors under signal, cond, cond_missing and obs_mask are stacked; meta values are gathered into lists. With a mask_generator, each sample without its own obs_mask gets one drawn from it.

Source code in src/ecgen/data/collate.py

class ECGCollate:
    """
    ``collate_fn`` that stacks dict samples into an :class:`ECGBatch`.

    Tensors under ``signal``, ``cond``, ``cond_missing`` and ``obs_mask`` are
    stacked; ``meta`` values are gathered into lists. With a ``mask_generator``,
    each sample without its own ``obs_mask`` gets one drawn from it.
    """

    def __init__(
        self,
        lead_names: Sequence[str],
        fs: int,
        mask_generator: Optional[MaskGenerator] = None,
    ) -> None:
        self.lead_names = list(lead_names)
        self.fs = int(fs)
        self.mask_generator = mask_generator

    def __call__(self, samples: List[Dict[str, Any]]) -> ECGBatch:
        if self.mask_generator is not None:
            for s in samples:
                if s.get("obs_mask") is None:
                    c, t = s["signal"].shape
                    s["obs_mask"] = self.mask_generator(c, t)
        first = samples[0]
        cond = {
            k: _stack([s["cond"][k] for s in samples]) for k in first.get("cond", {})
        }
        cond_missing = {
            k: torch.stack([s["cond_missing"][k] for s in samples])
            for k in first.get("cond_missing", {})
        }
        obs_mask = (
            torch.stack([s["obs_mask"] for s in samples])
            if first.get("obs_mask") is not None
            else None
        )
        meta = {k: [s["meta"][k] for s in samples] for k in first.get("meta", {})}
        return ECGBatch(
            signal=torch.stack([s["signal"] for s in samples]),
            lead_names=list(self.lead_names),
            fs=self.fs,
            cond=cond,
            cond_missing=cond_missing,
            obs_mask=obs_mask,
            meta=meta,
        )

ecgen.data.datamodule

Generic LightningDataModule over any registered dataset.

ECGDataModule

Bases: LightningDataModule

Builds train/val/test splits of one dataset and yields :class:~ecgen.types.ECGBatch.

Parameters:

Name	Type	Description	Default
dataset	Mapping[str, Any]	Dataset spec, {"name" | "target": ..., "params": {...}}; the split argument is filled in per split.	required
batch_size	int	Batch size for all splits.	32
num_workers	int	DataLoader workers.	4
pin_memory	Optional[bool]	Defaults to torch.cuda.is_available().	None
split_params	Optional[Mapping[str, Mapping[str, Any]]]	Per-split overrides of dataset.params, e.g. {"val": {"max_samples": 1000}}.	None
masks	Optional[Mapping[str, Any]]	Observation-mask spec (see :mod:ecgen.data.masks); adds obs_mask to every batch for partial-ECG conditioning.	None

The train split's feature_normalizer (if any) is passed to the val/test datasets that accept one, and saved in checkpoints via :meth:state_dict.

Source code in src/ecgen/data/datamodule.py

class ECGDataModule(pl.LightningDataModule):
    """
    Builds train/val/test splits of one dataset and yields :class:`~ecgen.types.ECGBatch`.

    Args:
        dataset: Dataset spec, ``{"name" | "target": ..., "params": {...}}``; the
            ``split`` argument is filled in per split.
        batch_size: Batch size for all splits.
        num_workers: DataLoader workers.
        pin_memory: Defaults to ``torch.cuda.is_available()``.
        split_params: Per-split overrides of ``dataset.params``, e.g.
            ``{"val": {"max_samples": 1000}}``.
        masks: Observation-mask spec (see :mod:`ecgen.data.masks`); adds
            ``obs_mask`` to every batch for partial-ECG conditioning.

    The train split's ``feature_normalizer`` (if any) is passed to the val/test
    datasets that accept one, and saved in checkpoints via :meth:`state_dict`.
    """

    def __init__(
        self,
        dataset: Mapping[str, Any],
        batch_size: int = 32,
        num_workers: int = 4,
        pin_memory: Optional[bool] = None,
        split_params: Optional[Mapping[str, Mapping[str, Any]]] = None,
        masks: Optional[Mapping[str, Any]] = None,
    ) -> None:
        super().__init__()
        self.dataset_spec = dict(dataset)
        self.batch_size = batch_size
        self.num_workers = num_workers
        self.pin_memory = (
            torch.cuda.is_available() if pin_memory is None else pin_memory
        )
        self.split_params = {k: dict(v) for k, v in (split_params or {}).items()}
        self.mask_spec = dict(masks) if masks else None

        self.datasets: Dict[str, BaseECGDataset] = {}

    def build_dataset(self, split: str) -> BaseECGDataset:
        cls = DATASETS.resolve(self.dataset_spec)
        params = {
            **dict(self.dataset_spec.get("params") or {}),
            **self.split_params.get(split, {}),
            "split": split,
        }
        train = self.datasets.get("train")
        if (
            split != "train"
            and "feature_normalizer" not in params
            and train is not None
            and getattr(train, "feature_normalizer", None) is not None
            and "feature_normalizer" in inspect.signature(cls).parameters
        ):
            params["feature_normalizer"] = train.feature_normalizer
        return cls(**params)

    def setup(self, stage: Optional[str] = None) -> None:
        splits = {
            "fit": ("train", "val"),
            "validate": ("val",),
            "test": ("test",),
            "predict": ("test",),
        }
        wanted = (
            splits.get(stage, ("train", "val", "test"))
            if stage
            else ("train", "val", "test")
        )
        for split in wanted:
            if split not in self.datasets:
                self.datasets[split] = self.build_dataset(split)

    def _loader(self, split: str, shuffle: bool) -> DataLoader:
        if split not in self.datasets:
            self.datasets[split] = self.build_dataset(split)
        ds = self.datasets[split]
        return DataLoader(
            ds,
            batch_size=self.batch_size,
            shuffle=shuffle,
            num_workers=self.num_workers,
            pin_memory=self.pin_memory,
            collate_fn=ECGCollate(
                ds.lead_names, ds.fs, build_mask(self.mask_spec, ds.lead_names)
            ),
        )

    def train_dataloader(self) -> DataLoader:
        return self._loader("train", shuffle=True)

    def val_dataloader(self) -> DataLoader:
        return self._loader("val", shuffle=False)

    def test_dataloader(self) -> DataLoader:
        return self._loader("test", shuffle=False)

    @property
    def reference_dataset(self) -> BaseECGDataset:
        """Any built split (train preferred), for lead names, fs, etc."""
        for split in ("train", "val", "test"):
            if split in self.datasets:
                return self.datasets[split]
        raise RuntimeError("No dataset built yet; call setup() first.")

    def state_dict(self) -> Dict[str, Any]:
        if not self.datasets:
            return {}
        ds = self.reference_dataset
        train = self.datasets.get("train", ds)
        norm = getattr(train, "feature_normalizer", None)
        return {
            "lead_names": list(ds.lead_names),
            "fs": ds.fs,
            "seq_length": ds.seq_length,
            "scale": getattr(ds, "scale", None),
            "feature_normalizer": norm.to_dict() if norm is not None else None,
            "normalizers": {f: n.to_dict() for f, n in train.field_normalizers.items()},
        }

reference_dataset property

reference_dataset: BaseECGDataset

Any built split (train preferred), for lead names, fs, etc.