ecg-plotter 0.1.0__py3-none-any.whl

ecg_plotter/__init__.py

@@ -0,0 +1,63 @@
+"""ECG Plotter — filter and render 12-lead ECG traces for research / visualisation.
+
+Not intended for diagnostic use. Dependencies such as matplotlib, scipy, wfdb,
+pandas, and Pillow are imported lazily when needed.
+"""
+
+from __future__ import annotations
+
+from importlib import import_module
+from typing import Any
+
+from ecg_plotter._version import __version__
+
+__all__ = [
+    "ANONYMOUS_ID",
+    "LEAD_NAMES",
+    "STANDARD_LAYOUT",
+    "__version__",
+    "apply_to_leads",
+    "as_lead_dict",
+    "bandpass",
+    "leads_to_array",
+    "load_jsonl",
+    "load_wfdb",
+    "median_baseline",
+    "notch",
+    "paper_calibration",
+    "plot",
+]
+
+# name -> (module, attribute)
+# Drawing lives in ``sheet``, not a module named ``plot``. A submodule called
+# ``plot`` would replace the public ``ecg.plot`` function as soon as that
+# submodule was imported.
+_LAZY_ATTRS: dict[str, tuple[str, str]] = {
+    "ANONYMOUS_ID": ("ecg_plotter.sheet", "ANONYMOUS_ID"),
+    "LEAD_NAMES": ("ecg_plotter.leads", "LEAD_NAMES"),
+    "STANDARD_LAYOUT": ("ecg_plotter.leads", "STANDARD_LAYOUT"),
+    "apply_to_leads": ("ecg_plotter.filters", "apply_to_leads"),
+    "as_lead_dict": ("ecg_plotter.leads", "as_lead_dict"),
+    "bandpass": ("ecg_plotter.filters", "bandpass"),
+    "leads_to_array": ("ecg_plotter.io", "leads_to_array"),
+    "load_jsonl": ("ecg_plotter.io", "load_jsonl"),
+    "load_wfdb": ("ecg_plotter.io", "load_wfdb"),
+    "median_baseline": ("ecg_plotter.filters", "median_baseline"),
+    "notch": ("ecg_plotter.filters", "notch"),
+    "paper_calibration": ("ecg_plotter.sheet", "paper_calibration"),
+    "plot": ("ecg_plotter.sheet", "plot"),
+}
+
+
+def __getattr__(name: str) -> Any:
+    try:
+        module_name, attr_name = _LAZY_ATTRS[name]
+    except KeyError as exc:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from exc
+    value = getattr(import_module(module_name), attr_name)
+    globals()[name] = value
+    return value
+
+
+def __dir__() -> list[str]:
+    return sorted(set(globals()) | set(__all__))

ecg_plotter/_version.py

@@ -0,0 +1,3 @@
+"""Package version — single source of truth for runtime and packaging."""
+
+__version__ = "0.1.0"

ecg_plotter/cli.py

