{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n\n# Principal Component Analysis - Optimal Basis Sets (PCA-OBS) removing cardiac artefact\n\nThis script shows an example of how to use an adaptation of PCA-OBS\n:footcite:`NiazyEtAl2005`. PCA-OBS was originally designed to remove\nthe ballistocardiographic artefact in simultaneous EEG-fMRI. Here, it\nhas been adapted to remove the delay between the detected R-peak and the\nballistocardiographic artefact such that the algorithm can be applied to\nremove the cardiac artefact in EEG (electroencephalography) and ESG\n(electrospinography) data. We will illustrate how it works by applying the\nalgorithm to ESG data, where the effect of removal is most pronounced.\n\nSee: https://www.biorxiv.org/content/10.1101/2024.09.05.611423v1\nfor more details on the dataset and application for ESG data.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Authors: Emma Bailey <bailey@cbs.mpg.de>,\n#          Steinn Hauser Magnusson <hausersteinn@gmail.com>\n# License: BSD-3-Clause\n# Copyright the MNE-Python contributors.\n\nimport glob\n\nimport numpy as np"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Download sample subject data from OpenNeuro if you haven't already.\nThis will download simultaneous EEG and ESG data from a single run of a\nsingle participant after median nerve stimulation of the left wrist.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import openneuro\nfrom matplotlib import pyplot as plt\n\nimport mne\nfrom mne import Epochs, events_from_annotations\nfrom mne.io import read_raw_eeglab\nfrom mne.preprocessing import find_ecg_events, fix_stim_artifact\n\n# add the path where you want the OpenNeuro data downloaded. Each run is ~2GB of data\nds = \"ds004388\"\ntarget_dir = mne.datasets.default_path() / ds\nrun_name = \"sub-001/eeg/*median_run-03_eeg*.set\"\nif not glob.glob(str(target_dir / run_name)):\n    target_dir.mkdir(exist_ok=True)\n    openneuro.download(dataset=ds, target_dir=target_dir, include=run_name[:-4])\nblock_files = glob.glob(str(target_dir / run_name))\nassert len(block_files) == 1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Define the esg channels (arranged in two patches over the neck and lower back).\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "esg_chans = [\n    \"S35\",\n    \"S24\",\n    \"S36\",\n    \"Iz\",\n    \"S17\",\n    \"S15\",\n    \"S32\",\n    \"S22\",\n    \"S19\",\n    \"S26\",\n    \"S28\",\n    \"S9\",\n    \"S13\",\n    \"S11\",\n    \"S7\",\n    \"SC1\",\n    \"S4\",\n    \"S18\",\n    \"S8\",\n    \"S31\",\n    \"SC6\",\n    \"S12\",\n    \"S16\",\n    \"S5\",\n    \"S30\",\n    \"S20\",\n    \"S34\",\n    \"S21\",\n    \"S25\",\n    \"L1\",\n    \"S29\",\n    \"S14\",\n    \"S33\",\n    \"S3\",\n    \"L4\",\n    \"S6\",\n    \"S23\",\n]\n\n# Interpolation window in seconds for ESG data to remove stimulation artefact\ntstart_esg = -7e-3\ntmax_esg = 7e-3\n\n# Define timing of heartbeat epochs in seconds relative to R-peaks\niv_baseline = [-400e-3, -300e-3]\niv_epoch = [-400e-3, 600e-3]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Next, we perform minimal preprocessing including removing the\nstimulation artefact, downsampling and filtering.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "raw = read_raw_eeglab(block_files[0], verbose=\"error\")\nraw.set_channel_types(dict(ECG=\"ecg\"))\n# Isolate the ESG channels (include the ECG channel for R-peak detection)\nraw.pick(esg_chans + [\"ECG\"])\n# Trim duration and downsample (from 10kHz) to improve example speed\nraw.crop(0, 60).load_data().resample(2000)\n\n# Find trigger timings to remove the stimulation artefact\nevents, event_dict = events_from_annotations(raw)\ntrigger_name = \"Median - Stimulation\"\n\nfix_stim_artifact(\n    raw,\n    events=events,\n    event_id=event_dict[trigger_name],\n    tmin=tstart_esg,\n    tmax=tmax_esg,\n    mode=\"linear\",\n    stim_channel=None,\n)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Find ECG events and add to the raw structure as event annotations.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "ecg_events, ch_ecg, average_pulse = find_ecg_events(raw, ch_name=\"ECG\")\necg_event_samples = np.asarray(\n    [[ecg_event[0] for ecg_event in ecg_events]]\n)  # Samples only\n\nqrs_event_time = [\n    x / raw.info[\"sfreq\"] for x in ecg_event_samples.reshape(-1)\n]  # Divide by sampling rate to make times\nduration = np.repeat(0.0, len(ecg_event_samples))\ndescription = [\"qrs\"] * len(ecg_event_samples)\n\nraw.annotations.append(\n    qrs_event_time, duration, description, ch_names=[esg_chans] * len(qrs_event_time)\n)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Create evoked response around the detected R-peaks\nbefore and after cardiac artefact correction.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "events, event_ids = events_from_annotations(raw)\nevent_id_dict = {key: value for key, value in event_ids.items() if key == \"qrs\"}\nepochs = Epochs(\n    raw,\n    events,\n    event_id=event_id_dict,\n    tmin=iv_epoch[0],\n    tmax=iv_epoch[1],\n    baseline=tuple(iv_baseline),\n)\nevoked_before = epochs.average()\n\n# Apply function - modifies the data in place. Optionally high-pass filter\n# the data before applying PCA-OBS to remove low frequency drifts\nraw = mne.preprocessing.apply_pca_obs(\n    raw, picks=esg_chans, n_jobs=5, qrs_times=raw.times[ecg_event_samples.reshape(-1)]\n)\n\nepochs = Epochs(\n    raw,\n    events,\n    event_id=event_id_dict,\n    tmin=iv_epoch[0],\n    tmax=iv_epoch[1],\n    baseline=tuple(iv_baseline),\n)\nevoked_after = epochs.average()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Compare evoked responses to assess completeness of artefact removal.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "fig, axes = plt.subplots(1, 1, layout=\"constrained\")\ndata_before = evoked_before.get_data(units=dict(eeg=\"uV\")).T\ndata_after = evoked_after.get_data(units=dict(eeg=\"uV\")).T\nhs = list()\nhs.append(axes.plot(epochs.times, data_before, color=\"k\")[0])\nhs.append(axes.plot(epochs.times, data_after, color=\"green\", label=\"after\")[0])\naxes.set(ylim=[-500, 1000], ylabel=\"Amplitude (\u00b5V)\", xlabel=\"Time (s)\")\naxes.set(title=\"ECG artefact removal using PCA-OBS\")\naxes.legend(hs, [\"before\", \"after\"])\nplt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## References\n.. footbibliography::\n\n"
      ]
    }
  ],
  "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.12.3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}