atriakit 0.1.0__py3-none-any.whl

atriakit/__init__.py

@@ -0,0 +1,65 @@
+"""Public package interface for the ECG feature library."""
+
+from importlib import import_module
+
+try:
+    from importlib.metadata import version, PackageNotFoundError
+    __version__ = version("atriakit")
+except PackageNotFoundError:
+    __version__ = "unknown"
+
+__all__ = [
+    "create_demo_data",
+    "AnnotationRow",
+    "Annotations",
+    "AnnotationsLoaderConfig",
+    "AnnotationsLoader",
+    "BaseLoader",
+    "DicomLoader",
+    "ECGData",
+    "ECGDataset",
+    "ECGLoader",
+    "FeatureCalculators",
+    "FeatureComputationConfig",
+    "Pipeline",
+    "PipelineConfig",
+    "SegmentConfig",
+    "SignalPreprocessor",
+    "SignalPreprocessorConfig",
+]
+
+_PUBLIC_IMPORTS = {
+    "create_demo_data": ("atriakit.datasets", "generate"),
+    "AnnotationRow": ("atriakit.models.annotation_row", "AnnotationRow"),
+    "Annotations": ("atriakit.models.annotations", "Annotations"),
+    "AnnotationsLoaderConfig": ("atriakit.configs.annotations_loader_config", "AnnotationsLoaderConfig"),
+    "AnnotationsLoader": ("atriakit.io", "AnnotationsLoader"),
+    "BaseLoader": ("atriakit.io", "BaseLoader"),
+    "DicomLoader": ("atriakit.io", "DicomLoader"),
+    "ECGData": ("atriakit.models.ecg_data", "ECGData"),
+    "ECGDataset": ("atriakit.processing.dataset", "ECGDataset"),
+    "ECGLoader": ("atriakit.io", "ECGLoader"),
+    "FeatureCalculators": ("atriakit.feature_calculator", "FeatureCalculators"),
+    "FeatureComputationConfig": (
+        "atriakit.configs.feature_computation_config",
+        "FeatureComputationConfig",
+    ),
+    "Pipeline": ("atriakit.processing.pipeline", "Pipeline"),
+    "PipelineConfig": ("atriakit.configs.pipeline_config", "PipelineConfig"),
+    "SegmentConfig": ("atriakit.configs.segment_config", "SegmentConfig"),
+    "SignalPreprocessor": ("atriakit.preprocessing", "SignalPreprocessor"),
+    "SignalPreprocessorConfig": (
+        "atriakit.configs.signal_preprocessor_config",
+        "SignalPreprocessorConfig",
+    ),
+}
+
+
+def __getattr__(name):
+    if name not in _PUBLIC_IMPORTS:
+        raise AttributeError(f"module 'atriakit' has no attribute {name!r}")
+
+    module_name, attr_name = _PUBLIC_IMPORTS[name]
+    value = getattr(import_module(module_name), attr_name)
+    globals()[name] = value
+    return value

atriakit/configs/annotations_loader_config.py

@@ -0,0 +1,24 @@
+from dataclasses import dataclass
+from typing import Literal
+
+
+@dataclass(slots=True)
+class AnnotationsLoaderConfig:
+    """Configuration for the AnnotationsLoader.
+
+    Parameters
+    ----------
+    boundary_mode : {"per_lead", "cross_lead"}
+        Controls which onset/offset values are treated as active after loading.
+
+        ``"cross_lead"`` (default)
+            ``onset`` and ``offset`` hold the multilead consensus values
+            (cross-lead minimum onset, cross-lead maximum offset).
+            No substitution is performed.
+
+        ``"per_lead"``
+            ``onset`` and ``offset`` are replaced by ``onset_original`` and
+            ``offset_original`` — the raw per-lead annotation values.
+    """
+
+    boundary_mode: Literal["per_lead", "cross_lead"] = "cross_lead"

atriakit/configs/feature_computation_config.py