@@ -0,0 +1,172 @@
+"""Command-line interface for ecg-plotter."""
+
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+
+from ecg_plotter._version import __version__
+
+
+def _build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        prog="ecg-plotter",
+        description="Plot a 12-lead ECG as PNG, PDF, or SVG.",
+    )
+    parser.add_argument(
+        "record",
+        nargs="?",
+        help="WFDB record path (with or without .hea/.dat), or omit when using --jsonl",
+    )
+    parser.add_argument(
+        "-o",
+        "--output",
+        required=True,
+        type=Path,
+        help="Output path (.png / .pdf / .svg)",
+    )
+    parser.add_argument(
+        "--jsonl",
+        type=Path,
+        help="JSON Lines file of lead arrays (uses --index)",
+    )
+    parser.add_argument(
+        "--index",
+        type=int,
+        default=0,
+        help="Record index within --jsonl (default: 0)",
+    )
+    parser.add_argument(
+        "--inches",
+        nargs=2,
+        type=float,
+        metavar=("WIDTH", "HEIGHT"),
+        help="Figure size in inches (mutually exclusive with --resolution)",
+    )
+    parser.add_argument(
+        "--resolution",
+        nargs=2,
+        type=int,
+        metavar=("WIDTH_PX", "HEIGHT_PX"),
+        help="Figure size in pixels (mutually exclusive with --inches)",
+    )
+    parser.add_argument("--dpi", type=int, default=300, help="DPI (default: 300)")
+    parser.add_argument(
+        "--fs",
+        type=float,
+        default=None,
+        help="Sampling frequency in Hz (overrides WFDB/JSONL metadata when set)",
+    )
+    parser.add_argument(
+        "--greyscale",
+        action="store_true",
+        help="Write greyscale output instead of pink colour paper",
+    )
+    parser.add_argument(
+        "--anonymise",
+        action="store_true",
+        help="Replace footer ID with the fixed anonymous placeholder",
+    )
+    parser.add_argument(
+        "--filter",
+        action="store_true",
+        help="Apply notch/bandpass using --notch-hz / --bandpass-hz",
+    )
+    parser.add_argument(
+        "--notch-hz",
+        type=float,
+        default=None,
+        help="Notch frequency in Hz (footer and optional filtering)",
+    )
+    parser.add_argument(
+        "--bandpass-hz",
+        nargs=2,
+        type=float,
+        metavar=("LOW", "HIGH"),
+        default=None,
+        help="Band-pass edges in Hz (footer and optional filtering)",
+    )
+    parser.add_argument(
+        "--id",
+        default=None,
+        help="Footer ID text (ignored when --anonymise is set)",
+    )
+    parser.add_argument(
+        "--format",
+        choices=("png", "pdf", "svg"),
+        default=None,
+        help="Output format (default: inferred from --output suffix)",
+    )
+    parser.add_argument(
+        "--version",
+        action="version",
+        version=f"%(prog)s {__version__}",
+    )
+    return parser
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = _build_parser()
+    args = parser.parse_args(argv)
+
+    if args.inches is None and args.resolution is None:
+        args.inches = [9.84, 6.3]
+    if args.inches is not None and args.resolution is not None:
+        parser.error("Provide only one of --inches or --resolution")
+    if args.record is None and args.jsonl is None:
+        parser.error("Provide a WFDB record path or --jsonl")
+    if args.record is not None and args.jsonl is not None:
+        parser.error("Provide either a WFDB record path or --jsonl, not both")
+
+    from ecg_plotter.io import load_jsonl, load_wfdb
+    from ecg_plotter.sheet import plot
+
+    if args.jsonl is not None:
+        if args.index < 0:
+            parser.error("--index must be >= 0")
+        # Stop once the requested record is found. A research JSONL file can
+        # be huge; there is no need to parse every later line.
+        chosen = None
+        n_seen = 0
+        for n_seen, item in enumerate(load_jsonl(args.jsonl), start=1):
+            if n_seen - 1 == args.index:
+                chosen = item
+                break
+        if chosen is None:
+            if n_seen == 0:
+                parser.error(f"No records found in {args.jsonl}")
+            parser.error(f"--index {args.index} out of range (0..{n_seen - 1})")
+        assert chosen is not None
+        leads = chosen["leads"]
+        fs = args.fs if args.fs is not None else chosen["fs"]
+        record_id = args.id if args.id is not None else (
+            str(chosen["id"]) if chosen.get("id") is not None else None
+        )
+    else:
+        data = load_wfdb(args.record)
+        leads = data["leads"]
+        fs = args.fs if args.fs is not None else data["fs"]
+        record_id = args.id if args.id is not None else data["record_name"]
+
+    path = plot(
+        leads,
+        args.output,
+        inches=tuple(args.inches) if args.inches is not None else None,
+        resolution=tuple(args.resolution) if args.resolution is not None else None,
+        dpi=args.dpi,
+        fs=fs,
+        colour=not args.greyscale,
+        id=record_id,
+        anonymise=args.anonymise,
+        filter=args.filter,
+        notch_hz=args.notch_hz,
+        bandpass_hz=tuple(args.bandpass_hz) if args.bandpass_hz is not None else None,
+        format=args.format,
+    )
+    print(path)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

ecg_plotter/filters.py

