Conditioning API¶

Auto-generated from docstrings.

`ecgen.conditioning.composite`¶

Combining several field encoders into one :class:~ecgen.types.ConditionOutput, with per-field condition dropout (for classifier-free guidance).

CompositeConditioner ¶

Bases: Module

Encodes condition fields into the vector / tokens / dense slots.

vector encoders are projected to vector_dim and summed;
token encoders are projected to token_dim and concatenated along N;
dense encoders are concatenated along channels.

Config (also stored in model hyperparameters)::

conditioning:
  vector_dim: 256
  token_dim: 256
  encoders:
    - {field: measurements, name: numeric, params: {num_features: 9}}
    - {field: obs, name: masked_signal, params: {num_leads: 12}}
  dropout: {measurements: 0.1}      # per-field drop probability in training

A dropped or absent field is replaced by a learned null embedding (vector: one [vector_dim] row per field; tokens: one null token; dense: everything unobserved). The same null is used for the unconditional pass of classifier-free guidance (drop="all").

Parameters:

Name	Type	Description	Default
`encoders`	`Sequence[Mapping[str, Any]]`	Specs `{field, name \| target, params}`.	required
`vector_dim`	`int`	Size of the combined global vector.	`256`
`token_dim`	`int`	Size of each token.	`256`
`dropout`	`Optional[Mapping[str, float]]`	Per-field drop probability during training.	`None`

Source code in src/ecgen/conditioning/composite.py

