Skip to content

Evaluation API

Auto-generated from docstrings.

ecgen.evaluation

Measurements on (generated) ECGs, e.g. to check that a model follows the heart rate it was asked for.

detect_r_peaks 

detect_r_peaks(signal: ArrayLike, fs: int, lead: Optional[int] = None, method: str = 'simple') -> np.ndarray

R-peak sample indices of one ECG [C, T] (or [T]).

Parameters:

Name Type Description Default
signal ArrayLike

One ECG.

 required
fs int

Sampling rate in Hz.

 required
lead Optional[int]

Lead to use; default = the lead with the largest QRS energy.

 None
method str

"simple" (band-pass + peak picking, no extra dependency) or "neurokit" (neurokit2.ecg_peaks; needs ecgen[eval]).

 'simple'
Source code in src/ecgen/evaluation/measure.py 
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
def detect_r_peaks(
    signal: ArrayLike,
    fs: int,
    lead: Optional[int] = None,
    method: str = "simple",
) -> np.ndarray:
    """
    R-peak sample indices of one ECG ``[C, T]`` (or ``[T]``).

    Args:
        signal: One ECG.
        fs: Sampling rate in Hz.
        lead: Lead to use; default = the lead with the largest QRS energy.
        method: ``"simple"`` (band-pass + peak picking, no extra dependency) or
            ``"neurokit"`` (``neurokit2.ecg_peaks``; needs ``ecgen[eval]``).
    """
    x = _as_numpy(signal)
    if x.ndim == 1:
        x = x[None]
    if x.shape[-1] < int(0.5 * fs):
        return np.array([], dtype=int)

    nyq = fs / 2
    b, a = butter(2, [5 / nyq, min(30 / nyq, 0.99)], btype="band")
    filtered = filtfilt(b, a, x, axis=-1)
    if lead is None:
        lead = int(np.argmax(np.percentile(np.abs(filtered), 99, axis=-1)))
    sig = filtered[lead]

    if method == "neurokit":
        import neurokit2 as nk

        _, info = nk.ecg_peaks(x[lead], sampling_rate=fs)
        return np.asarray(info["ECG_R_Peaks"], dtype=int)
    if method != "simple":
        raise ValueError(f"Unknown method {method!r}")

    # QRS polarity varies by lead: use the dominant sign
    if np.abs(sig.min()) > np.abs(sig.max()):
        sig = -sig
    height = 0.4 * np.percentile(sig, 99.5)
    peaks, _ = find_peaks(sig, height=height, distance=int(0.25 * fs))
    return peaks

heart_rate 

heart_rate(signals: ArrayLike, fs: int, method: str = 'simple', lead: Optional[int] = None) -> np.ndarray

Heart rate (bpm) per ECG from the median RR interval; NaN when fewer than two beats are found.

Parameters:

Name Type Description Default
signals ArrayLike

[N, C, T], [C, T], or [T].

 required
fs int

Sampling rate in Hz.

 required
Source code in src/ecgen/evaluation/measure.py 
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
def heart_rate(
    signals: ArrayLike,
    fs: int,
    method: str = "simple",
    lead: Optional[int] = None,
) -> np.ndarray:
    """
    Heart rate (bpm) per ECG from the median RR interval; NaN when fewer than
    two beats are found.

    Args:
        signals: ``[N, C, T]``, ``[C, T]``, or ``[T]``.
        fs: Sampling rate in Hz.
    """
    x = _as_numpy(signals)
    if x.ndim == 1:
        x = x[None, None]
    elif x.ndim == 2:
        x = x[None]
    rates: List[float] = []
    for ecg in x:
        peaks = detect_r_peaks(ecg, fs, lead=lead, method=method)
        if len(peaks) < 2:
            rates.append(np.nan)
        else:
            rates.append(60.0 * fs / float(np.median(np.diff(peaks))))
    return np.asarray(rates)

adherence 

adherence(requested: Sequence[float], measured: Sequence[float]) -> dict

How well measured values follow requested ones (NaNs ignored).

Returns {"mae", "pearson_r", "n_valid", "n"}.

Source code in src/ecgen/evaluation/measure.py 
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
def adherence(requested: Sequence[float], measured: Sequence[float]) -> dict:
    """
    How well measured values follow requested ones (NaNs ignored).

    Returns ``{"mae", "pearson_r", "n_valid", "n"}``.
    """
    req = np.asarray(requested, dtype=np.float64)
    mea = np.asarray(measured, dtype=np.float64)
    ok = ~(np.isnan(req) | np.isnan(mea))
    out = {
        "n": int(len(req)),
        "n_valid": int(ok.sum()),
        "mae": float("nan"),
        "pearson_r": float("nan"),
    }
    if ok.sum() >= 2:
        out["mae"] = float(np.mean(np.abs(req[ok] - mea[ok])))
        if np.std(req[ok]) > 0 and np.std(mea[ok]) > 0:
            out["pearson_r"] = float(np.corrcoef(req[ok], mea[ok])[0, 1])
    return out

Feature extractors for distribution metrics (Fréchet distance, MMD, precision/recall, nearest neighbours).

There is no single pretrained ECG encoder that every user has, so the default is a set of interpretable hand-crafted features; any torch.nn.Module that maps [B, C, T] to [B, D] (e.g. the penultimate layer of an ECG classifier, see :mod:ecgen.evaluation.classifier) can be used instead.

handcrafted_features 

handcrafted_features(x: ArrayLike, fs: int = 500) -> np.ndarray

Per-ECG feature vector [N, D]:

  • heart rate (bpm, NaN → batch median);
  • per lead: standard deviation, 1st/99th percentile, and the relative power in the bands 0.5–5, 5–15, 15–40, 40–100 Hz;
  • the upper triangle of the inter-lead correlation matrix.
Source code in src/ecgen/evaluation/features.py 
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
def handcrafted_features(x: ArrayLike, fs: int = 500) -> np.ndarray:
    """
    Per-ECG feature vector ``[N, D]``:

    - heart rate (bpm, NaN → batch median);
    - per lead: standard deviation, 1st/99th percentile, and the relative power
      in the bands 0.5–5, 5–15, 15–40, 40–100 Hz;
    - the upper triangle of the inter-lead correlation matrix.
    """
    x = _as_numpy(x)
    if x.ndim == 2:
        x = x[None]
    n, c, t = x.shape
    hr = heart_rate(x, fs)
    hr = np.where(np.isnan(hr), np.nanmedian(hr) if np.isfinite(hr).any() else 0.0, hr)

    std = x.std(axis=-1)
    p1, p99 = np.percentile(x, 1, axis=-1), np.percentile(x, 99, axis=-1)
    freqs, psd = welch(x, fs=fs, nperseg=min(t, 2 * fs), axis=-1)
    total = psd.sum(-1, keepdims=True) + 1e-12
    bands = (
        np.stack(
            [
                psd[..., (freqs >= lo) & (freqs < min(hi, fs / 2))].sum(-1)
                for lo, hi in BANDS
            ],
            axis=-1,
        )
        / total
    )
    iu = np.triu_indices(c, k=1)
    corr = (
        np.stack([np.nan_to_num(np.corrcoef(e))[iu] for e in x])
        if c > 1
        else np.zeros((n, 0))
    )
    feats = [hr[:, None], std, p1, p99, bands.reshape(n, -1), corr]
    return np.concatenate(feats, axis=1).astype(np.float64)

module_features 

module_features(model: Module, batch_size: int = 64, device: Optional[str] = None) -> Extractor

Wrap a [B, C, T] -> [B, D] network as an extractor.

