02_load_and_concatenate.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load and concatenate\n",
    "\n",
    "In this script we load the data from the original BrainVision format into MNE-Python, and then save it as a `.fif` file.\n",
    "\n",
    "For a given subject, load all available datafiles (corresponding to tasks),\n",
    "and do the following:\n",
    "- set EEG reference to None (=keep original reference)\n",
    "- set EEG, ECG, EOG channel types\n",
    "- add a 1020 standard montage to have coordinates for plotting\n",
    "- extract events and scale the event markers for each task (**A**ctive or **Y**oked, **F**ixed or **V**ariable, or description)\n",
    "    - AF          += 100\n",
    "    - AV          += 200\n",
    "    - YF          += 300\n",
    "    - YV          += 400\n",
    "    - description += 500\n",
    "- load bad segments and channels and add them to the data\n",
    "- concatenate the raw data across trials\n",
    "- save the concatenated raw data and corresponding events\n",
    "  (including their updated event markers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib qt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-06-29T12:54:45.910603Z",
     "start_time": "2020-06-29T12:54:45.897360Z"
    }
   },
   "outputs": [],
   "source": [
    "import itertools\n",
    "import multiprocessing\n",
    "import os\n",
    "import os.path as op\n",
    "\n",
    "import mne\n",
    "\n",
    "from utils import (\n",
    "    BIDS_ROOT,\n",
    "    TASK_BASED_INCREMENTS,\n",
    "    provide_trigger_dict,\n",
    "    task_not_present_for_subject,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# IO: Where to find the data\n",
    "eeg_path_template = op.join(\n",
    "    BIDS_ROOT, \"sub-{0:02}\", \"eeg\", \"sub-{0:02}_task-{1}_eeg.vhdr\"\n",
    ")\n",
    "\n",
    "# Where to find the annotations and bad channels\n",
    "fname_annots_template = op.join(\n",
    "    BIDS_ROOT, \"derivatives\", \"sub-{0:02}\", \"sub-{0:02}_task-{1}_annotations.txt\"\n",
    ")\n",
    "fname_channels_template = op.join(\n",
    "    BIDS_ROOT, \"derivatives\", \"sub-{0:02}\", \"sub-{0:02}_task-{1}_badchannels.txt\"\n",
    ")\n",
    "\n",
    "# Where to save concatenated raw data and events\n",
    "fname_rawconcat_template = op.join(\n",
    "    BIDS_ROOT, \"derivatives\", \"sub-{0:02}\", \"sub-{0:02}_concat_eeg-raw.fif.gz\"\n",
    ")\n",
    "fname_events_template = op.join(\n",
    "    BIDS_ROOT, \"derivatives\", \"sub-{0:02}\", \"sub-{0:02}_concat_eeg-eve.fif.gz\"\n",
    ")\n",
    "\n",
    "# Pack all names in a dict\n",
    "name_templates = dict()\n",
    "name_templates[\"eeg\"] = eeg_path_template\n",
    "name_templates[\"annots\"] = fname_annots_template\n",
    "name_templates[\"channels\"] = fname_channels_template\n",
    "name_templates[\"rawconcat\"] = fname_rawconcat_template\n",
    "name_templates[\"events\"] = fname_events_template"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "subjects = range(1, 41)\n",
    "tasks = [\"ActiveFixed\", \"ActiveVariable\", \"YokedFixed\", \"YokedVariable\", \"description\"]\n",
    "\n",
    "# How many subjects to run over in parallel\n",
    "NJOBS = max(2, multiprocessing.cpu_count() - 6)\n",
    "\n",
    "# Whether or not to overwrite existing files\n",
    "overwrite = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_and_concatenate(subj, tasks, name_templates, overwrite):\n",
    "    \"\"\"Load BrainVision raw data and save as fif.\n",
    "    \n",
    "    For a given subject, load all available datafiles (corresponding to tasks),\n",
    "    and do the following:\n",
    "    - set EEG reference to None (=keep original reference)\n",
    "    - set EEG, ECG, EOG channel types\n",
    "    - add a 1020 standard montage to have coordinates for plotting\n",
    "    - extract events and scale the event markers for each task (**A**ctive or\n",
    "      **Y**oked, **F**ixed or **V**ariable, or description)\n",
    "        - AF          += 100\n",
    "        - AV          += 200\n",
    "        - YF          += 300\n",
    "        - YV          += 400\n",
    "        - description += 500\n",
    "    - load bad segments and channels and add them to the data\n",
    "    - concatenate the raw data across trials\n",
    "    - save the concatenated raw data and corresponding events\n",
    "      (including their updated event markers)\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    subj : int\n",
    "        The subject identifier in the range(1, 41).\n",
    "    tasks : list of str\n",
    "        The task names.\n",
    "    name_templates : dict\n",
    "        A dictionary of string templates. Needs the following keys:\n",
    "        \"eeg\", \"annots\", \"channels\", \"rawconcat\", \"events\"\n",
    "    overwrite : bool\n",
    "        Whether to overwrite existing files.\n",
    "\n",
    "    \"\"\"\n",
    "    # Unpack places to load and save data from\n",
    "    eeg_path_template = name_templates[\"eeg\"]\n",
    "    fname_annots_template = name_templates[\"annots\"]\n",
    "    fname_channels_template = name_templates[\"channels\"]\n",
    "    fname_rawconcat_template = name_templates[\"rawconcat\"]\n",
    "    fname_events_template = name_templates[\"events\"]\n",
    "\n",
    "    # Handle existing files\n",
    "    skip_events = False\n",
    "    if op.exists(fname_events_template.format(subj)):\n",
    "        if overwrite:\n",
    "            os.remove(fname_events_template.format(subj))\n",
    "        else:\n",
    "            skip_events = True\n",
    "\n",
    "    skip_raw = False\n",
    "    if op.exists(fname_rawconcat_template.format(subj)):\n",
    "        if overwrite:\n",
    "            os.remove(fname_rawconcat_template.format(subj))\n",
    "        else:\n",
    "            skip_raw = True\n",
    "\n",
    "    # If we have all needed files already, return early\n",
    "    if skip_events and skip_raw:\n",
    "        return\n",
    "\n",
    "    # Else, start preparing the files ...\n",
    "    # Get the markers that wered used in the experiment\n",
    "    marker_ids = [ord(i) for i in provide_trigger_dict().values()]\n",
    "\n",
    "    all_events = list()\n",
    "    files = list()\n",
    "    bads = list()\n",
    "    for task in tasks:\n",
    "        if task_not_present_for_subject(subj, task):\n",
    "            continue\n",
    "\n",
    "        # Get the data\n",
    "        raw = mne.io.read_raw_brainvision(\n",
    "            eeg_path_template.format(subj, task), preload=True\n",
    "        )\n",
    "\n",
    "        # Suppress an automatic \"average reference\"\n",
    "        raw.set_eeg_reference(ref_channels=[])\n",
    "\n",
    "        # Set the EOG and ECG channels to their type\n",
    "        raw.set_channel_types({\"ECG\": \"ecg\", \"HEOG\": \"eog\", \"VEOG\": \"eog\"})\n",
    "\n",
    "        # Set a standard montage for plotting later\n",
    "        montage = mne.channels.make_standard_montage(\"standard_1020\")\n",
    "        raw.set_montage(montage)\n",
    "\n",
    "        # Extract events, incrementing event markers according to a\n",
    "        # task dependent mapping\n",
    "        event_id = dict()\n",
    "        for marker in marker_ids:\n",
    "            key = \"Stimulus/S{: >3}\".format(marker)\n",
    "            val = marker + TASK_BASED_INCREMENTS[task]\n",
    "            event_id[key] = val\n",
    "\n",
    "        events, event_id = mne.events_from_annotations(raw, event_id)\n",
    "\n",
    "        # Add bad segments and channels to the data\n",
    "        fname_annots = fname_annots_template.format(subj, task)\n",
    "        fname_channels = fname_channels_template.format(subj, task)\n",
    "\n",
    "        # annotations have orig_date \"anonymized\", so it will be set to\n",
    "        # None upon reading.\n",
    "        # For concatenating our annotations with raw.annotations however,\n",
    "        # they need to have the same orig_time property.\n",
    "        # Taking the raw.annotations.orig_time results in the expected\n",
    "        # alignment of raw and annotations\n",
    "        annots = mne.read_annotations(fname_annots)\n",
    "        annots = mne.Annotations(annots.onset, annots.duration,\n",
    "                                 annots.description, raw.annotations.orig_time)\n",
    "\n",
    "        raw.set_annotations(raw.annotations + annots)\n",
    "\n",
    "        # Collect bads across all tasks\n",
    "        # Later set concatenated raw.info[\"bads\"]\n",
    "        # See also: raw.load_bad_channels\n",
    "        with open(fname_channels, \"r\") as fin:\n",
    "            bads += [line.strip() for line in fin.readlines()]\n",
    "\n",
    "        # append task based datafiles to list for concatenation\n",
    "        all_events.append(events)\n",
    "        files.append(raw)\n",
    "\n",
    "    # Set the union of bad channels across tasks to each task\n",
    "    for raw in files:\n",
    "        raw.info[\"bads\"] = list(set(bads))\n",
    "\n",
    "    # Concatenate datafiles and events\n",
    "    raw, events = mne.concatenate_raws(files, events_list=all_events)\n",
    "\n",
    "    # Save\n",
    "    if not skip_events:\n",
    "        mne.write_events(fname_events_template.format(subj), events)\n",
    "\n",
    "    if not skip_raw:\n",
    "        raw.save(fname_rawconcat_template.format(subj))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run the pipeline in parallel over subjects\n",
    "pool_inputs = itertools.product(subjects, [tasks], [name_templates], [overwrite])\n",
    "\n",
    "with multiprocessing.Pool(NJOBS) as pool:\n",
    "    pool.starmap(load_and_concatenate, pool_inputs)"
   ]
  }
 ],
 "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.7"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}