class CompositeConditioner(nn.Module):
    """
    Encodes condition fields into the ``vector`` / ``tokens`` / ``dense`` slots.

    - vector encoders are projected to ``vector_dim`` and summed;
    - token encoders are projected to ``token_dim`` and concatenated along N;
    - dense encoders are concatenated along channels.

    Config (also stored in model hyperparameters)::

        conditioning:
          vector_dim: 256
          token_dim: 256
          encoders:
            - {field: measurements, name: numeric, params: {num_features: 9}}
            - {field: obs, name: masked_signal, params: {num_leads: 12}}
          dropout: {measurements: 0.1}      # per-field drop probability in training

    A dropped or absent field is replaced by a learned *null* embedding
    (vector: one ``[vector_dim]`` row per field; tokens: one null token;
    dense: everything unobserved). The same null is used for the
    unconditional pass of classifier-free guidance (``drop="all"``).

    Args:
        encoders: Specs ``{field, name | target, params}``.
        vector_dim: Size of the combined global vector.
        token_dim: Size of each token.
        dropout: Per-field drop probability during training.
    """

    def __init__(
        self,
        encoders: Sequence[Mapping[str, Any]],
        vector_dim: int = 256,
        token_dim: int = 256,
        dropout: Optional[Mapping[str, float]] = None,
    ) -> None:
        super().__init__()
        if not encoders:
            raise ValueError("conditioning.encoders is empty")
        self.vector_dim = vector_dim
        self.token_dim = token_dim
        self.dropout = dict(dropout or {})

        self.fields: List[str] = []
        self.encoders = nn.ModuleDict()
        self.proj = nn.ModuleDict()
        self.null_vector = nn.ParameterDict()
        self.null_token = nn.ParameterDict()
        dense_channels = 0
        for spec in encoders:
            field = spec["field"]
            if field in self.encoders:
                raise ValueError(f"Field {field!r} has two encoders")
            enc = CONDITION_ENCODERS.build(
                {k: v for k, v in spec.items() if k != "field"}
            )
            self.fields.append(field)
            self.encoders[field] = enc
            if enc.slot == "vector":
                self.proj[field] = nn.Linear(enc.out_dim, vector_dim)
                self.null_vector[field] = nn.Parameter(torch.zeros(vector_dim))
            elif enc.slot == "tokens":
                self.proj[field] = nn.Linear(enc.out_dim, token_dim)
                self.null_token[field] = nn.Parameter(torch.randn(token_dim) * 0.02)
            elif enc.slot == "dense":
                dense_channels += enc.out_dim
            else:
                raise ValueError(f"Encoder for {field!r} has unknown slot {enc.slot!r}")
        unknown = set(self.dropout) - set(self.fields)
        if unknown:
            raise ValueError(
                f"dropout names fields without encoders: {sorted(unknown)}"
            )
        self.dense_channels = dense_channels

    # ------------------------------------------------------------ properties

    @property
    def slots(self) -> set[str]:
        return {self.encoders[f].slot for f in self.fields}

    @property
    def uses_obs(self) -> bool:
        return OBS_FIELD in self.fields

    @property
    def cond_fields(self) -> List[str]:
        """Fields read from ``cond`` (everything except ``obs``)."""
        return [f for f in self.fields if f != OBS_FIELD]

    # --------------------------------------------------------------- forward

    def _drop_masks(
        self, batch_size: int, device: torch.device, drop: Drop
    ) -> Dict[str, Tensor]:
        """Per-field ``[B]`` bool masks, True = replace with null."""
        masks = {
            f: torch.zeros(batch_size, dtype=torch.bool, device=device)
            for f in self.fields
        }
        if self.training:
            for f, p in self.dropout.items():
                if p > 0:
                    masks[f] = torch.rand(batch_size, device=device) < p
        if drop is None:
            return masks
        if isinstance(drop, Mapping):
            for f, m in drop.items():
                masks[f] = masks[f] | m.to(device).bool()
            return masks
        names = (
            self.fields
            if drop == "all"
            else ([drop] if isinstance(drop, str) else list(drop))
        )
        for f in names:
            masks[f] = torch.ones(batch_size, dtype=torch.bool, device=device)
        return masks

    def forward(
        self,
        cond: Optional[Mapping[str, Tensor]] = None,
        cond_missing: Optional[Mapping[str, Tensor]] = None,
        obs: Optional[Tensor] = None,
        obs_mask: Optional[Tensor] = None,
        batch_size: Optional[int] = None,
        drop: Drop = None,
    ) -> ConditionOutput:
        """
        Args:
            cond: Normalized condition fields; a missing key means "not given".
            cond_missing: Per-field missing masks.
            obs: Partially observed ECG ``[B, C, T]``.
            obs_mask: Mask of ``obs`` (1 = observed).
            batch_size: Needed when no field is given at all.
            drop: ``"all"``, field name(s), or ``{field: [B] bool}`` to force nulls
                (e.g. the unconditional half of classifier-free guidance).
        """
        cond = dict(cond or {})
        cond_missing = dict(cond_missing or {})
        b = batch_size or _infer_batch(cond, obs)
        device = _infer_device(self, cond, obs)
        drops = self._drop_masks(b, device, drop)

        vector: Optional[Tensor] = None
        tokens: List[Tensor] = []
        token_masks: List[Tensor] = []
        dense: List[Tensor] = []
        for field in self.fields:
            enc = self.encoders[field]
            dropped = drops[field]
            if field == OBS_FIELD:
                given = obs is not None and obs_mask is not None
            else:
                given = field in cond
            if not given:
                dropped = torch.ones_like(dropped)

            if enc.slot == "dense":
                dense.append(self._dense(field, enc, obs, obs_mask, dropped, b, device))
                continue

            null_rows = dropped[:, None]
            # Given fields are always encoded, so a fully dropped batch keeps the
            # same token count as a conditional one (CFG can batch both passes).
            out = enc(cond[field], cond_missing.get(field)) if given else None

            if enc.slot == "vector":
                null = self.null_vector[field].expand(b, -1)
                v = (
                    null
                    if out is None
                    else torch.where(null_rows, null, self.proj[field](out))
                )
                vector = v if vector is None else vector + v
            else:
                null_tok = self.null_token[field].view(1, 1, -1).expand(b, 1, -1)
                if out is None:
                    tok, tmask = (
                        null_tok,
                        torch.ones(b, 1, dtype=torch.bool, device=device),
                    )
                else:
                    raw_tok, raw_mask = out
                    tok = self.proj[field](raw_tok)
                    # dropped rows: only the null token (in slot 0) is valid
                    tok = torch.where(
                        null_rows[..., None], null_tok.expand_as(tok), tok
                    )
                    tmask = raw_mask.bool() & ~null_rows
                    tmask[:, 0] |= dropped
                tokens.append(tok)
                token_masks.append(tmask)

        return ConditionOutput(
            vector=vector,
            tokens=torch.cat(tokens, dim=1) if tokens else None,
            tokens_mask=torch.cat(token_masks, dim=1) if token_masks else None,
            dense=torch.cat(dense, dim=1) if dense else None,
        )

    def _dense(self, field, enc, obs, obs_mask, dropped, b, device) -> Tensor:
        if obs is None or obs_mask is None:
            raise ValueError(
                f"Dense encoder for {field!r} needs obs/obs_mask; configure data.masks "
                "for training or pass obs= and obs_mask= when sampling"
            )
        keep = (~dropped).to(obs.dtype).view(-1, *([1] * (obs.dim() - 1)))
        return enc((obs * keep, obs_mask.to(obs.dtype) * keep))

cond_fields `property` ¶

cond_fields: List[str]

Fields read from cond (everything except obs).

forward ¶

forward(cond: Optional[Mapping[str, Tensor]] = None, cond_missing: Optional[Mapping[str, Tensor]] = None, obs: Optional[Tensor] = None, obs_mask: Optional[Tensor] = None, batch_size: Optional[int] = None, drop: Drop = None) -> ConditionOutput

Parameters:

Name	Type	Description	Default
`cond`	`Optional[Mapping[str, Tensor]]`	Normalized condition fields; a missing key means "not given".	`None`
`cond_missing`	`Optional[Mapping[str, Tensor]]`	Per-field missing masks.	`None`
`obs`	`Optional[Tensor]`	Partially observed ECG `[B, C, T]`.	`None`
`obs_mask`	`Optional[Tensor]`	Mask of `obs` (1 = observed).	`None`
`batch_size`	`Optional[int]`	Needed when no field is given at all.	`None`
`drop`	`Drop`	`"all"`, field name(s), or `{field: [B] bool}` to force nulls (e.g. the unconditional half of classifier-free guidance).	`None`