Source code in src/ecgen/evaluation/features.py 
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
def module_features(
    model: torch.nn.Module, batch_size: int = 64, device: Optional[str] = None
) -> Extractor:
    """Wrap a ``[B, C, T] -> [B, D]`` network as an extractor."""
    device = device or ("cuda" if torch.cuda.is_available() else "cpu")

    def extract(x: ArrayLike) -> np.ndarray:
        xt = torch.as_tensor(_as_numpy(x), dtype=torch.float32)
        model.eval().to(device)
        out = []
        with torch.no_grad():
            for i in range(0, len(xt), batch_size):
                out.append(
                    model(xt[i : i + batch_size].to(device)).float().cpu().numpy()
                )
        return np.concatenate(out)

    return extract

Metrics for generated ECGs (plan §8): distribution fidelity and diversity, condition adherence, partial-ECG reconstruction, lead consistency, and memorisation.

Arrays are [N, C, T] (NumPy or torch). Distribution metrics work on feature vectors from an extractor (:mod:ecgen.evaluation.features).

frechet_distance 

frechet_distance(feats_a: ndarray, feats_b: ndarray, eps: float = 1e-06) -> float

Fréchet distance between Gaussians fitted to two feature sets (FID-style).

Source code in src/ecgen/evaluation/metrics.py 
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
def frechet_distance(
    feats_a: np.ndarray, feats_b: np.ndarray, eps: float = 1e-6
) -> float:
    """Fréchet distance between Gaussians fitted to two feature sets (FID-style)."""
    mu_a, mu_b = feats_a.mean(0), feats_b.mean(0)
    sa, sb = (
        np.atleast_2d(np.cov(feats_a, rowvar=False)),
        np.atleast_2d(np.cov(feats_b, rowvar=False)),
    )
    covmean, _ = linalg.sqrtm(sa @ sb, disp=False)
    if not np.isfinite(covmean).all():
        off = np.eye(sa.shape[0]) * eps
        covmean = linalg.sqrtm((sa + off) @ (sb + off))
    covmean = covmean.real
    return float(
        ((mu_a - mu_b) ** 2).sum() + np.trace(sa) + np.trace(sb) - 2 * np.trace(covmean)
    )

mmd_rbf 

mmd_rbf(feats_a: ndarray, feats_b: ndarray, bandwidth: Optional[float] = None) -> float

Unbiased squared MMD with an RBF kernel (bandwidth = median heuristic).

Source code in src/ecgen/evaluation/metrics.py 
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
def mmd_rbf(
    feats_a: np.ndarray, feats_b: np.ndarray, bandwidth: Optional[float] = None
) -> float:
    """Unbiased squared MMD with an RBF kernel (bandwidth = median heuristic)."""
    x, y = np.asarray(feats_a, float), np.asarray(feats_b, float)
    z = np.concatenate([x, y])
    d2 = ((z[:, None] - z[None]) ** 2).sum(-1)
    if bandwidth is None:
        bandwidth = np.sqrt(0.5 * np.median(d2[d2 > 0])) if (d2 > 0).any() else 1.0
    k = np.exp(-d2 / (2 * bandwidth**2))
    n, m = len(x), len(y)
    kxx, kyy, kxy = k[:n, :n], k[n:, n:], k[:n, n:]
    return float(
        (kxx.sum() - np.trace(kxx)) / (n * (n - 1))
        + (kyy.sum() - np.trace(kyy)) / (m * (m - 1))
        - 2 * kxy.mean()
    )

precision_recall 

precision_recall(real: ndarray, fake: ndarray, k: int = 3) -> Dict[str, float]

k-NN manifold precision / recall (Kynkäänniemi et al., 2019): precision = share of generated samples inside the real manifold, recall = share of real samples inside the generated manifold.

Source code in src/ecgen/evaluation/metrics.py 
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
def precision_recall(
    real: np.ndarray, fake: np.ndarray, k: int = 3
) -> Dict[str, float]:
    """
    k-NN manifold precision / recall (Kynkäänniemi et al., 2019): precision = share
    of generated samples inside the real manifold, recall = share of real samples
    inside the generated manifold.
    """

    def radii(a: np.ndarray) -> np.ndarray:
        d = np.sqrt(((a[:, None] - a[None]) ** 2).sum(-1))
        return np.sort(d, axis=1)[:, min(k, len(a) - 1)]

    def inside(points: np.ndarray, ref: np.ndarray, r: np.ndarray) -> float:
        d = np.sqrt(((points[:, None] - ref[None]) ** 2).sum(-1))
        return float((d <= r[None]).any(axis=1).mean())

    return {
        "precision": inside(fake, real, radii(real)),
        "recall": inside(real, fake, radii(fake)),
    }

distribution_metrics 

distribution_metrics(real: ArrayLike, fake: ArrayLike, fs: int = 500, extractor: Optional[Callable[[ArrayLike], ndarray]] = None, k: int = 3) -> Dict[str, float]

Fréchet distance, MMD and precision/recall between real and generated ECGs on (standardized) features; default features: :func:handcrafted_features.

Source code in src/ecgen/evaluation/metrics.py 
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
def distribution_metrics(
    real: ArrayLike,
    fake: ArrayLike,
    fs: int = 500,
    extractor: Optional[Callable[[ArrayLike], np.ndarray]] = None,
    k: int = 3,
) -> Dict[str, float]:
    """
    Fréchet distance, MMD and precision/recall between real and generated ECGs on
    (standardized) features; default features: :func:`handcrafted_features`.
    """
    extract = extractor or (lambda x: handcrafted_features(x, fs))
    fr, ff = extract(real), extract(fake)
    fr, ff = _standardize(fr, ff)
    out = {"frechet_distance": frechet_distance(fr, ff), "mmd": mmd_rbf(fr, ff)}
    out.update(precision_recall(fr, ff, k))
    return out

inter_lead_correlation_error 

inter_lead_correlation_error(real: ArrayLike, fake: ArrayLike) -> float

Mean |corr_real − corr_fake| over lead pairs, using mean correlation matrices (MIDT-ECG).

Source code in src/ecgen/evaluation/metrics.py 
111
112
113
114
115
116
117
118
119
def inter_lead_correlation_error(real: ArrayLike, fake: ArrayLike) -> float:
    """Mean |corr_real − corr_fake| over lead pairs, using mean correlation matrices (MIDT-ECG)."""

    def mean_corr(x: np.ndarray) -> np.ndarray:
        return np.nanmean(np.stack([np.corrcoef(e) for e in x]), axis=0)

    r, f = mean_corr(_as_numpy(real)), mean_corr(_as_numpy(fake))
    iu = np.triu_indices(r.shape[0], k=1)
    return float(np.abs(r[iu] - f[iu]).mean())

einthoven_residual 

einthoven_residual(x: ArrayLike, lead_names: Sequence[str]) -> Dict[str, float]

Relative residuals of the limb-lead identities (III = II − I, aVR = −(I + II)/2, aVL = I − II/2, aVF = II − I/2) after per-lead centering; ~0 for real recordings, larger when a model generates all 12 leads independently.

Source code in src/ecgen/evaluation/metrics.py 
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
def einthoven_residual(x: ArrayLike, lead_names: Sequence[str]) -> Dict[str, float]:
    """
    Relative residuals of the limb-lead identities (III = II − I,
    aVR = −(I + II)/2, aVL = I − II/2, aVF = II − I/2) after per-lead centering;
    ~0 for real recordings, larger when a model generates all 12 leads independently.
    """
    x = _as_numpy(x)
    x = x - x.mean(-1, keepdims=True)
    idx = {n: i for i, n in enumerate(lead_names)}
    need = ("I", "II", "III", "aVR", "aVL", "aVF")
    if not all(n in idx for n in need):
        raise ValueError(
            f"einthoven_residual needs leads {need}; got {list(lead_names)}"
        )
    i, ii = x[:, idx["I"]], x[:, idx["II"]]
    expected = {
        "III": ii - i,
        "aVR": -(i + ii) / 2,
        "aVL": i - ii / 2,
        "aVF": ii - i / 2,
    }
    out = {}
    for name, exp in expected.items():
        actual = x[:, idx[name]]
        out[name] = float(
            np.linalg.norm(actual - exp) / (np.linalg.norm(actual) + 1e-12)
        )
    out["mean"] = float(np.mean(list(out.values())))
    return out