@@ -0,0 +1,142 @@
+"""Signal conditioning filters for ECG traces.
+
+Each function operates on a single 1-D lead. Use :func:`apply_to_leads` to
+run the same filter across a full 12-lead recording.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any
+
+from numpy import asarray, ndarray
+from scipy.signal import butter, filtfilt, iirnotch, medfilt
+
+from ecg_plotter.leads import LeadInput, _is_dataframe, as_lead_dict
+
+if TYPE_CHECKING:
+    from pandas import DataFrame
+
+
+def _require_filtfilt_length(x: ndarray, b: ndarray, a: ndarray, *, name: str) -> None:
+    """Raise a clear error before filtfilt's cryptic padlen failure."""
+    padlen = 3 * max(len(a), len(b))
+    if len(x) <= padlen:
+        raise ValueError(
+            f"{name}: signal length ({len(x)}) is too short for zero-phase filtering "
+            f"(needs more than {padlen} samples for this filter). "
+            "Use a longer recording, lower filter order, or skip filtering."
+        )
+
+
+def bandpass(
+    signal: ndarray,
+    fs: float,
+    lowcut: float = 0.05,
+    highcut: float = 150.0,
+    order: int = 4,
+) -> ndarray:
+    """Zero-phase Butterworth band-pass filter.
+
+    Parameters
+    ----------
+    signal
+        1-D ECG samples in millivolts (mV).
+    fs
+        Sampling frequency in Hz.
+    lowcut, highcut
+        Pass-band edges in Hz.
+    order
+        Filter order.
+    """
+    x = asarray(signal, dtype=float).reshape(-1)
+    if lowcut <= 0 or highcut <= lowcut:
+        raise ValueError("Require 0 < lowcut < highcut.")
+    nyq = 0.5 * fs
+    if highcut >= nyq:
+        raise ValueError(f"highcut ({highcut}) must be < Nyquist ({nyq}).")
+    low = lowcut / nyq
+    high = highcut / nyq
+    b, a = butter(order, [low, high], btype="band")
+    _require_filtfilt_length(x, b, a, name="bandpass")
+    return filtfilt(b, a, x)
+
+
+def notch(
+    signal: ndarray,
+    fs: float,
+    freq: float = 60.0,
+    quality: float = 30.0,
+) -> ndarray:
+    """Zero-phase IIR notch filter (mains rejection).
+
+    Parameters
+    ----------
+    signal
+        1-D ECG samples in millivolts (mV).
+    fs
+        Sampling frequency in Hz.
+    freq
+        Notch centre frequency in Hz (typically 50 or 60).
+    quality
+        Quality factor ``Q = freq / bandwidth``. Higher Q → narrower notch.
+    """
+    x = asarray(signal, dtype=float).reshape(-1)
+    nyq = 0.5 * fs
+    if not 0 < freq < nyq:
+        raise ValueError(f"Notch frequency must be in (0, {nyq}); got {freq}.")
+    b, a = iirnotch(freq, quality, fs=fs)
+    _require_filtfilt_length(x, b, a, name="notch")
+    return filtfilt(b, a, x)
+
+
+def median_baseline(
+    signal: ndarray,
+    kernel_size: int = 201,
+) -> ndarray:
+    """Remove baseline wander by subtracting a median-filtered estimate.
+
+    ``kernel_size`` must be odd. Larger kernels track slower drift.
+
+    Parameters
+    ----------
+    signal
+        1-D ECG samples in millivolts (mV).
+    kernel_size
+        Odd window length for the median baseline estimate.
+    """
+    x = asarray(signal, dtype=float).reshape(-1)
+    if kernel_size < 1 or kernel_size % 2 == 0:
+        raise ValueError("kernel_size must be a positive odd integer.")
+    if kernel_size > len(x):
+        raise ValueError(
+            f"kernel_size ({kernel_size}) cannot exceed signal length ({len(x)})."
+        )
+    baseline = medfilt(x, kernel_size=kernel_size)
+    return x - baseline
+
+
+def apply_to_leads(
+    leads: LeadInput,
+    filter_fn: Callable[..., ndarray],
+    **kwargs: Any,
+) -> dict[str, ndarray] | DataFrame:
+    """Apply a 1-D filter function to every lead.
+
+    Accepts a lead mapping, ``(12, n)`` / ``(n, 12)`` array, or DataFrame.
+    DataFrame inputs are returned as a DataFrame with the same lead columns;
+    other inputs return a lead dictionary.
+
+    Example
+    -------
+    >>> cleaned = apply_to_leads(leads, bandpass, fs=500, lowcut=0.05, highcut=150)
+    >>> cleaned = apply_to_leads(cleaned, notch, fs=500, freq=60)
+    """
+    was_df = _is_dataframe(leads)
+    lead_map = as_lead_dict(leads)
+    out = {name: filter_fn(sig, **kwargs) for name, sig in lead_map.items()}
+    if was_df:
+        from pandas import DataFrame
+
+        return DataFrame(out)
+    return out

ecg_plotter/io.py

