foof.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyMTSJ83ISI+gattGQih24Z9",
      "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/foof.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": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "F0Th5F1NcG9M",
        "outputId": "cb78fe5f-a1b8-491f-ec9a-398a4f5a1c2b"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Mounted at /content/drive\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "clean_folder = \"/content/drive/My Drive/drdo/cleanData\"\n",
        "results_foldername = \"/content/drive/My Drive/drdo/spectral_analysis\""
      ],
      "metadata": {
        "id": "2HhrWnBDcJm4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "pip install mne --upgrade\n",
        "pip install fooof\n",
        "pip install signal-processing\n",
        "pip install basic"
      ],
      "metadata": {
        "id": "_-H0-7-dgh1I"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import os, mne\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "import matplotlib.pyplot as plt\n",
        "import seaborn as sns\n",
        "from fooof import FOOOF\n",
        "from fooof.bands import Bands\n",
        "from fooof.analysis import get_band_peak_fm, get_band_peak_fg\n",
        "from fooof.plts.periodic import plot_peak_fits, plot_peak_params\n",
        "from fooof.plts.aperiodic import plot_aperiodic_params, plot_aperiodic_fits"
      ],
      "metadata": {
        "id": "T_Qor1MSc-Qw"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "read_datasetA(clean_data_dir)"
      ],
      "metadata": {
        "id": "ll5OnW0KiNJi"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zRV7XA3PcD_i",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 235
        },
        "outputId": "8438b864-31b1-48bf-ca59-cc24f6dd02d8"
      },
      "outputs": [
        {
          "output_type": "error",
          "ename": "NameError",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-1-19b193f46c5d>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0;31m# Define frequency bands of interest\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m bands = Bands({'delta' : [1, 4],\n\u001b[0m\u001b[1;32m      3\u001b[0m                \u001b[0;34m'theta'\u001b[0m \u001b[0;34m:\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;36m4\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m7.5\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m                \u001b[0;34m'alpha'\u001b[0m \u001b[0;34m:\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;36m7.5\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m13\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m                'beta' : [13, 30]})\n",
            "\u001b[0;31mNameError\u001b[0m: name 'Bands' is not defined"
          ]
        }
      ],
      "source": [
        "# Define frequency bands of interest\n",
        "bands = Bands({'delta' : [1, 4],\n",
        "               'theta' : [4, 7.5],\n",
        "               'alpha' : [7.5, 13],\n",
        "               'beta' : [13, 30]})\n",
        "clean_data_dir = \"/content/drive/My Drive/drdo/cleanData\"\n",
        "\n",
        "subject_names = ['VP001_nback1','VP001_nback2','VP001_nback3',\n",
        "                 'VP002_nback1','VP002_nback2','VP002_nback3',\n",
        "                 'VP003_nback1','VP003_nback2','VP003_nback3',\n",
        "                 'VP004_nback1','VP004_nback2','VP004_nback3',\n",
        "                 'VP005_nback1','VP005_nback2','VP005_nback3',\n",
        "                 'VP006_nback1','VP006_nback2','VP006_nback3',\n",
        "                 'VP007_nback1','VP007_nback2','VP007_nback3',\n",
        "                 'VP008_nback1','VP008_nback2','VP008_nback3',\n",
        "                 'VP009_nback1','VP009_nback2','VP009_nback3',]\n",
        "\n",
        "root, dirs, files = next(os.walk(clean_data_dir))\n",
        "\n",
        "brain_regions = {'Frontal':['AFF5','F1','FP1','FC5','FC1','AFF6','F2','FP2','FC6','FC2','AFz'],\n",
        "                 'Parietal':['CP5','CP1','P7','P3','O1','CP6','CP2','P8','P4','O2','Pz','POz']}\n",
        "\n",
        "# shape: [region][subject]\n",
        "deltas = [[None]*len(dirs)*3]*len(brain_regions.keys())\n",
        "thetas = [[None]*len(dirs)*3]*len(brain_regions.keys())\n",
        "alphas = [[None]*len(dirs)*3]*len(brain_regions.keys())\n",
        "betas = [[None]*len(dirs)*3]*len(brain_regions.keys())\n",
        "aps = [[None]*len(dirs)*3]*len(brain_regions.keys())\n",
        "exp_condition  = \"\"\n",
        "exp_folder = r\"\"\n",
        "\n",
        "freq_range = [1,30]\n",
        "# note: probably have to re-preprocess the data with higher high-cutoff frequency than the used freq in here\n",
        "# to avoid that part of the higher frequency spectra gets cut off which affects model fitting\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "def read_datasetA(subjectPath):\n",
        "    \n",
        "    try:\n",
        "        root, dirs, files = next(os.walk(subjectPath))\n",
        "\n",
        "        i=0\n",
        "        \n",
        "        for folder_name in dirs:\n",
        "          path_file_fif1 = subjectPath + '/' + folder_name + '/nback1.fif'\n",
        "          # data_vis_A(i,path_file_fif1, folder_name,'nback1')\n",
        "          file_path = path_file_fif1\n",
        "          subject_name = folder_name + \"_\" + \"nback1\"\n",
        "          epochs = mne.read_epochs(fname=file_path,verbose=False)\n",
        "          # Calculate Welch's power spectrum density (FFT) -> (epochs, channels, freq bins) shape\n",
        "          [psds,freqs] = calculate_psd(epochs,subject_name,fminmax=freq_range,window='hamming',\n",
        "                                 window_duration=2,window_overlap=0.5,plot=False)\n",
        "          # Average all epochs and channels together -> (freq bins,) shape\n",
        "          psds_allch = psds.mean(axis=(0,1))\n",
        "\n",
        "          psds = psds.mean(axis=(0))\n",
        "          df_psds_ch = array_to_df(subject_name,epochs,psds).reset_index().drop(columns='Subject')\n",
        "\n",
        "          df_psds_regions = df_channels_to_regions(df_psds_ch,brain_regions).reset_index().drop(columns='Subject')\n",
        "\n",
        "          for region in df_psds_regions.columns:\n",
        "            if i == 0:\n",
        "              globals()[\"df_fooof_\"+region] = pd.DataFrame(index=subject_names,columns=['Exponent','Offset','Peaks_CF','Peaks_PW','Peaks_BW','R_2','Error'])\n",
        "\n",
        "            psds_temp = df_psds_regions[region].to_numpy()\n",
        "            # Initialize power spectrum model objects and fit the power spectra\n",
        "            # note: aperiodic_mode should be 'knee' if the frequency range is > 40 Hz\n",
        "            max_n_peaks = float(\"inf\")\n",
        "            fm = FOOOF(peak_width_limits=[1,12],max_n_peaks=max_n_peaks,min_peak_height=0.225,\n",
        "                peak_threshold=2.0,aperiodic_mode='fixed',verbose=True)\n",
        "            fm.fit(freqs,psds_temp,freq_range)\n",
        "            # Plot the fitted model with power spectra \n",
        "            fig = plt.figure()\n",
        "            ax = fig.add_axes([0,0,1,1])\n",
        "            ax.set_title('PSD, {}'.format(subject_name))\n",
        "            if fm.get_params('peak_params').shape==(3,) or max_n_peaks == 0:\n",
        "              fm.plot(ax=ax)\n",
        "            else:\n",
        "              fm.plot(plot_peaks='shade', peak_kwargs={'color' : 'green'},ax=ax)\n",
        "\n",
        "            # Add model parameters to dataframe\n",
        "            globals()[\"df_fooof_\"+region].loc[globals()[\"df_fooof_\"+region].index[i],'Exponent'] = fm.get_params('aperiodic_params','exponent')\n",
        "            globals()[\"df_fooof_\"+region].loc[globals()[\"df_fooof_\"+region].index[i],'Offset'] = fm.get_params('aperiodic_params','offset')\n",
        "            globals()[\"df_fooof_\"+region].loc[globals()[\"df_fooof_\"+region].index[i],'Peaks_CF'] = fm.get_params('peak_params', 'CF')\n",
        "            globals()[\"df_fooof_\"+region].loc[globals()[\"df_fooof_\"+region].index[i],'Peaks_PW'] = fm.get_params('peak_params', 'PW')\n",
        "            globals()[\"df_fooof_\"+region].loc[globals()[\"df_fooof_\"+region].index[i],'Peaks_BW'] = fm.get_params('peak_params', 'BW')\n",
        "            globals()[\"df_fooof_\"+region].loc[globals()[\"df_fooof_\"+region].index[i],'R_2'] = fm.get_params('r_squared')\n",
        "            globals()[\"df_fooof_\"+region].loc[globals()[\"df_fooof_\"+region].index[i],'Error'] = fm.get_params('error')\n",
        "\n",
        "# ///////////////////////////////////////////////////////////////////////////////////////\n",
        "          path_file_fif2 = subjectPath + '/' + folder_name + '/nback2.fif'\n",
        "          # data_vis_A(i,path_file_fif1, folder_name,'nback2')\n",
        "\n",
        "          file_path = path_file_fif1\n",
        "          subject_name = folder_name + \"_\" + \"nback2\"\n",
        "          epochs = mne.read_epochs(fname=file_path,verbose=False)\n",
        "          # Calculate Welch's power spectrum density (FFT) -> (epochs, channels, freq bins) shape\n",
        "          [psds,freqs] = calculate_psd(epochs,subject_name,fminmax=freq_range,window='hamming',\n",
        "                                 window_duration=2,window_overlap=0.5,plot=False)\n",
        "          # Average all epochs and channels together -> (freq bins,) shape\n",
        "          psds_allch = psds.mean(axis=(0,1))\n",
        "\n",
        "          psds = psds.mean(axis=(0))\n",
        "          df_psds_ch = array_to_df(subject_name,epochs,psds).reset_index().drop(columns='Subject')\n",
        "\n",
        "          df_psds_regions = df_channels_to_regions(df_psds_ch,brain_regions).reset_index().drop(columns='Subject')\n",
        "\n",
        "          for region in df_psds_regions.columns:\n",
        "            if i == 0:\n",
        "              globals()[\"df_fooof_\"+region] = pd.DataFrame(index=subject_names,columns=['Exponent','Offset','Peaks_CF','Peaks_PW','Peaks_BW','R_2','Error'])\n",
        "\n",
        "            psds_temp = df_psds_regions[region].to_numpy()\n",
        "            # Initialize power spectrum model objects and fit the power spectra\n",
        "            # note: aperiodic_mode should be 'knee' if the frequency range is > 40 Hz\n",
        "            max_n_peaks = float(\"inf\")\n",
        "            fm = FOOOF(peak_width_limits=[1,12],max_n_peaks=max_n_peaks,min_peak_height=0.225,\n",
        "                peak_threshold=2.0,aperiodic_mode='fixed',verbose=True)\n",
        "            fm.fit(freqs,psds_temp,freq_range)\n",
        "            # Plot the fitted model with power spectra \n",
        "            fig = plt.figure()\n",
        "            ax = fig.add_axes([0,0,1,1])\n",
        "            ax.set_title('PSD, {}'.format(subject_name))\n",
        "            if fm.get_params('peak_params').shape==(3,) or max_n_peaks == 0:\n",
        "              fm.plot(ax=ax)\n",
        "            else:\n",
        "              fm.plot(plot_peaks='shade', peak_kwargs={'color' : 'green'},ax=ax)\n",
        "\n",
        "            # Add model parameters to dataframe\n",
        "            globals()[\"df_fooof_\"+region].loc[globals()[\"df_fooof_\"+region].index[i],'Exponent'] = fm.get_params('aperiodic_params','exponent')\n",
        "            globals()[\"df_fooof_\"+region].loc[globals()[\"df_fooof_\"+region].index[i],'Offset'] = fm.get_params('aperiodic_params','offset')\n",
        "            globals()[\"df_fooof_\"+region].loc[globals()[\"df_fooof_\"+region].index[i],'Peaks_CF'] = fm.get_params('peak_params', 'CF')\n",
        "            globals()[\"df_fooof_\"+region].loc[globals()[\"df_fooof_\"+region].index[i],'Peaks_PW'] = fm.get_params('peak_params', 'PW')\n",
        "            globals()[\"df_fooof_\"+region].loc[globals()[\"df_fooof_\"+region].index[i],'Peaks_BW'] = fm.get_params('peak_params', 'BW')\n",
        "            globals()[\"df_fooof_\"+region].loc[globals()[\"df_fooof_\"+region].index[i],'R_2'] = fm.get_params('r_squared')\n",
        "            globals()[\"df_fooof_\"+region].loc[globals()[\"df_fooof_\"+region].index[i],'Error'] = fm.get_params('error')\n",
        "\n",
        "# ///////////////////////////////////////////////////////////////////////////////////////\n",
        "\n",
        "          path_file_fif3 = subjectPath + '/' + folder_name + '/nback3.fif'\n",
        "          # data_vis_A(i,path_file_fif1, folder_name,'nback3')\n",
        "\n",
        "          file_path = path_file_fif1\n",
        "          subject_name = folder_name + \"_\" + \"nback3\"\n",
        "          epochs = mne.read_epochs(fname=file_path,verbose=False)\n",
        "          # Calculate Welch's power spectrum density (FFT) -> (epochs, channels, freq bins) shape\n",
        "          [psds,freqs] = calculate_psd(epochs,subject_name,fminmax=freq_range,window='hamming',\n",
        "                                 window_duration=2,window_overlap=0.5,plot=False)\n",
        "          # Average all epochs and channels together -> (freq bins,) shape\n",
        "          psds_allch = psds.mean(axis=(0,1))\n",
        "\n",
        "          psds = psds.mean(axis=(0))\n",
        "          df_psds_ch = array_to_df(subject_name,epochs,psds).reset_index().drop(columns='Subject')\n",
        "\n",
        "          df_psds_regions = df_channels_to_regions(df_psds_ch,brain_regions).reset_index().drop(columns='Subject')\n",
        "\n",
        "          for region in df_psds_regions.columns:\n",
        "            if i == 0:\n",
        "              globals()[\"df_fooof_\"+region] = pd.DataFrame(index=subject_names,columns=['Exponent','Offset','Peaks_CF','Peaks_PW','Peaks_BW','R_2','Error'])\n",
        "\n",
        "            psds_temp = df_psds_regions[region].to_numpy()\n",
        "            # Initialize power spectrum model objects and fit the power spectra\n",
        "            # note: aperiodic_mode should be 'knee' if the frequency range is > 40 Hz\n",
        "            max_n_peaks = float(\"inf\")\n",
        "            fm = FOOOF(peak_width_limits=[1,12],max_n_peaks=max_n_peaks,min_peak_height=0.225,\n",
        "                peak_threshold=2.0,aperiodic_mode='fixed',verbose=True)\n",
        "            fm.fit(freqs,psds_temp,freq_range)\n",
        "            # Plot the fitted model with power spectra \n",
        "            fig = plt.figure()\n",
        "            ax = fig.add_axes([0,0,1,1])\n",
        "            ax.set_title('PSD, {}'.format(subject_name))\n",
        "            if fm.get_params('peak_params').shape==(3,) or max_n_peaks == 0:\n",
        "              fm.plot(ax=ax)\n",
        "            else:\n",
        "              fm.plot(plot_peaks='shade', peak_kwargs={'color' : 'green'},ax=ax)\n",
        "\n",
        "            # Add model parameters to dataframe\n",
        "            globals()[\"df_fooof_\"+region].loc[globals()[\"df_fooof_\"+region].index[i],'Exponent'] = fm.get_params('aperiodic_params','exponent')\n",
        "            globals()[\"df_fooof_\"+region].loc[globals()[\"df_fooof_\"+region].index[i],'Offset'] = fm.get_params('aperiodic_params','offset')\n",
        "            globals()[\"df_fooof_\"+region].loc[globals()[\"df_fooof_\"+region].index[i],'Peaks_CF'] = fm.get_params('peak_params', 'CF')\n",
        "            globals()[\"df_fooof_\"+region].loc[globals()[\"df_fooof_\"+region].index[i],'Peaks_PW'] = fm.get_params('peak_params', 'PW')\n",
        "            globals()[\"df_fooof_\"+region].loc[globals()[\"df_fooof_\"+region].index[i],'Peaks_BW'] = fm.get_params('peak_params', 'BW')\n",
        "            globals()[\"df_fooof_\"+region].loc[globals()[\"df_fooof_\"+region].index[i],'R_2'] = fm.get_params('r_squared')\n",
        "            globals()[\"df_fooof_\"+region].loc[globals()[\"df_fooof_\"+region].index[i],'Error'] = fm.get_params('error')\n",
        "\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",
        "        # Export FOOOF data for all regions\n",
        "        for region in df_psds_regions.columns:\n",
        "          globals()[\"df_fooof_\"+region].to_excel(r\"{}/FOOOF/{}_{}_fooof.xlsx\".format(results_foldername,exp_folder,region))\n",
        "          display(globals()[\"df_fooof_\"+region])\n",
        "\n",
        "\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",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "idHT3HZSeRL_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def data_vis_A(i, file_path, subject_name, sesssion_name):\n",
        "  subject_name = subject_name + \"_\" + sesssion_name\n",
        "  epochs = mne.read_epochs(fname=file_path,verbose=False)\n",
        "  # Calculate Welch's power spectrum density (FFT) -> (epochs, channels, freq bins) shape\n",
        "  [psds,freqs] = calculate_psd(epochs,subject_name,fminmax=freq_range,window='hamming',\n",
        "                                 window_duration=2,window_overlap=0.5,plot=False)\n",
        "  # Average all epochs and channels together -> (freq bins,) shape\n",
        "  psds_allch = psds.mean(axis=(0,1))\n",
        "\n",
        "  psds = psds.mean(axis=(0))\n",
        "  df_psds_ch = array_to_df(subject_name,epochs,psds).reset_index().drop(columns='Subject')\n",
        "\n",
        "  df_psds_regions = df_channels_to_regions(df_psds_ch,brain_regions).reset_index().drop(columns='Subject')\n",
        "\n",
        "  for region in df_psds_regions.columns:\n",
        "    if i == 0:\n",
        "      globals()[\"df_fooof_\"+region] = pd.DataFrame(index=subject_names,columns=['Exponent','Offset','Peaks_CF','Peaks_PW','Peaks_BW','R_2','Error'])\n",
        "\n",
        "    psds_temp = df_psds_regions[region].to_numpy()\n",
        "    # Initialize power spectrum model objects and fit the power spectra\n",
        "    # note: aperiodic_mode should be 'knee' if the frequency range is > 40 Hz\n",
        "    max_n_peaks = float(\"inf\")\n",
        "    fm = FOOOF(peak_width_limits=[1,12],max_n_peaks=max_n_peaks,min_peak_height=0.225,\n",
        "                peak_threshold=2.0,aperiodic_mode='fixed',verbose=True)\n",
        "    fm.fit(freqs,psds_temp,freq_range)\n",
        "    # Plot the fitted model with power spectra \n",
        "    fig = plt.figure()\n",
        "    ax = fig.add_axes([0,0,1,1])\n",
        "    ax.set_title('PSD, {}'.format(subject_name))\n",
        "    if fm.get_params('peak_params').shape==(3,) or max_n_peaks == 0:\n",
        "      fm.plot(ax=ax)\n",
        "    else:\n",
        "      fm.plot(plot_peaks='shade', peak_kwargs={'color' : 'green'},ax=ax)\n",
        "\n",
        "    # Add model parameters to dataframe\n",
        "    globals()[\"df_fooof_\"+region].loc[globals()[\"df_fooof_\"+region].index[i],'Exponent'] = fm.get_params('aperiodic_params','exponent')\n",
        "    globals()[\"df_fooof_\"+region].loc[globals()[\"df_fooof_\"+region].index[i],'Offset'] = fm.get_params('aperiodic_params','offset')\n",
        "    globals()[\"df_fooof_\"+region].loc[globals()[\"df_fooof_\"+region].index[i],'Peaks_CF'] = fm.get_params('peak_params', 'CF')\n",
        "    globals()[\"df_fooof_\"+region].loc[globals()[\"df_fooof_\"+region].index[i],'Peaks_PW'] = fm.get_params('peak_params', 'PW')\n",
        "    globals()[\"df_fooof_\"+region].loc[globals()[\"df_fooof_\"+region].index[i],'Peaks_BW'] = fm.get_params('peak_params', 'BW')\n",
        "    globals()[\"df_fooof_\"+region].loc[globals()[\"df_fooof_\"+region].index[i],'R_2'] = fm.get_params('r_squared')\n",
        "    globals()[\"df_fooof_\"+region].loc[globals()[\"df_fooof_\"+region].index[i],'Error'] = fm.get_params('error')\n",
        "\n"
      ],
      "metadata": {
        "id": "0w_9c5hPelMe"
      },
      "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": "f9IKBTmrhYlS"
      },
      "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": "WuVXGhgwhelQ"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}