Source code in src/ecgen/conditioning/composite.py

def forward(
    self,
    cond: Optional[Mapping[str, Tensor]] = None,
    cond_missing: Optional[Mapping[str, Tensor]] = None,
    obs: Optional[Tensor] = None,
    obs_mask: Optional[Tensor] = None,
    batch_size: Optional[int] = None,
    drop: Drop = None,
) -> ConditionOutput:
    """
    Args:
        cond: Normalized condition fields; a missing key means "not given".
        cond_missing: Per-field missing masks.
        obs: Partially observed ECG ``[B, C, T]``.
        obs_mask: Mask of ``obs`` (1 = observed).
        batch_size: Needed when no field is given at all.
        drop: ``"all"``, field name(s), or ``{field: [B] bool}`` to force nulls
            (e.g. the unconditional half of classifier-free guidance).
    """
    cond = dict(cond or {})
    cond_missing = dict(cond_missing or {})
    b = batch_size or _infer_batch(cond, obs)
    device = _infer_device(self, cond, obs)
    drops = self._drop_masks(b, device, drop)

    vector: Optional[Tensor] = None
    tokens: List[Tensor] = []
    token_masks: List[Tensor] = []
    dense: List[Tensor] = []
    for field in self.fields:
        enc = self.encoders[field]
        dropped = drops[field]
        if field == OBS_FIELD:
            given = obs is not None and obs_mask is not None
        else:
            given = field in cond
        if not given:
            dropped = torch.ones_like(dropped)

        if enc.slot == "dense":
            dense.append(self._dense(field, enc, obs, obs_mask, dropped, b, device))
            continue

        null_rows = dropped[:, None]
        # Given fields are always encoded, so a fully dropped batch keeps the
        # same token count as a conditional one (CFG can batch both passes).
        out = enc(cond[field], cond_missing.get(field)) if given else None

        if enc.slot == "vector":
            null = self.null_vector[field].expand(b, -1)
            v = (
                null
                if out is None
                else torch.where(null_rows, null, self.proj[field](out))
            )
            vector = v if vector is None else vector + v
        else:
            null_tok = self.null_token[field].view(1, 1, -1).expand(b, 1, -1)
            if out is None:
                tok, tmask = (
                    null_tok,
                    torch.ones(b, 1, dtype=torch.bool, device=device),
                )
            else:
                raw_tok, raw_mask = out
                tok = self.proj[field](raw_tok)
                # dropped rows: only the null token (in slot 0) is valid
                tok = torch.where(
                    null_rows[..., None], null_tok.expand_as(tok), tok
                )
                tmask = raw_mask.bool() & ~null_rows
                tmask[:, 0] |= dropped
            tokens.append(tok)
            token_masks.append(tmask)

    return ConditionOutput(
        vector=vector,
        tokens=torch.cat(tokens, dim=1) if tokens else None,
        tokens_mask=torch.cat(token_masks, dim=1) if token_masks else None,
        dense=torch.cat(dense, dim=1) if dense else None,
    )

build_conditioner ¶

build_conditioner(spec: Optional[Mapping[str, Any]]) -> Optional[CompositeConditioner]

None / empty spec → None (unconditional).

Source code in src/ecgen/conditioning/composite.py

def build_conditioner(
    spec: Optional[Mapping[str, Any]],
) -> Optional[CompositeConditioner]:
    """``None`` / empty spec → ``None`` (unconditional)."""
    if not spec:
        return None
    return CompositeConditioner(**dict(spec))

`ecgen.conditioning.encoders`¶

Field encoders: one condition field → one :class:ConditionOutput slot.

Each encoder declares slot ("vector", "tokens" or "dense") and out_dim (feature size, or channels for dense). forward gets the field's value and its missing mask (None if the field has none) and returns:

vector: [B, out_dim]
tokens: ([B, N, out_dim], mask [B, N] bool)
dense: [B, out_dim, T]

CategoricalEncoder ¶

Bases: FieldEncoder

Single class label [B] (int) → embedding [B, dim].

A missing label (mask True, or value < 0) maps to a learned "unknown" embedding.

Source code in src/ecgen/conditioning/encoders.py

@CONDITION_ENCODERS.register("categorical")
class CategoricalEncoder(FieldEncoder):
    """
    Single class label ``[B]`` (int) → embedding ``[B, dim]``.

    A missing label (mask ``True``, or value ``< 0``) maps to a learned
    "unknown" embedding.
    """

    def __init__(self, num_classes: int, dim: int = 128) -> None:
        super().__init__()
        self.num_classes = num_classes
        self.out_dim = dim
        self.embed = nn.Embedding(num_classes + 1, dim)  # last row = unknown

    def forward(self, value: Tensor, missing: Optional[Tensor] = None) -> Tensor:
        idx = value.long().reshape(-1)
        unknown = idx < 0
        if missing is not None:
            unknown = unknown | missing.reshape(-1).bool()
        idx = torch.where(unknown, torch.full_like(idx, self.num_classes), idx)
        return self.embed(idx)

