electrotrace 1.9.0__py3-none-any.whl

electrotrace/__init__.py

@@ -0,0 +1,28 @@
+"""ElectroTrace: reproducible ECG/electrophysiology research and benchmarking tools."""
+
+__version__ = "1.9.0"
+
+from .candidate_suppressor import CandidateSuppressor
+from .io import load_recording
+from .provenance import DatasetManifest, manifest_from_dict
+from .qrs_delineation import delineate_qrs
+from .research_validation import build_validation_report, summarize_records_rigorous, write_validation_report
+from .signal import apply_pipeline
+from .validation import validate_record
+from .validation_detectors import detect_r_peaks, detect_r_peaks_two_stage
+
+__all__ = [
+    "__version__",
+    "CandidateSuppressor",
+    "DatasetManifest",
+    "manifest_from_dict",
+    "load_recording",
+    "apply_pipeline",
+    "detect_r_peaks",
+    "detect_r_peaks_two_stage",
+    "delineate_qrs",
+    "validate_record",
+    "build_validation_report",
+    "summarize_records_rigorous",
+    "write_validation_report",
+]

electrotrace/__main__.py

@@ -0,0 +1,3 @@
+from .cli import main
+
+raise SystemExit(main())

electrotrace/annotations.py

@@ -0,0 +1,194 @@
+"""Annotation model, validation, serialization, review state, and agreement metrics."""
+from __future__ import annotations
+
+import json
+import math
+import uuid
+from dataclasses import asdict, dataclass, field
+from typing import Literal, Optional
+
+AnnotationType = Literal["interval", "point"]
+ReviewStatus = Literal["unreviewed", "accepted", "flagged"]
+DEFAULT_LABELS = ["P Wave", "QRS", "T Wave", "Pacemaker Spike", "R Peak", "Artifact", "Abnormal Beat", "Graft Activation", "Arrhythmia", "Custom"]
+
+
+@dataclass
+class Annotation:
+    label: str
+    type: AnnotationType
+    channel: str
+    start: Optional[float] = None
+    end: Optional[float] = None
+    time: Optional[float] = None
+    confidence: float = 1.0
+    notes: str = ""
+    annotator: str = ""
+    status: ReviewStatus = "unreviewed"
+    reviewer: str = ""
+    review_notes: str = ""
+    id: str = field(default_factory=lambda: uuid.uuid4().hex[:10])
+
+    def validate(self, duration_s: float | None = None, start_time_s: float = 0.0, end_time_s: float | None = None) -> None:
+        if not self.label.strip():
+            raise ValueError("label must not be empty")
+        if not self.channel.strip():
+            raise ValueError("channel must not be empty")
+        if not math.isfinite(float(self.confidence)) or not 0.0 <= self.confidence <= 1.0:
+            raise ValueError("confidence must be between 0 and 1")
+        if self.status not in {"unreviewed", "accepted", "flagged"}:
+            raise ValueError("invalid review status")
+        start_bound = float(start_time_s)
+        end_bound = float(end_time_s) if end_time_s is not None else (start_bound + float(duration_s) if duration_s is not None else None)
+        if not math.isfinite(start_bound) or (end_bound is not None and not math.isfinite(end_bound)) or (end_bound is not None and end_bound < start_bound):
+            raise ValueError("invalid recording time bounds")
+        if self.type == "interval":
+            if self.start is None or self.end is None:
+                raise ValueError("interval annotations require start and end")
+            if not math.isfinite(self.start) or not math.isfinite(self.end):
+                raise ValueError("interval coordinates must be finite")
+            if self.end <= self.start:
+                raise ValueError("end must be greater than start")
+            if self.time is not None:
+                raise ValueError("interval annotations must not define time")
+            if end_bound is not None and (self.start < start_bound or self.end > end_bound):
+                raise ValueError("interval is outside the recording bounds")
+        elif self.type == "point":
+            if self.time is None:
+                raise ValueError("point annotations require time")
+            if not math.isfinite(self.time):
+                raise ValueError("point time must be finite")
+            if self.start is not None or self.end is not None:
+                raise ValueError("point annotations must not define start/end")
+            if end_bound is not None and not start_bound <= self.time <= end_bound:
+                raise ValueError("point is outside the recording bounds")
+        else:
+            raise ValueError(f"unknown annotation type: {self.type}")
+
+    @property
+    def position(self) -> float:
+        return self.start if self.type == "interval" else self.time  # type: ignore[return-value]
+
+    def to_dict(self) -> dict:
+        return asdict(self)
+
+
+class AnnotationStore:
+    def __init__(self, duration_s: float | None = None, start_time_s: float = 0.0, end_time_s: float | None = None):
+        self.duration_s = duration_s
+        self.start_time_s = float(start_time_s)
+        self.end_time_s = float(end_time_s) if end_time_s is not None else (self.start_time_s + float(duration_s) if duration_s is not None else None)
+        self._items: list[Annotation] = []
+
+    @property
+    def items(self) -> list[Annotation]:
+        return list(self._items)
+
+    def add(self, ann: Annotation) -> Annotation:
+        ann.validate(self.duration_s, self.start_time_s, self.end_time_s)
+        if any(a.id == ann.id for a in self._items):
+            raise ValueError(f"duplicate annotation id: {ann.id}")
+        self._items.append(ann)
+        self._sort()
+        return ann
+
+    def update(self, ann_id: str, **changes) -> Annotation:
+        for idx, existing in enumerate(self._items):
+            if existing.id == ann_id:
+                data = {**asdict(existing), **changes}
+                updated = Annotation(**data)
+                updated.validate(self.duration_s, self.start_time_s, self.end_time_s)
+                self._items[idx] = updated
+                self._sort()
+                return updated
+        raise KeyError(ann_id)
+
+    def delete(self, ann_id: str) -> bool:
+        before = len(self._items)
+        self._items = [a for a in self._items if a.id != ann_id]
+        return len(self._items) != before
+
+    def duplicate(self, ann_id: str) -> Annotation:
+        for existing in self._items:
+            if existing.id == ann_id:
+                data = asdict(existing)
+                data["id"] = uuid.uuid4().hex[:10]
+                data["status"] = "unreviewed"
+                data["reviewer"] = ""
+                data["review_notes"] = ""
+                return self.add(Annotation(**data))
+        raise KeyError(ann_id)
+
+    def clear(self) -> None:
+        self._items.clear()
+
+    def _sort(self) -> None:
+        self._items.sort(key=lambda a: a.position)
+
+    def to_dict(self, source_file: str = "", metadata: dict | None = None) -> dict:
+        meta = dict(metadata or {})
+        meta.setdefault("time_start_s", self.start_time_s)
+        if self.end_time_s is not None:
+            meta.setdefault("time_end_s", self.end_time_s)
+        return {"schema": "electrotrace.annotation/v2", "file": source_file, "metadata": meta, "annotations": [a.to_dict() for a in self._items]}
+
+    def to_json(self, source_file: str = "", metadata: dict | None = None) -> str:
+        return json.dumps(self.to_dict(source_file, metadata), indent=2)
+
+    @classmethod
+    def from_dict(cls, data: dict, duration_s: float | None = None, start_time_s: float | None = None, end_time_s: float | None = None) -> "AnnotationStore":
+        metadata = data.get("metadata") or {}
+        start = float(metadata.get("time_start_s", 0.0) if start_time_s is None else start_time_s)
+        end = metadata.get("time_end_s") if end_time_s is None else end_time_s
+        store = cls(duration_s=duration_s, start_time_s=start, end_time_s=float(end) if end is not None else None)
+        schema = data.get("schema", data.get("annotation_schema", ""))
+        if schema and schema not in {"electrotrace.annotation/v2", "v1"}:
+            raise ValueError(f"unsupported annotation schema: {schema}")
+        for raw in data.get("annotations", []):
+            store.add(Annotation(**raw))
+        return store
+
+    @classmethod
+    def from_json(cls, text: str, duration_s: float | None = None, start_time_s: float | None = None, end_time_s: float | None = None) -> "AnnotationStore":
+        try:
+            data = json.loads(text)
+        except json.JSONDecodeError as exc:
+            raise ValueError(f"invalid JSON: {exc}") from exc
+        return cls.from_dict(data, duration_s=duration_s, start_time_s=start_time_s, end_time_s=end_time_s)
+
+    def to_csv_rows(self, source_file: str = "") -> list[dict]:
+        return [{"file": source_file, "id": a.id, "type": a.type, "label": a.label, "channel": a.channel, "start": a.start, "end": a.end, "time": a.time, "confidence": a.confidence, "notes": a.notes, "annotator": a.annotator, "status": a.status, "reviewer": a.reviewer, "review_notes": a.review_notes} for a in self._items]
+
+
+def point_agreement(a: list[Annotation], b: list[Annotation], tolerance_s: float = 0.04) -> dict:
+    if tolerance_s <= 0 or not math.isfinite(tolerance_s):
+        raise ValueError("tolerance_s must be positive and finite")
+    left = [x for x in a if x.type == "point"]
+    right = [x for x in b if x.type == "point"]
+    used: set[str] = set()
+    errors = []
+    matches = 0
+    for x in left:
+        candidates = [
+            y for y in right
+            if y.id not in used
+            and y.label == x.label
+            and y.channel == x.channel
+            and y.time is not None
+            and x.time is not None
+            and abs(y.time - x.time) <= tolerance_s
+        ]
+        if candidates:
+            y = min(candidates, key=lambda z: abs(z.time - x.time))
+            used.add(y.id)
+            matches += 1
+            errors.append(abs(y.time - x.time))
+    total = max(len(left), len(right), 1)
+    return {"matches": matches, "agreement_rate": matches / total, "mean_absolute_error_s": sum(errors) / len(errors) if errors else None}
+
+
+def interval_iou(a: Annotation, b: Annotation) -> float:
+    if a.type != "interval" or b.type != "interval":
+        return 0.0
+    inter = max(0.0, min(a.end, b.end) - max(a.start, b.start))  # type: ignore[arg-type]
+    union = max(a.end, b.end) - min(a.start, b.start)  # type: ignore[arg-type]
+    return inter / union if union else 0.0

