04_EEG_Behavior_Models.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "initial_id",
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import warnings\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import seaborn as sns\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy.stats import zscore, f_oneway\n",
    "from patsy import dmatrices\n",
    "import statsmodels.api as sm\n",
    "import statsmodels.formula.api as smf\n",
    "from statsmodels.tools.sm_exceptions import ConvergenceWarning\n",
    "from statsmodels.stats.outliers_influence import variance_inflation_factor\n",
    "from statsmodels.stats.multitest import multipletests\n",
    "from matplotlib.lines import Line2D\n",
    "from matplotlib.colors import LinearSegmentedColormap\n",
    "from matplotlib.ticker import LinearLocator\n",
    "plt.rc('font', family='arial')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a40581f5141b84b1",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Function to fit mixed-effects models\n",
    "def fit_mixed_effects_model(df, dependent_var, independent_vars, fixed, interaction_term, func):\n",
    "    # Convert 'Timepoint' and 'Randomization' columns to categorical data type\n",
    "    df['Timepoint'] = df['Timepoint'].astype('category')\n",
    "    df['Randomization'] = df['Randomization'].astype('category')\n",
    "    \n",
    "    # Create a string of fixed effects by joining independent variables and fixed terms\n",
    "    fixed_effects = ' + '.join(independent_vars + fixed)\n",
    "    \n",
    "    # Construct the formula for the mixed-effects model\n",
    "    if interaction_term == '':\n",
    "        formula = f\"{dependent_var} ~ {fixed_effects}\"\n",
    "    else:\n",
    "        formula = f\"{dependent_var} ~ {fixed_effects} + {interaction_term}\"\n",
    "    \n",
    "    # Fit the mixed-effects model while suppressing convergence warnings\n",
    "    with warnings.catch_warnings():\n",
    "        warnings.filterwarnings(\"ignore\", category=ConvergenceWarning)\n",
    "        mixed_model = smf.mixedlm(formula, df, groups=df['Record ID'], re_formula=\"1\").fit(method=func, maxiter=1000, full_output=True)\n",
    "    \n",
    "    # Create design matrices for the model\n",
    "    y, X = dmatrices(formula, df, return_type='dataframe')\n",
    "    \n",
    "    # Calculate Variance Inflation Factor (VIF) for each feature\n",
    "    vif = pd.DataFrame()\n",
    "    vif[\"VIF Factor\"] = [variance_inflation_factor(X.values, i) for i in range(X.shape[1])]\n",
    "    vif[\"features\"] = X.columns\n",
    "    \n",
    "    # Return the fitted model and VIF dataframe\n",
    "    return mixed_model, vif"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "233cc40d4a62872f",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Function to run models and plot results\n",
    "def run_models_and_plot(data, dependent_vars, independent_vars, fixed, interaction_term, func, y_var_names):\n",
    "    models = {}\n",
    "    emmeans_data = {}\n",
    "    \n",
    "    # Fit models for each dependent variable\n",
    "    for dependent_var in dependent_vars:\n",
    "        # Fit the mixed-effects model\n",
    "        model, vif = fit_mixed_effects_model(data, dependent_var, independent_vars, fixed, interaction_term, func)\n",
    "        models[dependent_var] = {'model': model, 'VIF': vif}\n",
    "        \n",
    "        # Print model summary and VIF results\n",
    "        print(f\"Model summary for {dependent_var}:\")\n",
    "        print(model.summary())\n",
    "        print(\"VIF results:\")\n",
    "        print(vif)\n",
    "        print(\"\\n\" + \"=\"*80 + \"\\n\")\n",
    "        \n",
    "        # Add predicted values to the data\n",
    "        data[f'{dependent_var}_predicted'] = model.fittedvalues\n",
    "        emmeans_data[dependent_var] = data[['Timepoint', dependent_var, f'{dependent_var}_predicted'] + independent_vars]\n",
    "    \n",
    "    # Prepare coefficient data for plotting\n",
    "    coef_df = pd.DataFrame()\n",
    "    for model_name, model_details in models.items():\n",
    "        model = model_details['model']\n",
    "        err_series = model.params - model.conf_int()[0]\n",
    "        temp_df = pd.DataFrame({\n",
    "            'coef': model.params.values[1:], \n",
    "            'pval': model.pvalues[1:],\n",
    "            'err': err_series.values[1:], \n",
    "            'varname': err_series.index.values[1:], \n",
    "            'model': model_name  \n",
    "        })\n",
    "        coef_df = pd.concat([coef_df, temp_df], ignore_index=True)\n",
    "    \n",
    "    # Set up plot parameters\n",
    "    base_y = np.arange(len(coef_df['varname'].unique()[0:-1]))\n",
    "    width = 0.2\n",
    "    marker_list = ['o', 's', 'D', '^', 'v']\n",
    "    model_names = coef_df['model'].unique()\n",
    "    \n",
    "    # Create the plot\n",
    "    fig, ax = plt.subplots(figsize=(8, 7))\n",
    "    for i, mod in enumerate(model_names):\n",
    "        mod_df = coef_df[coef_df.model == mod]\n",
    "        mod_df = mod_df.set_index('varname').reindex(coef_df['varname'].unique()[0:-1])\n",
    "        Y = base_y + width*i\n",
    "        \n",
    "        # Plot error bars\n",
    "        ax.barh(Y, mod_df['coef'], height=width, color='none', edgecolor='none', xerr=mod_df['err'], capsize=5)\n",
    "        \n",
    "        # Set grid and title\n",
    "        ax.grid(False)\n",
    "        total_subjs = len(data['Record ID'].unique())\n",
    "        ax.set_title(f\"{total_subjs} subjects\")\n",
    "        \n",
    "        # Plot coefficients with color-coded p-values\n",
    "        colors = ['red' if p < 0.05 else 'green' if p < 0.07 else 'black' for p in mod_df['pval']]\n",
    "        for y, coef, color in zip(Y, mod_df['coef'], colors):\n",
    "            ax.scatter(y=y, x=coef, marker=marker_list[i % len(marker_list)], s=50, color=color, zorder=3)\n",
    "    \n",
    "    # Add vertical line at x=0\n",
    "    ax.axvline(x=0, linestyle='--', color='black', linewidth=1)\n",
    "    \n",
    "    # Set y-axis labels\n",
    "    ax.set_yticks(base_y + width)\n",
    "    print(y_var_names)\n",
    "    ax.set_yticklabels(y_var_names, rotation=0, fontsize=16)\n",
    "    \n",
    "    # Set x-axis label\n",
    "    ax.set_xlabel('Coefficient and 95% Confidence Interval', fontsize=16)\n",
    "    \n",
    "    # Add legend\n",
    "    legend_elements = [Line2D([0], [0], marker=marker_list[i % len(marker_list)], label=f'{model_names[i]}', color='k', markersize=8) for i in range(len(model_names))]\n",
    "    ax.legend(handles=legend_elements, loc='upper right', prop={'size': 16}, labelspacing=1.2)\n",
    "    \n",
    "    # Adjust layout and display plot\n",
    "    plt.tight_layout()\n",
    "    plt.show()        \n",
    "    \n",
    "    return models, emmeans_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b24430d6914ed6d5",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Function to calculate Cohen's f²\n",
    "def calculate_cohens_f2(model, X):\n",
    "    r2 = model.rsquared\n",
    "    f2 = r2 / (1 - r2)\n",
    "    return f2\n",
    "\n",
    "# Function to plot interaction effects\n",
    "def plot_interaction_effect(models, emmeans_data, y_var_names, plot_name):\n",
    "    for dependent_var, model_info in models.items():\n",
    "        model = model_info['model']\n",
    "        pvalues = model.pvalues\n",
    "        # Identify significant (and trending) variables (p < 0.07)\n",
    "        significant_vars = pvalues[pvalues < 0.07].index.tolist()\n",
    "        \n",
    "        for sig_var in significant_vars:\n",
    "            # Skip certain variables\n",
    "            if sig_var not in ['Intercept', 'Age', 'ADOS_Module', 'C(Timepoint)[T.Post_CBD]', 'C(Timepoint)[T.Post_Placebo]', 'C(Timepoint)[T.Post_Wash]']:\n",
    "                # Create a 1x2 subplot\n",
    "                fig, axs = plt.subplots(1, 2, figsize=(6, 6), sharey=False, sharex=False)\n",
    "                \n",
    "                # Plot for Post_CBD and Post_Placebo timepoints\n",
    "                for i, timepoint in enumerate(['Post_CBD', 'Post_Placebo']):\n",
    "                    subset_data = emmeans_data[dependent_var][emmeans_data[dependent_var]['Timepoint'] == timepoint]\n",
    "                    first_color = sns.color_palette(\"Set2\")[0]\n",
    "                    \n",
    "                    # Create regression plot\n",
    "                    sns.regplot(data=subset_data, x=sig_var, y=f'{dependent_var}_predicted', ax=axs[i], color=first_color, scatter_kws={'edgecolor': 'none', 's': 100})\n",
    "                    \n",
    "                    # Perform OLS regression\n",
    "                    X = sm.add_constant(subset_data[sig_var])\n",
    "                    y = subset_data[f'{dependent_var}_predicted']\n",
    "                    est = sm.OLS(y, X).fit()\n",
    "                    \n",
    "                    # Calculate statistics\n",
    "                    raw_p_value = est.pvalues.iloc[1]  # p-value of the slope\n",
    "                    bonf_p_value = min(raw_p_value * 3, 1.0)  # Bonferroni correction\n",
    "                    r2 = est.rsquared\n",
    "                    cohens_f2 = calculate_cohens_f2(est, X)\n",
    "                    \n",
    "                    # Set plot title and labels\n",
    "                    axs[i].set_title(f'{timepoint}', fontsize=20, fontweight='bold')\n",
    "                    axs[i].text(0.5, -0.3, f'p = {raw_p_value:.3f}', fontsize=16, ha='center', transform=axs[i].transAxes)\n",
    "                    axs[i].text(0.5, -0.4, f'$R^2$: {r2:.3f}', fontsize=16, ha='center', transform=axs[i].transAxes)\n",
    "                    axs[i].text(0.5, -0.5, f'f²: {cohens_f2:.3f}', fontsize=16, ha='center', transform=axs[i].transAxes)\n",
    "                    \n",
    "                    axs[i].set_xlabel(sig_var, fontsize=20, fontweight='bold')\n",
    "                    if i == 0:\n",
    "                        axs[i].set_ylabel(f'Predicted {dependent_var}', fontsize=20, fontweight='bold')\n",
    "                    else:\n",
    "                        axs[i].set_ylabel('')\n",
    "                    axs[i].tick_params(axis='both', which='major', labelsize=16)\n",
    "                \n",
    "                # Set overall title\n",
    "                model_p_value = model.pvalues[sig_var]\n",
    "                plt.suptitle(f'Effect of {sig_var} on {dependent_var} (Stratified by Study Timepoint)\\nModel p-value: {model_p_value:.3f}', fontsize=20)\n",
    "                \n",
    "                # Adjust layout\n",
    "                plt.subplots_adjust(bottom=0.5)\n",
    "                plt.tight_layout()\n",
    "                \n",
    "                # Save plot\n",
    "                directory = '05_Plots\\LMM'\n",
    "                file_format = 'pdf'\n",
    "                file_path = os.path.join(directory, f'{plot_name}_Effect_of_{sig_var}_on_{dependent_var}.{file_format}')\n",
    "                plt.savefig(file_path, format=file_format)\n",
    "                \n",
    "                plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c93ce14c1c98f1e4",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def plot_pvalue_heatmap(models, dependent_var_names, heatmap_names, plot_name):\n",
    "    # Extract z-values and p-values for each independent variable across different dependent variables\n",
    "    z_values = {}\n",
    "    p_values = {}\n",
    "    \n",
    "    # Iterate through models to extract z-values and p-values\n",
    "    for dep_var, model_info in models.items():\n",
    "        model = model_info['model']\n",
    "        print(model)\n",
    "        summary = model.summary()\n",
    "        # Extract z-values and p-values from the model summary\n",
    "        result_table = summary.tables[1]\n",
    "        z_values[dep_var] = result_table['z'].to_dict()\n",
    "        p_values[dep_var] = result_table['P>|z|'].to_dict()\n",
    "    \n",
    "    # Create DataFrames for z-values and p-values\n",
    "    indep_vars = list(z_values[next(iter(z_values))].keys())\n",
    "    dep_vars = list(z_values.keys())\n",
    "    \n",
    "    z_data = pd.DataFrame(index=indep_vars, columns=dep_vars)\n",
    "    p_data = pd.DataFrame(index=indep_vars, columns=dep_vars)\n",
    "    \n",
    "    # Populate DataFrames with extracted values\n",
    "    for dep_var in dep_vars:\n",
    "        z_data[dep_var] = z_data.index.map(z_values[dep_var])\n",
    "        p_data[dep_var] = p_data.index.map(p_values[dep_var])\n",
    "    \n",
    "    # Convert all values to numeric, forcing any non-numeric to NaN\n",
    "    z_data = z_data.apply(pd.to_numeric, errors='coerce')\n",
    "    p_data = p_data.apply(pd.to_numeric, errors='coerce')\n",
    "    \n",
    "    # Transpose the DataFrames to have dependent variables on the top\n",
    "    z_data = z_data.T\n",
    "    p_data = p_data.T\n",
    "    \n",
    "    # Create a custom colormap using variations of the first color in \"Set2\" palette\n",
    "    first_color = sns.color_palette(\"Set2\")[0]\n",
    "    colors = [(first_color[0], first_color[1], first_color[2], alpha) for alpha in np.linspace(1, 0, 256)]\n",
    "    cmap = LinearSegmentedColormap.from_list('custom_set2', colors, N=256)\n",
    "    \n",
    "    # Remove 'Group' from heatmap_names and adjust data accordingly\n",
    "    group_index = heatmap_names.index('Group')\n",
    "    heatmap_names = heatmap_names[:group_index]\n",
    "    p_data = p_data.iloc[:, :group_index]\n",
    "    z_data = z_data.iloc[:, :group_index]\n",
    "\n",
    "    # Create figure and axes\n",
    "    fig, ax = plt.subplots(figsize=(18, 16))  # Increased height for better visibility\n",
    "\n",
    "    # Plot heatmap with adjusted parameters and no grid\n",
    "    heatmap = sns.heatmap(p_data, annot=z_data, fmt=\".2f\", cmap=cmap, vmin=0, vmax=0.07,\n",
    "                cbar_kws={'label': 'p-value', 'use_gridspec': False, 'location': 'right'},\n",
    "                linewidths=0, linecolor='none',  # Remove grid lines\n",
    "                annot_kws={\"size\": 16, \"weight\": \"normal\"},\n",
    "                cbar=True, square=True, ax=ax)\n",
    "\n",
    "    # Adjust colorbar\n",
    "    cbar = heatmap.collections[0].colorbar\n",
    "    cbar.ax.set_ylabel('p-value', rotation=270, labelpad=20)\n",
    "    cbar.ax.yaxis.set_label_position(\"left\")\n",
    "    cbar.ax.set_position([0.92, ax.get_position().y0, 0.02, ax.get_position().height])\n",
    "\n",
    "    # Add stars for significant p-values\n",
    "    for i in range(p_data.shape[0]):\n",
    "        for j in range(p_data.shape[1]):\n",
    "            if p_data.iloc[i, j] <= 0.001:\n",
    "                ax.text(j + 0.5, i + 0.35, '***', ha='center', va='bottom', color='black', fontweight='bold', fontsize=20)\n",
    "            elif p_data.iloc[i, j] <= 0.01:\n",
    "                ax.text(j + 0.5, i + 0.35, '**', ha='center', va='bottom', color='black', fontweight='bold', fontsize=20)\n",
    "            elif p_data.iloc[i, j] <= 0.05:\n",
    "                ax.text(j + 0.5, i + 0.35, '*', ha='center', va='bottom', color='black', fontweight='bold', fontsize=20)\n",
    "            elif p_data.iloc[i, j] <= 0.07:\n",
    "                ax.text(j + 0.5, i + 0.35, '#', ha='center', va='bottom', color='black', fontweight='bold', fontsize=20)\n",
    "\n",
    "    # Set labels and title (empty in this case)\n",
    "    ax.set_title('', fontsize=16, color='black')\n",
    "    ax.set_xlabel('', fontsize=20, color='black')\n",
    "    ax.set_ylabel('', fontsize=20, color='black')\n",
    "    \n",
    "    # Set and format x-axis and y-axis labels\n",
    "    ax.set_xticks(np.arange(len(heatmap_names)) + 0.5)\n",
    "    ax.set_xticklabels(heatmap_names, rotation=45, ha='center', va='top', fontsize=14, color='black', fontweight='bold')\n",
    "    ax.set_yticklabels(dependent_var_names, rotation=0, va='center', fontsize=20, fontweight='bold')\n",
    "    \n",
    "    # Remove ticks\n",
    "    ax.tick_params(axis='both', which='both', length=0)\n",
    "    \n",
    "    # Adjust bottom margin to accommodate centered labels\n",
    "    plt.subplots_adjust(bottom=0.2)\n",
    "\n",
    "    # Save the plot as a PDF file\n",
    "    directory = '05_Plots\\LMM'\n",
    "    file_format = 'pdf'\n",
    "    file_path = os.path.join(directory, f'{plot_name}.{file_format}')\n",
    "    plt.savefig(file_path, format=file_format, bbox_inches='tight', dpi=300)\n",
    "\n",
    "    # Display the plot\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6b378f29a46eead",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def plot_pvalue_heatmap_beh(models, dependent_var_names, heatmap_names, plot_name):\n",
    "    # Extract z-values and p-values for each independent variable across different dependent variables\n",
    "    z_values = {}\n",
    "    p_values = {}\n",
    "    \n",
    "    # Iterate through models to extract z-values and p-values\n",
    "    for dep_var, model_info in models.items():\n",
    "        model = model_info['model']\n",
    "        summary = model.summary()\n",
    "        result_table = summary.tables[1]\n",
    "        z_values[dep_var] = result_table['z'].to_dict()\n",
    "        p_values[dep_var] = result_table['P>|z|'].to_dict()\n",
    "    \n",
    "    # Create DataFrames for z-values and p-values\n",
    "    indep_vars = list(z_values[next(iter(z_values))].keys())\n",
    "    dep_vars = list(z_values.keys())\n",
    "    \n",
    "    z_data = pd.DataFrame(index=indep_vars, columns=dep_vars)\n",
    "    p_data = pd.DataFrame(index=indep_vars, columns=dep_vars)\n",
    "    \n",
    "    # Populate DataFrames with extracted values\n",
    "    for dep_var in dep_vars:\n",
    "        z_data[dep_var] = z_data.index.map(z_values[dep_var])\n",
    "        p_data[dep_var] = p_data.index.map(p_values[dep_var])\n",
    "    \n",
    "    # Convert all values to numeric, forcing any non-numeric to NaN\n",
    "    z_data = z_data.apply(pd.to_numeric, errors='coerce')\n",
    "    p_data = p_data.apply(pd.to_numeric, errors='coerce')\n",
    "    \n",
    "    # Transpose the DataFrames to have dependent variables on the top\n",
    "    z_data = z_data.T\n",
    "    p_data = p_data.T\n",
    "    \n",
    "    # Create a custom colormap using variations of the first color in \"Set2\" palette\n",
    "    first_color = sns.color_palette(\"Set2\")[0]\n",
    "    colors = [(first_color[0], first_color[1], first_color[2], alpha) for alpha in np.linspace(1, 0, 256)]\n",
    "    cmap = LinearSegmentedColormap.from_list('custom_set2', colors, N=256)\n",
    "    \n",
    "    # Create figure and axes\n",
    "    fig, ax = plt.subplots(figsize=(18, 16))  # Increased height for better visibility\n",
    "\n",
    "    # Plot heatmap with adjusted parameters and no grid\n",
    "    heatmap = sns.heatmap(p_data, annot=z_data, fmt=\".2f\", cmap=cmap, vmin=0, vmax=0.07,\n",
    "                cbar_kws={'label': 'p-value', 'use_gridspec': False, 'location': 'right'},\n",
    "                linewidths=0, linecolor='none',\n",
    "                annot_kws={\"size\": 16, \"weight\": \"normal\"},\n",
    "                cbar=True, square=True, ax=ax)\n",
    "\n",
    "    # Adjust colorbar\n",
    "    cbar = heatmap.collections[0].colorbar\n",
    "    cbar.ax.set_ylabel('p-value', rotation=270, labelpad=20)\n",
    "    cbar.ax.yaxis.set_label_position(\"left\")\n",
    "    cbar.ax.set_position([0.92, ax.get_position().y0, 0.02, ax.get_position().height])\n",
    "\n",
    "    # Add stars for significant p-values\n",
    "    for i in range(p_data.shape[0]):\n",
    "        for j in range(p_data.shape[1]):\n",
    "            if p_data.iloc[i, j] <= 0.001:\n",
    "                ax.text(j + 0.5, i + 0.35, '***', ha='center', va='bottom', color='black', fontweight='bold', fontsize=20)\n",
    "            elif p_data.iloc[i, j] <= 0.01:\n",
    "                ax.text(j + 0.5, i + 0.35, '**', ha='center', va='bottom', color='black', fontweight='bold', fontsize=20)\n",
    "            elif p_data.iloc[i, j] <= 0.05:\n",
    "                ax.text(j + 0.5, i + 0.35, '*', ha='center', va='bottom', color='black', fontweight='bold', fontsize=20)\n",
    "            elif p_data.iloc[i, j] <= 0.07:\n",
    "                ax.text(j + 0.5, i + 0.35, '#', ha='center', va='bottom', color='black', fontweight='bold', fontsize=20)\n",
    "\n",
    "    # Set labels and title (empty in this case)\n",
    "    ax.set_title('', fontsize=16, color='black')\n",
    "    ax.set_xlabel('', fontsize=20, color='black')\n",
    "    ax.set_ylabel('', fontsize=20, color='black')\n",
    "    \n",
    "    # Set and format x-axis and y-axis labels\n",
    "    ax.set_xticks(np.arange(len(heatmap_names)) + 0.5)\n",
    "    ax.set_xticklabels(heatmap_names, rotation=45, ha='center', va='top', fontsize=14, color='black', fontweight='bold')\n",
    "    ax.set_yticklabels(dependent_var_names, rotation=0, va='center', fontsize=20, fontweight='bold')\n",
    "    \n",
    "    # Remove ticks\n",
    "    ax.tick_params(axis='both', which='both', length=0)\n",
    "    \n",
    "    # Adjust bottom margin to accommodate centered labels\n",
    "    plt.subplots_adjust(bottom=0.2)\n",
    "\n",
    "    # Save the plot as a PDF file\n",
    "    directory = '05_Plots\\LMM'\n",
    "    file_format = 'pdf'\n",
    "    file_path = os.path.join(directory, f'{plot_name}.{file_format}')\n",
    "    plt.savefig(file_path, format=file_format, bbox_inches='tight', dpi=300)\n",
    "\n",
    "    # Display the plot\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd8ee96e15d7aaa2",
   "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": null,
   "id": "e579e415f9622643",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def plot_single_violin_ANOVA(data, channels, y_vars, y_var_names, timepoint_col='Timepoint', subject_col='Record ID', figsize=(25, 10), font=20, x_tick_labels=None):\n",
    "    # Determine the number of plots\n",
    "    n_cols = len(y_vars)\n",
    "    if len(channels) != 0:\n",
    "        n_rows = len(channels)\n",
    "    else:\n",
    "        n_rows = 1\n",
    "         \n",
    "    # Create subplots\n",
    "    fig, axs = plt.subplots(n_rows, n_cols, figsize=figsize)\n",
    "    axs = axs.flatten()\n",
    "    p_values = []\n",
    "    y_vars_grouped = []\n",
    "\n",
    "    # Loop through each variable to plot\n",
    "    for idx, var_group in enumerate(y_vars):\n",
    "        if len(channels) != 0:\n",
    "            var = f\"{channels[0]}_{var_group}\"\n",
    "        else:\n",
    "            var = var_group\n",
    "    \n",
    "        # Filter out rows with missing data\n",
    "        plot_df = data[[subject_col, timepoint_col, var]].dropna()\n",
    "    \n",
    "        # Collect groups for ANOVA\n",
    "        groups = [group[var].dropna() for name, group in plot_df.groupby(timepoint_col)]\n",
    "    \n",
    "        # Perform one-way ANOVA across all unique timepoints\n",
    "        if len(groups) > 1:\n",
    "            f_stat, p_value = f_oneway(*groups)\n",
    "            p_values.append(p_value)\n",
    "        else:\n",
    "            p_values.append(1)  # Non-significant p-value if not enough groups\n",
    "    \n",
    "        # Generate violin plots\n",
    "        order = ['Baseline', 'Post_Placebo', 'Post_CBD', 'Post_Wash']\n",
    "        palette = sns.color_palette(\"Set2\")\n",
    "        \n",
    "        # Create box and strip plots\n",
    "        sns.boxplot(ax=axs[idx], data=plot_df, x=timepoint_col, y=var, order=order, gap=.001, fill=False, showfliers=False, palette=palette)\n",
    "        sns.stripplot(ax=axs[idx], data=plot_df, x=timepoint_col, y=var, dodge=False, edgecolor='white', size=10, alpha=.5, order=order, palette=palette)\n",
    "    \n",
    "        # Set labels and format axes\n",
    "        axs[idx].set_ylabel(y_var_names[idx], fontsize=font, fontweight='bold')\n",
    "        axs[idx].set_xlabel('', fontsize=font, fontweight='bold')\n",
    "        axs[idx].tick_params(axis='both', which='major', labelsize=font, labelrotation=45)\n",
    "        axs[idx].grid(False)\n",
    "        \n",
    "        # Set custom x-tick labels if provided\n",
    "        if x_tick_labels:\n",
    "            axs[idx].set_xticklabels(x_tick_labels)\n",
    "\n",
    "    # Apply Bonferroni correction if statistical tests were performed\n",
    "    if len(p_values) > 1:\n",
    "        corrected_p_values = multipletests(p_values, alpha=0.05, method='bonferroni')[1]\n",
    "        for idx, p in enumerate(corrected_p_values):\n",
    "            if idx < len(y_vars):\n",
    "                if len(channels) != 0:\n",
    "                    axs[idx].set_title(f'{channels[0]} Electrodes\\nANOVA p = {p_values[idx]:.4f}, Bonf. p = {p:.4f}', fontsize=font, fontweight='bold')\n",
    "                else:\n",
    "                    axs[idx].set_title(f'All Subjects \\nANOVA p = {p_values[idx]:.4f}', fontsize=font, fontweight='bold')            \n",
    "    \n",
    "    # Hide unused axes\n",
    "    for ax in axs[len(y_vars):]:\n",
    "        ax.set_visible(False)\n",
    "        \n",
    "    # Adjust layout and save figure\n",
    "    plt.tight_layout()\n",
    "    directory = '05_Plots'\n",
    "    file_format = 'pdf'\n",
    "    file_path = os.path.join(directory, f'Descriptive.{file_format}')\n",
    "    plt.savefig(file_path, format=file_format)\n",
    "    \n",
    "    plt.show()\n",
    "\n",
    "# Example usage\n",
    "channel_groups = []\n",
    "y_vars = ['Duration_Secs', 'CBD', 'OHCBD', 'COOHCBD']\n",
    "y_var_names = [\n",
    "    'rsEEG Duration\\n(seconds)', 'CBD\\n(ng/mL)', '7-OHCBD\\n(ng/mL)', '7-COOHCBD\\n(ng/mL)'\n",
    "]\n",
    "\n",
    "titles = ['All Subjects']\n",
    "\n",
    "# Custom x-tick labels\n",
    "x_tick_labels = ['Baseline', 'Post-Placebo', 'Post-CBD', 'Post-Wash']\n",
    "\n",
    "# Call the function with the specified parameters\n",
    "plot_single_violin_ANOVA(\n",
    "    data=filt_data, \n",
    "    channels=[],  # Pass an empty list for all electrodes\n",
    "    y_vars=y_vars, \n",
    "    y_var_names=y_var_names, \n",
    "    timepoint_col='Timepoint', \n",
    "    subject_col='Record ID',\n",
    "    x_tick_labels=x_tick_labels  # Use custom x-tick labels\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62b5adecffa033ce",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Define parameters for the statistical model and visualization\n",
    "\n",
    "# Fixed effects for the mixed-effects model\n",
    "fixed = ['Age', 'ADOS_Module', 'C(Timepoint)', 'C(Randomization)']\n",
    "\n",
    "# Optimization method for the mixed-effects model\n",
    "func = 'nm'  # Nelder-Mead optimization algorithm\n",
    "\n",
    "# No interaction term specified\n",
    "interaction_term = ''\n",
    "\n",
    "# Labels for y-axis in visualization\n",
    "y_var_names = [\n",
    "    'CBD Timepoint',\n",
    "    'Placebo Timepoint',\n",
    "    'Wash Timepoint',\n",
    "    'Randomization',\n",
    "    'Offset\\n(uV²)',  \n",
    "    'Exponent\\n(uV² / Hz)', \n",
    "    'Delta\\n(uV²)',\n",
    "    'Theta\\n(uV²)',\n",
    "    'Alpha\\n(uV²)', \n",
    "    'ADOS\\n(module)', \n",
    "    'Age\\n(years)'\n",
    "]\n",
    "\n",
    "# Labels for heatmap\n",
    "heatmap_names = [\n",
    "    'Intercept',\n",
    "    'Post-CBD',\n",
    "    'Post-Placebo',\n",
    "    'Post-Wash',\n",
    "    'Randomization',\n",
    "    'Offset\\n(uV²)',  \n",
    "    'Exponent\\n(uV²/Hz)', \n",
    "    'Delta SNR\\n(uV²/Hz)',\n",
    "    'Theta SNR\\n(uV²/Hz)',\n",
    "    'Alpha SNR\\n(uV²/Hz)', \n",
    "    'Age',\n",
    "    'ADOS',\n",
    "    'Group']\n",
    "\n",
    "# Dependent variable name (metabolite)\n",
    "dependent_var_name = ['7-COOH-CBD (z)']\n",
    "\n",
    "# Names of independent variables for visualization\n",
    "independent_var_name = [\n",
    "    'Int', \n",
    "    'Post-CBD',\n",
    "    'Post-Placebo',\n",
    "    'Post-Wash',\n",
    "    'Randomization', \n",
    "    'Offset\\n(uV²)',\n",
    "    'Exponent\\n(uV² / Hz)',\n",
    "    'Delta\\n(uV²)', \n",
    "    'Theta\\n(uV²)', \n",
    "    'Alpha\\n(uV²)', \n",
    "    'Age\\n(years)', \n",
    "    'ADOS\\n(module)'\n",
    "]\n",
    "\n",
    "# Dependent variables for the model\n",
    "dependent_vars = ['COOHCBD_Z_score']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c814aa736fe2b9e2",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Initialize list to store results\n",
    "all_results = []\n",
    "\n",
    "# Loop through each channel group\n",
    "for group_name, channel_indices in Channel_Groups.items():\n",
    "    independent_vars = []\n",
    "    \n",
    "    # Add exponent and offset variables for the current group\n",
    "    exponent_var = f\"{group_name}_Exponent\"\n",
    "    offset_var = f\"{group_name}_Offset\"\n",
    "    independent_vars.extend([offset_var, exponent_var])\n",
    "    \n",
    "    # Add band-specific variables for the current group\n",
    "    for band_name, freq_range in BANDS.items():\n",
    "        band_var = f\"{group_name}_{band_name}_SNR\"\n",
    "        independent_vars.append(band_var)\n",
    "    \n",
    "    # Run models and create plots for the current group\n",
    "    print(f\"Running models and plotting for group: {group_name}\")\n",
    "    models, emmeans_results = run_models_and_plot(filt_data, dependent_vars, independent_vars, fixed, interaction_term, func, y_var_names)\n",
    "    \n",
    "    # Generate interaction effect plot\n",
    "    plot_name = f\"{group_name}\"\n",
    "    plot_interaction_effect(models, emmeans_results, y_var_names, plot_name)\n",
    "    \n",
    "    # Generate heatmap plot\n",
    "    plot_name = f\"{group_name}_heatmap\"\n",
    "    plot_pvalue_heatmap(models, dependent_var_name, heatmap_names, plot_name)\n",
    "    \n",
    "    # Store results for the current group\n",
    "    all_results.append({\n",
    "        'group': group_name,\n",
    "        'models': models,\n",
    "        'emmeans_results': emmeans_results,\n",
    "        'independent_vars': independent_vars,\n",
    "        'dependent_vars': dependent_vars\n",
    "    })"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98bda939ad0724a3",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Model parameters\n",
    "func = 'nm'  # Optimization method (likely Nelder-Mead)\n",
    "interaction_term = ''  # No interaction term specified\n",
    "fixed = ['Age', 'ADOS_Module', 'C(Timepoint)', 'C(Randomization)']  # Fixed effects\n",
    "\n",
    "# Dependent variable (CBD metabolite)\n",
    "dependent_vars = ['COOHCBD_Z_score']  # Focusing on 7-COOH-CBD z-scores\n",
    "\n",
    "# Independent variables (various cognitive and behavioral assessments)\n",
    "independent_vars = [\n",
    "    'rbs_total_score', 'ppvt_raw_score', 'toni4_raw_score', 'eowpvt4_raw_score', \n",
    "    'beery_vmi_raw_score', 'beery_vp_raw_score', 'beery_mc_raw_score'\n",
    "]\n",
    "\n",
    "# Labels for y-axis in visualization\n",
    "y_var_names_constant = [\n",
    "    'CBD Timepoint', 'Placebo Timepoint', 'Wash Timepoint', 'Randomization', 'Age', 'ADOS'\n",
    "]\n",
    "\n",
    "y_var_names_append = [\n",
    "    'Total RBS\\nScore', 'PPVT\\nRaw Score', 'TONI-4\\nRaw Score', 'EOWPVT-4\\nRaw Score',\n",
    "    'Beery VMI\\nRaw Score', 'Beery VP\\nRaw Score', 'Beery MC\\nRaw Score',\n",
    "]\n",
    "\n",
    "# Labels for heatmap\n",
    "heatmap_names = [\n",
    "    'Intercept', 'Post-CBD', 'Post-Placebo', 'Post-Wash', 'Randomization',\n",
    "    'TONI-4\\nScore', 'Age', 'ADOS', 'Group'\n",
    "]\n",
    "\n",
    "# Dependent variable name for visualization\n",
    "dependent_var_name = ['7-COOH-CBD (z)']\n",
    "\n",
    "# Loop through each independent variable\n",
    "for ind_var, y_var_name in zip(independent_vars, y_var_names_append):\n",
    "    # Prepare y-axis labels\n",
    "    y_var_names = y_var_names_constant[:4] + [y_var_name] + y_var_names_constant[4:]\n",
    "    print(y_var_names)\n",
    "    \n",
    "    # Prepare data\n",
    "    df_copy = filt_data.copy()\n",
    "    df = df_copy.dropna(subset=[ind_var])\n",
    "    df_beh_filt = df.groupby('Record ID').filter(lambda x: len(x['Timepoint'].unique()) >= 3)\n",
    "    \n",
    "    # Calculate and print statistics for each timepoint\n",
    "    result = df_beh_filt.groupby('Timepoint').apply(calculate_stats).reset_index()\n",
    "    print(ind_var)\n",
    "    print(result)\n",
    "    \n",
    "    # Run models and generate plots\n",
    "    models, emmeans_results = run_models_and_plot(df_beh_filt, dependent_vars, [ind_var], fixed, interaction_term, func, y_var_names)\n",
    "    \n",
    "    # Generate interaction effect plot\n",
    "    plot_name = f\"{ind_var}_beh_reg\"\n",
    "    plot_interaction_effect(models, emmeans_results, y_var_names, plot_name)\n",
    "    \n",
    "    # Generate heatmap\n",
    "    plot_name = f\"{ind_var}_beh_heat\"\n",
    "    plot_pvalue_heatmap_beh(models, dependent_var_name, heatmap_names, plot_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40047b89840c767c",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Beery VMI, VP, MC analysis\n",
    "\n",
    "# Dependent variable (CBD metabolite)\n",
    "dependent_vars = ['COOHCBD_Z_score']  # Focusing on 7-COOH-CBD z-scores\n",
    "\n",
    "# Fixed effects for the model\n",
    "fixed = ['Age', 'ADOS_Module', 'C(Timepoint)', 'C(Randomization)']\n",
    "\n",
    "# Independent variables (Beery subtests)\n",
    "independent_vars = ['beery_vmi_raw_score', 'beery_vp_raw_score', 'beery_mc_raw_score']\n",
    "\n",
    "# Labels for y-axis in visualization\n",
    "y_var_names = [\n",
    "    'Session_CBD', 'Session_Placebo', 'Session_Wash', 'Randomization',\n",
    "    'Total Beery VMI\\nRaw Score', 'Total Beery VP\\nRaw Score', 'Total Beery MC\\nRaw Score',\n",
    "    'Age', 'ADOS'\n",
    "]\n",
    "\n",
    "# Labels for heatmap\n",
    "heatmap_names = [\n",
    "    'Intercept', 'Post-CBD', 'Post-Placebo', 'Post-Wash', 'Randomization',\n",
    "    'VMI\\nScore', 'VP\\nScore', 'MC\\nScore', 'Age', 'ADOS'\n",
    "]\n",
    "\n",
    "# Dependent variable name for visualization\n",
    "dependent_var_name = ['7-COOH-CBD (z)']\n",
    "\n",
    "# Data preparation\n",
    "df_copy = filt_data.copy()\n",
    "df = df_copy.dropna(subset=['beery_vmi_raw_score', 'beery_vp_raw_score', 'beery_mc_raw_score'])\n",
    "df = df.groupby('Record ID').filter(lambda x: len(x['Timepoint'].unique()) >= 3)\n",
    "\n",
    "# Calculate and print statistics for each timepoint\n",
    "result = df.groupby('Timepoint').apply(calculate_stats).reset_index()\n",
    "print(ind_var)  # Note: 'ind_var' is not defined in this snippet\n",
    "print(result)\n",
    "\n",
    "# Run models and generate plots\n",
    "models, emmeans_results = run_models_and_plot(df, dependent_vars, independent_vars, fixed, interaction_term, func, y_var_names)\n",
    "\n",
    "# Generate interaction effect plot\n",
    "plot_name = 'beery_beh_reg'\n",
    "plot_interaction_effect(models, emmeans_results, y_var_names, plot_name)\n",
    "\n",
    "# Generate heatmap\n",
    "plot_name = 'beery_beh_heat'\n",
    "plot_pvalue_heatmap_beh(models, dependent_var_name, heatmap_names, plot_name)\n",
    "\n",
    "# Print statistics again\n",
    "result = df.groupby('Timepoint').apply(calculate_stats).reset_index()\n",
    "print(result)"
   ]
  }
 ],
 "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
}