spectral_analysis.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyNHaafpt3KysWQABtlwiX/F",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/ayushi15092002/mental-health-workload/blob/main/spectral_analysis.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Load the Drive helper and mount\n",
        "from google.colab import drive\n",
        "\n",
        "# This will prompt for authorization.\n",
        "drive.mount('/content/drive')"
      ],
      "metadata": {
        "id": "LsVk6gRl9DaD",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "07d1a903-ca56-420f-d425-96184bece628"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Mounted at /content/drive\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "pip install mne --upgrade"
      ],
      "metadata": {
        "id": "awBwoWaL9LoZ",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "e753ea5b-9e46-4ebf-b0c3-1f0ac7443981"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Collecting mne\n",
            "  Downloading mne-1.2.2-py3-none-any.whl (7.6 MB)\n",
            "\u001b[K     |████████████████████████████████| 7.6 MB 3.8 MB/s \n",
            "\u001b[?25hRequirement already satisfied: numpy>=1.15.4 in /usr/local/lib/python3.7/dist-packages (from mne) (1.21.6)\n",
            "Requirement already satisfied: pooch>=1.5 in /usr/local/lib/python3.7/dist-packages (from mne) (1.6.0)\n",
            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.7/dist-packages (from mne) (2.11.3)\n",
            "Requirement already satisfied: matplotlib in /usr/local/lib/python3.7/dist-packages (from mne) (3.2.2)\n",
            "Requirement already satisfied: packaging in /usr/local/lib/python3.7/dist-packages (from mne) (21.3)\n",
            "Requirement already satisfied: scipy>=1.1.0 in /usr/local/lib/python3.7/dist-packages (from mne) (1.7.3)\n",
            "Requirement already satisfied: decorator in /usr/local/lib/python3.7/dist-packages (from mne) (4.4.2)\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.7/dist-packages (from mne) (4.64.1)\n",
            "Requirement already satisfied: requests>=2.19.0 in /usr/local/lib/python3.7/dist-packages (from pooch>=1.5->mne) (2.23.0)\n",
            "Requirement already satisfied: appdirs>=1.3.0 in /usr/local/lib/python3.7/dist-packages (from pooch>=1.5->mne) (1.4.4)\n",
            "Requirement already satisfied: pyparsing!=3.0.5,>=2.0.2 in /usr/local/lib/python3.7/dist-packages (from packaging->mne) (3.0.9)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.7/dist-packages (from requests>=2.19.0->pooch>=1.5->mne) (2022.9.24)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.7/dist-packages (from requests>=2.19.0->pooch>=1.5->mne) (2.10)\n",
            "Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.7/dist-packages (from requests>=2.19.0->pooch>=1.5->mne) (3.0.4)\n",
            "Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.7/dist-packages (from requests>=2.19.0->pooch>=1.5->mne) (1.24.3)\n",
            "Requirement already satisfied: MarkupSafe>=0.23 in /usr/local/lib/python3.7/dist-packages (from jinja2->mne) (2.0.1)\n",
            "Requirement already satisfied: kiwisolver>=1.0.1 in /usr/local/lib/python3.7/dist-packages (from matplotlib->mne) (1.4.4)\n",
            "Requirement already satisfied: cycler>=0.10 in /usr/local/lib/python3.7/dist-packages (from matplotlib->mne) (0.11.0)\n",
            "Requirement already satisfied: python-dateutil>=2.1 in /usr/local/lib/python3.7/dist-packages (from matplotlib->mne) (2.8.2)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.7/dist-packages (from kiwisolver>=1.0.1->matplotlib->mne) (4.1.1)\n",
            "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.7/dist-packages (from python-dateutil>=2.1->matplotlib->mne) (1.15.0)\n",
            "Installing collected packages: mne\n",
            "Successfully installed mne-1.2.2\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "pip install signal-processing"
      ],
      "metadata": {
        "id": "KtWBseV99PUS"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "pip install basic"
      ],
      "metadata": {
        "id": "CqCfjf759QDt"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "pip install sklearn.cross_validation"
      ],
      "metadata": {
        "id": "CRzjlyA7HSzW"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "pip install matplotlib==3.1.1"
      ],
      "metadata": {
        "id": "6WNuKCvKHY5r",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 398
        },
        "outputId": "a726487d-34c5-4d1f-f3d7-190761915e4c"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Collecting matplotlib==3.1.1\n",
            "  Downloading matplotlib-3.1.1-cp37-cp37m-manylinux1_x86_64.whl (13.1 MB)\n",
            "\u001b[K     |████████████████████████████████| 13.1 MB 3.3 MB/s \n",
            "\u001b[?25hRequirement already satisfied: numpy>=1.11 in /usr/local/lib/python3.7/dist-packages (from matplotlib==3.1.1) (1.21.6)\n",
            "Requirement already satisfied: kiwisolver>=1.0.1 in /usr/local/lib/python3.7/dist-packages (from matplotlib==3.1.1) (1.4.4)\n",
            "Requirement already satisfied: pyparsing!=2.0.4,!=2.1.2,!=2.1.6,>=2.0.1 in /usr/local/lib/python3.7/dist-packages (from matplotlib==3.1.1) (3.0.9)\n",
            "Requirement already satisfied: cycler>=0.10 in /usr/local/lib/python3.7/dist-packages (from matplotlib==3.1.1) (0.11.0)\n",
            "Requirement already satisfied: python-dateutil>=2.1 in /usr/local/lib/python3.7/dist-packages (from matplotlib==3.1.1) (2.8.2)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.7/dist-packages (from kiwisolver>=1.0.1->matplotlib==3.1.1) (4.1.1)\n",
            "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.7/dist-packages (from python-dateutil>=2.1->matplotlib==3.1.1) (1.15.0)\n",
            "Installing collected packages: matplotlib\n",
            "  Attempting uninstall: matplotlib\n",
            "    Found existing installation: matplotlib 3.2.2\n",
            "    Uninstalling matplotlib-3.2.2:\n",
            "      Successfully uninstalled matplotlib-3.2.2\n",
            "Successfully installed matplotlib-3.1.1\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.colab-display-data+json": {
              "pip_warning": {
                "packages": [
                  "matplotlib",
                  "mpl_toolkits"
                ]
              }
            }
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tbeugm8Jhijj"
      },
      "outputs": [],
      "source": [
        "import os, mne\n",
        "import pandas as pd\n",
        "import seaborn as sns\n",
        "import matplotlib.pyplot as plt\n",
        "# from signal_processing import *\n",
        "# from basic import *\n",
        "from scipy.signal import welch\n",
        "from scipy.integrate import simps\n",
        "from mne.time_frequency import psd_array_multitaper "
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "b_names = ['Delta','Theta','Alpha','Beta']\n",
        "b_freqs = [[1,3.9],[4,7.9],[8,12],[12.1,30]]\n",
        "brain_regions = {'Left frontal':['AFF5','F1','FP1','FC5','FC1'],\n",
        "                 'Right frontal':['AFF6','F2','FP2','FC6','FC2'],\n",
        "                 'Left posterior':['CP5','CP1','P7','P3','O1'],\n",
        "                 'Right posterior':['CP6','CP2','P8','P4','O2'],\n",
        "                 'Midline frontal':['AFz'],\n",
        "                 'Midline central':['Cz','T8','T7','C3','C4'],\n",
        "                 'Midline posterior':['Pz','POz']}"
      ],
      "metadata": {
        "id": "idL3qz3HELQ1"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "clean_data_dir = \"/content/drive/My Drive/drdo/cleanData\"\n",
        "read_datasetA(clean_data_dir)"
      ],
      "metadata": {
        "id": "uQnehcT39sOZ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# dir_inprogress = os.path.join(clean_folder,exp_folder)\n",
        "# file_dirs, subject_names = read_files(dir_inprogress,\"_clean-epo.fif\")\n",
        "df_faa = pd.DataFrame(columns=['Pre-frontal','Mid-frontal','Lateral-frontal'])\n",
        "results_foldername = \"/content/drive/My Drive/drdo/spectral_analysis\"\n",
        "exp_folder = r\"\"\n",
        "exp_condition = r\"\"\n",
        "\n",
        "def read_datasetA(subjectPath):\n",
        "    \n",
        "    try:\n",
        "        root, dirs, files = next(os.walk(clean_data_dir))\n",
        "\n",
        "        i=0\n",
        "        \n",
        "        for folder_name in dirs:\n",
        "          path_file_fif1 = subjectPath + '/' + folder_name + '/nback1.fif'\n",
        "          spectral_analysis_A(i,path_file_fif1, folder_name,'nback1')\n",
        "\n",
        "          path_file_fif2 = subjectPath + '/' + folder_name + '/nback2.fif'\n",
        "          spectral_analysis_A(i,path_file_fif1, folder_name,'nback2')\n",
        "\n",
        "          path_file_fif3 = subjectPath + '/' + folder_name + '/nback3.fif'\n",
        "          spectral_analysis_A(i,path_file_fif1, folder_name,'nback3')\n",
        "          i = i+1\n",
        "        # Calculate alpha asymmetry for three channel pairs\n",
        "\n",
        "        print(\"df_psd_Alpha >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>\")\n",
        "        print(df_psd_Alpha)\n",
        "        df_faa['Pre-frontal'] = calculate_asymmetry_ch(df_psd_Alpha,'FP1','FP2')\n",
        "        df_faa['Mid-frontal'] = calculate_asymmetry_ch(df_psd_Alpha,'F1','F2')\n",
        "        df_faa['Lateral-frontal'] = calculate_asymmetry_ch(df_psd_Alpha,'FC5','FC6')\n",
        "\n",
        "        # Pre-create results folders for spectral analysis data\n",
        "        create_results_folders(exp_folder=exp_folder,results_foldername=results_foldername)\n",
        "\n",
        "        # Export asymmetry data\n",
        "        df_faa.to_excel(r\"{}/Absolute PSD/{}_frontal_asymmetry.xlsx\".format(results_foldername,exp_condition))\n",
        "\n",
        "        # Export power spectra data for each band\n",
        "        for band in b_names:\n",
        "          export_psd_results(globals()[\"df_psd_\"+band],globals()[\"df_rel_psd_\"+band],exp_folder,exp_condition,band,brain_regions,results_foldername=results_foldername)\n",
        "            \n",
        "    except StopIteration:\n",
        "        pass\n",
        "        print(\"Error ocurred:\")\n",
        "        print(\"Directory with dataset does not found!\")\n",
        "        print(\"Program will be terminated\")\n",
        "        exit(1)\n"
      ],
      "metadata": {
        "id": "iqr4yRNVDmgz"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "\n",
        "def spectral_analysis_A(i, file_path, subject_name, sesssion_name):\n",
        "  subject_name = subject_name + \"_\" + sesssion_name\n",
        "  # Read the clean data from the disk\n",
        "  epochs = mne.read_epochs(fname=file_path,verbose=False)\n",
        "    \n",
        "  # Calculate Welch's power spectrum density (FFT)\n",
        "  [psds,freqs] = calculate_psd(epochs,subject_name,fminmax=[1,30],window='hamming',window_duration=2,window_overlap=0.5)\n",
        "\n",
        "  # Total PSD of chosen frequency range (for later relative PSD calculation)\n",
        "  psd_ch_allbands = bandpower_per_channel(psds,freqs,[b_freqs[0][0],b_freqs[-1][-1]],'All bands',subject_name,epochs)\n",
        "  # Find power for all bands and add to dataframe including the previous subjects\n",
        "  vmin = float('inf')\n",
        "  vmax = 0\n",
        "  for j in range(len(b_names)):\n",
        "    if i == 0:\n",
        "      globals()[\"df_psd_\"+b_names[j]] = pd.DataFrame()\n",
        "      globals()[\"df_rel_psd_\"+b_names[j]] = pd.DataFrame()\n",
        "        \n",
        "    # Devide the PSD to frequency band bins and calculate absolute and relative bandpowers\n",
        "    globals()[\"psd_ch_\"+b_names[j]] = bandpower_per_channel(psds,freqs,b_freqs[j],b_names[j],subject_name,epochs)\n",
        "    globals()[\"rel_psd_ch_\"+b_names[j]] = globals()[\"psd_ch_\"+b_names[j]] / psd_ch_allbands\n",
        "\n",
        "    # Convert the array to dataframe and concatenate it to dataframe including the previous subjects\n",
        "    globals()[\"temp_df_psd_\"+b_names[j]] = array_to_df(subject_name,epochs,globals()[\"psd_ch_\"+b_names[j]])\n",
        "    globals()[\"df_psd_\"+b_names[j]] = pd.concat([globals()[\"df_psd_\"+b_names[j]],globals()[\"temp_df_psd_\"+b_names[j]]])\n",
        "    globals()[\"temp_df_rel_psd_\"+b_names[j]] = array_to_df(subject_name,epochs,globals()[\"rel_psd_ch_\"+b_names[j]])\n",
        "    globals()[\"df_rel_psd_\"+b_names[j]] = pd.concat([globals()[\"df_rel_psd_\"+b_names[j]],globals()[\"temp_df_rel_psd_\"+b_names[j]]])\n",
        "\n",
        "    # Save the minimum and maximum PSD values as an integer for later colorbar use\n",
        "    vmin = min([vmin,min(globals()[\"psd_ch_\"+b_names[j]])])\n",
        "    vmax = max([vmax,max(globals()[\"psd_ch_\"+b_names[j]])])\n",
        "\n",
        "  # Plot topomaps for all bands\n",
        "  sns.set_style(\"white\",{'font.family': ['sans-serif']})\n",
        "  fig,axs = plt.subplots(nrows=1,ncols=len(b_names),figsize=(10,3))\n",
        "  fig.suptitle(\"Frequency topomaps ({})\".format(subject_name))\n",
        "  for topo in range(len(b_names)):\n",
        "      im,_ = mne.viz.plot_topomap(globals()[\"psd_ch_\"+b_names[topo]],epochs.info,axes=axs[topo],vlim=[vmin,vmax],show=False)\n",
        "      axs[topo].set_title(b_names[topo]+'\\n'+str(b_freqs[topo]))\n",
        "  cbar_ax = fig.add_axes([0.95,0.35,0.04,0.4])\n",
        "  clb = fig.colorbar(im, cax=cbar_ax)\n",
        "  clb.ax.set_ylabel('uV\\u00b2/Hz')\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "q0Bz8hFkDeTL"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# ========== Packages ==========\n",
        "import mne\n",
        "import matplotlib.pyplot as plt\n",
        "import seaborn as sns\n",
        "import numpy as np\n",
        "from scipy import stats\n",
        "\n",
        "# ========== Functions ==========\n",
        "def calculate_psd(epochs,subjectname,fminmax=[1,100],window='hamming',window_duration=2,window_overlap=0.5,plot=True):\n",
        "    \"\"\"\n",
        "    Calculate power spectrum density with FFT/Welch's method and plot the result.\n",
        "    Parameters\n",
        "    ----------\n",
        "    epochs: Epochs-type (MNE-Python) EEG file\n",
        "    subjectname: A string for subject's name\n",
        "    fminmax (optional): The minimum and maximum frequency range for estimating Welch's PSD\n",
        "    window (optional): The window type for estimating Welch's PSD\n",
        "    window_duration (optional): An integer for the length of that window\n",
        "    window_overlap (optional): A float for the percentage of overlap between the windows\n",
        "    Returns\n",
        "    -------\n",
        "    psds: An array for power spectrum density values\n",
        "    freqs: An array for the corresponding frequencies\n",
        "    \"\"\"\n",
        "    # Calculate PSD with Welch's method\n",
        "    window_size = int(epochs.info['sfreq']*window_duration)\n",
        "    psds, freqs = mne.time_frequency.psd_welch(epochs,n_fft=window_size,verbose=False,\n",
        "                                               n_overlap=int(window_size*window_overlap),\n",
        "                                               fmin=fminmax[0],fmax=fminmax[1],window=window)\n",
        "\n",
        "    # Unit conversion from V^2/Hz to uV^2/Hz\n",
        "    psds = psds*1e12\n",
        "\n",
        "    if plot == True:\n",
        "        # Plot average PSD for all epochs and channels (only for plot)\n",
        "        psds_mean_all = psds.mean(axis=(0, 1))\n",
        "\n",
        "        sns.set_style(\"darkgrid\",{'font.family': ['sans-serif']})\n",
        "        plt.figure()\n",
        "        plt.plot(freqs,psds_mean_all)\n",
        "        plt.fill_between(freqs,psds_mean_all)\n",
        "        plt.xlabel('Frequency (Hz)')\n",
        "        plt.ylabel('PSD (uV\\u00b2/Hz)')\n",
        "        plt.title(\"PSD ({})\".format(subjectname))\n",
        "        plt.xlim(1,40)\n",
        "        plt.ylim(0,1.1*max(psds_mean_all))\n",
        "\n",
        "    return [psds,freqs]\n",
        "\n",
        "def siqnal_quality_check(psds,freqs,band,b_name,subjectname,epochs):\n",
        "    \"\"\"\n",
        "    Plot topographically PSD values and calculate median absolute deviation\n",
        "    for the first and second half of the signal for signal reliability control.\n",
        "    Parameters\n",
        "    ----------\n",
        "    psds: An array for power spectrum density values\n",
        "    freqs: An array for corresponding frequencies\n",
        "    band: A list of lower and higher frequency for the frequency band in interest\n",
        "    b_name: A string for frequency band in interest\n",
        "    subjectname: A string for subject's name\n",
        "    epochs: Epochs-type (MNE-Python) EEG file\n",
        "    Returns\n",
        "    -------\n",
        "    psd_max_mad_error: A float for max median absolute deviation for PSD values\n",
        "    \"\"\"\n",
        "\n",
        "    # Divide the epochs into two\n",
        "    idx_mid_epoch = round(len(psds[:,0,0])/2)\n",
        "    psds_p1 = psds[:idx_mid_epoch,:,:]\n",
        "    psds_p2 = psds[idx_mid_epoch:,:,:]\n",
        "\n",
        "    # Define the bandpower indices\n",
        "    low, high = band\n",
        "    idx_band = np.logical_and(freqs >= low, freqs <= high)\n",
        "\n",
        "    # Get only the bandpower values that we need for both parts\n",
        "    psd_band_ch_p1 = psds_p1[:,:,idx_band]\n",
        "    psd_band_mean_epochs_p1 = psd_band_ch_p1.mean(axis=(2))\n",
        "    psd_band_ch_p2 = psds_p2[:,:,idx_band]\n",
        "    psd_band_mean_epochs_p2 = psd_band_ch_p2.mean(axis=(2))\n",
        "\n",
        "    # Find Median Absolute Deviation along all the epochs for each channel\n",
        "    psd_band_mean_epochs = np.concatenate((psd_band_mean_epochs_p1, psd_band_mean_epochs_p2))\n",
        "    psd_median = np.median(psd_band_mean_epochs,axis=0)\n",
        "    psd_mad = stats.median_abs_deviation(psd_band_mean_epochs,axis=0)\n",
        "\n",
        "    # For each epoch in all channels, calculate the Z-score (MAD error) using median and MAD\n",
        "    psd_mad_error_epochs_per_ch_p1 = np.zeros(shape=psd_band_mean_epochs_p1.shape)\n",
        "    psd_mad_error_epochs_per_ch_p2 = np.zeros(shape=psd_band_mean_epochs_p2.shape)\n",
        "    for i in range(len(psds[0,:,0])):\n",
        "        psd_mad_error_epochs_per_ch_p1[:,i] = abs(psd_band_mean_epochs_p1[:,i] - psd_median[i])/psd_mad[i]\n",
        "        psd_mad_error_epochs_per_ch_p2[:,i] = abs(psd_band_mean_epochs_p2[:,i] - psd_median[i])/psd_mad[i]\n",
        "\n",
        "    # Average the Z-scores (MAD errors) for all epochs for the two parts -> get score for 19 channels\n",
        "    psd_mad_error_per_ch_p1 = np.mean(psd_mad_error_epochs_per_ch_p1,axis=0)\n",
        "    psd_mad_error_per_ch_p2 = np.mean(psd_mad_error_epochs_per_ch_p2,axis=0)\n",
        "\n",
        "    # Find the average Z-score (MAD error) for the whole scalp (all channels together) for the two parts\n",
        "    psd_mad_error_avg = [np.round(np.mean(psd_mad_error_per_ch_p1),3),\n",
        "                         np.round(np.mean(psd_mad_error_per_ch_p2),3)]\n",
        "    \n",
        "    # Get the maximum of the two (for outputting it later)\n",
        "    psd_max_mad_error = np.max(psd_mad_error_avg)\n",
        "\n",
        "    # Find the average PSD values for each channel for both parts\n",
        "    psds_all_channels_p1 = psds_p1.mean(axis=(0))\n",
        "    psd_band_ch_p1 = psds_all_channels_p1[:,idx_band]\n",
        "    psd_band_mean_ch_p1 = psd_band_ch_p1.mean(axis=(1))\n",
        "    psds_all_channels_p2 = psds_p2.mean(axis=(0))\n",
        "    psd_band_ch_p2 = psds_all_channels_p2[:,idx_band]\n",
        "    psd_band_mean_ch_p2 = psd_band_ch_p2.mean(axis=(1))\n",
        "\n",
        "    \n",
        "    # Visually display the bandpower for both parts for visual inspection\n",
        "    vmin = min([min(psd_band_mean_ch_p1),min(psd_band_mean_ch_p2)])\n",
        "    vmax = max([max(psd_band_mean_ch_p1),max(psd_band_mean_ch_p2)])\n",
        "\n",
        "    if psd_max_mad_error >= 2:\n",
        "        sns.set_style(\"white\",{'font.family': ['sans-serif']})\n",
        "        fig,(ax1,ax2) = plt.subplots(ncols=2)\n",
        "        fig.suptitle(\"MAD >= 2 ! Quality control for {} ({})\".format(b_name,subjectname),y=1.1,x=0.575)\n",
        "        im,cm = mne.viz.plot_topomap(psd_band_mean_ch_p1,epochs.info,axes=ax1,vlim=[vmin,vmax],show=False)\n",
        "        im,cm = mne.viz.plot_topomap(psd_band_mean_ch_p2,epochs.info,axes=ax2,vlim=[vmin,vmax],show=False)\n",
        "        ax1.set_title(\"Epochs 0-{}\\nAvg MAD error = {}\".format(idx_mid_epoch-1,psd_mad_error_avg[0]))\n",
        "        ax2.set_title(\"Epochs {}-{}\\nAvg MAD error = {}\".format(idx_mid_epoch,len(psds[:,0,0]),psd_mad_error_avg[1]))\n",
        "        ax_x_start = 0.95\n",
        "        ax_x_width = 0.04\n",
        "        ax_y_start = 0.1\n",
        "        ax_y_height = 0.9\n",
        "        cbar_ax = fig.add_axes([ax_x_start, ax_y_start, ax_x_width, ax_y_height])\n",
        "        clb = fig.colorbar(im, cax=cbar_ax)\n",
        "        clb.ax.set_ylabel('uV\\u00b2/Hz'); \n",
        "\n",
        "    return psd_max_mad_error\n",
        "\n",
        "def bandpower_per_channel(psds,freqs,band,b_name,subjectname,epochs,ln_normalization=False):\n",
        "    \"\"\"\n",
        "    Find frequency band power in interest for all the channels.\n",
        "    Parameters\n",
        "    ----------\n",
        "    psds: An array for power spectrum density values\n",
        "    freqs: An array for corresponding frequencies\n",
        "    band: A list of lower and higher frequency for the frequency band in interest\n",
        "    b_name: A string for frequency band in interest\n",
        "    subjectname: A string for subject's name\n",
        "    epochs: Epochs-type (MNE-Python) EEG file\n",
        "    Returns\n",
        "    -------\n",
        "    psd_band_mean_ch: An array for a frequency band power values for all the channels.\n",
        "    \"\"\"\n",
        "    # Calculate the MAD error (z-score) of the bandpower to be sure of the quality\n",
        "    psd_max_mad_error = siqnal_quality_check(psds,freqs,band,b_name,subjectname,epochs)\n",
        "    \n",
        "    low, high = band\n",
        "    psds_all_channels = psds.mean(axis=(0))\n",
        "    idx_band = np.logical_and(freqs >= low, freqs <= high)\n",
        "    psd_band_ch = psds_all_channels[:,idx_band]\n",
        "    psd_band_mean_ch = psd_band_ch.mean(axis=(1))\n",
        "\n",
        "    #print(band,'before norm',psd_band_mean_ch)\n",
        "    if ln_normalization == True:\n",
        "        psd_band_mean_ch = np.log(psd_band_mean_ch)\n",
        "\n",
        "    #print(band,'after norm',psd_band_mean_ch)\n",
        "    # # If the error is larger than 2, print it as a result (visual inspection)\n",
        "    # if psd_max_mad_error < 2:\n",
        "    #     print(subjectname,b_name,\"MAD error is OK:\",psd_max_mad_error)\n",
        "    # else:\n",
        "    #     print(subjectname,b_name,\"MAD error is NOT OK:\",psd_max_mad_error)\n",
        "\n",
        "    return psd_band_mean_ch\n",
        "\n",
        "def calculate_asymmetry_ch(df_psd_band,left_ch,right_ch):\n",
        "    \"\"\"\n",
        "    Calculate asymmetry between brain hemispheres.\n",
        "    Parameters\n",
        "    ----------\n",
        "    df_psd_band: A dataframe with PSD values (for each region/channel) per subject for one band\n",
        "    left_ch: A string for the left channel (or region)\n",
        "    right_ch: A string for the right channel (or region)\n",
        "    Returns\n",
        "    -------\n",
        "    df_asymmetry: A dataframe for calculated asymmetry for all the subjects\n",
        "    \"\"\"\n",
        "\n",
        "    print(\"df_psd_band[left_ch] >>>>>>>>>>>>>>>>>>>>>>>\")\n",
        "    print(df_psd_band[left_ch])\n",
        "    df_asymmetry = (df_psd_band[left_ch] - df_psd_band[right_ch])/(df_psd_band[left_ch] + df_psd_band[right_ch])*100\n",
        "    return df_asymmetry"
      ],
      "metadata": {
        "id": "WI2tti8LOq35"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# ========== Packages ==========\n",
        "import os\n",
        "import pandas as pd\n",
        "\n",
        "# ========== Functions ==========\n",
        "def read_files(dir_inprogress,filetype,exclude_subjects=None,verbose=True):\n",
        "    \"\"\"\n",
        "    Get all the (EEG) file directories and subject names.\n",
        "    Parameters\n",
        "    ----------\n",
        "    dir_inprogress: A string with directory to look for files\n",
        "    filetype: A string with the ending of the files we are looking for (e.g. '.bdf')\n",
        "    Returns\n",
        "    -------\n",
        "    file_dirs: A list of strings with file directories for all the (EEG) files\n",
        "    subject_names: A list of strings with all the corresponding subject names\n",
        "    \"\"\"\n",
        "    file_dirs = []\n",
        "    subject_names = []\n",
        "    for file in os.listdir(dir_inprogress):\n",
        "        if file.endswith(filetype):\n",
        "            file_dirs.append(os.path.join(dir_inprogress, file))\n",
        "            subject_names.append(os.path.join(file).removesuffix(filetype))\n",
        "\n",
        "    for excl_sub in exclude_subjects:\n",
        "        for i in range(len(subject_names)):\n",
        "            if excl_sub in subject_names[i]:\n",
        "                if verbose == True:\n",
        "                    print('EXCLUDED SUBJECT: ',excl_sub,'in',subject_names[i],'at',file_dirs[i])\n",
        "                del subject_names[i]\n",
        "                del file_dirs[i]\n",
        "                break\n",
        "\n",
        "    if verbose == True:\n",
        "        print(\"Files in folder read in:\",len(file_dirs))\n",
        "\n",
        "    return [file_dirs, subject_names]\n",
        "\n",
        "def array_to_df(subjectname,epochs,array_channels):\n",
        "    \"\"\"\n",
        "    Convert channel-based array to Pandas dataframe with channels' and subjects' names.\n",
        "    Parameters\n",
        "    ----------\n",
        "    subjectname: A string for subject's name\n",
        "    epochs: Epochs-type (MNE-Python) EEG file\n",
        "    array_channels: An array with values for each channel\n",
        "    Returns\n",
        "    -------\n",
        "    df_channels: A dataframe with values for each channel\n",
        "    \"\"\"\n",
        "    df_channels = pd.DataFrame(array_channels).T\n",
        "    df_channels.columns = epochs.info.ch_names\n",
        "    df_channels['Subject'] = subjectname\n",
        "    df_channels.set_index('Subject', inplace=True)\n",
        "\n",
        "    return df_channels\n",
        "\n",
        "def df_channels_to_regions(df_psd_band,brain_regions):\n",
        "    \"\"\"\n",
        "    Average channels together based on the defined brain regions.\n",
        "    Parameters\n",
        "    ----------\n",
        "    df_psd_band: A dataframe with PSD values for each channel per subject\n",
        "    brain_regions: A dictionary of brain regions and EEG channels which they contain\n",
        "    Returns\n",
        "    -------\n",
        "    df_psd_reg_band: A dataframe with PSD values for each brain region per subject\n",
        "    \"\"\"\n",
        "    df_psd_reg_band = pd.DataFrame()\n",
        "    print(\"brain_regions >>>>>>>>>>>>>>>>>>>>>>>>>>>\")\n",
        "    print(brain_regions)\n",
        "    for reg_name in brain_regions:\n",
        "      print(\"df_psd_band[brain_regions[reg_name]] >>>>>>>>>>>>>>>>>>>>\")\n",
        "      print(df_psd_band[brain_regions[reg_name]])\n",
        "      df_temp = df_psd_band[brain_regions[reg_name]].copy().mean(axis=1)\n",
        "      df_psd_reg_band = pd.concat([df_psd_reg_band,df_temp],axis=1)\n",
        "        \n",
        "    df_psd_reg_band.columns = brain_regions.keys()\n",
        "    df_psd_reg_band.index.name = 'Subject'\n",
        "\n",
        "    return df_psd_reg_band\n",
        "\n",
        "def read_excel_psd(exp_folder,psd_folder,verbose=True):\n",
        "    \"\"\"\n",
        "    Get all PSD file directories and corresponding bands and experiment conditions.\n",
        "    Parameters\n",
        "    ----------\n",
        "    exp_folder: A string with a relative directory to experiment folder (e.g. 'Eyes Closed\\Baseline')\n",
        "    psd_folder: A string with a relative directory to the results folder (e.g. 'Results\\PSD\\regions')\n",
        "    Returns\n",
        "    -------\n",
        "    dir_inprogress: A string with directory to look for files\n",
        "    b_names: A list of strings for frequency bands of the files\n",
        "    condition: A list of strings for experiment conditions of the files\n",
        "    \"\"\"\n",
        "    dir_inprogress = os.path.join(psd_folder,exp_folder)\n",
        "    _, b_names = read_files(dir_inprogress,\".xlsx\",verbose=verbose)\n",
        "\n",
        "    condition = [None]*len(b_names)\n",
        "    for i in range(len(b_names)):\n",
        "        condition[i] = b_names[i].split(\"_psd_\", 1)\n",
        "    \n",
        "    return [dir_inprogress,b_names,condition]\n",
        "\n",
        "def create_results_folders(exp_folder,results_foldername=\"Results\"):\n",
        "    \"\"\"\n",
        "    Dummy way to try to pre-create folders for PSD results before exporting them\n",
        "    Parameters\n",
        "    ----------\n",
        "    exp_folder: A string with a relative directory to experiment folder (e.g. 'Eyes Closed\\Baseline')\n",
        "    \"\"\"\n",
        "    try:\n",
        "        os.makedirs(os.path.join(r\"{}/Absolute PSD/channels\".format(results_foldername),exp_folder))\n",
        "    except FileExistsError:\n",
        "        pass\n",
        "    try:\n",
        "        os.makedirs(os.path.join(r\"{}/Absolute PSD/regions\".format(results_foldername),exp_folder))\n",
        "    except FileExistsError:\n",
        "        pass\n",
        "    try:\n",
        "        os.makedirs(os.path.join(r\"{}/Relative PSD/channels\".format(results_foldername),exp_folder))\n",
        "    except FileExistsError:\n",
        "        pass\n",
        "    try:\n",
        "        os.makedirs(os.path.join(r\"{}/Relative PSD/regions\".format(results_foldername),exp_folder))\n",
        "    except FileExistsError:\n",
        "        pass\n",
        "    try:\n",
        "        os.makedirs(os.path.join(r\"{}/FOOOF\".format(results_foldername),exp_folder))\n",
        "    except FileExistsError:\n",
        "        pass\n",
        "\n",
        "def export_psd_results(df_psd_band,df_rel_psd_band,exp_folder,exp_condition,band,brain_regions,results_foldername='Results'):\n",
        "    \"\"\"\n",
        "    Export PSD results (for channels & regions) as Excel files\n",
        "    Parameters\n",
        "    ----------\n",
        "    df_psd_band: A dataframe with PSD values for each channel per subject\n",
        "    df_rel_psd_band: A dataframe with relative PSD values for each channel per subject\n",
        "    exp_folder: A string with a relative directory to experiment folder (e.g. 'Eyes Closed\\Baseline')\n",
        "    exp_condition: A string for experiment short code (e.g. 'EC_00')\n",
        "    band: The frequency band of the PSD values (e.g. 'Alpha')\n",
        "    brain_regions: A dictionary of brain regions and EEG channels which they contain\n",
        "    \"\"\"\n",
        "    # Save the PSD values for each channel for each band in Excel format\n",
        "    df_psd_band.to_excel(r\"{}/Absolute PSD/channels/{}/{}_psd_{}.xlsx\".format(results_foldername,exp_folder,exp_condition,band))\n",
        "    df_rel_psd_band.to_excel(r\"{}/Relative PSD/channels/{}/{}_rel_psd_{}.xlsx\".format(results_foldername,exp_folder,exp_condition,band))\n",
        "\n",
        "    # Find regional band powers and save them to Excel as well\n",
        "    df_psd_band_reg = df_channels_to_regions(df_psd_band,brain_regions)\n",
        "    df_psd_band_reg.to_excel(r\"{}/Absolute PSD/regions/{}/{}_psd_{}.xlsx\".format(results_foldername,exp_folder,exp_condition,band))\n",
        "    df_rel_psd_band_reg = df_channels_to_regions(df_rel_psd_band,brain_regions)\n",
        "    df_rel_psd_band_reg.to_excel(r\"{}/Relative PSD/regions/{}/{}_rel_psd_{}.xlsx\".format(results_foldername,exp_folder,exp_condition,band))"
      ],
      "metadata": {
        "id": "Y8VvFqPgsOMP"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}