reconstruction_metrics 

reconstruction_metrics(pred: ArrayLike, true: ArrayLike, unobserved: ArrayLike) -> Dict[str, float]

MAE, RMSE and mean Pearson r on the unobserved entries (unobserved: [N, C, T] or [C] / [N, C] bool, True = to evaluate).

Source code in src/ecgen/evaluation/metrics.py 
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
def reconstruction_metrics(
    pred: ArrayLike, true: ArrayLike, unobserved: ArrayLike
) -> Dict[str, float]:
    """
    MAE, RMSE and mean Pearson r on the unobserved entries (``unobserved``:
    ``[N, C, T]`` or ``[C]`` / ``[N, C]`` bool, True = to evaluate).
    """
    p, t = _as_numpy(pred), _as_numpy(true)
    m = np.asarray(_as_numpy(unobserved), dtype=bool)
    if m.ndim == 1:
        m = np.broadcast_to(m[None, :, None], p.shape)
    elif m.ndim == 2:
        m = np.broadcast_to(m[:, :, None], p.shape)
    err = (p - t)[m]
    rs = []
    for n in range(p.shape[0]):
        for c in range(p.shape[1]):
            sel = m[n, c]
            if sel.sum() > 2 and p[n, c][sel].std() > 0 and t[n, c][sel].std() > 0:
                rs.append(np.corrcoef(p[n, c][sel], t[n, c][sel])[0, 1])
    return {
        "mae": float(np.abs(err).mean()) if err.size else float("nan"),
        "rmse": float(np.sqrt((err**2).mean())) if err.size else float("nan"),
        "pearson_r": float(np.mean(rs)) if rs else float("nan"),
    }

nearest_neighbour_distances 

nearest_neighbour_distances(query: ArrayLike, reference: ArrayLike, batch: int = 256) -> np.ndarray

L2 distance from each query ECG to its nearest reference ECG (flattened signals).

Source code in src/ecgen/evaluation/metrics.py 
186
187
188
189
190
191
192
193
194
195
196
197
198
def nearest_neighbour_distances(
    query: ArrayLike, reference: ArrayLike, batch: int = 256
) -> np.ndarray:
    """L2 distance from each query ECG to its nearest reference ECG (flattened signals)."""
    q = _as_numpy(query).reshape(len(_as_numpy(query)), -1)
    r = _as_numpy(reference).reshape(len(_as_numpy(reference)), -1)
    r2 = (r**2).sum(1)
    out = []
    for i in range(0, len(q), batch):
        qb = q[i : i + batch]
        d2 = (qb**2).sum(1)[:, None] + r2[None] - 2 * qb @ r.T
        out.append(np.sqrt(np.clip(d2.min(1), 0, None)))
    return np.concatenate(out)

memorization_ratio 

memorization_ratio(fake: ArrayLike, train: ArrayLike, real_holdout: ArrayLike) -> Dict[str, float]

Median nearest-neighbour distance of generated ECGs to the training set, divided by that of held-out real ECGs. Values well below 1 suggest copying.

Source code in src/ecgen/evaluation/metrics.py 
201
202
203
204
205
206
207
208
209
210
211
212
213
214
def memorization_ratio(
    fake: ArrayLike, train: ArrayLike, real_holdout: ArrayLike
) -> Dict[str, float]:
    """
    Median nearest-neighbour distance of generated ECGs to the training set,
    divided by that of held-out real ECGs. Values well below 1 suggest copying.
    """
    d_fake = nearest_neighbour_distances(fake, train)
    d_real = nearest_neighbour_distances(real_holdout, train)
    return {
        "nn_fake_median": float(np.median(d_fake)),
        "nn_real_median": float(np.median(d_real)),
        "ratio": float(np.median(d_fake) / (np.median(d_real) + 1e-12)),
    }

A small 1D-ResNet ECG classifier for condition adherence ("does a classifier trained on real ECGs recognise the requested labels in generated ones?") and for TSTR / TRTR (train on synthetic / real, test on real).

Train it on data separate from the generator's training split, otherwise adherence scores are optimistic.

ResNet1dClassifier

Bases: Module

[B, C, T] -> [B, num_classes] logits; :meth:features gives the pooled embedding.

Source code in src/ecgen/evaluation/classifier.py 
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
class ResNet1dClassifier(nn.Module):
    """``[B, C, T] -> [B, num_classes]`` logits; :meth:`features` gives the pooled embedding."""

    def __init__(
        self,
        in_channels: int = 12,
        num_classes: int = 5,
        width: int = 64,
        depth: int = 4,
    ) -> None:
        super().__init__()
        self.stem = nn.Sequential(
            nn.Conv1d(in_channels, width, 15, 2, 7, bias=False),
            nn.BatchNorm1d(width),
            nn.ReLU(),
            nn.MaxPool1d(3, 2, 1),
        )
        layers, ch = [], width
        for i in range(depth):
            out = width * min(2**i, 4)
            layers.append(_Block(ch, out, 1 if i == 0 else 2))
            ch = out
        self.layers = nn.Sequential(*layers)
        self.head = nn.Linear(ch, num_classes)
        self.embed_dim = ch

    def features(self, x: Tensor) -> Tensor:
        return self.layers(self.stem(x)).mean(-1)

    def forward(self, x: Tensor) -> Tensor:
        return self.head(self.features(x))

fit_classifier 

fit_classifier(x: ArrayLike, y: ArrayLike, multilabel: bool = True, epochs: int = 20, lr: float = 0.001, batch_size: int = 64, width: int = 64, device: Optional[str] = None, seed: int = 0) -> ResNet1dClassifier

Train :class:ResNet1dClassifier on x [N, C, T] and labels y (multi-hot [N, K] with multilabel=True, else class indices [N]).