electrotrace/baseline_detectors.py

@@ -0,0 +1,122 @@
+"""Classical open R-peak detectors for locked baseline comparison.
+
+Implementations follow the classic literature at a research level:
+- Pan & Tompkins (1985): bandpass → derivative → square → moving integrate → adaptive threshold
+- Hamilton & Tompkins (1986): similar pipeline with Hamilton-style peak decision rules
+
+These are retrospective full-record detectors (same evaluation mode as ElectroTrace
+Stage-1) and are intended only for protocol-matched comparison on MIT-BIH.
+"""
+from __future__ import annotations
+
+import numpy as np
+from scipy import signal as sps
+
+
+def _validate(signal: np.ndarray, fs_hz: float) -> tuple[np.ndarray, float]:
+    x = np.asarray(signal, dtype=float)
+    fs = float(fs_hz)
+    if x.ndim != 1 or x.size < 32:
+        raise ValueError("signal must be one-dimensional with at least 32 samples")
+    if not np.isfinite(x).all():
+        raise ValueError("signal must contain only finite values")
+    if not np.isfinite(fs) or fs <= 0:
+        raise ValueError("fs_hz must be positive and finite")
+    return x, fs
+
+
+def _bandpass(x: np.ndarray, fs_hz: float, low: float = 5.0, high: float = 15.0) -> np.ndarray:
+    nyq = fs_hz / 2.0
+    low = min(low, nyq * 0.45)
+    high = min(high, nyq * 0.9)
+    if high <= low or low <= 0:
+        return x - np.median(x)
+    sos = sps.butter(2, [low, high], btype="bandpass", fs=fs_hz, output="sos")
+    try:
+        return sps.sosfiltfilt(sos, x - np.median(x))
+    except ValueError:
+        return sps.sosfilt(sos, x - np.median(x))
+
+
+def pan_tompkins_r_peaks(signal: np.ndarray, fs_hz: float) -> np.ndarray:
+    """Pan–Tompkins style R-peak detector (full-record retrospective)."""
+    x, fs = _validate(signal, fs_hz)
+    band = _bandpass(x, fs, 5.0, 15.0)
+    deriv = np.convolve(band, np.array([1, 2, 0, -2, -1], dtype=float) / 8.0, mode="same") * fs
+    squared = deriv ** 2
+    win = max(1, int(round(0.150 * fs)))
+    integrated = np.convolve(squared, np.ones(win, dtype=float) / win, mode="same")
+
+    init_n = min(len(integrated), int(2.0 * fs))
+    peak_level = float(np.max(integrated[:init_n])) if init_n else 0.0
+    noise_level = float(np.median(integrated[:init_n])) if init_n else 0.0
+    threshold = noise_level + 0.25 * (peak_level - noise_level)
+
+    min_distance = max(1, int(round(0.2 * fs)))
+    candidates, props = sps.find_peaks(integrated, distance=min_distance, height=threshold * 0.5)
+    if len(candidates) == 0:
+        return np.asarray([], dtype=int)
+
+    peaks: list[int] = []
+    last = -min_distance
+    for idx in candidates:
+        val = float(integrated[idx])
+        if idx - last < min_distance:
+            continue
+        if val >= threshold:
+            half = max(1, int(round(0.075 * fs)))
+            lo = max(0, int(idx) - half)
+            hi = min(len(band), int(idx) + half + 1)
+            local = band[lo:hi]
+            if local.size == 0:
+                continue
+            local_peak = lo + int(np.argmax(np.abs(local)))
+            peaks.append(local_peak)
+            last = idx
+            peak_level = 0.125 * val + 0.875 * peak_level
+            threshold = noise_level + 0.25 * (peak_level - noise_level)
+        else:
+            noise_level = 0.125 * val + 0.875 * noise_level
+            threshold = noise_level + 0.25 * (peak_level - noise_level)
+
+    return np.asarray(peaks, dtype=int)
+
+
+def hamilton_r_peaks(signal: np.ndarray, fs_hz: float) -> np.ndarray:
+    """Hamilton–Tompkins style R-peak detector (full-record retrospective)."""
+    x, fs = _validate(signal, fs_hz)
+    band = _bandpass(x, fs, 8.0, 16.0)
+    deriv = np.diff(band, prepend=band[0]) * fs
+    squared = deriv ** 2
+    win = max(1, int(round(0.080 * fs)))
+    integrated = np.convolve(squared, np.ones(win, dtype=float) / win, mode="same")
+
+    min_distance = max(1, int(round(0.2 * fs)))
+    med = float(np.median(integrated))
+    mad = float(np.median(np.abs(integrated - med)))
+    scale = 1.4826 * mad if mad > 1e-12 else float(np.std(integrated) + 1e-12)
+    height = med + 2.5 * scale
+
+    candidates, _ = sps.find_peaks(integrated, distance=min_distance, height=height)
+    if len(candidates) == 0:
+        candidates, _ = sps.find_peaks(integrated, distance=min_distance, height=med + 1.0 * scale)
+
+    peaks: list[int] = []
+    half = max(1, int(round(0.06 * fs)))
+    for idx in candidates:
+        lo = max(0, int(idx) - half)
+        hi = min(len(band), int(idx) + half + 1)
+        local = band[lo:hi]
+        if local.size == 0:
+            continue
+        local_peak = lo + int(np.argmax(np.abs(local)))
+        peaks.append(local_peak)
+    if not peaks:
+        return np.asarray([], dtype=int)
+    return np.asarray(sorted(set(peaks)), dtype=int)
+
+
+BASELINE_DETECTORS = {
+    "pan_tompkins": pan_tompkins_r_peaks,
+    "hamilton": hamilton_r_peaks,
+}

