ECG and physiological time-series models

import numpy as np

FS, SECONDS = 100, 10          # 100 samples per second (real ECGs are often 250-500 Hz), 10-second strips

def synthetic_strip(irregular, rng, noise):
    """A crude single-lead ECG-like strip: a spike per beat, baseline wander and noise.
    irregular=True mimics atrial fibrillation; regular strips include occasional ectopic beats."""
    t = np.arange(FS * SECONDS) / FS
    signal = 0.15 * np.sin(2 * np.pi * 0.3 * t) + rng.normal(0, noise, t.size)
    mean_rr, beat = rng.uniform(0.65, 1.0), rng.uniform(0, 0.5)      # 60-90 beats per minute
    while beat < SECONDS:
        signal += rng.uniform(0.6, 1.2) * np.exp(-((t - beat) ** 2) / (2 * 0.012 ** 2))
        if irregular:
            beat += mean_rr * rng.uniform(0.6, 1.4)
        elif rng.random() < 0.08:
            beat += mean_rr * 0.6                                     # occasional ectopic beat
        else:
            beat += mean_rr * rng.normal(1, 0.03)
    return signal

def make_dataset(n, noise, seed=0):
    rng = np.random.default_rng(seed)
    labels = rng.integers(0, 2, n)                                    # 1 = irregular rhythm
    return np.stack([synthetic_strip(bool(l), rng, noise) for l in labels]), labels

signals, labels = make_dataset(2000, noise=0.08)
print(f"{signals.shape[0]} strips of {signals.shape[1]} samples ({SECONDS} s at {FS} Hz); "
      f"{labels.mean():.0%} irregular")

# 2000 strips of 1000 samples (10 s at 100 Hz); 52% irregular

from scipy.signal import find_peaks
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import roc_auc_score
from sklearn.model_selection import train_test_split

def rr_features(signal):
    peaks, _ = find_peaks(signal, height=0.4, distance=int(0.3 * FS))   # detect beats
    rr = np.diff(peaks) / FS
    return [np.std(rr) / np.mean(rr), len(peaks)]                      # RR variability, beat count

def feature_model_auroc(signals, labels):
    feats = np.array([rr_features(s) for s in signals])
    tr, te = train_test_split(np.arange(len(labels)), test_size=0.3, random_state=0, stratify=labels)
    p = LogisticRegression().fit(feats[tr], labels[tr]).predict_proba(feats[te])[:, 1]
    return roc_auc_score(labels[te], p)

print(f"RR-variability features + logistic regression: AUROC {feature_model_auroc(signals, labels):.3f}")

# RR-variability features + logistic regression: AUROC 0.917

import torch
from torch import nn

def cnn_auroc(signals, labels, epochs=15):
    torch.manual_seed(0)
    model = nn.Sequential(
        nn.Conv1d(1, 8, kernel_size=7, padding=3), nn.ReLU(), nn.MaxPool1d(4),
        nn.Conv1d(8, 16, kernel_size=7, padding=3), nn.ReLU(), nn.MaxPool1d(4),
        nn.Conv1d(16, 16, kernel_size=7, padding=3), nn.ReLU(), nn.AdaptiveAvgPool1d(1),
        nn.Flatten(), nn.Linear(16, 2))
    X, y = torch.tensor(signals, dtype=torch.float32).unsqueeze(1), torch.tensor(labels)
    tr, te = train_test_split(np.arange(len(labels)), test_size=0.3, random_state=0, stratify=labels)
    tr = torch.tensor(tr)
    opt = torch.optim.Adam(model.parameters(), lr=3e-3)
    for _ in range(epochs):
        for batch in tr[torch.randperm(len(tr))].split(64):
            opt.zero_grad()
            nn.functional.cross_entropy(model(X[batch]), y[batch]).backward()
            opt.step()
    with torch.no_grad():
        p = model(X[te]).softmax(1)[:, 1].numpy()
    return roc_auc_score(labels[te], p)

print(f"1D CNN on the raw signal:                      AUROC {cnn_auroc(signals, labels):.3f}")

# 1D CNN on the raw signal:                      AUROC 0.945

noisy_signals, noisy_labels = make_dataset(2000, noise=0.25, seed=1)
print(f"Noisy strips, RR features: AUROC {feature_model_auroc(noisy_signals, noisy_labels):.3f}")
print(f"Noisy strips, 1D CNN:      AUROC {cnn_auroc(noisy_signals, noisy_labels):.3f}")

# Noisy strips, RR features: AUROC 0.601
# Noisy strips, 1D CNN:      AUROC 0.812