05_EEG_Mediation_Analysis.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c3e81898004b064",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from scipy.stats import zscore\n",
    "import pingouin as pg\n",
    "plt.rc('font', family = 'arial')\n",
    "from pingouin import mediation_analysis\n",
    "from scipy import stats\n",
    "from statsmodels.stats.multitest import fdrcorrection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa69968e4670f07f",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Frequency bands definition\n",
    "BANDS = {'Delta': (1, 4), 'Theta': (4, 8), 'Alpha': (8, 13)}\n",
    "\n",
    "# Preload channel names\n",
    "CHANNELS = ['F7', 'Fp1', 'Fp2', 'F8', 'F3', 'Fz', 'F4', 'C3', 'Cz', 'P8', 'P7', 'Pz', 'P4', 'T3', 'P3', 'O1', 'O2', 'C4', 'T4']\n",
    "\n",
    "# Find the index location of channels in anatomical groups\n",
    "All_Channel_indices = np.arange(len(CHANNELS))\n",
    "Occipital_Channel_indices = [CHANNELS.index('O1'), CHANNELS.index('O2')]\n",
    "Frontal_Channel_indices = [CHANNELS.index('F7'), CHANNELS.index('F3'), CHANNELS.index('Fz'), CHANNELS.index('F4'), CHANNELS.index('F8')]\n",
    "Central_Channel_indices = [CHANNELS.index('Cz'), CHANNELS.index('C3'), CHANNELS.index('C4')]\n",
    "\n",
    "# Name of anatomical channel groups and channel indices\n",
    "Channel_Groups = {\n",
    "    'All': All_Channel_indices,\n",
    "    'Occipital': Occipital_Channel_indices, \n",
    "    'Frontal': Frontal_Channel_indices,\n",
    "    'Central': Central_Channel_indices,\n",
    "}\n",
    "\n",
    "# Directory paths\n",
    "csv_path = '04_Features_EEG/Range_13_10_Sec_Epoch_Fixed_Merged.csv'\n",
    "\n",
    "# Load dataframe\n",
    "all_data = pd.read_csv(csv_path)\n",
    "filt_data = all_data.copy()\n",
    "\n",
    "# Filter dataframe for specific timepoints and Record IDs with at least 3 timepoints\n",
    "filt_data = all_data[(all_data['Timepoint'].isin(['Baseline', 'Post_Wash', 'Post_Placebo', 'Post_CBD']))]\n",
    "filt_data = filt_data.groupby('Record ID').filter(lambda x: len(x['Timepoint'].unique()) >= 3)\n",
    "\n",
    "# Z-score normalize metabolite levels across subjects\n",
    "for metabolite in ['CBD', 'OHCBD', 'COOHCBD', 'AEA']:\n",
    "    filt_data[f'{metabolite}_Z_score'] = zscore(filt_data[metabolite])\n",
    "\n",
    "# Calculate combined score\n",
    "filt_data['Combined_score'] = (zscore(filt_data['CBD']) + zscore(filt_data['OHCBD']) + zscore(filt_data['COOHCBD'])) / 3\n",
    "\n",
    "# Convert categorical variables to numeric codes\n",
    "for col in ['Timepoint', 'Randomization', 'ADOS_Module']:\n",
    "    filt_data[f'{col}_numeric'] = filt_data[col].astype('category').cat.codes\n",
    "\n",
    "# Sort the dataframe\n",
    "filt_data = filt_data.sort_values(by=['Record ID', 'Timepoint'], ascending=[True, False])\n",
    "\n",
    "# Print summary information\n",
    "print(f\"Number of unique 'Record ID' in the comparisons: n = {len(filt_data['Record ID'].unique())}\")\n",
    "print(filt_data['Record ID'].unique())\n",
    "\n",
    "# Function to calculate statistics for each group\n",
    "def calculate_stats(group):\n",
    "    return pd.Series({\n",
    "        'Mean Age': group['Age'].mean(),\n",
    "        'SEM Age': group['Age'].sem(),\n",
    "        'Min Age': group['Age'].min(),\n",
    "        'Max Age': group['Age'].max(),\n",
    "        'Unique Record ID': group['Record ID'].nunique()\n",
    "    })\n",
    "\n",
    "# Group by 'Timepoint' and calculate statistics\n",
    "result = filt_data.groupby('Timepoint').apply(calculate_stats).reset_index()\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "id": "initial_id",
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def plot_significant_mediation(results, x_label, m_label, y_label, title, data, np=np, plt=plt, stats=stats, pg=pg):\n",
    "    # Check for significant indirect effects\n",
    "    significant_indirect = results[(results['path'].str.contains('Indirect')) & (results['sig'] == 'Yes')]\n",
    "    \n",
    "    if significant_indirect.empty:\n",
    "        print(f\"No significant indirect effect found for {title}\")\n",
    "        return\n",
    "    \n",
    "    # Determine which column contains p-values\n",
    "    p_value_col = 'pval' if 'pval' in results.columns else 'p-val' if 'p-val' in results.columns else None\n",
    "    if p_value_col is None:\n",
    "        print(\"Warning: Could not find p-value column. Skipping p-value annotation.\")\n",
    "    \n",
    "    # Extract coefficients for paths a, b, and c\n",
    "    a = results.loc[results['path'].str.contains('~ X'), 'coef'].values\n",
    "    b = results.loc[results['path'].str.startswith('Y ~'), 'coef'].values\n",
    "    c = results.loc[results['path'] == 'Direct', 'coef'].values[0]\n",
    "    \n",
    "    # Generate data points for plotting\n",
    "    x = np.linspace(data[x_label].min(), data[x_label].max(), 100)\n",
    "    m = np.outer(x, a)\n",
    "    y_direct = c * x\n",
    "    y_indirect = np.dot(m, b)\n",
    "    y_total = y_direct + y_indirect\n",
    "\n",
    "    x_data = data[x_label]\n",
    "    y_data = data[y_label]\n",
    "    \n",
    "    # Calculate correlations and partial correlations\n",
    "    r_a = stats.pearsonr(data[x_label], data[m_label])[0]\n",
    "    r2_a = r_a ** 2\n",
    "    r_b = pg.partial_corr(data=data, x=m_label, y=y_label, covar=x_label)['r'].values[0]\n",
    "    r2_b = r_b ** 2\n",
    "    r_c = stats.pearsonr(data[x_label], data[y_label])[0]\n",
    "    r2_c = r_c ** 2\n",
    "    r_c_prime = pg.partial_corr(data=data, x=x_label, y=y_label, covar=m_label)['r'].values[0]\n",
    "    r2_c_prime = r_c_prime ** 2\n",
    "    \n",
    "    # Plot 1: Relationships between X, mediators, and Y\n",
    "    fig1, ax1 = plt.subplots(figsize=(12, 8))\n",
    "    for i, m_path in enumerate(results['path'][results['path'].str.contains('~ X')]):\n",
    "        mediator = m_path.split('~')[0].strip()\n",
    "        color = 'red' if f'Indirect {mediator}' in significant_indirect['path'].values else 'gray'\n",
    "        ax1.plot(x, m[:, i], label=f'{x_label} -> {mediator}', color=color)\n",
    "    ax1.plot(x, y_direct, label=f'{x_label} -> {y_label} (Direct)', linestyle='--', color='black')\n",
    "    ax1.set_xlabel(x_label, fontsize=14)\n",
    "    ax1.set_ylabel('Estimated Value', fontsize=14)\n",
    "    ax1.legend(bbox_to_anchor=(1.05, 1), loc='upper left', fontsize=12)\n",
    "    ax1.set_title(f'Relationships: {x_label} to mediators and {y_label}', fontsize=16)\n",
    "    ax1.tick_params(axis='both', which='major', labelsize=12)\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "    # Plot 2: Indirect and Total Effects\n",
    "    fig2, ax2 = plt.subplots(figsize=(12, 8))\n",
    "    ax2.plot(x, y_indirect, label=f'Indirect Effect on {y_label}', color='blue')\n",
    "    ax2.plot(x, y_total, label=f'Total Effect on {y_label}', color='green')\n",
    "    ax2.set_xlabel(x_label, fontsize=14)\n",
    "    ax2.set_ylabel(y_label, fontsize=14)\n",
    "    ax2.legend(bbox_to_anchor=(1.05, 1), loc='upper left', fontsize=12)\n",
    "    ax2.set_title(f'Indirect and Total Effects on {y_label}', fontsize=16)\n",
    "    ax2.tick_params(axis='both', which='major', labelsize=12)\n",
    "    \n",
    "    # Add effect size annotation\n",
    "    for _, row in significant_indirect.iterrows():\n",
    "        annotation_text = f\"{row['path']}\\nEffect: {row['coef']:.3f}\"\n",
    "        if p_value_col:\n",
    "            annotation_text += f\"\\np-value: {row[p_value_col]:.3f}\"\n",
    "        ax2.annotate(annotation_text, xy=(0.05, 0.95), xycoords='axes fraction',\n",
    "                     va='top', ha='left', bbox=dict(boxstyle='round', fc='white', ec='gray', alpha=0.8),\n",
    "                     fontsize=12)\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "    # Plot 3: Scatter plot with regression lines\n",
    "    fig3, ax3 = plt.subplots(figsize=(12, 8))\n",
    "    post_cbd_mask = data['Timepoint'] == 'Post_CBD'\n",
    "    ax3.scatter(x_data[post_cbd_mask], y_data[post_cbd_mask], color='red', label='Post_CBD', alpha=0.7)\n",
    "    ax3.scatter(x_data[~post_cbd_mask], y_data[~post_cbd_mask], color='black', label='Other Timepoints', alpha=0.7)\n",
    "    \n",
    "    # Post_CBD regression\n",
    "    post_cbd_fit = np.polyfit(x_data[post_cbd_mask], y_data[post_cbd_mask], 1)\n",
    "    post_cbd_line = np.poly1d(post_cbd_fit)\n",
    "    ax3.plot(x_data, post_cbd_line(x_data), color='red', linestyle='--', label='Post_CBD Regression')\n",
    "    \n",
    "    # Other Timepoints regression\n",
    "    other_fit = np.polyfit(x_data[~post_cbd_mask], y_data[~post_cbd_mask], 1)\n",
    "    other_line = np.poly1d(other_fit)\n",
    "    ax3.plot(x_data, other_line(x_data), color='black', linestyle='--', label='Other Timepoints Regression')\n",
    "    \n",
    "    ax3.set_xlabel(x_label, fontsize=14)\n",
    "    ax3.set_ylabel(y_label, fontsize=14)\n",
    "    ax3.legend(bbox_to_anchor=(1.05, 1), loc='upper left', fontsize=12)\n",
    "    ax3.set_title(f'{y_label} with Linear Regression', fontsize=16)\n",
    "    ax3.tick_params(axis='both', which='major', labelsize=12)\n",
    "\n",
    "    # Calculate the difference between regressions\n",
    "    slope_diff = post_cbd_fit[0] - other_fit[0]\n",
    "    intercept_diff = post_cbd_fit[1] - other_fit[1]\n",
    "    \n",
    "    # Perform statistical test for slope difference\n",
    "    from scipy import stats\n",
    "    t_stat, p_value = stats.ttest_ind(y_data[post_cbd_mask] - post_cbd_line(x_data[post_cbd_mask]),\n",
    "                                      y_data[~post_cbd_mask] - other_line(x_data[~post_cbd_mask]))\n",
    "    \n",
    "    # Add difference statistics annotation\n",
    "    ax3.text(0.05, 0.95, f\"Slope difference: {slope_diff:.3f}\", transform=ax3.transAxes, \n",
    "             verticalalignment='top', fontsize=12)\n",
    "    ax3.text(0.05, 0.90, f\"Intercept difference: {intercept_diff:.3f}\", transform=ax3.transAxes, \n",
    "             verticalalignment='top', fontsize=12)\n",
    "    ax3.text(0.05, 0.85, f\"Difference p-value: {p_value:.3f}\", transform=ax3.transAxes, \n",
    "             verticalalignment='top', fontsize=12)\n",
    "\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "    # Plot 4: Mediation Triangle with Full Statistics\n",
    "    fig4, ax4 = plt.subplots(figsize=(12, 10))\n",
    "    ax4.axis('off')\n",
    "    ax4.set_title('Mediation Triangle with Full Statistics', fontsize=16)\n",
    "    \n",
    "    # Draw the triangle\n",
    "    triangle = plt.Polygon([(0,0), (1,0), (0.5,0.866)], fill=False)\n",
    "    ax4.add_patch(triangle)\n",
    "    \n",
    "    # Add labels\n",
    "    ax4.text(0, -0.1, x_label, ha='center', va='center', fontsize=12)\n",
    "    ax4.text(1, -0.1, y_label, ha='center', va='center', fontsize=12)\n",
    "    ax4.text(0.5, 0.4, m_label, ha='center', va='center', fontsize=12)\n",
    "    \n",
    "    # Add arrows and statistics\n",
    "    ax4.annotate('', xy=(0.9, 0.05), xytext=(0.1, 0.05), arrowprops=dict(arrowstyle='<->', color='r'))\n",
    "    direct_effect = results.loc[results['path'] == 'Direct']\n",
    "    ax4.text(0.5, -0.05, f\"Direct Effect (c'):\\nCoef = {direct_effect['coef'].values[0]:.3f}\\np-value = {direct_effect[p_value_col].values[0]:.3f}\\nr = {r_c_prime:.3f}, r² = {r2_c_prime:.3f}\", ha='center', va='center', fontsize=10)\n",
    "    \n",
    "    ax4.annotate('', xy=(0.25, 0.433), xytext=(0.05, 0.05), arrowprops=dict(arrowstyle='<->', color='b'))\n",
    "    a_effect = results.loc[results['path'].str.contains('~ X')]\n",
    "    ax4.text(0.2, 0.25, f\"X -> M (a):\\nCoef = {a_effect['coef'].values[0]:.3f}\\np-value = {a_effect[p_value_col].values[0]:.3f}\\nr = {r_a:.3f}, r² = {r2_a:.3f}\", ha='center', va='center', fontsize=10)\n",
    "    \n",
    "    ax4.annotate('', xy=(0.95, 0.05), xytext=(0.75, 0.433), arrowprops=dict(arrowstyle='<->', color='g'))\n",
    "    b_effect = results.loc[results['path'].str.startswith('Y ~')]\n",
    "    ax4.text(0.8, 0.25, f\"M -> Y (b):\\nCoef = {b_effect['coef'].values[0]:.3f}\\np-value = {b_effect[p_value_col].values[0]:.3f}\\nr = {r_b:.3f}, r² = {r2_b:.3f}\", ha='center', va='center', fontsize=10)\n",
    "\n",
    "    # Add full model statistics below the diagram\n",
    "    stats_text = \"Full Model Statistics:\\n\\n\"\n",
    "    for _, row in results.iterrows():\n",
    "        if row['path'] == 'Total':\n",
    "            stats_text += f\"{row['path']}:\\nCoefficient = {row['coef']:.3f}, p-value = {row[p_value_col]:.3f}\\nr = {r_c:.3f}, r² = {r2_c:.3f}\\n\\n\"\n",
    "        elif row['path'] == 'Direct':\n",
    "            stats_text += f\"{row['path']}:\\nCoefficient = {row['coef']:.3f}, p-value = {row[p_value_col]:.3f}\\nr = {r_c_prime:.3f}, r² = {r2_c_prime:.3f}\\n\\n\"\n",
    "        else:\n",
    "            stats_text += f\"{row['path']}:\\nCoefficient = {row['coef']:.3f}, p-value = {row[p_value_col]:.3f}\\n\\n\"\n",
    "    \n",
    "    ax4.text(0.5, -0.5, stats_text, ha='center', va='center', fontsize=10, bbox=dict(facecolor='white', edgecolor='black', alpha=0.8))\n",
    "\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "    # Print overall title and significant indirect effects\n",
    "    print(f\"\\n{title}\\nSignificant Indirect Effect(s):\\n{', '.join(significant_indirect['path'])}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae14b6a9e46b7132",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Initialize dictionaries and lists to store results\n",
    "mediation_results = {}\n",
    "all_p_values = []\n",
    "all_result_keys = []\n",
    "\n",
    "# Define lists of metabolites, behavioral tests, and electrode groups to analyze\n",
    "metabolite_list = ['COOHCBD_Z_score']\n",
    "beh_test_list = ['rbs_total_score', 'ppvt_raw_score', 'toni4_raw_score', 'eowpvt4_raw_score', 'beery_vmi_raw_score', 'beery_vp_raw_score', 'beery_mc_raw_score']\n",
    "electrode_group = ['Occipital', 'Frontal', 'Central']\n",
    "\n",
    "# Iterate through each behavioral test\n",
    "for beh_test in beh_test_list:\n",
    "    \n",
    "    # Drop rows with missing values in the relevant behavioral test column\n",
    "    df = filt_data.dropna(subset=[beh_test])\n",
    "    # Filter to include only subjects with at least 3 timepoints\n",
    "    df_beh_filt = df.groupby('Record ID').filter(lambda x: len(x['Timepoint'].unique()) >= 3)\n",
    "     \n",
    "    # Calculate statistics for each timepoint\n",
    "    result = df_beh_filt.groupby('Timepoint').apply(calculate_stats).reset_index()\n",
    "    print(beh_test)\n",
    "    print(result)\n",
    "    \n",
    "    # Iterate through each metabolite\n",
    "    for metabolite in metabolite_list:\n",
    "        # Iterate through each electrode group\n",
    "        for electrodes in electrode_group:\n",
    "            # Define fixed covariates\n",
    "            fixed_covariates = ['Age', 'Timepoint_numeric', 'Randomization_numeric', 'ADOS_numeric']\n",
    "            # Define EEG covariates for the current electrode group\n",
    "            eeg_covariates = [f'{electrodes}_Exponent', f'{electrodes}_Alpha_SNR', f'{electrodes}_Offset', f'{electrodes}_Delta_SNR', f'{electrodes}_Theta_SNR']\n",
    "            \n",
    "            # Iterate through each EEG covariate\n",
    "            for eeg_covariate in eeg_covariates:\n",
    "                # Create a unique key for the current analysis\n",
    "                result_key = f'{beh_test}_{metabolite}_{electrodes}_{eeg_covariate}'\n",
    "                \n",
    "                # Perform mediation analysis\n",
    "                results = mediation_analysis(data=df_beh_filt, x=eeg_covariate, m=metabolite, y=beh_test, covar=fixed_covariates, n_boot=1000, seed=42).round(3)\n",
    "                \n",
    "                # Print results\n",
    "                print(result_key)\n",
    "                print('\\n')\n",
    "                print(results)\n",
    "                print('\\n')\n",
    "                print('\\n')\n",
    "                \n",
    "                # Store results in the mediation_results dictionary\n",
    "                mediation_results[result_key] = results\n",
    "                \n",
    "                # Store p-values and result keys for FDR correction\n",
    "                all_p_values.extend(results['pval'].tolist())\n",
    "                all_result_keys.extend([result_key] * len(results))\n",
    "            \n",
    "                # Plot significant mediation relationships\n",
    "                plot_significant_mediation(\n",
    "                    results=results, \n",
    "                    x_label=f'{eeg_covariate}',\n",
    "                    m_label=metabolite,\n",
    "                    y_label=beh_test,\n",
    "                    title=result_key,\n",
    "                    data=df_beh_filt\n",
    "                )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "755a84b49446cc30",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def hierarchical_fdr_correction(mediation_results):\n",
    "    # Dictionary to store FDR-corrected results\n",
    "    fdr_corrected_results = {}\n",
    "    \n",
    "    # Identify unique electrode groups from the result keys\n",
    "    electrode_groups = set()\n",
    "    for key in mediation_results.keys():\n",
    "        parts = key.split('_')\n",
    "        electrode_group = parts[-2]  # Assuming electrode group is always second to last\n",
    "        electrode_groups.add(electrode_group)\n",
    "    \n",
    "    # First level: FDR correction across electrode groups\n",
    "    group_min_p_values = []\n",
    "    group_keys = []\n",
    "    \n",
    "    # Find the minimum p-value for each electrode group\n",
    "    for group in electrode_groups:\n",
    "        group_p_values = []\n",
    "        for key, results in mediation_results.items():\n",
    "            if group in key:\n",
    "                indirect_effect = results[results['path'].str.contains('Indirect')]\n",
    "                if not indirect_effect.empty and indirect_effect['sig'].values[0] == 'Yes':\n",
    "                    p_value = indirect_effect['pval'].values[0]\n",
    "                    group_p_values.append(p_value)\n",
    "        \n",
    "        if group_p_values:\n",
    "            min_p_value = min(group_p_values)\n",
    "            group_min_p_values.append(min_p_value)\n",
    "            group_keys.append(group)\n",
    "    \n",
    "    # Perform FDR correction on the minimum p-values across groups\n",
    "    rejected_groups, corrected_group_p_values = fdrcorrection(group_min_p_values, alpha=0.05, method='indep')\n",
    "    \n",
    "    # Second level: FDR correction within significant electrode groups\n",
    "    for group, is_rejected in zip(group_keys, rejected_groups):\n",
    "        if is_rejected:\n",
    "            group_p_values = []\n",
    "            group_result_keys = []\n",
    "            \n",
    "            # Collect p-values for significant indirect effects within the group\n",
    "            for key, results in mediation_results.items():\n",
    "                if group in key:\n",
    "                    indirect_effect = results[results['path'].str.contains('Indirect')]\n",
    "                    if not indirect_effect.empty and indirect_effect['sig'].values[0] == 'Yes':\n",
    "                        p_value = indirect_effect['pval'].values[0]\n",
    "                        group_p_values.append(p_value)\n",
    "                        group_result_keys.append(key)\n",
    "            \n",
    "            # Perform FDR correction within the group\n",
    "            rejected, corrected_p_values = fdrcorrection(group_p_values, alpha=0.05, method='indep')\n",
    "            \n",
    "            # Update results with corrected p-values and significance\n",
    "            for key, p_value, is_rejected in zip(group_result_keys, corrected_p_values, rejected):\n",
    "                if key not in fdr_corrected_results:\n",
    "                    fdr_corrected_results[key] = mediation_results[key].copy()\n",
    "                \n",
    "                indirect_index = fdr_corrected_results[key]['path'].str.contains('Indirect')\n",
    "                fdr_corrected_results[key].loc[indirect_index, 'corrected_pval'] = p_value\n",
    "                fdr_corrected_results[key].loc[indirect_index, 'significant_after_correction'] = is_rejected\n",
    "    \n",
    "    return fdr_corrected_results\n",
    "\n",
    "# Apply hierarchical FDR correction to the mediation results\n",
    "fdr_corrected_results = hierarchical_fdr_correction(mediation_results)\n",
    "\n",
    "# Print FDR-corrected results\n",
    "for key, results in fdr_corrected_results.items():\n",
    "    print(f\"Results for {key}:\")\n",
    "    print(results)\n",
    "    print(\"\\n\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}