MultiLabelEncoder ¶

Bases: FieldEncoder

Multi-hot labels [B, L] → [B, dim] (linear projection, SSSD-ECG style), or with as_tokens=True one token per present label for cross-attention. Missing labels (mask [B, L]) are treated as absent.

Source code in src/ecgen/conditioning/encoders.py

@CONDITION_ENCODERS.register("multilabel")
class MultiLabelEncoder(FieldEncoder):
    """
    Multi-hot labels ``[B, L]`` → ``[B, dim]`` (linear projection, SSSD-ECG style),
    or with ``as_tokens=True`` one token per *present* label for cross-attention.
    Missing labels (mask ``[B, L]``) are treated as absent.
    """

    def __init__(
        self, num_labels: int, dim: int = 128, as_tokens: bool = False
    ) -> None:
        super().__init__()
        self.num_labels = num_labels
        self.out_dim = dim
        self.as_tokens = as_tokens
        self.slot = "tokens" if as_tokens else "vector"
        if as_tokens:
            self.label_emb = nn.Parameter(torch.randn(num_labels, dim) * 0.02)
        else:
            self.proj = _mlp(num_labels, dim, None)

    def forward(self, value: Tensor, missing: Optional[Tensor] = None):
        y = value.float()
        if missing is not None:
            y = y * (~missing.bool()).float()
        if not self.as_tokens:
            return self.proj(y)
        tokens = self.label_emb[None] * y[..., None]  # [B, L, D]
        return tokens, y > 0.5

NumericEncoder ¶

Bases: FieldEncoder

Normalized numeric features [B, F] with a missing mask → [B, dim].

The MLP sees (value * observed, missing), so missingness is explicit. With as_tokens=True each feature becomes one token value * w_f + b_f (or a learned missing token), for cross-attention.

Source code in src/ecgen/conditioning/encoders.py

@CONDITION_ENCODERS.register("numeric")
class NumericEncoder(FieldEncoder):
    """
    Normalized numeric features ``[B, F]`` with a missing mask → ``[B, dim]``.

    The MLP sees ``(value * observed, missing)``, so missingness is explicit.
    With ``as_tokens=True`` each feature becomes one token
    ``value * w_f + b_f`` (or a learned missing token), for cross-attention.
    """

    def __init__(
        self,
        num_features: int,
        dim: int = 128,
        hidden: Optional[int] = None,
        as_tokens: bool = False,
    ) -> None:
        super().__init__()
        self.num_features = num_features
        self.out_dim = dim
        self.as_tokens = as_tokens
        self.slot = "tokens" if as_tokens else "vector"
        if as_tokens:
            self.weight = nn.Parameter(torch.randn(num_features, dim) * 0.02)
            self.bias = nn.Parameter(torch.randn(num_features, dim) * 0.02)
            self.missing_token = nn.Parameter(torch.randn(num_features, dim) * 0.02)
        else:
            self.mlp = _mlp(2 * num_features, dim, hidden)

    def forward(self, value: Tensor, missing: Optional[Tensor] = None):
        x = value.float()
        m = torch.zeros_like(x, dtype=torch.bool) if missing is None else missing.bool()
        x = torch.where(m, torch.zeros_like(x), x)
        if not self.as_tokens:
            return self.mlp(torch.cat([x, m.float()], dim=-1))
        tokens = x[..., None] * self.weight + self.bias
        tokens = torch.where(m[..., None], self.missing_token.expand_as(tokens), tokens)
        return tokens, torch.ones_like(m)

FourierScalarEncoder ¶

Bases: FieldEncoder

Numeric features [B, F] → random Fourier features → MLP [B, dim].

Source code in src/ecgen/conditioning/encoders.py

@CONDITION_ENCODERS.register("fourier")
class FourierScalarEncoder(FieldEncoder):
    """Numeric features ``[B, F]`` → random Fourier features → MLP ``[B, dim]``."""

    def __init__(
        self, num_features: int, dim: int = 128, n_freqs: int = 16, sigma: float = 1.0
    ) -> None:
        super().__init__()
        self.out_dim = dim
        self.fourier = FourierFeatures(num_features, n_freqs, sigma)
        self.mlp = _mlp(self.fourier.out_dim + num_features, dim, None)

    def forward(self, value: Tensor, missing: Optional[Tensor] = None) -> Tensor:
        x = value.float()
        m = torch.zeros_like(x, dtype=torch.bool) if missing is None else missing.bool()
        x = torch.where(m, torch.zeros_like(x), x)
        feats = self.fourier(x) * (~m).float().repeat_interleave(
            self.fourier.out_dim // x.shape[1], dim=1
        )
        return self.mlp(torch.cat([feats, m.float()], dim=-1))

