{ "cells": [ { "cell_type": "markdown", "id": "da4881be", "metadata": {}, "source": [ "# Optimal Transport in linear ICA: EEG Artifact Removal\n", "#### In this application, we apply OT-ICA to a real-world clinical EEG dataset to demonstrate its ability to isolate sparse, super-Gaussian biological artifacts (like eye-blinks) from continuous brain wave recordings.\n", "\n", "### The Challenge of Volume Conduction\n", "In EEG recordings, the human skull acts as a linear volume conductor. This means the electrical signals generated by various brain regions, as well as non-brain muscles (like the eyes and jaw), mix instantaneously and linearly before reaching the scalp electrodes. This perfectly satisfies the foundational linear mixing model: $X = AS$.\n", "\n", "Eye blinks (ocular artifacts) are notoriously problematic in neuroscience because their high amplitude completely masks the underlying cognitive brain waves. However, because eye blinks are mathematically highly sparse and \"super-Gaussian\", ICA is the gold-standard technique for isolating and removing them.\n", "\n", "In this notebook, we use the `mne` library to fetch a real clinical MEG/EEG sample dataset, and deploy OT-ICA to blindly separate the blink artifact into a single independent component." ] }, { "cell_type": "code", "execution_count": 1, "id": "92550e87", "metadata": {}, "outputs": [], "source": [ "import numpy as np\n", "import torch\n", "import matplotlib.pyplot as plt\n", "import mne\n", "import warnings\n", "warnings.filterwarnings('ignore')\n", "\n", "from wasserstein_ica import WassersteinICA\n", "\n", "# ==========================================\n", "# 1. Thesis Theme Configuration\n", "# ==========================================\n", "def set_thesis_theme():\n", " plt.rcParams.update({\n", " 'figure.figsize': (14, 8),\n", " 'figure.dpi': 150,\n", " 'axes.grid': True,\n", " 'grid.alpha': 0.3,\n", " 'axes.spines.top': False,\n", " 'axes.spines.right': False,\n", " 'font.size': 12,\n", " 'lines.linewidth': 1.0\n", " })\n", "set_thesis_theme()" ] }, { "cell_type": "code", "execution_count": 2, "id": "e741a55a", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Fetching MNE Sample Dataset...\n", "Using default location ~/mne_data for sample...\n", "Creating /home/ashutoshjha/mne_data\n", "Fetching 1 file for the sample dataset ...\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "Downloading file 'MNE-sample-data-processed.tar.gz' from 'https://osf.io/download/86qa2?version=6' to '/home/ashutoshjha/mne_data'.\n" ] }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "fd6b093a045d4d6aa6045599594c188e", "version_major": 2, "version_minor": 0 }, "text/plain": [ " 0%| | 0.00/1.65G [00:00" ] }, "metadata": {}, "output_type": "display_data" }, { "name": "stdout", "output_type": "stream", "text": [ "Takeaway: The massive, synchronized spikes spanning all frontal channels (the eye-blinks) have been perfectly isolated into a single independent component, successfully cleaning the underlying continuous brain waves in the other components.\n" ] } ], "source": [ "# ==========================================\n", "# 4. Visualization: Raw Data vs OT-ICA Components\n", "# ==========================================\n", "fig, axes = plt.subplots(2, 1, figsize=(14, 10), sharex=True)\n", "colors = ['#0173B2', '#DE8F05', '#029E73', '#D55E00', '#CC78BC']\n", "\n", "# Plot A: Real Scalp EEG Electrodes (Corrupted by Blinks)\n", "for i in range(dim):\n", " # Offset by i*5 to stack the signals vertically\n", " axes[0].plot(time, X_real[i] - i*5, color='#555555')\n", "axes[0].set_title(\"Real Clinical EEG (Frontal Electrodes) Corrupted by Eye Blinks\", fontsize=14, fontweight='bold')\n", "axes[0].set_yticks([-20, -15, -10, -5, 0])\n", "axes[0].set_yticklabels(frontal_channels[::-1]) \n", "\n", "# Plot B: OT-ICA Recovered Components\n", "for i in range(dim):\n", " # Offset by i*5 to stack the signals vertically\n", " axes[1].plot(time, Z_hat[i] - i*5, color=colors[i])\n", "axes[1].set_title(\"OT-ICA Recovered Independent Components\", fontsize=14, fontweight='bold')\n", "axes[1].set_yticks([-20, -15, -10, -5, 0])\n", "axes[1].set_yticklabels([f'Comp {i+1}' for i in range(dim)][::-1])\n", "axes[1].set_xlabel(\"Time (Seconds)\", fontsize=12)\n", "\n", "plt.tight_layout()\n", "plt.show()\n", "\n", "print(\"Takeaway: The massive, synchronized spikes spanning all frontal channels (the eye-blinks) have been perfectly isolated into a single independent component, successfully cleaning the underlying continuous brain waves in the other components.\")" ] }, { "cell_type": "code", "execution_count": null, "id": "4d107284", "metadata": {}, "outputs": [], "source": [] } ], "metadata": { "kernelspec": { "display_name": "ot_ica", "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.4" } }, "nbformat": 4, "nbformat_minor": 5 }