@@ -0,0 +1,45 @@
+from dataclasses import dataclass
+
+
+@dataclass(slots=True)
+class FeatureComputationConfig:
+    """Feature-computation parameters passed to FeatureCalculators.compute_all()."""
+
+    # Entropy
+    shannon_entropy_n_bins: int = 32  # number of histogram bins for Shannon entropy
+    shannon_entropy_bin_range: tuple[float, float] | None = (
+        None  # fixed amplitude range for the histogram; None uses per-signal min/max. For comparable results should be set to the same value across the datasets.
+    )
+    sample_entropy_m: int = (
+        2  # embedding dimension (template length) for sample entropy
+    )
+    sample_entropy_r_factor: float = (
+        0.25  # tolerance as fraction of segment SD; controls how strictly templates must match
+    )
+
+    # Complexity / extrema
+    extrema_threshold_multiplier: float = (
+        0.1  # peak must exceed this fraction of max amplitude to be counted (complexity feature)
+    )
+
+    # Noise & morphology
+    noise_sd_multiplier: float = (
+        3.0  # a signal deflection must exceed this many noise SDs to be counted as a fragment or extremum
+    )
+    morphology_min_phase_fraction: float = (
+        0.1  # minimum fraction of segment length a sign group must span to count as a distinct morphology phase (0.1 = at least 10% of P-wave)
+    )
+    morphology_noise_sd_multiplier: float = (
+        3.0  # same as noise_sd_multiplier but applied only inside the morphology classifier
+    )
+
+    # Fragmentation
+    fragment_noise_multiplier: float = (
+        3.0  # a fragment's amplitude change must exceed this many noise SDs to count as a valid fragment boundary
+    )
+    min_fragment_length_ms: float = (
+        0.0  # fragments shorter than this (ms) are discarded
+    )
+    normalize_by_duration: bool = (
+        False  # normalise fragment metrics per 100ms of P-wave duration; True matches the original paper
+    )

atriakit/configs/pipeline_config.py

@@ -0,0 +1,32 @@
+from dataclasses import dataclass, field
+
+from atriakit.configs.feature_computation_config import FeatureComputationConfig
+from atriakit.configs.segment_config import SegmentConfig
+from atriakit.configs.signal_preprocessor_config import (
+    SignalPreprocessorConfig,
+    default_morphology_preprocessor_config,
+    default_signal_preprocessor_config,
+)
+
+
+@dataclass(slots=True)
+class PipelineConfig:
+    """Batch pipeline parameters."""
+
+    feature_computation: FeatureComputationConfig = field(
+        default_factory=FeatureComputationConfig  # settings passed to FeatureCalculators.compute_all(); see FeatureComputationConfig
+    )
+    segment_config: SegmentConfig = field(
+        default_factory=SegmentConfig  # segment-boundary and baseline-correction settings; see SegmentConfig
+    )
+    signal_preprocessor_config: SignalPreprocessorConfig = field(
+        default_factory=default_signal_preprocessor_config  # preprocessing for amplitude and duration features
+    )
+    morphology_preprocessor_config: SignalPreprocessorConfig = field(
+        default_factory=default_morphology_preprocessor_config  # preprocessing for shape-based features
+    )
+
+    # Grouping
+    group_tolerance_ms: int = (
+        200  # max time gap (ms) between P-wave peaks across leads to be grouped as the same beat
+    )

atriakit/configs/segment_config.py

@@ -0,0 +1,11 @@
+from typing import Literal
+
+from pydantic.dataclasses import dataclass
+
+
+@dataclass
+class SegmentConfig:
+    """Segment-boundary and baseline-correction settings for ``SegmentProcessor``."""
+
+    baseline_correction_type: Literal["none", "linear", "onset"] = "none"  # Baseline correction applied to extracted segments before feature calculation.
+    skip_first_ms: int = 0  # Drops the first samples of each annotated segment before feature extraction.

atriakit/configs/signal_preprocessor_config.py

@@ -0,0 +1,48 @@
+from dataclasses import dataclass
+from typing import Literal
+
+import numpy as np
+
+
+@dataclass(slots=True)
+class SignalPreprocessorConfig:
+    """Default parameters for ECG signal preprocessing."""
+
+    lowcut: float | None = (
+        None  # Low cutoff for the bandpass filter; `None` disables bandpass filtering.
+    )
+    highcut: float | None = (
+        None  # High cutoff for the bandpass filter; `None` disables bandpass filtering.
+    )
+    notch_freq: float | list[float] | None = (
+        50.0  # Line-noise notch frequency/frequencies in Hz; set to `None` to skip the notch filter.
+    )
+    order: int = 2  # Butterworth bandpass filter order.
+    quality_factor: float = (
+        20.0  # Q factor for the notch filter; higher values make the notch narrower.
+    )
+    global_baseline_mode: Literal["none", "p_onset_spline"] = (
+        "p_onset_spline"  # Full-recording baseline correction mode applied before segment extraction.
+    )
+    global_baseline_smoothing_factor: float = (
+        0  # Dimensionless smoothing strength for the SciPy spline fit; 0 keeps exact interpolation, larger values follow only slower baseline drift.
+    )
+    normalization_type: Literal["none", "zscore"] = (
+        "none"  # Dataset-level normalization to apply after filtering; use `"none"` to disable it.
+    )
+    mean: np.ndarray | None = (
+        None  # Per-lead mean used for zscore normalization; Either provide or compute with ECGDataset. Required if `normalization_type` is `"zscore"`.
+    )
+    std: np.ndarray | None = (
+        None  # Per-lead standard deviation used for zscore normalization;  Either provide or compute with ECGDataset. Required if `normalization_type` is `"zscore"`.
+    )
+
+
+def default_signal_preprocessor_config() -> SignalPreprocessorConfig:
+    """Low-pass at 120 Hz with 50/100 Hz notch. Used for amplitude and duration features."""
+    return SignalPreprocessorConfig(lowcut=None, highcut=120, notch_freq=[50.0, 100.0])
+
+
+def default_morphology_preprocessor_config() -> SignalPreprocessorConfig:
+    """Low-pass at 30 Hz with 50/100 Hz notch. Used for P-wave shape classification and inflection detection."""
+    return SignalPreprocessorConfig(lowcut=None, highcut=30, notch_freq=[50.0, 100.0])

