<filename>src/main.py
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
__author__ = "<NAME>"
__version__ = "0.2.0"
__license__ = "MIT"
import pyxdf  # Open xdf files
import argparse  # Read arguments from command line
import numpy as np  # Numeric library for python
import pandas as pd  # Tabular DataFrames
from scipy import signal  # Digital Signal Processing
from biosppy.signals import ecg as ecg_analysis  # ECG algorithms
from scipy.stats import mannwhitneyu as utest  # Mann-Whitney test


def bandpass_filter(data, lowcut, highcut, fs, order=4):
    nyq = 0.5 * fs
    low = lowcut / nyq
    high = highcut / nyq
    b, a = signal.butter(order, [low, high], btype='band')
    y = signal.lfilter(b, a, data)
    return y


def notch_filter(data, fs, f0=50, Q=30):
    b, a = signal.iirnotch(f0, Q, fs)
    y = signal.lfilter(b, a, data)
    return y


def detrend_filter(data):
    return data - data.mean()


def rolling_std(vec, window_size):
    """Rolling window for SDNN"""
    shape = vec.shape[:-1] + (vec.shape[-1] - window_size + 1, window_size)
    strides = vec.strides + (vec.strides[-1],)
    windows = np.lib.stride_tricks.as_strided(vec,
                                              shape=shape,
                                              strides=strides)
    rstd = np.hstack((np.zeros(window_size // 2) + np.nan,
                      np.std(windows, 1),
                      np.zeros((window_size // 2) - 1) + np.nan))
    return rstd


def main(args):
    file_path = args.file_path
    data, header = pyxdf.load_xdf(file_path)

    # Read the biosignals from the first stream
    stream = data[0]
    y = stream['time_series']
    fs = int(float(stream["info"]["nominal_srate"][0]))  # Store sample frequency (fs)
    # Store two biosignals, and a reference -1 value
    ecg, gsr, _ = np.split(y, [1, 2], axis=1)
    t = stream["time_stamps"]  # Extract time array

    # Read the markers from the second stream
    stream = data[1]
    mark_l = []
    mark_m = []
    mark_h = []
    for marker, ts in zip(stream["time_series"], stream["time_stamps"]):
        if marker[0].startswith("Low"):
            mark_l.append(ts)
        elif marker[0].startswith("Mid"):
            mark_m.append(ts)
        elif marker[0].startswith("High"):
            mark_h.append(ts)

    # load raw ECG signal
    out = ecg_analysis.ecg(signal=ecg.flatten(), sampling_rate=fs, show=False)

    # HRV
    rr_intervals = np.diff(out["ts"][out["rpeaks"]], prepend=[np.nan])
    rr_intervals *= 1000  # Convert from s to ms
    n = 8  # Must be even number for now
    hrv = rolling_std(rr_intervals, n)

    # Place markers for gsr
    markers_lb = np.array(data[1]["time_series"])
    markers_ts = np.array(data[1]["time_stamps"])
    # Select the closest sample point to each marker timestamp. Error +/- 2* fs
    markers_ts = np.array([min(t, key=lambda x:abs(x-i)) for i in markers_ts])
    markers_idx = np.array([np.where(t == i)[0] for i in markers_ts])
    assert len(markers_idx) == len(markers_lb)

    # Create dataframe for statistical analysis
    label = markers_lb.squeeze()
    gsr_vals = gsr[markers_idx].squeeze()
    table = np.vstack([label, gsr_vals])
    df = pd.DataFrame(table.transpose(), columns=["Marker", "GSR"])

    # Place markers for HR and HRV
    t = (out["ts"][out["rpeaks"]] + stream["time_stamps"].min())[1:]
    markers_lb = np.array(data[1]["time_series"])
    markers_ts = np.array(data[1]["time_stamps"])
    # Select the closest sample point to each marker timestamp. Error +/- 2* fs
    markers_ts = np.array([min(t, key=lambda x:abs(x-i)) for i in markers_ts])
    markers_idx = np.array([np.where(t == i)[0] for i in markers_ts])
    assert len(markers_idx) == len(markers_lb)
    markers_idx = markers_idx.squeeze()
    df["HR"] = out["heart_rate"][markers_idx]
    df["HRV"] = hrv[1:][markers_idx]

    # Create condition column from markers
    l = ["Low" if i else "" for i in df["Marker"].str.startswith('Low', na=False)]
    m = ["Mid" if i else "" for i in df["Marker"].str.startswith('Mid', na=False)]
    h = ["High" if i else "" for i in df["Marker"].str.startswith('High', na=False)]
    cond = np.vstack([l,m,h]).transpose()
    cond = ["".join(i) for i in cond]
    df["Condition"] = cond

    # Statistical Analysis
    print("Results:")
    for var in ["GSR", "HR", "HRV"]:
        print(var)
        # Three independant tests
        for (a, b) in [("Low", "Mid"), ("Low", "High"), ("Mid", "High")]:
            i = df[df['Condition'] == a][var].to_numpy()
            j = df[df['Condition'] == b][var].to_numpy()
            i = [float(w) for w in i]
            j = [float(w) for w in j]
            u = utest(i, j)
            if u[1] < 0.05:
                print(f"Test {a}-{b}: {u[0]} p={u[1]:.4f} *")
            else:
                print(f"Test {a}-{b}: {u[0]} p={u[1]:.4f}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("file_path", help="XDF file to be processed")
    parser.add_argument(
        "--version",
        action="version",
        version="%(prog)s (version {version})".format(version=__version__))
    args = parser.parse_args()
    main(args)
