learn_graph.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "b7ba7a90-ad88-4e14-a6b7-a8006af52cd9",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import pickle\n",
    "import os\n",
    "import numpy as np\n",
    "import gc\n",
    "\n",
    "#import mne\n",
    "#mne.set_log_level(verbose='WARNING')\n",
    "\n",
    "import torch\n",
    "from sklearn.metrics import mean_absolute_error, mean_squared_error,r2_score,explained_variance_score\n",
    "import torch.nn as nn\n",
    "#import pandas as pd\n",
    "import warnings\n",
    "#warnings.simplefilter(action='ignore', category=FutureWarning)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "2c7d032e-8a3d-4dbf-8ef4-c830ac1caa14",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class LearnGraph(nn.Module):\n",
    "    \n",
    "    def __init__(self,):\n",
    "        super(LearnGraph, self).__init__()\n",
    "        ##Adj Matrix to learn\n",
    "        self.adj_matrix = torch.nn.Parameter(torch.ones(76,76))\n",
    "        ##You can also initialize using a pre-built graph, with knowledge domain or a spatial graph\n",
    "        #self.adj_matrix = torch.nn.Parameter(torch.load('graph_adj_init.pt'))   \n",
    "        \n",
    "    def forward(self,x,channels_to_keep,channels_to_remove):\n",
    "        ## Solve the reconstruction problem using the closed form\n",
    "        adj = self.adj_matrix\n",
    "        D = torch.diag(adj.sum(dim=0))\n",
    "        L = D - adj\n",
    "        L_e = torch.stack([L[x,channels_to_remove] for x in channels_to_remove]).float()\n",
    "        L_e_inv = torch.linalg.inv(L_e)\n",
    "        L_eeb = torch.stack([L[x,channels_to_keep] for x in channels_to_remove]).float()\n",
    "        prod = torch.matmul(L_e_inv,L_eeb)\n",
    "        pred_se = -torch.einsum('ef,bft->bet',prod,x)\n",
    "        return pred_se"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "bd3a4012-18de-440c-87d5-bd819a2d7770",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "dset_name_path = 'Schirrmeister'\n",
    "path = '/users/local/eeg_datasets/'\n",
    "\n",
    "# Load the dataset you want to learn from\n",
    "arr = torch.from_numpy(np.load(os.path.join(path,dset_name_path+'/'+dset_name_path+'.npy'))).float()\n",
    "\n",
    "# You can try to see if the code works with a random array\n",
    "# arr = torch.rand((6000,70,1000))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "c7b5b6d4-f1ea-4fa6-a609-49967d606308",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "NxChannelsxTimesteps\n",
      "torch.Size([13484, 76, 2001])\n"
     ]
    }
   ],
   "source": [
    "print('NxChannelsxTimesteps')\n",
    "print(arr.shape) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "c9f95014-522f-4c33-9649-c3ab8293a25c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "lr = 0.1\n",
    "n_epoch = 2000\n",
    "n_iter = 10\n",
    "n_samples = 5000 # number of sample to use to solve reconstruction problem in parallel (related to how many RAM you have)\n",
    "device = 'cuda' if torch.cuda.is_available() else cpu\n",
    "model = LearnGraph().to(device)\n",
    "optimizer = torch.optim.Adam(model.parameters(),lr=lr)  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "5d03a2e2-c4c0-4695-8442-8ebd5a735aef",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "R²: \n",
      "  0:0.569   1:0.676   2:0.716   3:0.775   4:0.775   5:0.817   6:0.802   7:0.822   8:0.809   9:0.796  10:0.830  11:0.774  12:0.786  13:0.764  14:0.828  15:0.830  16:0.837  17:0.857  18:0.858  19:0.836  20:0.864  21:0.854  22:0.867  23:0.864  24:0.870  25:0.863  26:0.855  27:0.879  28:0.883  29:0.886  30:0.846  31:0.795  32:0.808  33:0.826  34:0.836  35:0.845  36:0.879  37:0.881  38:0.850  39:0.873  40:0.852  41:0.878 "
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[23], line 11\u001b[0m\n\u001b[1;32m      9\u001b[0m seb \u001b[38;5;241m=\u001b[39m arr[idx][:,channels_to_keep]  \u001b[38;5;66;03m# observed part of the signal\u001b[39;00m\n\u001b[1;32m     10\u001b[0m se \u001b[38;5;241m=\u001b[39m arr[idx][:,channels_to_remove] \u001b[38;5;66;03m# missing part of the signal\u001b[39;00m\n\u001b[0;32m---> 11\u001b[0m data, target \u001b[38;5;241m=\u001b[39m \u001b[43mseb\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mto\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdevice\u001b[49m\u001b[43m)\u001b[49m, se\u001b[38;5;241m.\u001b[39mto(device)\n\u001b[1;32m     12\u001b[0m optimizer\u001b[38;5;241m.\u001b[39mzero_grad()\n\u001b[1;32m     13\u001b[0m interpo \u001b[38;5;241m=\u001b[39m model(data,channels_to_keep,channels_to_remove) \u001b[38;5;66;03m# interpolate using the closed form\u001b[39;00m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "print(\"R²: \")\n",
    "for j in range(n_epoch):\n",
    "    list_r2 = [] \n",
    "    for i in range(n_iter):\n",
    "        idx = sorted(list(np.random.choice(np.arange(arr.shape[0]), n_samples, replace=False))) # randomly select samples \n",
    "        K = 38 \n",
    "        channels_to_remove = sorted(list(np.random.choice(np.arange(76), K, replace=False))) # randomly select K electrodes to remove \n",
    "        channels_to_keep = [x for x in np.arange(76) if x not in channels_to_remove]\n",
    "        seb = arr[idx][:,channels_to_keep]  # observed part of the signal\n",
    "        se = arr[idx][:,channels_to_remove] # missing part of the signal\n",
    "        data, target = seb.to(device), se.to(device)\n",
    "        optimizer.zero_grad()\n",
    "        interpo = model(data,channels_to_keep,channels_to_remove) # interpolate using the closed form\n",
    "        true_channel = target\n",
    "        residu = (true_channel - interpo)**2 \n",
    "        r2 = 1 - (residu.sum())/(((true_channel - true_channel.mean())**2).sum())\n",
    "        loss = 1 - r2\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "    print(\"{:3d}:{:.3f}\".format(j,1-loss.detach().to('cpu').numpy()),end = ' ')\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "61bc9116-dc28-447b-ad92-bd67f12cc789",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['AF3', 'AF4', 'AF7', 'AF8', 'AFF5h', 'AFF6h', 'AFz', 'C1', 'C2', 'C3', 'C4', 'C5', 'C6', 'CCP3h', 'CCP4h', 'CCP5h', 'CCP6h', 'CP1', 'CP2', 'CP3', 'CP4', 'CP5', 'CP6', 'CPz', 'Cz', 'F1', 'F2', 'F3', 'F4', 'F5', 'F6', 'F7', 'F8', 'FC1', 'FC2', 'FC3', 'FC4', 'FC5', 'FC6', 'FCC3h', 'FCC4h', 'FCC5h', 'FCC6h', 'FCz', 'FT7', 'FT8', 'Fp1', 'Fp2', 'Fpz', 'Fz', 'Iz', 'O1', 'O2', 'Oz', 'P1', 'P10', 'P2', 'P3', 'P4', 'P5', 'P6', 'P7', 'P8', 'P9', 'PO3', 'PO4', 'PO5', 'PO6', 'PO7', 'PO8', 'POz', 'Pz', 'T7', 'T8', 'TP7', 'TP8']\n"
     ]
    }
   ],
   "source": [
    "dset_name_path = 'Schirrmeister'\n",
    "\n",
    "with open(os.path.join(path,dset_name_path+'/Schirrmeister_channels.pkl'), \"rb\") as input_file:\n",
    "    channels = pickle.load(input_file)    \n",
    "print(channels)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "env_yass",
   "language": "python",
   "name": "env_yass"
  },
  "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}