electrotrace/beats.py

@@ -0,0 +1,47 @@
+"""Beat-level ECG segmentation from detected R peaks."""
+from __future__ import annotations
+
+from dataclasses import asdict, dataclass
+import numpy as np
+
+
+@dataclass(frozen=True)
+class Beat:
+    index: int
+    r_index: int
+    r_time: float
+    start: float
+    end: float
+    rr_prev_s: float | None
+    rr_next_s: float | None
+    heart_rate_bpm: float | None
+
+    def to_dict(self) -> dict:
+        return asdict(self)
+
+
+def segment_beats(time: np.ndarray, peaks: np.ndarray, pre_s: float = 0.35, post_s: float = 0.55) -> list[Beat]:
+    time = np.asarray(time, dtype=float)
+    peaks = np.asarray(peaks, dtype=int)
+    if time.ndim != 1 or len(time) < 2:
+        raise ValueError("time must be a one-dimensional array with at least two samples")
+    if not np.isfinite(time).all() or np.any(np.diff(time) <= 0):
+        raise ValueError("time must be finite and strictly increasing")
+    if len(peaks) == 0:
+        return []
+    if not np.isfinite(pre_s) or not np.isfinite(post_s) or pre_s <= 0 or post_s <= 0:
+        raise ValueError("beat windows must be positive and finite")
+    peaks = np.unique(peaks[(peaks >= 0) & (peaks < len(time))])
+    if len(peaks) == 0:
+        return []
+    r_times = time[peaks]
+    beats: list[Beat] = []
+    for i, (idx, rt) in enumerate(zip(peaks, r_times)):
+        prev_rr = float(rt - r_times[i - 1]) if i else None
+        next_rr = float(r_times[i + 1] - rt) if i + 1 < len(r_times) else None
+        rr = prev_rr if prev_rr is not None and prev_rr > 0 else next_rr
+        hr = 60.0 / rr if rr is not None and rr > 0 else None
+        start = max(float(time[0]), float(rt - pre_s))
+        end = min(float(time[-1]), float(rt + post_s))
+        beats.append(Beat(i, int(idx), float(rt), start, end, prev_rr, next_rr, hr))
+    return beats

