Source code for neurokit2.ecg.ecg_intervalrelated

# -*- coding: utf-8 -*-
import numpy as np
import pandas as pd

from ..hrv import hrv




def ecg_intervalrelated(data, sampling_rate=1000):
    """**Interval-related analysis of ECG**

    Performs ECG analysis on longer periods of data (typically > 10 seconds), such as resting-state
    data.

    Parameters
    ----------
    data : Union[dict, pd.DataFrame]
        A DataFrame containing the different processed signal(s) as different columns, typically
        generated by ``ecg_process()`` or ``bio_process()``. Can also take a dict containing sets of
        separately processed DataFrames.
    sampling_rate : int
        The sampling frequency of the signal (in Hz, i.e., samples/second). Defaults to 1000Hz.

    Returns
    -------
    DataFrame
        A dataframe containing the analyzed ECG features. The analyzed features consist of the following:

        .. codebookadd::
            ECG_Rate_Mean|The mean heart rate.

        * ``ECG_HRV``: the different heart rate variability metrices.

        See :func:`.hrv_summary()` docstrings for details.

    See Also
    --------
    .bio_process, ecg_eventrelated

    Examples
    ----------
    .. ipython:: python

      import neurokit2 as nk

      # Download data
      data = nk.data("bio_resting_5min_100hz")

      # Process the data
      df, info = nk.ecg_process(data["ECG"], sampling_rate=100)

      # Single dataframe is passed
      nk.ecg_intervalrelated(df, sampling_rate=100)

      epochs = nk.epochs_create(df, events=[0, 15000], sampling_rate=100, epochs_end=150)
      nk.ecg_intervalrelated(epochs)

    """
    intervals = {}

    # Format input
    if isinstance(data, pd.DataFrame):
        rate_cols = [col for col in data.columns if "ECG_Rate" in col]
        if len(rate_cols) == 1:
            intervals.update(_ecg_intervalrelated_formatinput(data))
            intervals.update(_ecg_intervalrelated_hrv(data, sampling_rate))
        else:
            raise ValueError(
                "NeuroKit error: ecg_intervalrelated(): Wrong input,"
                "we couldn't extract heart rate. Please make sure"
                "your DataFrame contains an `ECG_Rate` column."
            )
        ecg_intervals = pd.DataFrame.from_dict(intervals, orient="index").T

    elif isinstance(data, dict):
        for index in data:
            intervals[index] = {}  # Initialize empty container

            # Add label info
            intervals[index]["Label"] = data[index]["Label"].iloc[0]

            # Rate
            intervals[index] = _ecg_intervalrelated_formatinput(data[index], intervals[index])

            # HRV
            intervals[index] = _ecg_intervalrelated_hrv(
                data[index], sampling_rate, intervals[index]
            )

        ecg_intervals = pd.DataFrame.from_dict(intervals, orient="index")

    return ecg_intervals



# =============================================================================
# Internals
# =============================================================================


def _ecg_intervalrelated_formatinput(data, output={}):

    # Sanitize input
    colnames = data.columns.values
    if len([i for i in colnames if "ECG_Rate" in i]) == 0:
        raise ValueError(
            "NeuroKit error: ecg_intervalrelated(): Wrong input,"
            "we couldn't extract heart rate. Please make sure"
            "your DataFrame contains an `ECG_Rate` column."
        )
    signal = data["ECG_Rate"].values
    output["ECG_Rate_Mean"] = np.mean(signal)

    return output


def _ecg_intervalrelated_hrv(data, sampling_rate, output={}):

    # Sanitize input
    colnames = data.columns.values
    if len([i for i in colnames if "ECG_R_Peaks" in i]) == 0:
        raise ValueError(
            "NeuroKit error: ecg_intervalrelated(): Wrong input,"
            "we couldn't extract R-peaks. Please make sure"
            "your DataFrame contains an `ECG_R_Peaks` column."
        )

    # Transform rpeaks from "signal" format to "info" format.
    rpeaks = np.where(data["ECG_R_Peaks"].values)[0]
    rpeaks = {"ECG_R_Peaks": rpeaks}

    results = hrv(rpeaks, sampling_rate=sampling_rate)
    for column in results.columns:
        # Add and convert to float
        output[column] = results[[column]].values

    return output