Source code in src/ecgen/evaluation/classifier.py 
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
def fit_classifier(
    x: ArrayLike,
    y: ArrayLike,
    multilabel: bool = True,
    epochs: int = 20,
    lr: float = 1e-3,
    batch_size: int = 64,
    width: int = 64,
    device: Optional[str] = None,
    seed: int = 0,
) -> ResNet1dClassifier:
    """
    Train :class:`ResNet1dClassifier` on ``x [N, C, T]`` and labels ``y``
    (multi-hot ``[N, K]`` with ``multilabel=True``, else class indices ``[N]``).
    """
    device = device or ("cuda" if torch.cuda.is_available() else "cpu")
    torch.manual_seed(seed)
    xt = torch.as_tensor(_as_numpy(x), dtype=torch.float32)
    yt = torch.as_tensor(np.asarray(_as_numpy(y)))
    num_classes = yt.shape[1] if multilabel else int(yt.max()) + 1
    model = ResNet1dClassifier(xt.shape[1], num_classes, width).to(device)
    opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-4)
    sched = torch.optim.lr_scheduler.OneCycleLR(
        opt, lr, total_steps=epochs * ((len(xt) + batch_size - 1) // batch_size)
    )
    model.train()
    for _ in range(epochs):
        perm = torch.randperm(len(xt))
        for i in range(0, len(xt), batch_size):
            idx = perm[i : i + batch_size]
            logits = model(xt[idx].to(device))
            target = yt[idx].to(device)
            if multilabel:
                loss = F.binary_cross_entropy_with_logits(logits, target.float())
            else:
                loss = F.cross_entropy(logits, target.long())
            opt.zero_grad()
            loss.backward()
            opt.step()
            sched.step()
    return model.eval()

predict 

predict(model: Module, x: ArrayLike, batch_size: int = 128, multilabel: bool = True) -> np.ndarray

Probabilities: sigmoid (multi-label) or softmax (single-label).

Source code in src/ecgen/evaluation/classifier.py 
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
@torch.no_grad()
def predict(
    model: nn.Module, x: ArrayLike, batch_size: int = 128, multilabel: bool = True
) -> np.ndarray:
    """Probabilities: sigmoid (multi-label) or softmax (single-label)."""
    device = next(model.parameters()).device
    xt = torch.as_tensor(_as_numpy(x), dtype=torch.float32)
    out = []
    model.eval()
    for i in range(0, len(xt), batch_size):
        logits = model(xt[i : i + batch_size].to(device)).float()
        out.append(
            (torch.sigmoid(logits) if multilabel else torch.softmax(logits, -1))
            .cpu()
            .numpy()
        )
    return np.concatenate(out)

label_metrics 

label_metrics(probs: ndarray, y: ndarray, multilabel: bool = True) -> Dict[str, float]

Macro AUROC over labels present with both values (multi-label), or accuracy + macro AUROC.

Source code in src/ecgen/evaluation/classifier.py 
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
def label_metrics(
    probs: np.ndarray, y: np.ndarray, multilabel: bool = True
) -> Dict[str, float]:
    """Macro AUROC over labels present with both values (multi-label), or accuracy + macro AUROC."""
    from sklearn.metrics import roc_auc_score

    y = np.asarray(y)
    if multilabel:
        aucs = [
            roc_auc_score(y[:, k], probs[:, k])
            for k in range(y.shape[1])
            if 0 < y[:, k].sum() < len(y)
        ]
        return {
            "macro_auroc": float(np.mean(aucs)) if aucs else float("nan"),
            "n_labels_scored": len(aucs),
        }
    acc = float((probs.argmax(1) == y).mean())
    try:
        auc = float(roc_auc_score(y, probs, multi_class="ovr"))
    except ValueError:
        auc = float("nan")
    return {"accuracy": acc, "macro_auroc": auc}

label_adherence 

label_adherence(model: Module, generated: ArrayLike, requested: ArrayLike, multilabel: bool = True) -> Dict[str, float]

Score generated ECGs against the labels they were generated for.

Source code in src/ecgen/evaluation/classifier.py 
164
165
166
167
168
169
170
171
172
173
174
175
def label_adherence(
    model: nn.Module,
    generated: ArrayLike,
    requested: ArrayLike,
    multilabel: bool = True,
) -> Dict[str, float]:
    """Score generated ECGs against the labels they were generated for."""
    return label_metrics(
        predict(model, generated, multilabel=multilabel),
        _as_numpy(requested),
        multilabel,
    )

tstr 

tstr(train_x: ArrayLike, train_y: ArrayLike, test_x: ArrayLike, test_y: ArrayLike, multilabel: bool = True, **fit_kwargs) -> Dict[str, float]

Train on train_* (synthetic for TSTR, real for TRTR), test on real test_*.

Source code in src/ecgen/evaluation/classifier.py 
178
179
180
181
182
183
184
185
186
187
188
189
190
def tstr(
    train_x: ArrayLike,
    train_y: ArrayLike,
    test_x: ArrayLike,
    test_y: ArrayLike,
    multilabel: bool = True,
    **fit_kwargs,
) -> Dict[str, float]:
    """Train on ``train_*`` (synthetic for TSTR, real for TRTR), test on real ``test_*``."""
    model = fit_classifier(train_x, train_y, multilabel=multilabel, **fit_kwargs)
    return label_metrics(
        predict(model, test_x, multilabel=multilabel), _as_numpy(test_y), multilabel
    )

ecgen.api

Public sampling API: ecgen.sample(checkpoint, n, ...) for any checkpoint trained with ecgen train.

load_model 

load_model(checkpoint: PathLike, family: Optional[str] = None, device: Optional[str] = None)

Load a generative model from an ecgen train checkpoint.

The family is read from the checkpoint (ecgen_family); pass family for checkpoints written before it was recorded. hf://<repo>/<path> checkpoints are downloaded (see :func:ecgen.generate).

Source code in src/ecgen/api.py 
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
def load_model(
    checkpoint: PathLike, family: Optional[str] = None, device: Optional[str] = None
):
    """
    Load a generative model from an ``ecgen train`` checkpoint.

    The family is read from the checkpoint (``ecgen_family``); pass ``family``
    for checkpoints written before it was recorded. ``hf://<repo>/<path>``
    checkpoints are downloaded (see :func:`ecgen.generate`).
    """
    load_builtins()
    from ecgen.generate import _resolve_model_path

    path = _resolve_model_path(checkpoint)
    device = device or ("cuda" if torch.cuda.is_available() else "cpu")
    if family is None:
        meta = torch.load(path, map_location="cpu", weights_only=False)
        family = meta.get("ecgen_family")
        if family is None:
            raise ValueError(
                f"{path} does not record its family; pass family= (one of {', '.join(FAMILIES)})"
            )
    model = FAMILIES.get(family).load_from_checkpoint(str(path), map_location=device)
    return model.to(device).eval()

sample 

sample(checkpoint: Union[PathLike, Any], n: int = 16, *, family: Optional[str] = None, cond: Optional[Mapping[str, Any]] = None, obs: Optional[Any] = None, obs_mask: Optional[Any] = None, guidance: Optional[float] = None, steps: Optional[int] = None, seed: Optional[int] = None, batch_size: int = 32, device: Optional[str] = None, derive_12: bool = False, output_dir: Optional[PathLike] = None, format: Optional[str] = 'npy', plot: int = 0, **sample_kwargs: Any) -> np.ndarray

Generate n ECGs from a checkpoint (or an already loaded model).

Parameters:

Name Type Description Default
checkpoint Union[PathLike, Any]

.ckpt path, hf:// path, or a loaded GenerativeModule.

 required
n int

Number of ECGs.

 16
family Optional[str]

Registry name, for checkpoints that do not record it.

 None
cond Optional[Mapping[str, Any]]

Raw condition values (numeric fields in raw units; NaN = unknown; strings for text fields). A value with a leading dimension of n (or a list of n strings) gives one condition per sample.

 None
obs Optional[Any]

Partially observed ECGs [n, C, T].

 None
obs_mask Optional[Any]

Masks of obs (1 = observed).

 None
guidance Optional[float]

Classifier-free guidance scale (conditional diffusion / flow / prior).

 None
steps Optional[int]

Sampling steps (diffusion / flow).

 None
seed Optional[int]

Seed for reproducible sampling.

 None
batch_size int

ECGs generated per forward pass.

 32
device Optional[str]

Default: CUDA if available.

 None
derive_12 bool

Complete 8 independent leads (I, II, V1–V6) to the standard 12.

 False
output_dir Optional[PathLike]

Write the samples here (format: npy, csv, wfdb, or None to skip) plus samples.json metadata.

 None
plot int

Render the first plot samples as interactive HTML (ecgen.plot_ecg).

 0
sample_kwargs Any

Passed to the model's sample (e.g. eta, solver, inpaint, temperature, seq_length).

 {}

Returns:

Type Description
ndarray

[n, C, T] float32 array.
Source code in src/ecgen/api.py 
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
def sample(
    checkpoint: Union[PathLike, Any],
    n: int = 16,
    *,
    family: Optional[str] = None,
    cond: Optional[Mapping[str, Any]] = None,
    obs: Optional[Any] = None,
    obs_mask: Optional[Any] = None,
    guidance: Optional[float] = None,
    steps: Optional[int] = None,
    seed: Optional[int] = None,
    batch_size: int = 32,
    device: Optional[str] = None,
    derive_12: bool = False,
    output_dir: Optional[PathLike] = None,
    format: Optional[str] = "npy",
    plot: int = 0,
    **sample_kwargs: Any,
) -> np.ndarray:
    """
    Generate ``n`` ECGs from a checkpoint (or an already loaded model).

    Args:
        checkpoint: ``.ckpt`` path, ``hf://`` path, or a loaded ``GenerativeModule``.
        n: Number of ECGs.
        family: Registry name, for checkpoints that do not record it.
        cond: Raw condition values (numeric fields in raw units; NaN = unknown;
            strings for text fields). A value with a leading dimension of ``n``
            (or a list of ``n`` strings) gives one condition per sample.
        obs: Partially observed ECGs ``[n, C, T]``.
        obs_mask: Masks of ``obs`` (1 = observed).
        guidance: Classifier-free guidance scale (conditional diffusion / flow / prior).
        steps: Sampling steps (diffusion / flow).
        seed: Seed for reproducible sampling.
        batch_size: ECGs generated per forward pass.
        device: Default: CUDA if available.
        derive_12: Complete 8 independent leads (I, II, V1–V6) to the standard 12.
        output_dir: Write the samples here (``format``: ``npy``, ``csv``, ``wfdb``, or
            ``None`` to skip) plus ``samples.json`` metadata.
        plot: Render the first ``plot`` samples as interactive HTML (``ecgen.plot_ecg``).
        sample_kwargs: Passed to the model's ``sample`` (e.g. ``eta``, ``solver``,
            ``inpaint``, ``temperature``, ``seq_length``).

    Returns:
        ``[n, C, T]`` float32 array.
    """
    model = (
        checkpoint
        if hasattr(checkpoint, "sample") and not isinstance(checkpoint, (str, Path))
        else load_model(checkpoint, family, device)
    )
    info: Dict[str, Any] = dict(getattr(model, "data_info", {}) or {})
    params = inspect.signature(model.sample).parameters
    extra: Dict[str, Any] = dict(sample_kwargs)
    if guidance is not None:
        extra["guidance"] = guidance
    if steps is not None:
        extra["steps"] = steps
    if "seq_length" in params and "seq_length" not in extra and info.get("seq_length"):
        extra["seq_length"] = info["seq_length"]
    if seed is not None:
        torch.manual_seed(seed)

    out: List[np.ndarray] = []
    for start in range(0, n, batch_size):
        stop = min(n, start + batch_size)
        kw = dict(extra)
        if cond:
            kw["cond"] = {k: _slice(v, start, stop, n) for k, v in cond.items()}
        if obs is not None:
            kw["obs"] = torch.as_tensor(
                np.asarray(obs) if not isinstance(obs, torch.Tensor) else obs
            )[start:stop]
            if obs_mask is not None:
                m = (
                    obs_mask
                    if isinstance(obs_mask, torch.Tensor)
                    else torch.as_tensor(np.asarray(obs_mask))
                )
                kw["obs_mask"] = m[start:stop]
        if seed is not None and "seed" in params:
            kw["seed"] = seed + start
        with torch.no_grad():
            x = model.sample(stop - start, **kw)
        out.append(x.detach().float().cpu().numpy())
    samples = np.concatenate(out).astype(np.float32)

    lead_names = list(
        info.get("lead_names") or [f"lead{i}" for i in range(samples.shape[1])]
    )
    if derive_12:
        from ecgen.data.leads import STANDARD_12, derive_standard12

        samples = derive_standard12(samples, lead_names).astype(np.float32)
        lead_names = list(STANDARD_12)

    if output_dir is not None:
        write_samples(
            samples,
            output_dir,
            lead_names,
            int(info.get("fs") or 500),
            format,
            scale=info.get("scale"),
            cond=cond,
            checkpoint=checkpoint,
            plot=plot,
        )
    return samples

write_samples 

write_samples(samples: ndarray, output_dir: PathLike, lead_names: List[str], fs: int, format: Optional[str] = 'npy', scale: Any = None, cond: Optional[Mapping[str, Any]] = None, checkpoint: Any = None, plot: int = 0) -> List[Path]

Write [n, C, T] samples as npy / csv / wfdb (+ samples.json, plots).

Source code in src/ecgen/api.py 
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
def write_samples(
    samples: np.ndarray,
    output_dir: PathLike,
    lead_names: List[str],
    fs: int,
    format: Optional[str] = "npy",
    scale: Any = None,
    cond: Optional[Mapping[str, Any]] = None,
    checkpoint: Any = None,
    plot: int = 0,
) -> List[Path]:
    """Write ``[n, C, T]`` samples as ``npy`` / ``csv`` / ``wfdb`` (+ ``samples.json``, plots)."""
    out = Path(output_dir)
    out.mkdir(parents=True, exist_ok=True)
    created: List[Path] = []
    digits = len(str(len(samples)))
    units = "mV" if scale == "none" else "NU"  # z-scored data has no physical unit
    if format == "npy":
        created.append(out / "samples.npy")
        np.save(created[-1], samples)
    elif format == "csv":
        for i, ecg in enumerate(samples):
            path = out / f"sample_{i + 1:0{digits}d}.csv"
            np.savetxt(
                path,
                ecg.T,
                delimiter=",",
                header=",".join(lead_names),
                comments="",
                fmt="%.6f",
            )
            created.append(path)
    elif format == "wfdb":
        import wfdb

        for i, ecg in enumerate(samples):
            name = f"sample_{i + 1:0{digits}d}"
            wfdb.wrsamp(
                name,
                fs=fs,
                units=[units] * len(lead_names),
                sig_name=list(lead_names),
                p_signal=ecg.T.astype(np.float64),
                fmt=["16"] * len(lead_names),
                write_dir=str(out),
            )
            created.append(out / f"{name}.hea")
    elif format is not None:
        raise ValueError(f"format must be npy, csv, wfdb or None; got {format!r}")

    if plot:
        from ecgen.plot import plot_ecg

        for i in range(min(plot, len(samples))):
            path = out / f"sample_{i + 1:0{digits}d}.html"
            plot_ecg(
                samples[i],
                sample_rate=fs,
                lead_names=lead_names,
                title=f"Generated ECG {i + 1}",
                output_path=path,
            )
            created.append(path)

    meta = {
        "n": int(len(samples)),
        "shape": list(samples.shape),
        "lead_names": list(lead_names),
        "fs": fs,
        "units": units,
        "format": format,
        "checkpoint": str(checkpoint) if isinstance(checkpoint, (str, Path)) else None,
        "cond": {
            k: (v.tolist() if hasattr(v, "tolist") else v)
            for k, v in (cond or {}).items()
        },
    }
    (out / "samples.json").write_text(json.dumps(meta, indent=2, default=str))
    return created

evaluation.metrics.signal_quality

Signal quality metrics for ECG evaluation.

signal_to_noise_ratio 

signal_to_noise_ratio(ecg: Tensor | ndarray) -> float

Calculate signal-to-noise ratio.

Parameters:

Name Type Description Default
ecg Tensor | ndarray

ECG signal

 required

Returns:

Type Description
float

SNR in dB
Source code in evaluation/metrics/signal_quality.py 
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
def signal_to_noise_ratio(ecg: Tensor | np.ndarray) -> float:
    """
    Calculate signal-to-noise ratio.

    Args:
        ecg: ECG signal

    Returns:
        SNR in dB
    """
    if isinstance(ecg, Tensor):
        ecg = ecg.cpu().numpy()

    signal_power = np.mean(ecg ** 2)
    noise_estimate = np.std(np.diff(ecg))
    noise_power = noise_estimate ** 2

    snr = 10 * np.log10(signal_power / (noise_power + 1e-8))

    return float(snr)

baseline_wander 

baseline_wander(ecg: Tensor | ndarray, sampling_rate: int = 500) -> float

Measure baseline wander in ECG signal.

Parameters:

Name Type Description Default
ecg Tensor | ndarray

ECG signal

 required
sampling_rate int

Sampling rate in Hz

 500

Returns:

Type Description
float

Baseline wander metric
Source code in evaluation/metrics/signal_quality.py 
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
def baseline_wander(ecg: Tensor | np.ndarray, sampling_rate: int = 500) -> float:
    """
    Measure baseline wander in ECG signal.

    Args:
        ecg: ECG signal
        sampling_rate: Sampling rate in Hz

    Returns:
        Baseline wander metric
    """
    if isinstance(ecg, Tensor):
        ecg = ecg.cpu().numpy()

    # High-pass filter to remove baseline
    nyquist = sampling_rate / 2
    cutoff = 0.5 / nyquist
    b, a = signal.butter(4, cutoff, btype='high')
    filtered = signal.filtfilt(b, a, ecg)

    # Baseline is the difference
    baseline = ecg - filtered
    wander = np.std(baseline)

    return float(wander)

powerline_interference 

powerline_interference(ecg: Tensor | ndarray, sampling_rate: int = 500) -> float

Measure 50/60 Hz powerline interference.

Parameters:

Name Type Description Default
ecg Tensor | ndarray

ECG signal

 required
sampling_rate int

Sampling rate in Hz

 500

Returns:

Type Description
float

Powerline interference metric
Source code in evaluation/metrics/signal_quality.py 
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
def powerline_interference(ecg: Tensor | np.ndarray, sampling_rate: int = 500) -> float:
    """
    Measure 50/60 Hz powerline interference.

    Args:
        ecg: ECG signal
        sampling_rate: Sampling rate in Hz

    Returns:
        Powerline interference metric
    """
    if isinstance(ecg, Tensor):
        ecg = ecg.cpu().numpy()

    # FFT
    fft = np.fft.fft(ecg)
    freqs = np.fft.fftfreq(len(ecg), 1/sampling_rate)
    power = np.abs(fft) ** 2

    # Power at 50 Hz and 60 Hz
    idx_50 = np.argmin(np.abs(freqs - 50))
    idx_60 = np.argmin(np.abs(freqs - 60))

    powerline_power = power[idx_50] + power[idx_60]
    total_power = np.sum(power)

    interference = powerline_power / (total_power + 1e-8)

    return float(interference)

calculate_quality_metrics 

calculate_quality_metrics(ecg: Tensor | ndarray, sampling_rate: int = 500) -> dict

Calculate comprehensive signal quality metrics.

Parameters:

Name Type Description Default
ecg Tensor | ndarray

ECG signal

 required
sampling_rate int

Sampling rate in Hz

 500

Returns:

Type Description
dict

Dictionary of quality metrics
Source code in evaluation/metrics/signal_quality.py 
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
def calculate_quality_metrics(ecg: Tensor | np.ndarray, sampling_rate: int = 500) -> dict:
    """
    Calculate comprehensive signal quality metrics.

    Args:
        ecg: ECG signal
        sampling_rate: Sampling rate in Hz

    Returns:
        Dictionary of quality metrics
    """
    return {
        "snr": signal_to_noise_ratio(ecg),
        "baseline_wander": baseline_wander(ecg, sampling_rate),
        "powerline_interference": powerline_interference(ecg, sampling_rate),
    }

evaluation.metrics.diversity

Diversity metrics for generated ECG samples.

pairwise_distance_diversity 

pairwise_distance_diversity(samples: Tensor | ndarray, metric: str = 'euclidean') -> float

Calculate diversity using pairwise distances.

Parameters:

Name Type Description Default
samples Tensor | ndarray

Generated samples of shape (n_samples, ...)

 required
metric str

Distance metric ("euclidean", "cosine", etc.)

 'euclidean'

Returns:

Type Description
float

Mean pairwise distance
Source code in evaluation/metrics/diversity.py 
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
def pairwise_distance_diversity(samples: Tensor | np.ndarray, metric: str = "euclidean") -> float:
    """
    Calculate diversity using pairwise distances.

    Args:
        samples: Generated samples of shape (n_samples, ...)
        metric: Distance metric ("euclidean", "cosine", etc.)

    Returns:
        Mean pairwise distance
    """
    if isinstance(samples, Tensor):
        samples = samples.cpu().numpy()

    # Flatten samples
    n_samples = samples.shape[0]
    samples_flat = samples.reshape(n_samples, -1)

    # Calculate pairwise distances
    distances = pdist(samples_flat, metric=metric)

    return float(np.mean(distances))

intra_class_diversity 

intra_class_diversity(samples: Tensor | ndarray, labels: ndarray) -> dict

Calculate diversity within each class.

Parameters:

Name Type Description Default
samples Tensor | ndarray

Generated samples

 required
labels ndarray

Class labels for each sample

 required

Returns:

Type Description
dict

Dictionary of diversity per class
Source code in evaluation/metrics/diversity.py 
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
def intra_class_diversity(samples: Tensor | np.ndarray, labels: np.ndarray) -> dict:
    """
    Calculate diversity within each class.

    Args:
        samples: Generated samples
        labels: Class labels for each sample

    Returns:
        Dictionary of diversity per class
    """
    if isinstance(samples, Tensor):
        samples = samples.cpu().numpy()

    unique_labels = np.unique(labels)
    diversity_per_class = {}

    for label in unique_labels:
        class_samples = samples[labels == label]
        if len(class_samples) > 1:
            diversity = pairwise_distance_diversity(class_samples)
            diversity_per_class[str(label)] = diversity

    return diversity_per_class

mode_coverage 

mode_coverage(samples: Tensor | ndarray, real_samples: Tensor | ndarray, n_clusters: int = 10) -> float

Estimate mode coverage using clustering.

Parameters:

Name Type Description Default
samples Tensor | ndarray

Generated samples

 required
real_samples Tensor | ndarray

Real samples

 required
n_clusters int

Number of clusters

 10

Returns:

Type Description
float

Mode coverage score (0-1)
Source code in evaluation/metrics/diversity.py 
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
def mode_coverage(samples: Tensor | np.ndarray, real_samples: Tensor | np.ndarray, 
                  n_clusters: int = 10) -> float:
    """
    Estimate mode coverage using clustering.

    Args:
        samples: Generated samples
        real_samples: Real samples
        n_clusters: Number of clusters

    Returns:
        Mode coverage score (0-1)
    """
    if isinstance(samples, Tensor):
        samples = samples.cpu().numpy()
    if isinstance(real_samples, Tensor):
        real_samples = real_samples.cpu().numpy()

    from sklearn.cluster import KMeans

    # Flatten
    samples_flat = samples.reshape(len(samples), -1)
    real_flat = real_samples.reshape(len(real_samples), -1)

    # Cluster real samples
    kmeans = KMeans(n_clusters=n_clusters, random_state=42)
    kmeans.fit(real_flat)

    # Check which clusters are covered by generated samples
    generated_labels = kmeans.predict(samples_flat)
    covered_clusters = len(np.unique(generated_labels))

    coverage = covered_clusters / n_clusters

    return float(coverage)

evaluation.metrics.fidelity

Fidelity metrics for generated ECG samples.

frechet_distance 

frechet_distance(mu1: ndarray, sigma1: ndarray, mu2: ndarray, sigma2: ndarray, eps: float = 1e-06) -> float

Calculate Fréchet Distance between two Gaussian distributions.

Parameters:

Name Type Description Default
mu1 ndarray

Mean of first distribution

 required
sigma1 ndarray

Covariance of first distribution

 required
mu2 ndarray

Mean of second distribution

 required
sigma2 ndarray

Covariance of second distribution

 required
eps float

Small value for numerical stability

 1e-06

Returns:

Type Description
float

Fréchet distance
Source code in evaluation/metrics/fidelity.py 
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
def frechet_distance(mu1: np.ndarray, sigma1: np.ndarray, 
                    mu2: np.ndarray, sigma2: np.ndarray, eps: float = 1e-6) -> float:
    """
    Calculate Fréchet Distance between two Gaussian distributions.

    Args:
        mu1: Mean of first distribution
        sigma1: Covariance of first distribution
        mu2: Mean of second distribution
        sigma2: Covariance of second distribution
        eps: Small value for numerical stability

    Returns:
        Fréchet distance
    """
    mu1 = np.atleast_1d(mu1)
    mu2 = np.atleast_1d(mu2)

    sigma1 = np.atleast_2d(sigma1)
    sigma2 = np.atleast_2d(sigma2)

    diff = mu1 - mu2

    # Product might be almost singular
    covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False)
    if not np.isfinite(covmean).all():
        offset = np.eye(sigma1.shape[0]) * eps
        covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset))

    # Numerical error might give slight imaginary component
    if np.iscomplexobj(covmean):
        if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3):
            m = np.max(np.abs(covmean.imag))
            raise ValueError(f"Imaginary component {m}")
        covmean = covmean.real

    tr_covmean = np.trace(covmean)

    return float(diff.dot(diff) + np.trace(sigma1) + np.trace(sigma2) - 2 * tr_covmean)