atriakit/constants.py

@@ -0,0 +1,18 @@
+"""Package-level constants shared across modules."""
+
+# Leads required by the Kors regression matrix to compute a 3-lead VCG.
+VCG_LEADS: list[str] = ["I", "II", "V1", "V2", "V3", "V4", "V5", "V6"]
+
+# Output columns produced by VCG feature computation.
+VCG_FEATURE_COLUMNS: list[str] = [
+    "vcg_eigenvalues_1", "vcg_eigenvalues_2", "vcg_eigenvalues_3",
+    "vcg_roundness", "vcg_flatness", "vcg_area",
+    "vcg_axis_elevation", "vcg_axis_azimuth",
+    "vcg_sum_fragment_count", "vcg_sum_fragment_width", "vcg_sum_fragment_height",
+    "vcg_x_fragment_count", "vcg_x_fragment_width", "vcg_x_fragment_height",
+    "vcg_y_fragment_count", "vcg_y_fragment_width", "vcg_y_fragment_height",
+    "vcg_z_fragment_count", "vcg_z_fragment_width", "vcg_z_fragment_height",
+]
+
+# Leads required to compute the frontal P-wave axis.
+AXIS_LEADS: list[str] = ["I", "aVF"]

atriakit/datasets.py

@@ -0,0 +1,286 @@
+"""Synthetic demo ECG data for AtriaKit.
+
+Requires neurokit2, which is not installed by default::
+
+    pip install neurokit2
+    from atriakit import create_demo_data
+    create_demo_data()                       # writes to ./demo_data/
+    create_demo_data(output_dir="./my_data")
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+import pydicom
+from pydicom.dataset import Dataset, FileDataset, FileMetaDataset
+from pydicom.sequence import Sequence
+from pydicom.uid import generate_uid
+
+from atriakit.models.annotation_schema import AnnotationSchema
+
+LEADS = ["I", "II", "III", "aVR", "aVL", "aVF", "V1", "V2", "V3", "V4", "V5", "V6"]
+FS = 500
+DURATION = 10  # seconds
+
+_LEAD_SCALES = {
+    "I":   0.70, "II":  1.00, "III": 0.40,
+    "aVR": -0.80, "aVL": 0.30, "aVF": 0.70,
+    "V1":  -0.30, "V2":  0.20, "V3":  0.80,
+    "V4":  1.20,  "V5":  1.10, "V6":  0.90,
+}
+_MV_TO_RAW = 1000  # 1 mV → 1000 int16 units; sensitivity = 0.001 mV/unit
+
+DEMO_RECORDINGS = [
+    {
+        "heart_rate": 60,
+        AnnotationSchema.PATIENT_ID: "DEMO001",
+        "seed": 0,
+        "date": "20240101",
+        "time": "090000",
+    },
+    {
+        "heart_rate": 75,
+        AnnotationSchema.PATIENT_ID: "DEMO002",
+        "seed": 1,
+        "date": "20240615",
+        "time": "143000",
+    },
+    {
+        "heart_rate": 90,
+        AnnotationSchema.PATIENT_ID: "DEMO003",
+        "seed": 2,
+        "date": "20241201",
+        "time": "111500",
+    },
+]
+
+
+def _require_neurokit():
+    try:
+        import neurokit2 as nk
+
+        return nk
+    except ImportError:
+        raise SystemExit(
+            "Demo data generation requires neurokit2.\n"
+            "Install it with:  pip install neurokit2"
+        )
+
+
+def simulate_12lead(heart_rate: int, noise: float = 0.02, seed: int = 0) -> np.ndarray:
+    """Simulate a 12-lead ECG and return a ``(12, n_samples)`` array in mV."""
+    nk = _require_neurokit()
+    rng = np.random.default_rng(seed)
+    lead_ii = np.array(
+        nk.ecg_simulate(
+            duration=DURATION, sampling_rate=FS, heart_rate=heart_rate, noise=noise
+        )
+    )
+    ecg = np.zeros((12, len(lead_ii)))
+    for i, lead in enumerate(LEADS):
+        per_lead_noise = rng.normal(0, noise * 0.05, len(lead_ii))
+        ecg[i] = _LEAD_SCALES[lead] * lead_ii + per_lead_noise
+    return ecg
+
+
+def to_dicom(
+    ecg: np.ndarray,
+    patient_id: str,
+    date: str = "20240101",
+    time: str = "120000",
+    leads: list[str] | None = None,
+) -> FileDataset:
+    """Pack an ECG array into a DICOM FileDataset.
+
+    Args:
+        ecg: Signal array with shape ``(n_leads, n_samples)`` in mV.
+        patient_id: Patient identifier written into DICOM metadata.
+        date: Study date in ``YYYYMMDD`` format.
+        time: Study time in ``HHMMSS`` format.
+        leads: Channel labels matching the row order of ``ecg``.
+            Defaults to the standard 12-lead ``LEADS`` list.
+
+    Returns:
+        A ``pydicom.FileDataset`` ready to save with ``ds.save_as(path)``.
+    """
+    if leads is None:
+        leads = LEADS
+    sop_uid = generate_uid()
+
+    file_meta = FileMetaDataset()
+    file_meta.MediaStorageSOPClassUID = "1.2.840.10008.5.1.4.1.1.9.1.1"
+    file_meta.MediaStorageSOPInstanceUID = sop_uid
+    file_meta.TransferSyntaxUID = pydicom.uid.ExplicitVRLittleEndian
+
+    ds = FileDataset(None, {}, file_meta=file_meta, preamble=b"\0" * 128)
+    ds.is_implicit_VR = False
+    ds.is_little_endian = True
+
+    ds.SOPClassUID = file_meta.MediaStorageSOPClassUID
+    ds.SOPInstanceUID = sop_uid
+    ds.PatientID = patient_id
+    ds.PatientName = f"Demo^{patient_id}"
+    ds.ContentDate = date
+    ds.ContentTime = time
+    ds.StudyDate = date
+    ds.StudyTime = time
+
+    wf = Dataset()
+    wf.MultiplexGroupLabel = "ECG"
+    wf.SamplingFrequency = float(FS)
+
+    n_ch, n_samples = ecg.shape
+    wf.add_new(pydicom.tag.Tag(0x003A, 0x0005), "US", n_ch)  # NumberOfChannels
+    wf.add_new(
+        pydicom.tag.Tag(0x003A, 0x0010), "UL", n_samples
+    )  # NumberOfWaveformSamples
+    wf.WaveformBitsAllocated = 16
+    wf.WaveformSampleInterpretation = "SS"
+
+    channels = []
+    for lead in leads:
+        ch = Dataset()
+        ch.ChannelLabel = lead
+        units = Dataset()
+        units.CodeValue = "mV"
+        units.CodingSchemeDesignator = "UCUM"
+        units.CodeMeaning = "millivolt"
+        ch.ChannelSensitivityUnitsSequence = Sequence([units])
+        ch.ChannelSensitivity = str(1.0 / _MV_TO_RAW)
+        ch.ChannelSensitivityCorrectionFactor = "1.0"
+        ch.WaveformBitsStored = 16
+        channels.append(ch)
+    wf.ChannelDefinitionSequence = Sequence(channels)
+
+    raw = np.clip((ecg * _MV_TO_RAW).T, -32768, 32767).astype(np.int16)
+    wf.WaveformData = raw.tobytes()
+
+    ds.WaveformSequence = Sequence([wf])
+    return ds
+
+
+def detect_p_waves(ecg_12lead: np.ndarray) -> list[dict]:
+    """Detect P-wave onsets and offsets on Lead II using NeuroKit2.
+
+    Returns a list of dicts with keys ``onset``, ``offset``, and ``qrs_onset``
+    (sample indices at the pipeline's sampling rate).
+    """
+    nk = _require_neurokit()
+    lead_ii = ecg_12lead[1]
+    signals, _ = nk.ecg_process(lead_ii, sampling_rate=FS)
+
+    p_onset_mask = (
+        signals.get("ECG_P_Onsets", pd.Series(0, index=signals.index))
+        .fillna(0)
+        .astype(int)
+    )
+    p_offset_mask = (
+        signals.get("ECG_P_Offsets", pd.Series(0, index=signals.index))
+        .fillna(0)
+        .astype(int)
+    )
+    r_peak_mask = (
+        signals.get("ECG_R_Peaks", pd.Series(0, index=signals.index))
+        .fillna(0)
+        .astype(int)
+    )
+
+    onsets = np.where(p_onset_mask == 1)[0]
+    offsets = np.where(p_offset_mask == 1)[0]
+    r_peaks = np.where(r_peak_mask == 1)[0]
+
+    records: list[dict] = []
+    used_offsets: set[int] = set()
+    for onset in onsets:
+        future_offsets = offsets[offsets > onset]
+        if not len(future_offsets):
+            continue
+        offset = int(future_offsets[0])
+        if offset in used_offsets:
+            continue
+        used_offsets.add(offset)
+        following_r = r_peaks[r_peaks > offset]
+        qrs_onset = int(following_r[0]) - 20 if len(following_r) else -1
+        records.append(
+            {
+                AnnotationSchema.ONSET: int(onset),
+                AnnotationSchema.OFFSET: int(offset),
+                AnnotationSchema.QRS_ONSET: qrs_onset,
+            }
+        )
+
+    return records
+
+
+def build_annotations(
+    file_path: str,
+    patient_id: str,
+    p_waves: list[dict],
+    seed: int = 42,
+) -> pd.DataFrame:
+    """Expand Lead II P-wave detections into a per-lead annotations DataFrame.
+
+    Adds small random jitter to onset/offset per lead to simulate realistic
+    inter-lead variation.
+    """
+    rng = np.random.default_rng(seed)
+    rows = []
+    for p_wave_id, pw in enumerate(p_waves):
+        for lead in LEADS:
+            rows.append(
+                {
+                    AnnotationSchema.PATIENT_ID: patient_id,
+                    AnnotationSchema.LEAD: lead,
+                    AnnotationSchema.P_WAVE_ID: p_wave_id,
+                    AnnotationSchema.ONSET: pw[AnnotationSchema.ONSET]
+                    + int(rng.integers(-2, 3)),
+                    AnnotationSchema.OFFSET: pw[AnnotationSchema.OFFSET]
+                    + int(rng.integers(-2, 3)),
+                    AnnotationSchema.QRS_ONSET: pw[AnnotationSchema.QRS_ONSET],
+                    AnnotationSchema.IGNORE: False,
+                    AnnotationSchema.FILE_PATH: file_path,
+                }
+            )
+    return pd.DataFrame(rows)
+
+
+def generate(output_dir: str | Path) -> Path:
+    """Generate demo DICOM files and an annotations CSV into ``output_dir``.
+
+    Creates one DICOM per entry in ``DEMO_RECORDINGS`` and a combined
+    ``demo_annotations.csv``. Returns the output directory path.
+    """
+    out = Path(output_dir)
+    out.mkdir(parents=True, exist_ok=True)
+
+    all_annotations: list[pd.DataFrame] = []
+    for cfg in DEMO_RECORDINGS:
+        patient_id = cfg[AnnotationSchema.PATIENT_ID]
+        print(f"[{patient_id}] Simulating 12-lead ECG at {cfg['heart_rate']} bpm …")
+
+        ecg = simulate_12lead(cfg["heart_rate"], seed=cfg["seed"])
+
+        dcm = to_dicom(ecg, patient_id=patient_id, date=cfg["date"], time=cfg["time"])
+        dcm_path = out / f"{patient_id}.dcm"
+        dcm.save_as(str(dcm_path))
+        print(f"  Saved {dcm_path}")
+
+        p_waves = detect_p_waves(ecg)
+        print(f"  Detected {len(p_waves)} P-waves")
+
+        anns = build_annotations(
+            f"{patient_id}.dcm", patient_id, p_waves, seed=cfg["seed"]
+        )
+        all_annotations.append(anns)
+
+    combined = pd.concat(all_annotations, ignore_index=True)
+    ann_path = out / "demo_annotations.csv"
+    combined.to_csv(ann_path, index=False)
+    print(
+        f"\nSaved {ann_path} "
+        f"({len(combined)} rows across {combined[AnnotationSchema.PATIENT_ID].nunique()} patients)"
+    )
+    return out