MaskedSignalEncoder ¶

Bases: FieldEncoder

Partially observed ECG → dense [B, 2C, T] = [x * mask, mask].

Reads obs / obs_mask (not a cond field). With channels set, a 1x1 conv projects to that many channels.

Source code in src/ecgen/conditioning/encoders.py

@CONDITION_ENCODERS.register("masked_signal")
class MaskedSignalEncoder(FieldEncoder):
    """
    Partially observed ECG → dense ``[B, 2C, T]`` = ``[x * mask, mask]``.

    Reads ``obs`` / ``obs_mask`` (not a ``cond`` field). With ``channels`` set,
    a 1x1 conv projects to that many channels.
    """

    slot = "dense"

    def __init__(self, num_leads: int, channels: Optional[int] = None) -> None:
        super().__init__()
        self.num_leads = num_leads
        self.proj = nn.Conv1d(2 * num_leads, channels, 1) if channels else None
        self.out_dim = channels or 2 * num_leads

    def forward(self, value: Tensor, missing: Optional[Tensor] = None) -> Tensor:
        obs, mask = value
        mask = mask.to(obs.dtype)
        dense = torch.cat([obs * mask, mask], dim=1)
        return self.proj(dense) if self.proj is not None else dense

`ecgen.nn.injection`¶

Layers that inject a condition into a backbone.

All modulation layers are zero-initialised, so a freshly built conditional network computes exactly what its unconditional counterpart would; the condition's influence is learned from there.

Shapes: signals/feature maps are [B, C, T], global conditions [B, D], token conditions [B, N, D].

ConcatInput ¶

Bases: Module

Module form of :func:concat_dense; out_channels = in_channels + cond_channels.

Source code in src/ecgen/nn/injection.py

class ConcatInput(nn.Module):
    """Module form of :func:`concat_dense`; ``out_channels = in_channels + cond_channels``."""

    def __init__(self, in_channels: int, cond_channels: int) -> None:
        super().__init__()
        self.out_channels = in_channels + cond_channels

    def forward(self, x: Tensor, dense: Optional[Tensor]) -> Tensor:
        return concat_dense(x, dense)

FiLM ¶

Bases: Module

h * (1 + gamma(c)) + beta(c) per channel.

Source code in src/ecgen/nn/injection.py

class FiLM(nn.Module):
    """``h * (1 + gamma(c)) + beta(c)`` per channel."""

    def __init__(self, cond_dim: int, channels: int) -> None:
        super().__init__()
        self.proj = _zero(nn.Linear(cond_dim, 2 * channels))

    def forward(self, h: Tensor, c: Optional[Tensor]) -> Tensor:
        if c is None:
            return h
        gamma, beta = self.proj(c).chunk(2, dim=-1)
        shape = (h.shape[0], -1) + (1,) * (h.dim() - 2)
        return h * (1 + gamma.view(shape)) + beta.view(shape)

AdaGN ¶

Bases: Module

GroupNorm, then FiLM from the condition (zero-initialised, so it starts as plain GroupNorm).

Source code in src/ecgen/nn/injection.py

class AdaGN(nn.Module):
    """GroupNorm, then FiLM from the condition (zero-initialised, so it starts as plain GroupNorm)."""

    def __init__(self, cond_dim: int, channels: int, groups: int = 8) -> None:
        super().__init__()
        self.norm = nn.GroupNorm(groups, channels, affine=True)
        self.film = FiLM(cond_dim, channels)

    def forward(self, h: Tensor, c: Optional[Tensor]) -> Tensor:
        return self.film(self.norm(h), c)

AdaLNZero ¶

Bases: Module

DiT-style adaptive LayerNorm for [B, N, D] tokens.

forward returns (modulated_x, gate); the block adds gate * f(modulated_x) to its residual. Gates start at 0.

Source code in src/ecgen/nn/injection.py

class AdaLNZero(nn.Module):
    """
    DiT-style adaptive LayerNorm for ``[B, N, D]`` tokens.

    ``forward`` returns ``(modulated_x, gate)``; the block adds
    ``gate * f(modulated_x)`` to its residual. Gates start at 0.
    """

    def __init__(self, cond_dim: int, dim: int) -> None:
        super().__init__()
        self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
        self.proj = _zero(nn.Linear(cond_dim, 3 * dim))

    def forward(self, x: Tensor, c: Tensor) -> tuple[Tensor, Tensor]:
        shift, scale, gate = self.proj(F.silu(c)).unsqueeze(1).chunk(3, dim=-1)
        return self.norm(x) * (1 + scale) + shift, gate

AddToTimeEmb ¶

Bases: Module

Project a global condition and add it to the time embedding.

Source code in src/ecgen/nn/injection.py

class AddToTimeEmb(nn.Module):
    """Project a global condition and add it to the time embedding."""

    def __init__(self, cond_dim: int, temb_dim: int) -> None:
        super().__init__()
        self.proj = _zero(nn.Linear(cond_dim, temb_dim))

    def forward(self, temb: Tensor, c: Optional[Tensor]) -> Tensor:
        return temb if c is None else temb + self.proj(c)