calculate_fid 

calculate_fid(real_features: Tensor | ndarray, generated_features: Tensor | ndarray) -> float

Calculate Fréchet Inception Distance (FID) between real and generated features.

Parameters:

Name Type Description Default
real_features Tensor | ndarray

Features from real samples

 required
generated_features Tensor | ndarray

Features from generated samples

 required

Returns:

Type Description
float

FID score
Source code in evaluation/metrics/fidelity.py 
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
def calculate_fid(real_features: Tensor | np.ndarray, 
                 generated_features: Tensor | np.ndarray) -> float:
    """
    Calculate Fréchet Inception Distance (FID) between real and generated features.

    Args:
        real_features: Features from real samples
        generated_features: Features from generated samples

    Returns:
        FID score
    """
    if isinstance(real_features, Tensor):
        real_features = real_features.cpu().numpy()
    if isinstance(generated_features, Tensor):
        generated_features = generated_features.cpu().numpy()

    # Calculate statistics
    mu_real = np.mean(real_features, axis=0)
    sigma_real = np.cov(real_features, rowvar=False)

    mu_gen = np.mean(generated_features, axis=0)
    sigma_gen = np.cov(generated_features, rowvar=False)

    # Calculate FID
    fid = frechet_distance(mu_real, sigma_real, mu_gen, sigma_gen)

    return fid

