neuromorph_plotting.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "5505dd52-7bad-471b-bb31-bb1d491261f2",
   "metadata": {},
   "source": [
    "### Imports and main params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1f5d8af-aa0d-4922-a923-93fe11b492ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import scipy.stats\n",
    "from mpl_toolkits.axes_grid1 import make_axes_locatable\n",
    "\n",
    "import mne\n",
    "from mne.channels import find_ch_adjacency\n",
    "from mne.datasets import sample\n",
    "from mne.stats import combine_adjacency, spatio_temporal_cluster_test\n",
    "from mne.viz import plot_compare_evokeds\n",
    "\n",
    "import os\n",
    "import pickle\n",
    "import pandas as pd\n",
    "from itertools import combinations\n",
    "from scipy.stats import ttest_ind\n",
    "from mne.stats import f_mway_rm, f_threshold_mway_rm\n",
    "\n",
    "from tqdm import tqdm\n",
    "\n",
    "mne.set_log_level('warning')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a06ec41-72ab-4439-9900-b842ba24be7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "ROOT = '/scratch/alr664/multiple_affix'\n",
    "os.chdir(ROOT)\n",
    "\n",
    "mri_dir = os.path.join(ROOT, 'mri')\n",
    "#stc_dir = os.path.join(ROOT, 'stc_signed')\n",
    "#stats_dir = os.path.join(ROOT, 'stats')\n",
    "\n",
    "#subjects = ['R1830', 'R2283', 'R2356', 'R2340', 'R2341', 'R2342', 'R2343', 'R2345', 'R2346', 'R2347', 'R2352', 'R2353', 'R2354', 'R2358', 'R2359', 'R2360', 'R2366', 'R2367', 'R2368', 'R2371']\n",
    "subjects = [\n",
    "    \"A0394\", \"A0421\", \"A0446\", \"A0451\", \"A0468\", \"A0484\", \"A0495\", \"A0502\", \"A0503\", \"A0508\", \n",
    "    \"A0509\", \"A0512\", \"A0513\", \"A0514\", \"A0516\", \"A0517\", \"A0518\", \"A0519\", \"A0520\", \"A0521\", \n",
    "    \"A0522\", \"A0523\", \"A0524\", \"A0525\"\n",
    "]\n",
    "print(len(subjects))\n",
    "#pres_type = 'rsvp' # subset all data to this pres type only\n",
    "# surprisals = ['high_high','high_low','low_high','low_low']\n",
    "# colors = {str('high_high'+\"_\"+pres_type):'tab:blue',\n",
    "#          str('high_low'+\"_\"+pres_type):'tab:green',\n",
    "#          str('low_high'+\"_\"+pres_type):'tab:red',\n",
    "#          str('low_low'+\"_\"+pres_type):'tab:orange'}\n",
    "\n",
    "\n",
    "# permutation test vars \n",
    "tmin = 0\n",
    "tmax = 0.5\n",
    "tail = 1\n",
    "p_thresh = 0.05\n",
    "n_permutations = 10000\n",
    "\n",
    "lex_event_ids = {\n",
    "    '0Suff w/ Lat.': 1,\n",
    "    '0Suff w/o Lat.': 2,\n",
    "    '0Suff NW': 4,\n",
    "    '1Suff w/ Lat.': 11,\n",
    "    '1Suff w/o Lat.': 12,\n",
    "    '1Suff PseudoStemNW': 14,\n",
    "    '1Suff RealStemNW': 15,\n",
    "    '2Suff w/ Lat.': 21,\n",
    "    '2Suff w/o Lat.': 22,\n",
    "    '2Suff Composite': 23,\n",
    "    '2Suff PseudoStemNW': 24,\n",
    "    '2Suff RealStemNW': 25,\n",
    "    'Practice': 99\n",
    "}\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24e6875b-5320-4270-9c44-80337a7ec93a",
   "metadata": {},
   "source": [
    "### Import epochs and subset by pres_type"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23442482-a59e-4294-bfb9-2027e28cf6b6",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Read data\n",
    "for i, subject in enumerate(tqdm(subjects)):\n",
    "                    \n",
    "    epochs_dir = f'meg/{subject}/epochs/'\n",
    "    epochs = mne.read_epochs(os.path.join(ROOT, 'meg', subject, f'{subject}_rejection-epo.fif'))\n",
    "    #epochs = epochs[epochs.metadata.pres_type == pres_type] # subsetting by pres_type\n",
    "\n",
    "    num_events = len(epochs.events)\n",
    "    epochs.metadata = pd.DataFrame(index=range(num_events))\n",
    "\n",
    "    # Step 2: Create the 'Trigger' column\n",
    "    epochs.metadata['Trigger'] = epochs.events[:, 2]\n",
    "    \n",
    "    # Step 3: Map the event IDs to condition labels and create the 'Condition' column\n",
    "    epochs.metadata['Condition'] = epochs.metadata['Trigger'].map(event_to_label)\n",
    "    \n",
    "    #epochs.metadata['full_cond'] = epochs.metadata['cond'] + '_' + epochs.metadata['pres_type'] # create merged condition names\n",
    "\n",
    "    if i == 0: # initialize variables on the first subject only\n",
    "        times = epochs.times # get the epochs times\n",
    "        event_to_label = {v: k for k, v in lex_event_ids.items()}\n",
    "        \n",
    "        #print(epochs.metadata)\n",
    "        conditions = epochs.metadata.Condition.unique()\n",
    "        epochs_data = np.zeros((len(subjects), len(conditions), 207,  len(times))) # create container for all subj epoch data\n",
    "        info = epochs[0].info\n",
    "        \n",
    "        print('times = ', times)\n",
    "        print('conditions = ', conditions)\n",
    "    \n",
    "    for j, condition in enumerate(conditions):\n",
    "        epochs_data[i, j] = epochs[epochs.metadata['Condition'] == condition].get_data(copy=True).mean(axis=0) # for each subj, for each cond: get average over trials  (axis 0)\n",
    "print(epochs_data.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "deaedca3-41fa-4ab2-bc2c-5feed88f1152",
   "metadata": {},
   "outputs": [],
   "source": [
    "average_waveforms = epochs_data.mean(axis=0)  # average across the subject axis\n",
    "\n",
    "print(average_waveforms.shape)\n",
    "\n",
    "evoked_dict = {}\n",
    "for i, condition in enumerate(conditions):\n",
    "    evoked_dict[condition] = mne.EvokedArray(average_waveforms[i], info, tmin=times[0])\n",
    "\n",
    "mne.viz.plot_compare_evokeds(evoked_dict, combine='gfp', title='Evoked Response') # need to set custom colors for all conds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4f057cd-ce86-496e-a5e8-cfe6d7672615",
   "metadata": {},
   "outputs": [],
   "source": [
    "# for simple plotting\n",
    "import matplotlib.pyplot as plt\n",
    "suffix_groups = {\n",
    "    '0Suff': [],\n",
    "    '1Suff': [],\n",
    "    '2Suff': []\n",
    "}\n",
    "\n",
    "for i, condition in enumerate(conditions):\n",
    "    if condition.startswith('0Suff'):\n",
    "        suffix_groups['0Suff'].append(i)\n",
    "    elif condition.startswith('1Suff'):\n",
    "        suffix_groups['1Suff'].append(i)\n",
    "    elif condition.startswith('2Suff'):\n",
    "        suffix_groups['2Suff'].append(i)\n",
    "\n",
    "average_waveforms = {}\n",
    "for suffix, indices in suffix_groups.items():\n",
    "    group_data = epochs_data[indices].mean(axis=0)  # average across the condition axis\n",
    "    average_waveforms[suffix] = group_data.mean(axis=0)  # average across the epoch axis\n",
    "\n",
    "evoked_dict = {}\n",
    "for suffix, data in average_waveforms.items():\n",
    "    evoked_dict[suffix] = mne.EvokedArray(data, info, tmin=times[0])\n",
    "\n",
    "# Plot the evoked responses\n",
    "fig = mne.viz.plot_compare_evokeds(evoked_dict, combine='gfp', title='Evoked Response by Number of Suffixes')\n",
    "\n",
    "fig[0].savefig('suffixation.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15222843-1194-4411-9622-4d1cf16a2e54",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "nw_group = []\n",
    "words_group = []\n",
    "\n",
    "# Categorize conditions\n",
    "for i, condition in enumerate(conditions):\n",
    "    if 'NW' in condition:\n",
    "        nw_group.append(i)\n",
    "    else:\n",
    "        words_group.append(i)\n",
    "\n",
    "# Calculate average waveforms for NW and Words groups\n",
    "average_waveforms = {}\n",
    "average_waveforms['NW'] = epochs_data[nw_group].mean(axis=0).mean(axis=0)  # average across the condition and epoch axes\n",
    "average_waveforms['Words'] = epochs_data[words_group].mean(axis=0).mean(axis=0)  # average across the condition and epoch axes\n",
    "\n",
    "# Create EvokedArray for each group\n",
    "evoked_dict = {}\n",
    "for group, data in average_waveforms.items():\n",
    "    evoked_dict[group] = mne.EvokedArray(data, info, tmin=times[0])\n",
    "\n",
    "fig = mne.viz.plot_compare_evokeds(evoked_dict, combine='gfp', title='Evoked Response by Nonwords and Words')\n",
    "fig[0].savefig('wordness.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3120a397-3a38-4ea8-801f-653f27a4d421",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "composite_group = []\n",
    "with_lat_group = []\n",
    "without_lat_group = []\n",
    "\n",
    "for i, condition in enumerate(conditions):\n",
    "    if 'Composite' in condition:\n",
    "        composite_group.append(i)\n",
    "    elif 'w/ Lat.' in condition:\n",
    "        with_lat_group.append(i)\n",
    "    elif 'w/o Lat.' in condition:\n",
    "        without_lat_group.append(i)\n",
    "\n",
    "average_waveforms = {}\n",
    "average_waveforms['Composite'] = epochs_data[composite_group].mean(axis=0).mean(axis=0)  # average across the condition and epoch axes\n",
    "average_waveforms['Productive'] = epochs_data[with_lat_group].mean(axis=0).mean(axis=0)  # average across the condition and epoch axes\n",
    "average_waveforms['Non-productive'] = epochs_data[without_lat_group].mean(axis=0).mean(axis=0)\n",
    "\n",
    "evoked_dict = {}\n",
    "for group, data in average_waveforms.items():\n",
    "    evoked_dict[group] = mne.EvokedArray(data, info, tmin=times[0])\n",
    "\n",
    "fig = mne.viz.plot_compare_evokeds(evoked_dict, combine='gfp', title='Evoked Response by Productivity')\n",
    "fig[0].savefig('productivity.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a33204df-ee69-4e84-8e24-dfebbd8f05c9",
   "metadata": {},
   "source": [
    "### Create X"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34805c22-deef-417f-97fe-ea13b18f42f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "X = epochs_data.copy()\n",
    "print(X.shape)\n",
    "\n",
    "# extract time window of interest\n",
    "idx_tmin = np.where(times == tmin)[0][0]\n",
    "idx_tmax = np.where(times == tmax)[0][0]\n",
    "X = X[:, :, :, idx_tmin : idx_tmax + 1]\n",
    "print(X.shape)\n",
    "\n",
    "# create array of times for the search window: needed to get cluster times later\n",
    "search_window = np.arange(tmin, tmax+1, 0.001)\n",
    "# print(search_window)\n",
    "\n",
    "# transpose for permutation test later\n",
    "X = X.transpose(0, 1, 3, 2)\n",
    "print(X.shape)\n",
    "\n",
    "# # average over space\n",
    "# X = X.mean(axis=3)\n",
    "# print(X.shape)\n",
    "\n",
    "# # convert to a list of len = num_conditions, each element of the list is an array of shape (nsubj x ntimes)\n",
    "X = [np.squeeze(x) for x in np.split(X, 4, axis=1)]\n",
    "print(len(X))\n",
    "print(X[0].shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b62cbac-8b38-4e29-b84c-945b69b98ad7",
   "metadata": {},
   "source": [
    "### Set up test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "963be3ac-4d9b-4644-bd6c-035b42f41937",
   "metadata": {},
   "outputs": [],
   "source": [
    "### if doing spatio-temp search, add adjacency\n",
    "def stat_fun(*args):\n",
    "    # Return only the F-values\n",
    "    return f_mway_rm(np.swapaxes(args, 1, 0), factor_levels=factor_levels,\n",
    "                     effects=effects, return_pvals=False)[0]\n",
    "\n",
    "#adjacency, ch_names = mne.channels.find_ch_adjacency(info, ch_type='mag')\n",
    "hemi = 'lh'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2915fbad-ce09-40de-957f-cb5d94329769",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Estimating F-test threshold\n",
    "factor_levels = [] # condition contrasts\n",
    "effects = ['A:B'] # A - main effect of A, B - main effect of B, A:B - interaction effect\n",
    "f_thresh = f_threshold_mway_rm(len(subjects), factor_levels, effects, p_thresh)\n",
    "\n",
    "print(\"Conditions: \", conditions)\n",
    "print(\"Time window of analysis: \", tmin, \" to \", tmax)\n",
    "print(\"Launching clustering test for effect:\", effects)\n",
    "print(\"Threshold:\", f_thresh)\n",
    "\n",
    "# Perform the clustering test\n",
    "F_obs, clusters, clusters_pvals, h0 = clu = mne.stats.permutation_cluster_test(X, #spatio_temporal_cluster_test\n",
    "                                            tail=tail,                                   \n",
    "                                            threshold=f_thresh,\n",
    "                                            stat_fun = stat_fun,\n",
    "                                            n_permutations=n_permutations,\n",
    "                                            # adjacency = adjacency,\n",
    "                                            out_type='indices',\n",
    "                                            n_jobs=-1)\n",
    "\n",
    "# Check outputs\n",
    "print(\"Cluster p-values:\", clusters_pvals)\n",
    "for cluster, pval in zip(clusters, clusters_pvals):\n",
    "    #print(cluster)\n",
    "    if pval < p_thresh:\n",
    "        cluster_start_time = search_window[cluster[0][0]]\n",
    "        cluster_end_time = search_window[cluster[0][-1]]\n",
    "        if len(search_window[cluster[0]]) > 30:\n",
    "            print(\"%s - %s ms, p-value: %s\" % (cluster_start_time, cluster_end_time, pval))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4762d321-925b-4daf-8ab8-793aa6706703",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c7d5b81c-c2db-4c09-b48e-9be9377a5b56",
   "metadata": {},
   "source": [
    "### plotting code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d54b813-a312-4116-bce2-a347e0014fde",
   "metadata": {},
   "outputs": [],
   "source": [
    "average_waveforms = epochs_data.mean(axis=0)  # average across the subject axis\n",
    "\n",
    "print(average_waveforms.shape)\n",
    "\n",
    "evoked_dict = {}\n",
    "for i, condition in enumerate(conditions):\n",
    "    evoked_dict[condition] = mne.EvokedArray(average_waveforms[i], info, tmin=times[0])\n",
    "\n",
    "# break down compare evokeds plot\n",
    "fig, ax = plt.subplots(1, 1, figsize=(10, 4))\n",
    "mne.viz.plot_compare_evokeds(evoked_dict, combine='gfp', title='', axes = ax, show=False) #, colors=colors\n",
    "for cluster, pval in zip(clusters, clusters_pvals):\n",
    "    if pval < 0.35:\n",
    "        cluster_start_time = search_window[cluster[0][0]]\n",
    "        cluster_end_time = search_window[cluster[0][-1]]\n",
    "        if len(search_window[cluster[0]]) > 1:\n",
    "            ax.axvspan(cluster_start_time, cluster_end_time,color='lightblue', alpha=0.5)\n",
    "            print(\"%s - %s ms, p-value: %s\" % (cluster_start_time, cluster_end_time, pval))\n",
    "\n",
    "plt.show()\n",
    "# fig.savefig(\"rsvp_0-0.5_B.pdf\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d82c599c-30e6-4d32-89dc-a571b2ef2031",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73e58f13-e6e5-44e2-9845-ef0acd8caa3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "for ind in cluster_end_time:\n",
    "    print(ind)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96ead12b-2239-4a1c-b75b-7ce30b325512",
   "metadata": {},
   "outputs": [],
   "source": [
    "X_avg = np.mean(X, axis=(0,1))\n",
    "X_avg.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae7e00c7-c149-4e17-8c72-7d03a1b32c8f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from scipy.stats import ttest_rel\n",
    "\n",
    "group1 = X_avg[:, :2]  \n",
    "group2 = X_avg[:, 2:]  \n",
    "\n",
    "t_stat, p_value = ttest_rel(group1, group2)\n",
    "print(f\"T-statistic: {t_stat}, P-value: {p_value}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d96861c8-de3a-4714-9120-d064410b50fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def stat_fun(*args):\n",
    "    return f_mway_rm(np.swapaxes(args, 1, 0), factor_levels=factor_levels, effects=effects, return_pvals=False)[0]"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}