Main.py

# -*- coding: utf-8 -*-
"""
Updated Aug 7 2024

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

This is the main Python file containing the code that support the findings of
https://doi.org/10.1101/2024.05.06.592342 

The data used in this analysis was previously described and preprocessed
by Zimmermann, M., Lomoriello, A. S., and Konvalinka, I.
Intra-individual behavioural and neural signatures of audience effects and
interactions in a mirror-game paradigm. Royal Society Open Science, 9(2) 2022
"""

# %% Load libraries
import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import mne
import pickle
import mat73
import time
import seaborn as sns
import nolds
from tqdm import tqdm # progress bar

# import Python script for microstates [von Wegner & Lauf, 2018]
# originally downloaded from https://github.com/Frederic-vW/eeg_microstates
# I modified the script for estimating two-brain microstates
# by defining kmeans_return_all and kmeans_dualmicro
from eeg_microstates3 import (kmeans_return_all, kmeans_dualmicro)
# import helper functions
from helper import (numpy_arr_to_pandas_df, time_now)

from dualmicro_functions import (load_epoch_from_fieldtrip, prepare_1P_micro_arr,
                                 plot_microstates, reorder_microstate_results,
                                 single_micro_fit_all_feature_computation,
                                 interbrain_microstate_feature_computation,
                                 prepare_2P_micro_arr_collapsed_events,
                                 plot_dualmicro, sign_swap_microstates,
                                 dualmicro_fit_all_feature_computation,
                                 load_microstate_arrays, 
                                 get_synch_events_from_pseudo_pairs,
                                 combine_two_person_microstate_arrays,
                                 pseudo_pair_dualmicro_backfitting,
                                 dualmicro_fit_all_pseudo_pair_feature_computation,
                                 compute_dualmicro_DFA, compute_dualmicro_DFA_pseudo,
                                 shifted_interbrain_microstate_feature_computation)

# Style for matplotlib/seaborn
plt.style.use('default')

# Root for project
os.chdir("C:/Users/glia/Documents/MirrorGame")

# Paths
data_path = "./data/external/EEG/"
mov_data_path = "./data/external/movement/"
fig_save_path = "./reports/figures/"
feat_save_path = "./data/features/"
microstate_save_path = "./data/features/microstates2/"
mov_save_path = "./data/features/movement/"


# %% Load preprocessed EEG data
# The data was originally preprocessed in Fieldtrip by Marius Zimmermann
# Get filenames for the EEG data
files = []
for r, d, f in os.walk(data_path):
    for file in f:
        if (".mat" in file) & ("ppn" in file):
            files.append(os.path.join(r, file))
# Sort the filenames
files.sort()

n_subjects = len(files)
# Get Subject_id
Subject_id = [0]*n_subjects
for i in range(n_subjects):
    id_number = files[i].split("/")[-1].split(".")[0].split("pair")[-1].split("pair")[-1].replace("_ppn","")
    Subject_id[i] = int(id_number)+1000 # add 1000 to keep the first 0

# There are data from 23 pairs
# Pair 21 and 25 were excluded in the original analysis
# After looking at the data, it seems pair 21, participant 1 and pair 25
# participant 2 only had 1254 and 1440 epochs respectively.
# Their data also do not end with resting-state condition
# All the other EEG data have around 2400 1s epochs and start and ends with rest
bad_subjects = [1211, 1212, 1251, 1252] # the whole pair is dropped
good_subject_idx = [not i in bad_subjects for i in Subject_id]
# Update Subject_id and files
Subject_id = list(np.array(Subject_id)[good_subject_idx])
n_subjects = len(Subject_id)
files = list(np.array(files)[good_subject_idx])