reconstruction_error 

reconstruction_error(real: Tensor | ndarray, reconstructed: Tensor | ndarray, metric: str = 'mse') -> float

Calculate reconstruction error.

Parameters:

Name Type Description Default
real Tensor | ndarray

Real samples

 required
reconstructed Tensor | ndarray

Reconstructed samples

 required
metric str

Error metric ("mse", "mae", "rmse")

 'mse'

Returns:

Type Description
float

Reconstruction error
Source code in evaluation/metrics/fidelity.py 
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
def reconstruction_error(real: Tensor | np.ndarray, 
                        reconstructed: Tensor | np.ndarray,
                        metric: str = "mse") -> float:
    """
    Calculate reconstruction error.

    Args:
        real: Real samples
        reconstructed: Reconstructed samples
        metric: Error metric ("mse", "mae", "rmse")

    Returns:
        Reconstruction error
    """
    if isinstance(real, Tensor):
        real = real.cpu().numpy()
    if isinstance(reconstructed, Tensor):
        reconstructed = reconstructed.cpu().numpy()

    if metric == "mse":
        error = np.mean((real - reconstructed) ** 2)
    elif metric == "mae":
        error = np.mean(np.abs(real - reconstructed))
    elif metric == "rmse":
        error = np.sqrt(np.mean((real - reconstructed) ** 2))
    else:
        raise ValueError(f"Unknown metric: {metric}")

    return float(error)

