FeatureEstimation.py

# -*- coding: utf-8 -*-
"""
Updated Oct 18 2022

@author: Qianliang Li (glia@dtu.dk)

This script contains the code to estimate the following EEG features:
    1. Power Spectral Density
    2. Frontal Theta/Beta Ratio
    3. Asymmetry
    4. Peak Alpha Frequency
    5. 1/f Exponents
    6. Microstates
    7. Long-Range Temporal Correlation (DFA Exponent)
Source localization and functional connectivity
    8. Imaginary part of Coherence
    9. Weighted Phase Lag Index
    10. (Orthogonalized) Power Envelope Correlations
    11. Granger Causality

It should be run after Preprocessing.py

All features are saved in pandas DataFrame format for the machine learning
pipeline

Note that the code has not been changed to fit the demonstration dataset,
thus just running it might introduce some errors.
The code is provided to show how we performed the feature estimation
and if you are adapting the code, you should check if it fits your dataset
e.g. the questionnaire data, sensors and source parcellation etc

The script was written in Spyder. The outline panel can be used to navigate
the different parts easier (Default shortcut: Ctrl + Shift + O)
"""

# Set working directory
import os
wkdir = "/home/glia/EEG"
os.chdir(wkdir)

# Load all libraries from the Preamble
from Preamble import *

# %% Load preprocessed epochs and questionnaire data
load_path = "./PreprocessedData"

# Get filenames
files = []
for r, d, f in os.walk(load_path):
    for file in f:
        if ".fif" in file:
            files.append(os.path.join(r, file))
files.sort()

# Subject IDs
Subject_id = [0] * len(files)
for i in range(len(files)):
    temp = files[i].split("\\")
    temp = temp[-1].split("_")
    Subject_id[i] = int(temp[0])

n_subjects = len(Subject_id)

# Load Final epochs
final_epochs = [0]*n_subjects
for n in range(n_subjects):
    final_epochs[n] = mne.read_epochs(fname = os.path.join(files[n]),
                    verbose=0)

# Load dropped epochs - used for gap idx in microstates
bad_subjects = [12345] # list with subjects that were excluded due to too many dropped epochs/chs
Drop_epochs_df = pd.read_pickle("./Preprocessing/dropped_epochs.pkl")

good_subject_df_idx = [not i in bad_subjects for i in Drop_epochs_df["Subject_ID"]]
Drop_epochs_df = Drop_epochs_df.loc[good_subject_df_idx,:]
Drop_epochs_df = Drop_epochs_df.sort_values(by=["Subject_ID"]).reset_index(drop=True)

### Load questionnaire data
# For the purposes of this demonstration I will make a dummy dataframe
# The original code imported csv files with questionnaire data and group status
final_qdf = pd.DataFrame({"Subject_ID":Subject_id,
                          "Age":[23,26],
                          "Gender":[0,0],
                          "Group_status":[0,1],
                          "PCL_total":[33,56],
                          "Q1":[1.2, 2.3],
                          "Q2":[1.7, 1.5],
                          "Qn":[2.1,1.0]})

# Define cases as >= 44 total PCL
# Type: numpy array with subject id
cases = np.array(final_qdf["Subject_ID"][final_qdf["PCL_total"]>=44])
n_groups = 2
Groups = ["CTRL", "PTSD"]

# Define folder for saving features
Feature_savepath = "./Features/"
Stat_savepath = "./Statistics/"
Model_savepath = "./Model/"

# %% Power spectrum features
Freq_Bands = {"delta": [1.25, 4.0],
              "theta": [4.0, 8.0],
              "alpha": [8.0, 13.0],
              "beta": [13.0, 30.0],
              "gamma": [30.0, 49.0]}
ch_names = final_epochs[0].info["ch_names"]
n_channels = final_epochs[0].info["nchan"]

# Pre-allocate memory
power_bands = [0]*len(final_epochs)

