method_plots.py

import ipdb
import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl
import seaborn as sns
from joblib import dump, load
import time
from tqdm import tqdm
from sys import platform
import glob
import re
import pandas as pd
import json

from tsfresh import extract_features
from tsfresh.feature_selection import relevance as tsfresh_relevance
from tsfresh.feature_extraction import settings as tsfresh_settings
from sklearn.decomposition import PCA, FastICA
from sklearn.preprocessing import StandardScaler

from datetime import datetime, timedelta
from os import makedirs, listdir
from os.path import isdir, getsize
from os.path import join as path_join
from os.path import exists as path_exists
from pprint import PrettyPrinter

from multiprocessing import Pool, cpu_count
from functools import partial
from ast import literal_eval

from modules.animationplotter import AnimationPlotter, AnimationPlotter2D
from modules.digitalsignalprocessing import *
from modules.datapipeline import SubjectData, datetime_to_ms\
        ,ms_to_datetime, DataSynchronizer
from modules.datapipeline import get_file_list, load_files_conditions
from modules.datapipeline import get_windowed_data
from modules.evaluations import Evaluation
from config import DEBUG, NROWS, MARKER_FS, BR_FS, N_MARKERS , WINDOW_SIZE \
        ,WINDOW_SHIFT, ACC_THOLD, FS_RESAMPLE, MIN_RESP_RATE, MAX_RESP_RATE \
        ,TRAIN_VAL_TEST_SPLIT, IMU_FS, ACC_FS, DPI, FIG_FMT, DATA_DIR

''' We're looking only at meditation and rest '''

WINDOW_SIZE = 15 # seconds
WINDOW_SHIFT = 1 # seconds
# mpl.rc('font', serif='Times New Roman')

sns.set_theme(style='ticks')
# sns.axes_style('white')