Pair_id = [0]*(n_subjects//2)
for i in range(n_subjects//2):
    Pair_id[i] = int(str(Subject_id[2*i])[1:-1])
# Add 100 to pair_id to fix sorting for 1 digit numbers, e.g. 03
Pair_id = [ele+100 for ele in Pair_id]
n_pairs = len(Pair_id)

# Save the IDs as environmental variables to be used in functions
# from dualmicro_functions.py
os.environ["Subject_id"] = Subject_id
os.environ["Pair_id"] = Pair_id

event_id = {"rest":1, "uncoupled":2, "coupled": 3, "observe, actor": 4,
            "observe, observer": 6, "imitate, leader": 5, "imitate, follower": 7,
            "control": 8}

# Clarification of the labels
# Cond4: ppn1 is observer, ppn2 is actor
# Cond6: ppn1 is actor, ppn2 is observer
# Cond5: ppn1 is follower, ppn2 is leader
# Cond5: ppn1 is leader, ppn2 is follower

event_id_inv = {v: k for k, v in event_id.items()}

# We collapsed condition 4 and 6 & 5 and 7 for two-brain microstates
# By swapping the EEG of ppn1 and ppn2 so ppn1 is always observer/follower and
# ppn2 actor/leader
collapsed_event_id = {"rest":1, "uncoupled":2, "coupled": 3,
                      "observer_actor": 4, "follower_leader": 5, "control": 8}
collapsed_event_id_inv = {v: k for k, v in collapsed_event_id.items()}

# Load the first EEG to get info about sfreq and n_channels
i = 0
epoch, trialinfo = load_epoch_from_fieldtrip(0, files, event_id)
n_channels = epoch.info["nchan"]
sfreq = int(epoch.info["sfreq"])

# Visualize the data
# epoch.plot(scalings=40e-6, n_channels=32)
# mne.viz.plot_events(epoch.events, sfreq = 1, event_id = event_id, first_samp=-3) # sfreq set to epoch length in s to reflect experiment time

# We compute microstates for the three frequency ranges
alpha_range = [8.0, 13.0]
beta_range = [13.0, 30.0]
broadband_range = None # Data is already 1 to 40 Hz broadband filtered
freq_names = ["alpha","beta","broadband"]
all_freq_ranges = [alpha_range, beta_range, broadband_range]

# %% Intrabrain microstates fit all data
# All subjects from all pairs are concatenated to find common microstates
single_brain_event_id = {"rest":1, "uncoupled":2, "coupled": 3, "observer": 4,
                         "actor": 6, "follower": 5, "leader": 7, "control": 8}
ppn2_correction = {6:4, 4:6, 7:5, 5:7}

# Loop over frequencies
for f in len(all_freq_ranges):
    ff = freq_names[f]
    freq_range0 = all_freq_ranges[f]
    # =========================================================================
    # First the microstate topographies are determined
    # It might be an advantage to run the estimation of microstates on a HPC
    # =========================================================================
    # Get data from all pairs before performing kmeans
    np.random.seed(1234)
    n_clusters=[3, 4, 5, 6, 7, 8, 9, 10]
    n_runs = 100 # increased to 100 runs!
    # Get current time
    c_time1 = time_now(); print(c_time1)
    # Save RAM by appending directly to array instead of making list and then array
    sub_arr_indices = [0]
    trialinfo_list = []
    for i in range(n_subjects):
        tmp_data, trialinfo = prepare_1P_micro_arr(i, ppn2_correction, sfreq,
                                                   freq_range=freq_range0, standardize=True)
        sub_arr_indices.append(len(tmp_data))
        trialinfo_list.append([Subject_id[i],trialinfo])
        if i == 0: # first run initiation
            micro_data_all = tmp_data
        else:
            micro_data_all = np.append(micro_data_all, tmp_data, axis=0)
        del tmp_data # clear up space
        print(f"Finished preparing microstate data for pair {Subject_id[i]}")
    
    # Use cumulative sum to determine indices for each subjects's data
    subject_indices = np.cumsum(sub_arr_indices)
    
    # Save the trialinfos from all subjects, for easier access in later steps
    with open(f"{microstate_save_path}Intrabrain_microstate_fit_all_{ff}_trialinfos.pkl", "wb") as filehandle:
        pickle.dump(trialinfo_list, filehandle)
    
    # # with args parser in hpc
    # n_maps = n_clusters[(args.map_idx-1)]
    # print(f"Running analysis for maps: {n_maps}")
    # print("Memory used by the micro data array (GB):",micro_data_all.nbytes*9.31e-10)
    
    # Run Kmeans
    for n_maps in n_clusters: # Don't use for loop on the HPC!
        # Run the 100 runs in batches of 10 to save underway in case the job script terminates
        best_cv_crit = 9999 # initialize unreasonably high value
        for r in range(10):
            microstate_results = list(kmeans_return_all(micro_data_all, n_maps,
                                                        n_runs=int(n_runs/10),maxiter=1000))
            # Overwrite the maps if a lower CV criterion was found for the initiation