muse_quality.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "electronic-adjustment",
   "metadata": {},
   "source": [
    "# Muse EEG data quality\n",
    "This notebook aggregates quality indicators of EEG signal (returned by the muse headband), for each participant.\n",
    "These values represent the percentage of the time the EEG headband was on the head, and the percentages of a good, medium, and bad quality signal separately for each electrode (obtained from the device).\n",
    "The data is available for each participant, each stimulus, and study phases.\n",
    "\n",
    "The quality of EEG signal is reported using [Horse Shoe Indicator (HSI)](https://mind-monitor.com/Technical_Manual.php#help_horseshoe) values provided by the device, which represent how well electrodes fit the participant's head.\n",
    "\n",
    "Data on signal quality is stored in `headband_dict`, which can be exported to a .json file."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "beneficial-dining",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "dependent-yesterday",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import pandas as pd\n",
    "import os\n",
    "import configparser\n",
    "from copy import deepcopy\n",
    "from tqdm.notebook import tqdm\n",
    "from pprint import pprint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "based-aging",
   "metadata": {},
   "outputs": [],
   "source": [
    "config = configparser.ConfigParser()\n",
    "config.read(\"config.ini\")\n",
    "\n",
    "if not os.path.exists(config['DataDirectories']['unzipped_dataset']):\n",
    "    raise Exception(\"Please set path for unzipped dataset in config.ini\")\n",
    "\n",
    "DATA_DIR = config['DataDirectories']['unzipped_dataset']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "digital-advice",
   "metadata": {},
   "outputs": [],
   "source": [
    "# generate paths\n",
    "paths = []\n",
    "for part in os.listdir(DATA_DIR):\n",
    "    for f_name in os.listdir(DATA_DIR + '/' + part):\n",
    "        if '_MUSE' in f_name:\n",
    "            paths.append((part, f_name.split('_')[1], f_name.split('_')[2], DATA_DIR + '/' + part + '/' + f_name))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "caroline-medication",
   "metadata": {},
   "outputs": [],
   "source": [
    "def update_emotion_stats(emotion_dict, new_ok_len, new_len):\n",
    "    \"\"\"\n",
    "    Update Avg_HeadBandOn and Num_samples in provided emotion_dict\n",
    "    \"\"\"\n",
    "    old_ok = emotion_dict.get('Avg_HeadBandOn')\n",
    "    old_len = emotion_dict.get('Num_samples')\n",
    "    total_len = old_len + new_len\n",
    "    emotion_dict['Avg_HeadBandOn'] = (old_ok*old_len + new_ok_len)/total_len\n",
    "    emotion_dict['Num_samples'] = total_len\n",
    "    \n",
    "def get_hsi_scores(df, HSI_COLS=['HSI_TP9', 'HSI_AF7', 'HSI_AF8', 'HSI_TP10']):\n",
    "    \"\"\"\n",
    "    Get HSI scores for each channel.\n",
    "    1 - Good, 2 - Medium, 4 - Bad\n",
    "    \"\"\"\n",
    "    ones = (df[HSI_COLS]==1).sum()\n",
    "    twos = (df[HSI_COLS]==2).sum()\n",
    "    fours = (df[HSI_COLS]==4).sum()\n",
    "    return {col: {1: ones[col]/len(df), 2: twos[col]/len(df), 4: fours[col]/len(df)} for col in HSI_COLS}\n",
    "\n",
    "def fix_headband_dict(dict_to_fix):\n",
    "    \"\"\"\n",
    "    Delete washout from baseline fields.\n",
    "    :param dict_to_fix:\n",
    "    :return:\n",
    "    \"\"\"\n",
    "    for part_dict in dict_to_fix.values():\n",
    "        for emotion, emo_dict in part_dict.items():\n",
    "            if emotion == 'BASELINE':\n",
    "                emo_dict.pop('WASHOUT')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "retired-missile",
   "metadata": {},
   "source": [
    "## Get the EEG data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "weird-emergency",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "defd15c13e7d47faac9a9f1eba26b2b0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1312 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "EMOTION_DICT_TEMPLATE = {\n",
    "    'Avg_HeadBandOn': 0,\n",
    "    'Num_samples': 0, \n",
    "    'STIMULUS': dict(), \n",
    "    'QUESTIONNAIRES': dict(), \n",
    "    'WASHOUT': dict()\n",
    "}\n",
    "HSI_COLS = ['HSI_TP9', 'HSI_AF7', 'HSI_AF8', 'HSI_TP10']\n",
    "HeadBandOn_COLS = ['HeadBandOn']\n",
    "\n",
    "headband_dict = dict()\n",
    "\n",
    "\n",
    "for part, emotion, phase, path in tqdm(paths):\n",
    "    # load data from one phase\n",
    "    df = pd.read_json(path)[HeadBandOn_COLS+HSI_COLS]\n",
    "    df.dropna(inplace=True)\n",
    "    # add loaded data to headband_dict\n",
    "    phase_len_ok = len(df[df['HeadBandOn'] == 1])\n",
    "    phase_len = len(df)\n",
    "    headband_dict.setdefault(part, dict())\n",
    "    headband_dict[part].setdefault(\n",
    "        emotion, \n",
    "        deepcopy(EMOTION_DICT_TEMPLATE)\n",
    "    )\n",
    "    update_emotion_stats(headband_dict[part][emotion], phase_len_ok, phase_len)    \n",
    "    headband_dict[part][emotion][phase].setdefault('HeadBandOn', phase_len_ok/phase_len)\n",
    "    headband_dict[part][emotion][phase].update(get_hsi_scores(df))\n",
    "\n",
    "# delete unnecessary fields\n",
    "fix_headband_dict(headband_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "inside-niagara",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_headband_off_instances(headband_dict):\n",
    "    \"\"\"\n",
    "    Return dict: instances with < 100% HeadBandOn time.\n",
    "        Participant: Emotion: HeadBandOn time % (value between 0-1)\n",
    "    \"\"\"\n",
    "    ret_dict = dict()\n",
    "    for part, part_dict in headband_dict.items():\n",
    "        for emotion, emo_dict in part_dict.items():\n",
    "            if emo_dict.get('Avg_HeadBandOn') < 1:\n",
    "                ret_dict.setdefault(part, dict())\n",
    "                ret_dict[part].setdefault(emotion, emo_dict.get('Avg_HeadBandOn'))\n",
    "    return ret_dict"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "previous-earthquake",
   "metadata": {},
   "source": [
    "## Print and save quality indicators"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "protecting-military",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'23': {'ANGER': 0.9930110493329372},\n",
      " '26': {'DISGUST': 0.9908630192727407, 'LIKING': 0.9925431711145997},\n",
      " '30': {'ANGER': 0.9925554199255542, 'LIKING': 0.9927327291985406},\n",
      " '33': {'NEUTRAL': 0.9932857628270473},\n",
      " '34': {'SURPRISE': 0.9905743178782676},\n",
      " '38': {'FEAR': 0.9928443131119995},\n",
      " '39': {'AMUSEMENT': 0.0,\n",
      "        'ANGER': 0.0,\n",
      "        'AWE': 0.0,\n",
      "        'BASELINE': 0.2779818266740266,\n",
      "        'DISGUST': 0.0,\n",
      "        'ENTHUSIASM': 0.0,\n",
      "        'LIKING': 0.738687376548752,\n",
      "        'NEUTRAL': 0.0,\n",
      "        'SADNESS': 0.0,\n",
      "        'SURPRISE': 0.0},\n",
      " '40': {'LIKING': 0.9927198058614897},\n",
      " '41': {'AMUSEMENT': 0.9929977745963798},\n",
      " '44': {'BASELINE': 0.9935078589076382},\n",
      " '47': {'ENTHUSIASM': 0.9929960868825498},\n",
      " '49': {'AWE': 0.9930354851898324, 'SURPRISE': 0.9909755206430398},\n",
      " '50': {'AMUSEMENT': 0.9933898894746702,\n",
      "        'ANGER': 0.0,\n",
      "        'AWE': 0.002656717284631689,\n",
      "        'DISGUST': 0.0,\n",
      "        'ENTHUSIASM': 0.0,\n",
      "        'FEAR': 0.0,\n",
      "        'LIKING': 0.0,\n",
      "        'SURPRISE': 0.0},\n",
      " '53': {'AWE': 0.9940053978544207},\n",
      " '54': {'ANGER': 0.9928849200633731, 'SADNESS': 0.992578570999339},\n",
      " '55': {'BASELINE': 0.9949899596357071, 'DISGUST': 0.9913095489409612},\n",
      " '56': {'NEUTRAL': 0.9927631771703144},\n",
      " '57': {'AMUSEMENT': 0.9927472398402631, 'LIKING': 0.9927893738140418},\n",
      " '58': {'AWE': 0.9904173961116742,\n",
      "        'DISGUST': 0.47722319093286836,\n",
      "        'SURPRISE': 0.6952954150389712},\n",
      " '60': {'BASELINE': 0.990509033203125},\n",
      " '61': {'AWE': 0.5339260837498976,\n",
      "        'BASELINE': 0.9827886215766379,\n",
      "        'DISGUST': 0.4864754329553788,\n",
      "        'NEUTRAL': 0.9867381637457494,\n",
      "        'SURPRISE': 0.8751204919395047},\n",
      " '62': {'BASELINE': 0.7876404252571773, 'NEUTRAL': 0.627879719223306},\n",
      " '63': {'BASELINE': 0.9941068403597929, 'FEAR': 0.9927926350651435},\n",
      " '64': {'FEAR': 0.9931456480415146, 'NEUTRAL': 0.993034404963339}}\n"
     ]
    }
   ],
   "source": [
    "# print instances with < 100% HeadBandOn time\n",
    "pprint(get_headband_off_instances(headband_dict))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "mounted-intensity",
   "metadata": {},
   "outputs": [],
   "source": [
    "# save whole headband_dict to .json file \n",
    "with open('muse_quality.json', 'w') as f:\n",
    "    json.dump(headband_dict, f)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}