structural_similarity 

structural_similarity(real: Tensor | ndarray, generated: Tensor | ndarray) -> float

Calculate structural similarity between real and generated samples.

Parameters:

Name Type Description Default
real Tensor | ndarray

Real samples

 required
generated Tensor | ndarray

Generated samples

 required

Returns:

Type Description
float

Structural similarity score
Source code in evaluation/metrics/fidelity.py 
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
def structural_similarity(real: Tensor | np.ndarray, 
                         generated: Tensor | np.ndarray) -> float:
    """
    Calculate structural similarity between real and generated samples.

    Args:
        real: Real samples
        generated: Generated samples

    Returns:
        Structural similarity score
    """
    if isinstance(real, Tensor):
        real = real.cpu().numpy()
    if isinstance(generated, Tensor):
        generated = generated.cpu().numpy()

    # Normalize
    real_norm = (real - real.mean()) / (real.std() + 1e-8)
    gen_norm = (generated - generated.mean()) / (generated.std() + 1e-8)

    # Cross-correlation
    correlation = np.mean(real_norm * gen_norm)

    return float(correlation)

evaluation.benchmark

Benchmarking utilities for model evaluation.

ModelBenchmark

Benchmark model performance.

Source code in evaluation/benchmark.py 
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
class ModelBenchmark:
    """Benchmark model performance."""

    def __init__(self, model: Any, device: str = "cuda"):
        self.model = model
        self.device = device
        self.results = {}

    def measure_inference_time(self, input_shape: tuple, n_runs: int = 100, warmup: int = 10):
        """
        Measure inference time.

        Args:
            input_shape: Input tensor shape
            n_runs: Number of runs for averaging
            warmup: Number of warmup runs
        """
        self.model.eval()
        dummy_input = torch.randn(input_shape).to(self.device)

        # Warmup
        with torch.no_grad():
            for _ in range(warmup):
                _ = self.model(dummy_input)

        # Measure
        if self.device == "cuda":
            torch.cuda.synchronize()

        start_time = time.time()
        with torch.no_grad():
            for _ in range(n_runs):
                _ = self.model(dummy_input)

        if self.device == "cuda":
            torch.cuda.synchronize()

        end_time = time.time()
        avg_time = (end_time - start_time) / n_runs

        self.results["inference_time_ms"] = avg_time * 1000
        self.results["throughput_samples_per_sec"] = input_shape[0] / avg_time

    def measure_memory_usage(self, input_shape: tuple):
        """Measure GPU memory usage."""
        if self.device != "cuda":
            return

        torch.cuda.reset_peak_memory_stats()
        dummy_input = torch.randn(input_shape).to(self.device)

        with torch.no_grad():
            _ = self.model(dummy_input)

        memory_allocated = torch.cuda.max_memory_allocated() / 1024**2  # MB
        memory_reserved = torch.cuda.max_memory_reserved() / 1024**2  # MB

        self.results["memory_allocated_mb"] = memory_allocated
        self.results["memory_reserved_mb"] = memory_reserved

    def count_parameters(self):
        """Count model parameters."""
        total_params = sum(p.numel() for p in self.model.parameters())
        trainable_params = sum(p.numel() for p in self.model.parameters() if p.requires_grad)

        self.results["total_parameters"] = total_params
        self.results["trainable_parameters"] = trainable_params
        self.results["total_parameters_millions"] = total_params / 1e6

    def run_full_benchmark(self, input_shape: tuple, save_path: Path = None):
        """Run full benchmark suite."""
        print("Running benchmark...")

        print("  - Counting parameters...")
        self.count_parameters()

        print("  - Measuring inference time...")
        self.measure_inference_time(input_shape)

        if self.device == "cuda":
            print("  - Measuring memory usage...")
            self.measure_memory_usage(input_shape)

        print("Benchmark complete!")

        if save_path:
            save_path = Path(save_path)
            save_path.parent.mkdir(parents=True, exist_ok=True)
            with open(save_path, 'w') as f:
                json.dump(self.results, f, indent=2)
            print(f"Results saved to {save_path}")

        return self.results

    def print_results(self):
        """Print benchmark results."""
        print("\n" + "="*50)
        print("BENCHMARK RESULTS")
        print("="*50)
        for key, value in self.results.items():
            if isinstance(value, float):
                print(f"{key}: {value:.4f}")
            else:
                print(f"{key}: {value}")
        print("="*50 + "\n")

measure_inference_time 

measure_inference_time(input_shape: tuple, n_runs: int = 100, warmup: int = 10)

Measure inference time.

Parameters:

Name Type Description Default
input_shape tuple