electrotrace/benchmark.py

@@ -0,0 +1,154 @@
+"""Leakage-safe model benchmarking with subject-level, stratified splits."""
+from __future__ import annotations
+
+from dataclasses import asdict, dataclass
+
+import numpy as np
+from scipy.stats import t as student_t
+from sklearn.ensemble import RandomForestClassifier
+from sklearn.linear_model import LogisticRegression
+from sklearn.metrics import accuracy_score, balanced_accuracy_score, f1_score, roc_auc_score, confusion_matrix
+from sklearn.model_selection import StratifiedGroupKFold
+from sklearn.pipeline import make_pipeline
+from sklearn.preprocessing import StandardScaler
+
+
+@dataclass
+class FoldMetrics:
+    fold: int
+    n_train: int
+    n_test: int
+    accuracy: float
+    balanced_accuracy: float
+    macro_f1: float
+    weighted_f1: float
+    roc_auc: float | None
+
+    def to_dict(self) -> dict:
+        return asdict(self)
+
+
+def _model_classes(model) -> np.ndarray:
+    classes = getattr(model, "classes_", None)
+    if classes is not None:
+        return np.asarray(classes)
+    steps = getattr(model, "steps", None)
+    if steps:
+        estimator = steps[-1][1]
+        classes = getattr(estimator, "classes_", None)
+        if classes is not None:
+            return np.asarray(classes)
+    raise ValueError("trained model does not expose class labels")
+
+
+def _auc(y_true: np.ndarray, proba: np.ndarray, classes: np.ndarray) -> float | None:
+    try:
+        if len(classes) == 2:
+            return float(roc_auc_score(y_true, proba[:, 1]))
+        return float(roc_auc_score(y_true, proba, multi_class="ovr", labels=classes))
+    except ValueError:
+        return None
+
+
+def _summary(values: list[float | None]) -> dict[str, float | None]:
+    """Describe cross-validation fold variability.
+
+    The t-based interval is explicitly descriptive across CV folds. Because CV
+    training sets overlap, it must not be interpreted as an independent-sample
+    confidence interval for future-subject performance.
+    """
+    x = np.asarray([v for v in values if v is not None and np.isfinite(v)], dtype=float)
+    if x.size == 0:
+        return {"n": 0, "mean": None, "std": None, "ci95_low": None, "ci95_high": None}
+    mean = float(np.mean(x))
+    if x.size == 1:
+        return {"n": 1, "mean": mean, "std": None, "ci95_low": None, "ci95_high": None}
+    sd = float(np.std(x, ddof=1))
+    critical = float(student_t.ppf(0.975, df=x.size - 1))
+    margin = critical * sd / np.sqrt(x.size)
+    return {"n": int(x.size), "mean": mean, "std": sd, "ci95_low": mean - margin, "ci95_high": mean + margin}
+
+
+def benchmark_models(X: np.ndarray, y: np.ndarray, groups: np.ndarray, folds: int = 5, seed: int = 42) -> dict:
+    """Benchmark baseline classifiers using leakage-safe group-stratified folds.
+
+    ``groups`` are the experimental units for splitting. Samples from one group
+    never occur in both train and test within a fold.
+    """
+    X = np.asarray(X, dtype=float)
+    y = np.asarray(y)
+    groups = np.asarray(groups)
+    if X.ndim != 2:
+        raise ValueError("X must be a two-dimensional feature matrix")
+    if not np.isfinite(X).all():
+        raise ValueError("X contains NaN or infinite values")
+    if len(X) != len(y) or len(y) != len(groups):
+        raise ValueError("X, y, and groups must have equal length")
+    labels = np.unique(y)
+    unique_groups = np.unique(groups)
+    if len(labels) < 2:
+        raise ValueError("Need at least two outcome classes for classification benchmarking")
+    if len(unique_groups) < 2:
+        raise ValueError("Need at least two subjects/groups for leakage-safe benchmarking")
+    requested_folds = int(folds)
+    if requested_folds < 2:
+        raise ValueError("folds must be at least 2")
+    groups_per_class = {str(cls): len(np.unique(groups[y == cls])) for cls in labels}
+    max_folds = min(len(unique_groups), min(groups_per_class.values()))
+    n_splits = min(requested_folds, max_folds)
+    if n_splits < 2:
+        raise ValueError(f"Need at least two distinct subjects in every class; observed {groups_per_class}")
+
+    models = {
+        "logistic_regression": make_pipeline(StandardScaler(), LogisticRegression(max_iter=2000, class_weight="balanced", random_state=seed)),
+        "random_forest": RandomForestClassifier(n_estimators=300, class_weight="balanced", random_state=seed, n_jobs=-1),
+    }
+    splitter = StratifiedGroupKFold(n_splits=n_splits, shuffle=True, random_state=seed)
+    result: dict[str, object] = {
+        "n_samples": int(len(y)),
+        "n_subjects": int(len(unique_groups)),
+        "experimental_unit": "group",
+        "folds": n_splits,
+        "splitter": "StratifiedGroupKFold",
+        "seed": int(seed),
+        "groups_per_class": groups_per_class,
+        "models": {},
+    }
+    metric_names = ("accuracy", "balanced_accuracy", "macro_f1", "weighted_f1", "roc_auc")
+    for name, model in models.items():
+        metrics: list[FoldMetrics] = []
+        confusion = None
+        test_groups_per_fold: list[int] = []
+        for fold, (train_idx, test_idx) in enumerate(splitter.split(X, y, groups), start=1):
+            if len(np.unique(y[train_idx])) < 2:
+                raise ValueError(f"Fold {fold} training set contains fewer than two outcome classes")
+            if set(groups[train_idx]).intersection(set(groups[test_idx])):
+                raise ValueError(f"Fold {fold} leaks experimental units between train and test")
+            model.fit(X[train_idx], y[train_idx])
+            pred = model.predict(X[test_idx])
+            proba = model.predict_proba(X[test_idx])
+            classes = _model_classes(model)
+            metrics.append(FoldMetrics(
+                fold=fold,
+                n_train=len(train_idx), n_test=len(test_idx),
+                accuracy=float(accuracy_score(y[test_idx], pred)),
+                balanced_accuracy=float(balanced_accuracy_score(y[test_idx], pred)),
+                macro_f1=float(f1_score(y[test_idx], pred, average="macro", zero_division=0)),
+                weighted_f1=float(f1_score(y[test_idx], pred, average="weighted", zero_division=0)),
+                roc_auc=_auc(y[test_idx], proba, classes),
+            ))
+            test_groups_per_fold.append(int(len(np.unique(groups[test_idx]))))
+            cm = confusion_matrix(y[test_idx], pred, labels=labels)
+            confusion = cm if confusion is None else confusion + cm
+        fold_values = {k: [getattr(m, k) for m in metrics] for k in metric_names}
+        summaries = {k: _summary(v) for k, v in fold_values.items()}
+        result["models"][name] = {
+            "folds": [m.to_dict() for m in metrics],
+            "summary": summaries,
+            "summary_interval_interpretation": "descriptive_t_interval_across_CV_folds; not an independent-subject generalization CI",
+            "mean": {k: s["mean"] for k, s in summaries.items()},
+            "confusion_matrix": confusion.tolist() if confusion is not None else [],
+            "labels": [str(x) for x in labels],
+            "test_groups_per_fold": test_groups_per_fold,
+        }
+    return result