def marker_main(subject='S13', condition='R', sens='imu_gyr', method='roddiger'):
    window_size = WINDOW_SIZE
    window_shift = WINDOW_SHIFT
    fs = MARKER_FS
    marker_sel = 0
    if sens == 'marker':
        fs = MARKER_FS
        title = 'Marker Raw'
    elif 'imu' in sens:
        fs = IMU_FS
        if 'acc' in sens: title = 'Head Worn Accelerometer Raw'
        if 'gyr' in sens: title = 'Head Worn Gyroscope Raw'
    elif sens=='h_acc':
        if int(subject[-2:]) < 27: fs=BR_FS
        else: fs=ACC_FS
        title = 'Harness Accelerometer Raw'

    # set glob based on sensor type and subject and condition
    if 'imu' in sens: data_glob = f"*{condition}_imu_df"
    elif sens == 'h_acc': data_glob = f"*{condition}_accel_df"
    else: data_glob = f"*{condition}_{sens}_df"
    marker_glob = f"*{condition}_marker_df"

    data_list = get_file_list(data_glob, sbj=subject)
    lbl_list = get_file_list(f'*{condition}_summary_df', sbj=subject)
    pss_list = get_file_list(f'*{condition}_pressure_df', sbj=subject)
    marker_flist = get_file_list(marker_glob, sbj=subject)

    # load the files in the glob
    marker_df = load_files_conditions(marker_flist, skip_ratio=0.0)
    data_df = load_files_conditions(data_list, skip_ratio=0.0)
    lbl_df = load_files_conditions(lbl_list, skip_ratio=0.0)
    pss_df = load_files_conditions(pss_list, skip_ratio=0.0)

    # bin to windows
    x_time = data_df['ms'].values
    marker_time = marker_df['ms'].values
    vsw = vectorized_slide_win

    data_inds = np.arange(0, len(x_time))
    wins = vsw(data_inds, len(data_inds),
                    sub_window_size=window_size*fs,
                    stride_size=window_shift*fs)

    is_marker = False
    # if sens == 'marker':
    #     is_marker = True
    if sens == 'imu_acc':
        data = np.array(data_df['accelerometer'].map(literal_eval).tolist())
    elif sens == 'imu_gyr':
        data = np.array(data_df['gyroscope'].map(literal_eval).tolist())
    elif sens == 'h_acc':
        sbj_data = SubjectData(condition=condition, subject=subject)
        sbj_data.accel_df = data_df
        data = sbj_data.get_accel_data()

    data_wins = get_windowed_data(x_time, data, wins)
    data_list = [win for win in data_wins]

    br_col = [col for col in pss_df.columns.values if\
              'breathing' in col.lower()][0]
    pss_dsp = pressure_signal_processing
    pss_data = pss_df[br_col].values
    pss_data = pss_dsp(pss_data)
    pss_data = np.interp(x_time, pss_df['ms'].values, pss_data)

    pss_data_wins = get_windowed_data(x_time, pss_data, wins)
    pss_list = [win for win in pss_data_wins]

    # Sync marker data
    sbj_data = SubjectData(condition=condition, subject=subject)
    sbj_data.marker_df = marker_df
    marker_data = sbj_data.get_marker_data()
    marker_data = marker_data[:, marker_sel, :]
    tmp = []
    for i in range(3):
        tmp.append(np.interp(x_time, marker_df['ms'].values,
                             marker_data[:,i]))
    marker_data = np.array(tmp).T
    marker_data = second_order_diff(marker_data, MARKER_FS)
    marker_wins = get_windowed_data(x_time, marker_data, wins)
    marker_list = [win for win in marker_wins]

    # Perform rejection
    std_scaler = StandardScaler()
    rejects = [i for i, win in enumerate(data_list) if
               reject_artefact((win-np.mean(win,axis=0))/np.std(win,axis=0))]
    inds_to_keep = np.delete(np.arange(0, len(data_list)), rejects)

    marker_list = [marker_list[i] for i in inds_to_keep]
    data_list = [data_list[i] for i in inds_to_keep]
    pss_list = [pss_list[i] for i in inds_to_keep]

    # plot the raw signal of a select window
    fig, axs = plt.subplots(3,2, figsize=(12, 6.5), dpi=DPI)
    mpl.rcParams['axes.titleweight'] = 'bold'

    axs = axs.flatten()
    win_ind = 22
    if len(data_list) < win_ind:
        data_win = data_list[int(len(data_list)//2)]
        marker_win = marker_list[int(len(data_list)//2)]
    else:
        data_win = data_list[win_ind]
        marker_win = marker_list[win_ind]

    fontsize = 10
    piter = 0.03
    x_win = np.linspace(0, window_size, len(data_win))
    axs[0].plot(x_win, StandardScaler().fit_transform(data_win), linewidth=1)
    axs[0].set_title(title, fontsize=fontsize)
    axs[0].set_xticklabels([])
    # axs[0].set_xlabel('Time (s)')
    if sens == 'marker' or 'acc' in sens:
        axs[0].set_ylabel('Acceleration (m/s${^2}$)', fontsize=fontsize)
    else:
        axs[0].set_ylabel('Angular Velocity (deg/s)', fontsize=fontsize)
    axs[0].legend(['X','Y', 'Z'], ncol=3, loc='lower left', fontsize=fontsize)

    x_win = np.linspace(0, window_size, len(marker_win))
    axs[1].plot(x_win, StandardScaler().fit_transform(marker_win), linewidth=1)
    axs[1].set_title(f'Marker {marker_sel+1} Acceleration Raw', fontsize=fontsize)
    axs[1].set_xticklabels([])
    # axs[1].set_xlabel('Time (s)')
    axs[1].set_ylabel('Acceleration (m/s${^2}$)', fontsize=fontsize)
    axs[1].legend(['X','Y', 'Z'], ncol=3, loc='lower left', fontsize=fontsize)

    # plot the processed signal with the pressure signal
    if method == 'roddiger': sig_dsp = roddiger_sp
    elif method == 'hernandez': sig_dsp = hernandez_sp

    interp, smth = sig_dsp(data=data_win, fs=fs, is_marker=False)
    marker_interp, marker_smth = sig_dsp(data=marker_win, fs=MARKER_FS,
                                         is_marker=False)

    pca = PCA(n_components=1, random_state=3)
    pca_out = StandardScaler().fit_transform(pca.fit_transform(smth))
    marker_out = StandardScaler().fit_transform(pca.fit_transform(marker_smth))
    pss_out = StandardScaler().fit_transform(pss_list[win_ind].reshape(-1, 1))\
            .squeeze()

    x_smth = np.linspace(0, window_size, len(pca_out))
    x_pss = np.linspace(0, window_size, len(pss_list[win_ind]))
    axs[2].plot(x_pss, pss_out)
    axs[2].plot(x_smth, pca_out, color='tab:red')
    axs[2].set_ylabel("Normalised Units", fontsize=fontsize)
    axs[2].set_xticks(np.arange(0, WINDOW_SIZE, 2))
    axs[2].set_yticks(np.arange(-4, 4+2, 2))
    axs[2].set_xlabel('Time (s)', fontsize=fontsize)
    if 'acc' in sens:
        axs[2].set_title('Accelerometer Signal vs Pressure', fontsize=fontsize)

        axs[2].legend(['PSS', 'ACC'], ncol=3, loc='lower left',
                      fontsize=fontsize)
    else:
        axs[2].set_title('Gyroscope Signal vs Pressure', fontsize=fontsize)
        axs[2].legend(['PSS', 'GYR'], ncol=3, loc='lower left',
                      fontsize=fontsize)

    x_marker = np.linspace(0, window_size, len(marker_out))
    axs[3].plot(x_pss, pss_out)
    axs[3].plot(x_marker, marker_out, color='tab:green')
    axs[3].set_xticks(np.arange(0, WINDOW_SIZE, 2))
    axs[3].set_yticks(np.arange(-4, 4+2, 2))
    axs[3].set_xlabel('Time (s)', fontsize=fontsize)
    axs[3].set_title(f'Marker {marker_sel+1} Signal vs Pressure', fontsize=fontsize)
    axs[3].legend(['PSS', 'MKR'], ncol=3, loc='lower left', fontsize=fontsize)

    for a in axs[:4]:
        a.tick_params(labelsize=fontsize)
        box = a.get_position()
        box.y0 += piter
        box.y1 += piter
        a.set_position(box)

    # plot the FFT w/o padding
    pad_len = npads_frequency_resolution(len(pca_out), fs=FS_RESAMPLE)
    data_pad = np.pad(pca_out.squeeze(), (0, pad_len), 'constant',
                      constant_values=0)
    data_xf, data_yf = run_fft(data_pad, FS_RESAMPLE)
    data_ind = data_yf.argmax(axis=0)

    pad_len = npads_frequency_resolution(len(pss_out), fs=fs)
    pss_pad = np.pad(pss_out, (0, pad_len), 'constant', constant_values=0)
    pss_xf, pss_yf = run_fft(pss_pad, fs)
    pss_ind = pss_yf.argmax(axis=0)

    pad_len = npads_frequency_resolution(len(marker_out), fs=FS_RESAMPLE)
    marker_pad = np.pad(marker_out.squeeze(), (0, pad_len), 'constant',
                      constant_values=0)
    marker_xf, marker_yf = run_fft(marker_pad, FS_RESAMPLE)
    marker_ind = marker_yf.argmax(axis=0)

    fig.delaxes(axs[-1])
    fig.delaxes(axs[-2])
    axs = np.delete(axs, [-2,-1])
    axs = np.insert(axs, len(axs), fig.add_subplot(3, 3, 7))
    axs = np.insert(axs, len(axs), fig.add_subplot(3, 3, 8))
    axs = np.insert(axs, len(axs), fig.add_subplot(3, 3, 9))

    ymax = 0.5

    axs[-3].plot(data_xf, data_yf, color='tab:red')
    axs[-3].plot(data_xf[data_ind], data_yf[data_ind], 'rx')
    axs[-3].set_xticks(np.arange(0, 1, 0.2))
    axs[-3].set_title('GYR Spectrum', fontsize=fontsize)
    axs[-3].set_xlim([0, 1])
    axs[-3].set_ylim([0, ymax])
    axs[-3].set_xlabel('Frequency (Hz)', fontsize=fontsize)
    axs[-3].set_ylabel('$|Y(f)|$', fontsize=fontsize)

    axs[-2].plot(pss_xf,  pss_yf)
    axs[-2].plot(pss_xf[pss_ind], pss_yf[pss_ind], 'rx')
    axs[-2].set_xticks(np.arange(0, 1, 0.2))
    axs[-2].set_title('PSS Spectrum', fontsize=fontsize)
    axs[-2].set_xlim([0, 1])
    axs[-2].set_ylim([0, ymax])
    axs[-2].set_xlabel('Frequency (Hz)', fontsize=fontsize)

    axs[-1].plot(marker_xf,  marker_yf, color='tab:green')
    axs[-1].plot(marker_xf[marker_ind], marker_yf[marker_ind], 'rx')
    axs[-1].set_xticks(np.arange(0, 1, 0.2))
    axs[-1].set_title(f'MKR{marker_sel+1} Spectrum', fontsize=fontsize)
    axs[-1].set_xlabel('Frequency (Hz)', fontsize=fontsize)
    axs[-1].set_xlim([0, 1])
    axs[-1].set_ylim([0, ymax])
    tmp = axs[-3:]
    for a in tmp:
        box = a.get_position()
        box.y0 -= piter
        box.y1 -= piter
        a.set_position(box)

    fmt = FIG_FMT
    fig.savefig(f"./graphs/methods/{method}_methods.{fmt}")

def main(subject='S14', condition='R', sens='imu_gyr', method='roddiger'):
    window_size = WINDOW_SIZE
    window_shift = WINDOW_SHIFT
    fs = MARKER_FS
    marker_sel = 6
    if sens == 'marker':
        fs = MARKER_FS
        title = f'Marker {marker_sel+1} Raw'
    elif 'imu' in sens:
        fs = IMU_FS
        if 'acc' in sens: title = 'Head Worn Accelerometer Raw'
        if 'gyr' in sens: title = 'Head Worn Gyroscope Raw'
    elif sens=='h_acc':
        if int(subject[-2:]) < 27: fs=BR_FS
        else: fs=ACC_FS
        title = 'Harness Accelerometer Raw'

    # set glob based on sensor type and subject and condition
    if 'imu' in sens: data_glob = f"*{condition}_imu_df"
    elif sens == 'h_acc': data_glob = f"*{condition}_accel_df"
    else: data_glob = f"*{condition}_{sens}_df"

    data_list = get_file_list(data_glob, sbj=subject)
    lbl_list = get_file_list(f'*{condition}_summary_df', sbj=subject)
    pss_list = get_file_list(f'*{condition}_pressure_df', sbj=subject)

    # load the files in the glob
    data_df = load_files_conditions(data_list, skip_ratio=0.0)
    lbl_df = load_files_conditions(lbl_list, skip_ratio=0.0)
    pss_df = load_files_conditions(pss_list, skip_ratio=0.0)

    # bin to windows
    x_time = data_df['ms'].values
    vsw = vectorized_slide_win

    data_inds = np.arange(0, len(x_time))
    wins = vsw(data_inds, len(data_inds),
                    sub_window_size=window_size*fs,
                    stride_size=window_shift*fs)

    is_marker = False
    if sens == 'marker':
        sbj_data = SubjectData(condition=condition, subject=subject)
        sbj_data.marker_df = data_df
        data = sbj_data.get_marker_data()
        data = second_order_diff(data, fs)
        is_marker = True
    elif sens == 'imu_acc':
        data = np.array(data_df['accelerometer'].map(literal_eval).tolist())
    elif sens == 'imu_gyr':
        data = np.array(data_df['gyroscope'].map(literal_eval).tolist())
    elif sens == 'h_acc':
        sbj_data = SubjectData(condition=condition, subject=subject)
        sbj_data.accel_df = data_df
        data = sbj_data.get_accel_data()

    data_wins = get_windowed_data(x_time, data, wins)
    data_list = [win for win in data_wins]

    br_col = [col for col in pss_df.columns.values if\
              'breathing' in col.lower()][0]
    pss_dsp = pressure_signal_processing
    pss_data = pss_df[br_col].values
    pss_data = pss_dsp(pss_data)
    pss_data = np.interp(x_time, pss_df['ms'].values, pss_data)

    pss_data_wins = get_windowed_data(x_time, pss_data, wins)
    pss_list = [win for win in pss_data_wins]

    # Perform rejection
    std_scaler = StandardScaler()
    rejects = [i for i, win in enumerate(data_list) if
               reject_artefact((win-np.mean(win,axis=0))/np.std(win,axis=0))]
    inds_to_keep = np.delete(np.arange(0, len(data_list)), rejects)
    data_list = [data_list[i] for i in inds_to_keep]
    pss_list = [pss_list[i] for i in inds_to_keep]

    # plot the raw signal of a select window
    fig, axs = plt.subplots(3,1)
    win_ind = 60
    if len(data_list) < win_ind:
        data_win = data_list[int(len(data_list)//2)]
    else:
        data_win = data_list[win_ind]

    if is_marker:
        data_win = data_win[:, marker_sel, :]
    x_win = np.linspace(0, window_size, len(data_win))
    axs[0].plot(x_win, StandardScaler().fit_transform(data_win), linewidth=1)
    axs[0].set_title(title)
    axs[0].set_xticklabels([])
    axs[0].set_xlabel('Time (s)')
    if sens == 'marker' or 'acc' in sens:
        axs[0].set_ylabel('Acceleration (m/s${^2}$)')
    else:
        axs[0].set_ylabel('Angular Velocity (deg/s)')

    axs[0].legend(['X','Y', 'Z'], ncol=3, loc='lower left')

    # plot the processed signal with the pressure signal
    if method == 'roddiger': sig_dsp = roddiger_sp
    elif method == 'hernandez': sig_dsp = hernandez_sp

    interp, smth = sig_dsp(data=data_win, fs=fs, is_marker=False)

    pca = PCA(n_components=1, random_state=3)
    pca_out = StandardScaler().fit_transform(pca.fit_transform(smth))

    pss_out = StandardScaler().fit_transform(pss_list[win_ind].reshape(-1, 1))\
            .squeeze()

    x_smth = np.linspace(0, window_size, len(pca_out))
    x_pss = np.linspace(0, window_size, len(pss_list[win_ind]))
    axs[1].plot(x_pss, pss_out)
    axs[1].plot(x_smth, pca_out, color='tab:red')
    axs[1].set_ylabel("Normalised Units")
    axs[1].set_xticks(np.arange(0, WINDOW_SIZE, 2))
    axs[1].set_xlabel('Time (s)')
    axs[1].set_title('Sensor Signal vs Pressure')
    if sens == 'marker' or 'acc' in sens:
        axs[1].legend(['PSS', 'ACC'], ncol=3, loc='lower left')
    else:
        axs[1].legend(['PSS', 'GYR'], ncol=3, loc='lower left')

    # plot the FFT w/o padding
    pad_len = npads_frequency_resolution(len(pca_out), fs=FS_RESAMPLE)
    data_pad = np.pad(pca_out.squeeze(), (0, pad_len), 'constant', constant_values=0)
    data_xf, data_yf = run_fft(data_pad, FS_RESAMPLE)
    data_ind = data_yf.argmax(axis=0)

    pad_len = npads_frequency_resolution(len(pss_out), fs=fs)
    pss_pad = np.pad(pss_out, (0, pad_len), 'constant', constant_values=0)
    pss_xf, pss_yf = run_fft(pss_pad, fs)
    pss_ind = pss_yf.argmax(axis=0)

    fig.delaxes(axs[-1])
    axs = np.insert(axs, -1, fig.add_subplot(3, 2, 5))
    axs = np.insert(axs, -1, fig.add_subplot(3, 2, 6))

    axs[2].plot(data_xf, data_yf)
    axs[2].plot(data_xf[data_ind], data_yf[data_ind], 'rx')
    axs[2].set_xticks(np.arange(0, 1, 0.2))
    axs[2].set_xlim([0, 1])
    axs[2].set_ylim([0, 0.4])
    axs[2].set_xlabel('Frequency (Hz)')
    axs[2].set_ylabel('$|Y(f)|$')

    axs[3].plot(pss_xf,  pss_yf)
    axs[3].plot(pss_xf[pss_ind], pss_yf[pss_ind], 'rx')
    axs[3].set_xticks(np.arange(0, 1, 0.2))
    axs[3].set_xlim([0, 1])
    axs[3].set_ylim([0, 0.4])
    axs[3].set_xlabel('Frequency (Hz)')
    axs[3].set_ylabel('$|Y(f)|$')

    fig.set_size_inches((5.5, 5))
    plt.tight_layout()
    plt.show()

if __name__ == '__main__':
    # Make sure to run sync_data before any of these operations
    marker_main(method='hernandez')
    marker_main(method='roddiger')
    plt.show()