demo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# BCGNet Demo"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following python toolbox trains a neural network intended for BCG artifact removal in EEG-fMRI datasets. More detail about our method can be found in the paper McIntosh et al. IEEE Trans Biomed Engi at https://ieeexplore.ieee.org/document/9124646"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import commands\n",
    "\n",
    "import os\n",
    "\n",
    "from pathlib import Path\n",
    "from config import get_config\n",
    "from session import Session"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Path setup"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The first step is to set up all the relevant path. Here for the purpose of the demo we will define all path here; however, for custom use it is recommended to set up all path in the yaml file."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Option 1: In YAML File (Recommended)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It is recommended for the user to set up the path in the yaml file\n",
    "\n",
    "\n",
    "Variables that needs to be set up:\n",
    "\n",
    "| Variable Name | Type | Description                                                                                                           |\n",
    "|---------------|------|-----------------------------------------------------------------------------------------------------------------------|\n",
    "| d_root        | str  | the absolute path to the root directory of this package, e.g. \"/home/jsmith/BCGNet/\"                                  |\n",
    "| d_data        | str  | the absolute path to the directory containing all raw datasets, e.g. \"/home/jsmith/data/\"                             |\n",
    "| d_model       | str  | the absolute path to the directory to save trained models                                                             |\n",
    "| d_output      | str  | the absolute path to the directory to save cleaned datasets                                                           |\n",
    "| d_eval        | str  | (Optional) the absolute path to directory containing all the metric datasets used for comparing performance of BCGNet |\n",
    "| str_eval      | str  | (Optional) must be specified if d_eval is given, the name of the alternative method used for comparison               |\n",
    "\n",
    "Once the user has successfully set up all these variable in the yaml file, it's only needed to execute the following command"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# cfg = get_config(d_root / 'config' / 'default_config.yaml')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Option 2: In Python"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For the purpose of this demo, we will set all the variables listed above in the Jupyter notebook. Additionally, here we will set them as pathlib objects instead of strings for convenience."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get the absolute path to the root directory of the package\n",
    "d_root = Path(os.getcwd())\n",
    "\n",
    "# get the absolute path to the directory containing all data\n",
    "# all dataset should be in EEGLAB formats\n",
    "# here the structure of directory is presumed to be\n",
    "# d_data / subXX / input_file_naming_format\n",
    "# where input_file_naming_format is defined in the yaml file\n",
    "d_data = d_root / 'example_data' / 'raw_data'\n",
    "\n",
    "# get the absolute path to the directory to save all trained models\n",
    "# structure of the directory will be\n",
    "# d_model / model_type / subXX / {model_type}_{time_stamp} / {model_type}_{time_stamp}.index\n",
    "\n",
    "# (note: depending on TF version, either save in the new TF checkpoint format or old h5 format)\n",
    "d_model = d_root / 'trained_model' / 'non_cv_model'\n",
    "\n",
    "# get the absolute path to the directory to save all cleaned dataset\n",
    "# structure of the directory will be\n",
    "# d_output / subXX / output_file_naming_format\n",
    "d_output = d_root / 'cleaned_data' / 'non_cv_data'\n",
    "\n",
    "# (Optional)\n",
    "# if the users wish, a dataset used to compare the performance of\n",
    "# BCGNet can be provided, here a OBS-cleaned dataset is used\n",
    "# convention is same as the d_data and all dataset\n",
    "# should be in EEGLAB format\n",
    "\n",
    "# get the absolute path to the directory containing all data\n",
    "# cleaned by the alternative method\n",
    "# here the structure of the directory is also presumed to be\n",
    "# d_eval / subXX / eval_file_naming_format\n",
    "d_eval = d_root / 'example_data' / 'obs_cleaned_data'\n",
    "\n",
    "# (Optional - relevant only if  d_eval is provided)\n",
    "# define the name of the alternative method\n",
    "str_eval = 'OBS'\n",
    "\n",
    "# generate a config (cfg) object from the yaml file\n",
    "# all hyperparameters are from the paper\n",
    "cfg = get_config(filename=d_root / 'config' / 'default_config.yaml')\n",
    "\n",
    "# change all the path (recommended to set these in the yaml file directory)\n",
    "cfg.d_root = d_root\n",
    "cfg.d_data = d_data\n",
    "cfg.d_model = d_model\n",
    "cfg.d_output = d_output\n",
    "cfg.d_eval = d_eval\n",
    "cfg.str_eval = str_eval"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Additional note"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If the user wants to conduct a quick test, the following line can be used to set the maximum number of training iterations to be used."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# If the user just want a quick test, can set the number of maximum epochs\n",
    "# to be few so training will be over quickly via the line below\n",
    "# cfg.num_epochs = 5"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initialize training session"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "All key hyperparamters relevant to preprocessing and training are set in the yaml file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# provide the name of the subject\n",
    "str_sub = 'sub34'\n",
    "\n",
    "# provide the index of the runs to be used for training\n",
    "# if just a single run, then [1] or [2]\n",
    "# if multiple runs then [1, 2]\n",
    "\n",
    "# for a run from sub11 and run index 1\n",
    "# filename is presumed to be\n",
    "# subXX_r0X_\n",
    "vec_idx_run = [1, 2]\n",
    "\n",
    "\n",
    "# str_arch specifies the type of the model to be used\n",
    "# if str_arch is not provided then the default model (same as paper)\n",
    "# is used. If user wants to define their own model, example on how to do it\n",
    "# can be found in models/gru_arch_000.py, the only caveat is that \n",
    "# the name of the file and class name has to be same as the type of the model\n",
    "# e.g. gru_arch_000\n",
    "\n",
    "# random_seed is set to ensure that the splitting of entire dataset into\n",
    "# training, validation and test sets is always the same, useful for model\n",
    "# selection\n",
    "\n",
    "# verbose sets the verbosity of Keras during model training\n",
    "# 0=silent, 1=progress bar, 2=one line per epoch\n",
    "\n",
    "# overwrite specifies whether or not to overwrite existing cleaned data\n",
    "\n",
    "# cv_mode specifies whether or not to use cross validation mode\n",
    "# more on this later\n",
    "s1 = Session(str_sub=str_sub, vec_idx_run=vec_idx_run, str_arch='default_rnn_model',\n",
    "             random_seed=1997, verbose=2, overwrite=False, cv_mode=False, num_fold=5, cfg=cfg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare for training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# loads all dataset\n",
    "s1.load_all_dataset()\n",
    "\n",
    "# preform preprocessing of all dataset and initialize model\n",
    "s1.prepare_training()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model training and generating cleaned dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# train the model\n",
    "s1.train()\n",
    "\n",
    "# generate cleaned dataset\n",
    "s1.clean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot the training history\n",
    "s1.plot_training_history()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluating the performance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Evaluate the performance of the model in terms of RMS and\n",
    "# ratio of band power of cleaned dataset in delta, theta \n",
    "# and alpha bands compared to the raw data\n",
    "\n",
    "# mode specifies which set to evaluate the performance on\n",
    "# mode='train' evaluates on training set\n",
    "# mode='valid' evaluates on validation set\n",
    "# mode='test' evaluates on test set\n",
    "s1.evaluate(mode='test')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot a random epoch from a specified channel and a set\n",
    "# str_ch_eeg should be set to standard EEG channel names, e.g. Pz, Fz, Oz etc.\n",
    "# mode='train' evaluates on training set\n",
    "# mode='valid' evaluates on validation set\n",
    "# mode='test' evaluates on test set\n",
    "s1.plot_random_epoch(str_ch_eeg='T8', mode='test')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot the power spectral density (PSD) from the mean/specified channel\n",
    "# mode='train' evaluates on training set\n",
    "# mode='valid' evaluates on validation set\n",
    "# mode='test' evaluates on test set\n",
    "\n",
    "# str_ch_eeg='avg' plots the mean PSD across all channels\n",
    "# str_ch_eeg could also be set to standard EEG channel names, e.g. Pz, Fz, Oz etc.\n",
    "s1.plot_psd(str_ch_eeg='avg', mode='test')\n",
    "s1.plot_psd(str_ch_eeg='T8', mode='test')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Saving trained model and cleaned dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save trained model\n",
    "s1.save_model()\n",
    "\n",
    "# save cleaned data in .mat files\n",
    "# the saved .mat file has one field 'data' which contains the \n",
    "# n_channel by n_time_stamp matrix holding all cleaned data\n",
    "\n",
    "# note that the unit of the data saved in the mat file \n",
    "# is in Volts instead of in microVolts\n",
    "s1.save_data()\n",
    "\n",
    "# alternatively, save cleaned data in Neuromag .fif format \n",
    "# (note that EEEGLAB support for .fif format is limited)\n",
    "# s1.save_dataset()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Cross validation mode"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Alternatively, if cross validation is deemed necessary, the users can set up a cross validation style session via the following command"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# first change the output and model directory\n",
    "d_model = d_root / 'trained_model' / 'cv_model'\n",
    "d_output = d_root / 'cleaned_data' / 'cv_data'\n",
    "cfg.d_model = d_model\n",
    "cfg.d_output = d_output\n",
    "\n",
    "# it is recommended for user to set the num_fold argument,\n",
    "# which specifies the number of cross validation folds\n",
    "# in which case, percentage of test set and validation set data\n",
    "# will be set to 1/num_fold and remaining data will be the training set\n",
    "# e.g.\n",
    "s2 = Session(str_sub=str_sub, vec_idx_run=vec_idx_run, str_arch='default_rnn_model',\n",
    "             random_seed=1997, verbose=2, overwrite=True, cv_mode=True, num_fold=5, cfg=cfg)\n",
    "\n",
    "# otherwise the number of cross validation folds will be inferred from\n",
    "# percentage of test set data set in the config yaml file via 1/per_test\n",
    "# s2 = Session(str_sub=str_sub, vec_idx_run=vec_idx_run, str_arch='default_rnn_model',\n",
    "#                     random_seed=1997, verbose=2, overwrite=True,\n",
    "#                     cv_mode=True, cfg=cfg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Remaining commands are the same"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s2.load_all_dataset()\n",
    "s2.prepare_training()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s2.train()\n",
    "s2.clean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s2.plot_training_history()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s2.evaluate(mode='test')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Additionally, in the cross validation mode, the user can\n",
    "# specify the fold (0-indexing) from which the figures are\n",
    "# to be plotted\n",
    "\n",
    "# For the demo, plot using the 3rd fold (note the 0-indexing)\n",
    "idx_fold = 2\n",
    "\n",
    "s2.plot_random_epoch(str_ch_eeg='T8', mode='test', idx_fold=idx_fold)\n",
    "\n",
    "s2.plot_psd(str_ch_eeg='avg', mode='test', idx_fold=idx_fold)\n",
    "s2.plot_psd(str_ch_eeg='T8', mode='test', idx_fold=idx_fold)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s2.save_model()\n",
    "s2.save_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.8.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}