CrossAttention1d ¶

Bases: Module

Feature map [B, C, T] attends to condition tokens [B, N, D].

Residual: x + out(attn(norm(x), tokens)), with out zero-initialised.

Source code in src/ecgen/nn/injection.py

class CrossAttention1d(nn.Module):
    """
    Feature map ``[B, C, T]`` attends to condition tokens ``[B, N, D]``.

    Residual: ``x + out(attn(norm(x), tokens))``, with ``out`` zero-initialised.
    """

    def __init__(
        self, channels: int, context_dim: int, heads: int = 4, groups: int = 8
    ) -> None:
        super().__init__()
        if channels % heads:
            raise ValueError(f"channels={channels} must be divisible by heads={heads}")
        self.heads = heads
        self.norm = nn.GroupNorm(groups, channels)
        self.q = nn.Linear(channels, channels)
        self.kv = nn.Linear(context_dim, 2 * channels)
        self.out = _zero(nn.Linear(channels, channels))

    def forward(
        self, x: Tensor, tokens: Optional[Tensor], mask: Optional[Tensor] = None
    ) -> Tensor:
        if tokens is None:
            return x
        b, c, t = x.shape
        h = self.norm(x).transpose(1, 2)  # [B, T, C]
        q = self.q(h).view(b, t, self.heads, -1).transpose(1, 2)
        k, v = self.kv(tokens).chunk(2, dim=-1)
        k = k.view(b, -1, self.heads, c // self.heads).transpose(1, 2)
        v = v.view(b, -1, self.heads, c // self.heads).transpose(1, 2)
        attn_mask = None
        if mask is not None:
            attn_mask = mask[:, None, None, :].to(torch.bool)  # True = attend
            # rows with no valid token would give NaN; let them attend to everything
            attn_mask = attn_mask | ~attn_mask.any(dim=-1, keepdim=True)
        out = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
        out = out.transpose(1, 2).reshape(b, t, c)
        return x + self.out(out).transpose(1, 2)

concat_dense ¶

concat_dense(x: Tensor, dense: Optional[Tensor]) -> Tensor

ConcatInput: stack time-aligned conditions onto the input channels.

Source code in src/ecgen/nn/injection.py

def concat_dense(x: Tensor, dense: Optional[Tensor]) -> Tensor:
    """``ConcatInput``: stack time-aligned conditions onto the input channels."""
    if dense is None:
        return x
    if dense.shape[-1] != x.shape[-1]:
        dense = F.interpolate(
            dense, size=x.shape[-1], mode="linear", align_corners=False
        )
    return torch.cat([x, dense], dim=1)

`ecgen.nn.embeddings`¶

Embeddings for diffusion time steps and continuous scalars.

TimestepEmbedding ¶

Bases: Module

Sinusoidal embedding followed by a 2-layer MLP (DDPM / ADM style).

Source code in src/ecgen/nn/embeddings.py

class TimestepEmbedding(nn.Module):
    """Sinusoidal embedding followed by a 2-layer MLP (DDPM / ADM style)."""

    def __init__(self, dim: int, hidden: int | None = None, scale: float = 1.0) -> None:
        super().__init__()
        hidden = hidden or 4 * dim
        self.dim = dim
        self.scale = scale
        self.mlp = nn.Sequential(
            nn.Linear(dim, hidden), nn.SiLU(), nn.Linear(hidden, hidden)
        )
        self.out_dim = hidden

    def forward(self, t: Tensor) -> Tensor:
        return self.mlp(sinusoidal_embedding(t * self.scale, self.dim))

FourierFeatures ¶

Bases: Module

Random Fourier features of scalars: [B, F] -> [B, F * 2 * n_freqs].

Frequencies are fixed (a buffer), drawn once with a fixed seed so models built twice from the same config are identical.

Source code in src/ecgen/nn/embeddings.py

class FourierFeatures(nn.Module):
    """
    Random Fourier features of scalars: ``[B, F] -> [B, F * 2 * n_freqs]``.

    Frequencies are fixed (a buffer), drawn once with a fixed seed so models
    built twice from the same config are identical.
    """

    def __init__(
        self, num_features: int, n_freqs: int = 16, sigma: float = 1.0
    ) -> None:
        super().__init__()
        gen = torch.Generator().manual_seed(0)
        self.register_buffer(
            "freqs", torch.randn(num_features, n_freqs, generator=gen) * sigma
        )
        self.out_dim = num_features * 2 * n_freqs

    def forward(self, x: Tensor) -> Tensor:
        args = 2 * math.pi * x[..., None] * self.freqs  # [B, F, n]
        return torch.cat([torch.cos(args), torch.sin(args)], dim=-1).flatten(1)

sinusoidal_embedding ¶

sinusoidal_embedding(t: Tensor, dim: int, max_period: float = 10000.0) -> Tensor

Transformer-style sinusoidal embedding of (possibly fractional) positions.

Parameters:

Name	Type	Description	Default
`t`	`Tensor`	`[B]` time steps or noise levels.	required
`dim`	`int`	Embedding size.	required

Returns:

Type	Description
`Tensor`	`[B, dim]`

Source code in src/ecgen/nn/embeddings.py

def sinusoidal_embedding(t: Tensor, dim: int, max_period: float = 10000.0) -> Tensor:
    """
    Transformer-style sinusoidal embedding of (possibly fractional) positions.

    Args:
        t: ``[B]`` time steps or noise levels.
        dim: Embedding size.

    Returns:
        ``[B, dim]``
    """
    half = dim // 2
    freqs = torch.exp(
        -math.log(max_period)
        * torch.arange(half, device=t.device, dtype=torch.float32)
        / half
    )
    args = t.float()[:, None] * freqs[None]
    emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
    if dim % 2:
        emb = torch.cat([emb, torch.zeros_like(emb[:, :1])], dim=-1)
    return emb

`ecgen.data.masks`¶

Observation masks for partial-ECG conditioning (lead reconstruction, imputation, forecasting).

A mask generator is called once per sample with the signal shape and returns an obs_mask [C, T] float tensor, 1 = observed. The data module attaches it to the batch (ECGBatch.obs_mask); models condition on signal * obs_mask and (optionally) compute the loss on the unobserved part.

Config (data.masks)::

masks: {name: lead_subset, params: {keep: [[I], [I, II], [II, V1]]}}
masks:
  name: mixture
  params:
    components:
      - {weight: 0.5, name: lead_subset, params: {keep: [[I]]}}
      - {weight: 0.3, name: time_span, params: {min_frac: 0.1, max_frac: 0.5}}
      - {weight: 0.2, name: none}

MaskGenerator ¶

Base class; subclasses implement :meth:__call__.

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

class MaskGenerator:
    """Base class; subclasses implement :meth:`__call__`."""

    def __init__(self, lead_names: Sequence[str]) -> None:
        self.lead_names = list(lead_names)

    def __call__(
        self, num_leads: int, length: int, generator: Optional[torch.Generator] = None
    ) -> Tensor:
        raise NotImplementedError

FullyObserved ¶

Bases: MaskGenerator

Everything observed (also: the "unconditional" branch of a mixture).

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

@MASKS.register("none")
class FullyObserved(MaskGenerator):
    """Everything observed (also: the "unconditional" branch of a mixture)."""

    def __call__(self, num_leads, length, generator=None):
        return torch.ones(num_leads, length)

FullyUnobserved ¶

Bases: MaskGenerator

Nothing observed: plain generation, when mixed with partial masks.

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

@MASKS.register("unobserved")
class FullyUnobserved(MaskGenerator):
    """Nothing observed: plain generation, when mixed with partial masks."""

    def __call__(self, num_leads, length, generator=None):
        return torch.zeros(num_leads, length)

LeadSubset ¶

Bases: MaskGenerator

Observe whole leads.

Parameters:

Name	Type	Description	Default
`keep`	`Optional[Sequence[Sequence[str]]]`	List of lead-name lists; one is chosen uniformly per sample (e.g. `[[I], [I, II]]` = "single lead I" or "leads I and II"). If omitted, a random subset of `min_leads..max_leads` leads is kept.	`None`
`min_leads`	`int`	Smallest random subset.	`1`
`max_leads`	`int`	Largest random subset.	`3`

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

@MASKS.register("lead_subset")
class LeadSubset(MaskGenerator):
    """
    Observe whole leads.

    Args:
        keep: List of lead-name lists; one is chosen uniformly per sample
            (e.g. ``[[I], [I, II]]`` = "single lead I" or "leads I and II").
            If omitted, a random subset of ``min_leads..max_leads`` leads is kept.
        min_leads: Smallest random subset.
        max_leads: Largest random subset.
    """

    def __init__(
        self,
        lead_names: Sequence[str],
        keep: Optional[Sequence[Sequence[str]]] = None,
        min_leads: int = 1,
        max_leads: int = 3,
    ) -> None:
        super().__init__(lead_names)
        self.keep: Optional[List[List[int]]] = None
        if keep is not None:
            self.keep = []
            for names in keep:
                unknown = [n for n in names if n not in self.lead_names]
                if unknown:
                    raise ValueError(
                        f"lead_subset keeps unknown leads {unknown}; data has {self.lead_names}"
                    )
                self.keep.append([self.lead_names.index(n) for n in names])
        self.min_leads, self.max_leads = min_leads, max_leads

    def __call__(self, num_leads, length, generator=None):
        mask = torch.zeros(num_leads, length)
        if self.keep is not None:
            idx = self.keep[_randint(0, len(self.keep) - 1, generator)]
        else:
            k = _randint(self.min_leads, min(self.max_leads, num_leads), generator)
            idx = torch.randperm(num_leads, generator=generator)[:k].tolist()
        mask[idx] = 1.0
        return mask

TimeSpan ¶

Bases: MaskGenerator

Hide n_spans contiguous segments (all leads) covering min_frac..max_frac of the record in total — imputation.

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

@MASKS.register("time_span")
class TimeSpan(MaskGenerator):
    """
    Hide ``n_spans`` contiguous segments (all leads) covering ``min_frac..max_frac``
    of the record in total — imputation.
    """

    def __init__(
        self,
        lead_names: Sequence[str],
        min_frac: float = 0.1,
        max_frac: float = 0.5,
        n_spans: int = 1,
    ) -> None:
        super().__init__(lead_names)
        self.min_frac, self.max_frac, self.n_spans = min_frac, max_frac, n_spans

    def __call__(self, num_leads, length, generator=None):
        mask = torch.ones(num_leads, length)
        frac = self.min_frac + (self.max_frac - self.min_frac) * _rand(generator)
        span = max(1, int(round(frac * length / self.n_spans)))
        for _ in range(self.n_spans):
            start = _randint(0, length - span, generator)
            mask[:, start : start + span] = 0.0
        return mask

Forecast ¶

Bases: MaskGenerator

Observe the first min_frac..max_frac of the record; predict the rest.

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

@MASKS.register("forecast")
class Forecast(MaskGenerator):
    """Observe the first ``min_frac..max_frac`` of the record; predict the rest."""

    def __init__(
        self, lead_names: Sequence[str], min_frac: float = 0.25, max_frac: float = 0.75
    ) -> None:
        super().__init__(lead_names)
        self.min_frac, self.max_frac = min_frac, max_frac

    def __call__(self, num_leads, length, generator=None):
        frac = self.min_frac + (self.max_frac - self.min_frac) * _rand(generator)
        mask = torch.zeros(num_leads, length)
        mask[:, : int(round(frac * length))] = 1.0
        return mask

RandomPoints ¶

Bases: MaskGenerator

Each sample point observed independently with probability p_observed.

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

@MASKS.register("random_points")
class RandomPoints(MaskGenerator):
    """Each sample point observed independently with probability ``p_observed``."""

    def __init__(self, lead_names: Sequence[str], p_observed: float = 0.5) -> None:
        super().__init__(lead_names)
        self.p = p_observed

    def __call__(self, num_leads, length, generator=None):
        return (torch.rand(num_leads, length, generator=generator) < self.p).float()

Mixture ¶

Bases: MaskGenerator

Pick one component per sample with probability proportional to weight.

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

@MASKS.register("mixture")
class Mixture(MaskGenerator):
    """Pick one component per sample with probability proportional to ``weight``."""

    def __init__(
        self, lead_names: Sequence[str], components: Sequence[Mapping[str, Any]]
    ) -> None:
        super().__init__(lead_names)
        if not components:
            raise ValueError("mixture needs at least one component")
        self.components = [build_mask(c, lead_names) for c in components]
        w = torch.tensor([float(c.get("weight", 1.0)) for c in components])
        self.probs = w / w.sum()

    def __call__(self, num_leads, length, generator=None):
        i = int(torch.multinomial(self.probs, 1, generator=generator))
        return self.components[i](num_leads, length, generator)

build_mask ¶

build_mask(spec: Optional[Mapping[str, Any]], lead_names: Sequence[str]) -> Optional[MaskGenerator]

Build a mask generator from {name | target, params} (weight is ignored).

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

def build_mask(
    spec: Optional[Mapping[str, Any]], lead_names: Sequence[str]
) -> Optional[MaskGenerator]:
    """Build a mask generator from ``{name | target, params}`` (``weight`` is ignored)."""
    if not spec:
        return None
    spec = {k: v for k, v in spec.items() if k != "weight"}
    return MASKS.build(spec, lead_names=lead_names)

Conditioning API¶

ecgen.conditioning.composite¶

CompositeConditioner ¶

cond_fields property ¶

forward ¶

build_conditioner ¶

ecgen.conditioning.encoders¶

CategoricalEncoder ¶

MultiLabelEncoder ¶

NumericEncoder ¶

FourierScalarEncoder ¶

MaskedSignalEncoder ¶

ecgen.nn.injection¶

ConcatInput ¶

FiLM ¶

AdaGN ¶

AdaLNZero ¶

AddToTimeEmb ¶

CrossAttention1d ¶

concat_dense ¶

ecgen.nn.embeddings¶

TimestepEmbedding ¶

FourierFeatures ¶

sinusoidal_embedding ¶

ecgen.data.masks¶

MaskGenerator ¶

FullyObserved ¶

FullyUnobserved ¶

LeadSubset ¶

TimeSpan ¶

Forecast ¶

RandomPoints ¶

Mixture ¶

build_mask ¶

`ecgen.conditioning.composite`¶

cond_fields `property` ¶

`ecgen.conditioning.encoders`¶

`ecgen.nn.injection`¶

`ecgen.nn.embeddings`¶

`ecgen.data.masks`¶