@@ -0,0 +1,169 @@
+"""Load ECG recordings from WFDB paths or JSON Lines files.
+
+Lead samples are expected in millivolts (mV). ``load_wfdb`` reads
+``record.p_signal``, which WFDB already scales to physical units using the
+header gain; a warning is issued if those units are not mV.
+"""
+
+from __future__ import annotations
+
+import json
+import warnings
+from collections.abc import Iterator, Mapping, Sequence
+from pathlib import Path
+from typing import Any
+
+from numpy import asarray, ndarray, vstack
+
+from ecg_plotter.leads import LEAD_NAMES, as_lead_dict
+
+_MV_ALIASES = frozenset({"mv", "millivolt", "millivolts"})
+
+
+def _warn_if_units_not_mv(units: Sequence[Any] | None, record_path: Path) -> None:
+    if not units:
+        return
+    unexpected = [
+        str(unit)
+        for unit in units
+        if str(unit).strip().lower() not in _MV_ALIASES
+    ]
+    if unexpected:
+        warnings.warn(
+            f"WFDB record {record_path} has non-mV units {unexpected}; "
+            "ecg-plotter expects lead amplitudes in millivolts (mV).",
+            UserWarning,
+            stacklevel=3,
+        )
+
+
+def leads_to_array(
+    leads: Mapping[str, ndarray],
+    lead_names: Sequence[str] = LEAD_NAMES,
+) -> ndarray:
+    """Stack a lead mapping into an array of shape ``(n_leads, n_samples)``.
+
+    Values are treated as millivolts (mV).
+    """
+    missing = [name for name in lead_names if name not in leads]
+    if missing:
+        raise KeyError(f"Missing leads: {missing}")
+    return vstack([asarray(leads[name], dtype=float) for name in lead_names])
+
+
+def load_wfdb(
+    record_path: str | Path,
+    *,
+    lead_names: Sequence[str] | None = None,
+) -> dict[str, Any]:
+    """Load a WFDB record into a lead dictionary plus metadata.
+
+    Lead arrays come from ``record.p_signal`` and are expected in millivolts
+    (mV). If the header ``units`` are not mV, a :class:`UserWarning` is
+    emitted; values are still returned unchanged.
+
+    Parameters
+    ----------
+    record_path
+        Path to the record *without* extension, or with ``.hea`` / ``.dat``.
+    lead_names
+        Optional override for signal names (must match the number of channels).
+        If omitted, names are taken from the WFDB header (``record.sig_name``),
+        so any 12-lead WFDB dataset works regardless of channel order
+        (e.g. ``aVL`` before or after ``aVF``).
+
+    Returns
+    -------
+    dict
+        ``leads`` (name → array in mV), ``fs``, ``sig_name``, ``record_name``,
+        ``n_samples``, ``units``, and the raw ``wfdb.Record`` under ``record``.
+    """
+    from wfdb import rdrecord
+
+    path = Path(record_path)
+    if path.suffix.lower() in {".hea", ".dat", ".mat"}:
+        path = path.with_suffix("")
+
+    record = rdrecord(str(path))
+    signal = asarray(record.p_signal, dtype=float)
+    names = list(record.sig_name) if lead_names is None else list(lead_names)
+    units = list(record.units) if getattr(record, "units", None) is not None else None
+    _warn_if_units_not_mv(units, path)
+
+    if signal.ndim != 2 or signal.shape[1] != len(names):
+        raise ValueError(
+            f"Signal shape {signal.shape} does not match {len(names)} lead names."
+        )
+
+    leads = {name: signal[:, i].copy() for i, name in enumerate(names)}
+    return {
+        "leads": leads,
+        "fs": float(record.fs),
+        "sig_name": names,
+        "record_name": record.record_name,
+        "n_samples": int(record.sig_len),
+        "units": units,
+        "record": record,
+    }
+
+
+def load_jsonl(
+    path: str | Path,
+    *,
+    lead_names: Sequence[str] | None = None,
+) -> Iterator[dict[str, Any]]:
+    """Yield records from a JSON Lines file of lead arrays.
+
+    Each line must be a JSON object with:
+
+    - lead arrays in millivolts (mV), e.g. ``"I": [...]``
+    - required ``fs`` (sampling frequency in Hz; must be positive)
+    - optional metadata such as ``id`` / ``label``
+
+    Non-lead keys are returned alongside ``leads``.
+    """
+    names = set(lead_names) if lead_names is not None else set(LEAD_NAMES)
+    path = Path(path)
+    with path.open("r", encoding="utf-8") as fh:
+        for line_no, line in enumerate(fh, start=1):
+            line = line.strip()
+            if not line:
+                continue
+            try:
+                row = json.loads(line)
+            except json.JSONDecodeError as exc:
+                raise ValueError(f"Invalid JSON on line {line_no} of {path}") from exc
+            if not isinstance(row, dict):
+                raise ValueError(f"Expected object on line {line_no}; got {type(row)}.")
+
+            lead_map = {
+                key: asarray(val, dtype=float).reshape(-1)
+                for key, val in row.items()
+                if key in names and not isinstance(val, str)
+            }
+            if not lead_map:
+                raise ValueError(f"No lead arrays found on line {line_no} of {path}.")
+
+            meta = {key: val for key, val in row.items() if key not in lead_map}
+            if "fs" not in meta or meta["fs"] is None:
+                raise ValueError(
+                    f"Missing required field 'fs' on line {line_no} of {path}."
+                )
+            try:
+                fs = float(meta["fs"])
+            except (TypeError, ValueError) as exc:
+                raise ValueError(
+                    f"Invalid 'fs' on line {line_no} of {path}: {meta['fs']!r}."
+                ) from exc
+            if fs <= 0:
+                raise ValueError(
+                    f"'fs' must be positive on line {line_no} of {path}; got {fs}."
+                )
+
+            yield {
+                "leads": as_lead_dict(lead_map),
+                "fs": fs,
+                "id": meta.get("id", meta.get("study_id")),
+                "label": meta.get("label", meta.get("condition")),
+                "meta": meta,
+            }