Input tensor shape

 required
n_runs int

Number of runs for averaging

 100
warmup int

Number of warmup runs

 10
Source code in evaluation/benchmark.py 
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
def measure_inference_time(self, input_shape: tuple, n_runs: int = 100, warmup: int = 10):
    """
    Measure inference time.

    Args:
        input_shape: Input tensor shape
        n_runs: Number of runs for averaging
        warmup: Number of warmup runs
    """
    self.model.eval()
    dummy_input = torch.randn(input_shape).to(self.device)

    # Warmup
    with torch.no_grad():
        for _ in range(warmup):
            _ = self.model(dummy_input)

    # Measure
    if self.device == "cuda":
        torch.cuda.synchronize()

    start_time = time.time()
    with torch.no_grad():
        for _ in range(n_runs):
            _ = self.model(dummy_input)

    if self.device == "cuda":
        torch.cuda.synchronize()

    end_time = time.time()
    avg_time = (end_time - start_time) / n_runs

    self.results["inference_time_ms"] = avg_time * 1000
    self.results["throughput_samples_per_sec"] = input_shape[0] / avg_time

measure_memory_usage 

measure_memory_usage(input_shape: tuple)

Measure GPU memory usage.

Source code in evaluation/benchmark.py 
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
def measure_memory_usage(self, input_shape: tuple):
    """Measure GPU memory usage."""
    if self.device != "cuda":
        return

    torch.cuda.reset_peak_memory_stats()
    dummy_input = torch.randn(input_shape).to(self.device)

    with torch.no_grad():
        _ = self.model(dummy_input)

    memory_allocated = torch.cuda.max_memory_allocated() / 1024**2  # MB
    memory_reserved = torch.cuda.max_memory_reserved() / 1024**2  # MB

    self.results["memory_allocated_mb"] = memory_allocated
    self.results["memory_reserved_mb"] = memory_reserved

count_parameters 

count_parameters()

Count model parameters.

Source code in evaluation/benchmark.py 
70
71
72
73
74
75
76
77
def count_parameters(self):
    """Count model parameters."""
    total_params = sum(p.numel() for p in self.model.parameters())
    trainable_params = sum(p.numel() for p in self.model.parameters() if p.requires_grad)

    self.results["total_parameters"] = total_params
    self.results["trainable_parameters"] = trainable_params
    self.results["total_parameters_millions"] = total_params / 1e6

run_full_benchmark 

run_full_benchmark(input_shape: tuple, save_path: Path = None)

Run full benchmark suite.

Source code in evaluation/benchmark.py 
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
def run_full_benchmark(self, input_shape: tuple, save_path: Path = None):
    """Run full benchmark suite."""
    print("Running benchmark...")

    print("  - Counting parameters...")
    self.count_parameters()

    print("  - Measuring inference time...")
    self.measure_inference_time(input_shape)

    if self.device == "cuda":
        print("  - Measuring memory usage...")
        self.measure_memory_usage(input_shape)

    print("Benchmark complete!")

    if save_path:
        save_path = Path(save_path)
        save_path.parent.mkdir(parents=True, exist_ok=True)
        with open(save_path, 'w') as f:
            json.dump(self.results, f, indent=2)
        print(f"Results saved to {save_path}")

    return self.results

print_results 

print_results()

Print benchmark results.

Source code in evaluation/benchmark.py 
104
105
106
107
108
109
110
111
112
113
114
def print_results(self):
    """Print benchmark results."""
    print("\n" + "="*50)
    print("BENCHMARK RESULTS")
    print("="*50)
    for key, value in self.results.items():
        if isinstance(value, float):
            print(f"{key}: {value:.4f}")
        else:
            print(f"{key}: {value}")
    print("="*50 + "\n")

evaluation.visualize

Evaluation visualization utilities.

plot_evaluation_results 

plot_evaluation_results(metrics: dict, save_path: str | Path = None, figsize: tuple = (12, 8)) -> plt.Figure

Plot evaluation results as bar charts.

Parameters:

Name Type Description Default
metrics dict

Dictionary of metric names and values

 required
save_path str | Path

Optional path to save figure

 None
figsize tuple

Figure size

 (12, 8)

Returns:

Type Description
Figure

matplotlib Figure object
Source code in evaluation/visualize.py 
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
def plot_evaluation_results(
    metrics: dict,
    save_path: str | Path = None,
    figsize: tuple = (12, 8),
) -> plt.Figure:
    """
    Plot evaluation results as bar charts.

    Args:
        metrics: Dictionary of metric names and values
        save_path: Optional path to save figure
        figsize: Figure size

    Returns:
        matplotlib Figure object
    """
    fig, ax = plt.subplots(figsize=figsize)

    metric_names = list(metrics.keys())
    metric_values = list(metrics.values())

    bars = ax.bar(range(len(metric_names)), metric_values, alpha=0.7, edgecolor='black')

    # Color bars based on value (assuming lower is better for most metrics)
    colors = plt.cm.RdYlGn_r(np.linspace(0.2, 0.8, len(bars)))
    for bar, color in zip(bars, colors):
        bar.set_color(color)

    ax.set_xticks(range(len(metric_names)))
    ax.set_xticklabels(metric_names, rotation=45, ha='right')
    ax.set_ylabel("Metric Value")
    ax.set_title("Evaluation Metrics")
    ax.grid(True, alpha=0.3, axis='y')

    # Add value labels on bars
    for i, (bar, value) in enumerate(zip(bars, metric_values)):
        height = bar.get_height()
        ax.text(bar.get_x() + bar.get_width()/2., height,
               f'{value:.4f}',
               ha='center', va='bottom', fontsize=9)

    plt.tight_layout()

    if save_path:
        fig.savefig(save_path, dpi=300, bbox_inches='tight')

    return fig

plot_comparison_matrix 

plot_comparison_matrix(models: list[str], metrics: dict[str, list[float]], save_path: str | Path = None, figsize: tuple = (12, 8)) -> plt.Figure

Plot comparison matrix of multiple models across metrics.

Parameters:

Name Type Description Default
models list[str]

List of model names

 required
metrics dict[str, list[float]]

Dictionary of metric names to lists of values (one per model)

 required
save_path str | Path

Optional path to save figure

 None
figsize tuple

Figure size

 (12, 8)

Returns:

Type Description
Figure

matplotlib Figure object
Source code in evaluation/visualize.py 
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
def plot_comparison_matrix(
    models: list[str],
    metrics: dict[str, list[float]],
    save_path: str | Path = None,
    figsize: tuple = (12, 8),
) -> plt.Figure:
    """
    Plot comparison matrix of multiple models across metrics.

    Args:
        models: List of model names
        metrics: Dictionary of metric names to lists of values (one per model)
        save_path: Optional path to save figure
        figsize: Figure size

    Returns:
        matplotlib Figure object
    """
    metric_names = list(metrics.keys())
    n_models = len(models)
    n_metrics = len(metric_names)

    # Create matrix
    matrix = np.array([metrics[m] for m in metric_names])

    fig, ax = plt.subplots(figsize=figsize)

    im = ax.imshow(matrix, cmap='RdYlGn_r', aspect='auto')

    # Set ticks
    ax.set_xticks(range(n_models))
    ax.set_yticks(range(n_metrics))
    ax.set_xticklabels(models, rotation=45, ha='right')
    ax.set_yticklabels(metric_names)

    # Add colorbar
    cbar = plt.colorbar(im, ax=ax)
    cbar.set_label("Metric Value", rotation=270, labelpad=20)

    # Add text annotations
    for i in range(n_metrics):
        for j in range(n_models):
            text = ax.text(j, i, f'{matrix[i, j]:.3f}',
                          ha="center", va="center", color="black", fontsize=9)

    ax.set_title("Model Comparison Matrix")
    plt.tight_layout()

    if save_path:
        fig.savefig(save_path, dpi=300, bbox_inches='tight')

    return fig