def power_band_estimation(n):
    # Get index for eyes open and eyes closed
    EC_index = final_epochs[n].events[:,2] == 1
    EO_index = final_epochs[n].events[:,2] == 2
    
    # Calculate the power spectral density
    psds, freqs = psd_multitaper(final_epochs[n], fmin = 1, fmax = 50) # output (epochs, channels, freqs)
    
    temp_power_band = []
    
    for fmin, fmax in Freq_Bands.values():
        # Calculate the power each frequency band
        psds_band = psds[:, :, (freqs >= fmin) & (freqs < fmax)].sum(axis=-1)
        # Calculate the mean for each eye status
        psds_band_eye = np.array([psds_band[EC_index,:].mean(axis=0),
                                      psds_band[EO_index,:].mean(axis=0)])
        # Append for each freq band
        temp_power_band.append(psds_band_eye)
        # Output: List with the 5 bands, and each element is a 2D array with eye status as 1st dimension and channels as 2nd dimension
    
    # The list is reshaped and absolute and relative power are calculated
    abs_power_band = np.reshape(temp_power_band, (5, 2, n_channels))
    abs_power_band = 10.*np.log10(abs_power_band) # Convert to decibel scale
    
    rel_power_band = np.reshape(temp_power_band, (5, 2, n_channels))
    rel_power_band = rel_power_band/np.sum(rel_power_band, axis=0, keepdims=True)
    # each eye condition and channel have been normalized to power in all 5 frequencies for that given eye condition and channel
    
    # Make one list in 1 dimension
    abs_power_values = np.concatenate(np.concatenate(abs_power_band, axis=0), axis=0)
    rel_power_values = np.concatenate(np.concatenate(rel_power_band, axis=0), axis=0)
    ## Output: First the channels, then the eye status and finally the frequency bands are concatenated
    ## E.g. element 26 is 3rd channel, eyes open, first frequency band
    #assert abs_power_values[26] == abs_power_band[0,1,2]
    #assert abs_power_values[47] == abs_power_band[0,1,23] # +21 channels to last
    #assert abs_power_values[50] == abs_power_band[1,0,2] # once all channels have been changed then the freq is changed and eye status
    
    # Get result
    res = np.concatenate([abs_power_values,rel_power_values],axis=0)
    return n, res

with concurrent.futures.ProcessPoolExecutor() as executor:
    for n, result in executor.map(power_band_estimation, range(len(final_epochs))): # Function and arguments
        power_bands[n] = result

# Combine all data into one dataframe
# First the columns are prepared
n_subjects = len(Subject_id)

# The group status (PTSD/CTRL) is made using the information about the cases
Group_status = np.array(["CTRL"]*n_subjects)
Group_status[np.array([i in cases for i in Subject_id])] = "PTSD"

# A variable that codes the channels based on A/P localization is also made
Frontal_chs = ["Fp1", "Fpz", "Fp2", "AFz", "Fz", "F3", "F4", "F7", "F8"]
Central_chs = ["Cz", "C3", "C4", "T7", "T8", "FT7", "FC3", "FCz", "FC4", "FT8", "TP7", "CP3", "CPz", "CP4", "TP8"]
Posterior_chs = ["Pz", "P3", "P4", "P7", "P8", "POz", "O1", "O2", "Oz"]

Brain_region_labels = ["Frontal","Central","Posterior"]
Brain_region = np.array(ch_names, dtype = "<U9")
Brain_region[np.array([i in Frontal_chs for i in ch_names])] = Brain_region_labels[0]
Brain_region[np.array([i in Central_chs for i in ch_names])] = Brain_region_labels[1]
Brain_region[np.array([i in Posterior_chs for i in ch_names])] = Brain_region_labels[2]

# A variable that codes the channels based on M/L localization
Left_chs = ["Fp1", "F3", "F7", "FC3", "FT7", "C3", "T7", "CP3", "TP7", "P3", "P7", "O1"]
Right_chs = ["Fp2", "F4", "F8", "FC4", "FT8", "C4", "T8", "CP4", "TP8", "P4", "P8", "O2"]
Mid_chs = ["Fpz", "AFz", "Fz", "FCz", "Cz", "CPz", "Pz", "POz", "Oz"]

Brain_side = np.array(ch_names, dtype = "<U5")
Brain_side[np.array([i in Left_chs for i in ch_names])] = "Left"
Brain_side[np.array([i in Right_chs for i in ch_names])] = "Right"
Brain_side[np.array([i in Mid_chs for i in ch_names])] = "Mid"

# Eye status is added
eye_status = list(final_epochs[0].event_id.keys())
n_eye_status = len(eye_status)

# Frequency bands
freq_bands_name = list(Freq_Bands.keys())
n_freq_bands = len(freq_bands_name)

# Quantification (Abs/Rel)
quant_status = ["Absolute", "Relative"]
n_quant_status = len(quant_status)

# The dataframe is made by combining all the unlisted pds values
# Each row correspond to a different channel. It is reset after all channel names have been used