{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "import os\n",
    "from scipy.stats import mannwhitneyu, wilcoxon\n",
    "from statsmodels.stats.multitest import multipletests\n",
    "\n",
    "from statannotations.Annotator import Annotator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "pd.options.display.float_format = '{:.3f}'.format\n",
    "pd.set_option('display.max_columns', None)\n",
    "\n",
    "plt.rcParams['figure.dpi']=170\n",
    "\n",
    "DPI=250"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from list_vars import LIST_PROFILERS, DIR_FIGURES, RESULTS_DIR, POOLS, CONTROLS"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Biological sample analysis (part 2)\n",
    "\n",
    "In this notebook we are going to do the analysis on the *biological* samples (POOL samples + controls).\n",
    "\n",
    "One thing we want to check is how S and mode affect the robustness of our results. So, to do that we are going to compare the results using a comparative plot. The comparativ plot shows how many new variable are shown compared to the previous case. So, we can do a comparative increase on mode by fixing S, or on S by fixing the mode.\n",
    "\n",
    "This analysis can be performed with many variables, and we are going to choose the following:\n",
    "1) The number of detected species, across all samples and per individual sample (after NORM+ and cutting of species with more than 65% NaNs).\n",
    "2) The number of detected differentially abundant species across the 4 comparisons.\n",
    "\n",
    "With this in mind, we can later select one S and one mode and do the following analyses.\n",
    "\n",
    "1) The importance of including the biological control samples to ensure that false positives are not considered.\n",
    "2) The importance of normalizing the reads considering the biogical samples.\n",
    "3) Plot the significantly differentially abundant species. \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Observing the effect of $S$ on species retention for biological controls"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "rename_mapping = {\n",
    "        'POOL1': 'RR1', 'POOL2': 'RR2', 'POOL3': 'RR3', 'POOL4': 'RR4',\n",
    "        'POOL5': 'SP1', 'POOL6': 'SP2', 'POOL7': 'SP3', 'POOL8': 'SP4',\n",
    "        'POOL9': 'HC1', 'POOL10': 'HC2', 'POOL11': 'HC3', 'POOL12': 'HC4',\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "dict_retained_discarded_species = {'mode': [], 'S': [], 'sample': [], 'retained': [], 'discarded': []}\n",
    "\n",
    "for mode in [3, 5, 7]:\n",
    "    for S in [0, 1, 2, 3, 4, 5, 6, 7, 10, 15]:\n",
    "        for sample in [f'POOL{i}' for i in range(1, 13)] + ['ACIDOLA', 'BLACTIS']:\n",
    "            flags_file =  f'{RESULTS_DIR}/summary/{sample}_pass2_mode{mode}_taxgenus_S{S}_NORM+.flags.tsv'\n",
    "            df_flags = pd.read_csv(flags_file, sep='\\t').set_index('taxonomy_id')[['name', 'lineage', 'mean_norm']]\n",
    "            \n",
    "            dict_retained_discarded_species['mode'].append(mode)\n",
    "            dict_retained_discarded_species['S'].append(S)\n",
    "            dict_retained_discarded_species['sample'].append(rename_mapping[sample] if 'POOL' in sample else sample)\n",
    "            dict_retained_discarded_species['retained'].append((df_flags['mean_norm'] == False).sum())\n",
    "            dict_retained_discarded_species['discarded'].append((df_flags['mean_norm'] == True).sum())          \n",
    "\n",
    "df_retained_discarded_species = pd.DataFrame(dict_retained_discarded_species)\n",
    "df_retained_discarded_species['total'] = df_retained_discarded_species['retained'] + df_retained_discarded_species['discarded']\n",
    "\n",
    "df_retained_discarded_species['percentage'] = 100 * df_retained_discarded_species['retained'] / df_retained_discarded_species['total']\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_retained_discarded_species"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_total_counts = df_retained_discarded_species.groupby(['mode', 'sample'])['total'].max().unstack(level=0)\n",
    "df_total_counts['increment_3_5'] = df_total_counts[5] / df_total_counts[3]\n",
    "df_total_counts['increment_3_7'] = df_total_counts[7] / df_total_counts[3]\n",
    "\n",
    "df_total_counts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_retained_discarded_species"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for column, name in zip(['retained', 'percentage'], ['Retained genera',  'Retained percentage']):\n",
    "    # Set up the grid for 14 samples in a 4x4 layout\n",
    "    samples = df_retained_discarded_species['sample'].unique()  # Unique samples\n",
    "    n_cols, n_rows = 7, 2  # Define the grid layout\n",
    "\n",
    "    # Initialize the grid layout\n",
    "    fig, axes = plt.subplots(n_rows, n_cols, figsize=(2*n_cols, 2*n_rows), sharex=True)\n",
    "\n",
    "    # Placeholder to collect handles and labels for the legend\n",
    "    handles, labels = None, None\n",
    "\n",
    "\n",
    "    # Loop through each sample and create its subplot\n",
    "    for i, sample in enumerate(samples):\n",
    "        row, col = divmod(i, n_cols)  # Get the row and column index for the subplot\n",
    "        ax = axes[row, col]\n",
    "\n",
    "        # Filter the DataFrame for the current sample\n",
    "        sample_data = df_retained_discarded_species[df_retained_discarded_species['sample'] == sample]\n",
    "\n",
    "        # Plot retained species as a function of S with mode as hue\n",
    "        lineplot = sns.lineplot(\n",
    "            data=sample_data,\n",
    "            x='S',\n",
    "            y=column,\n",
    "            hue='mode',\n",
    "            marker='o',\n",
    "            ax=ax,\n",
    "            legend=(i == 6)  # Enable legend only for the first subplot\n",
    "        )\n",
    "        ax.set_title(sample)  # Title of the subplot is the sample name\n",
    "        ax.set_xlabel(\"S\")\n",
    "        if col == 0:\n",
    "            ax.set_ylabel(name)\n",
    "        else:\n",
    "            ax.set_ylabel(\"\")\n",
    "\n",
    "        ax.set_xticks([0,  2,  4,  6,  10, 15])\n",
    "\n",
    "        if i==6:\n",
    "            ax.legend(frameon=False, bbox_to_anchor=(1.05, 1), title='Mode')\n",
    "\n",
    "        # Set a specific Y-axis limit for the last two samples\n",
    "        if column == 'retained':\n",
    "            if i in [len(samples) - 2, len(samples) - 1]:  # Last two samples\n",
    "                ax.set_ylim(0, 4)\n",
    "\n",
    "    # Remove extra subplots (for grid cells not needed)\n",
    "    for j in range(i + 1, n_rows * n_cols):\n",
    "        fig.delaxes(axes[j // n_cols, j % n_cols])\n",
    "\n",
    "\n",
    "    # Adjust the main plot area to leave space for the legend\n",
    "    plt.tight_layout()\n",
    "\n",
    "    for format in ['png', 'tiff']: \n",
    "        plt.savefig(f'{RESULTS_DIR}/figures/paper/s_effect_{column}.{format}', dpi=DPI)\n",
    "\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plotting sample detection and statistacally differential species across modes and S"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize an empty list to collect data rows\n",
    "data_rows = []\n",
    "\n",
    "# Iterate through modes and S values\n",
    "for mode in [3, 5, 7]:\n",
    "    for S in [0, 1, 2, 3, 4, 5, 6, 7, 10, 15]:\n",
    "        # Read the data\n",
    "        df_normpipe_retained_species = pd.read_csv(f'{RESULTS_DIR}/merged_counts/mode{mode}_S{S}_NORM+_retained.tsv', sep='\\t')\n",
    "        df_normpipe_discarded_normplus = pd.read_csv(f'{RESULTS_DIR}/merged_counts/mode{mode}_S{S}_NORM+_discarded_norm+.tsv', sep='\\t')\n",
    "        df_normpipe_discarded_common = pd.read_csv(f'{RESULTS_DIR}/merged_counts/mode{mode}_S{S}_NORM+_discarded_common.tsv', sep='\\t')\n",
    "\n",
    "        # Get the list of samples (assuming they are column names)\n",
    "        samples = df_normpipe_retained_species.columns[3:]\n",
    "\n",
    "        data_rows.append({\n",
    "                'mode': mode,\n",
    "                'S': S,\n",
    "                'sample': 'ALL',\n",
    "                'count_retained': len(df_normpipe_retained_species),\n",
    "                'count_discarded_norm': len(df_normpipe_discarded_normplus),\n",
    "                'count_discarded_all': len(df_normpipe_discarded_common)\n",
    "            })\n",
    "        \n",
    "        for sample in samples:\n",
    "            # Count non-NaN species for each dataframe and sample\n",
    "            count_retained = df_normpipe_retained_species[sample].notna().sum()\n",
    "            count_discarded_norm = df_normpipe_discarded_normplus[sample].notna().sum()\n",
    "            count_discarded_all = df_normpipe_discarded_common[sample].notna().sum()\n",
    "\n",
    "            # Collect the data as a dictionary\n",
    "            data_rows.append({\n",
    "                'mode': mode,\n",
    "                'S': S,\n",
    "                'sample': sample,\n",
    "                'count_retained': count_retained,\n",
    "                'count_discarded_norm': count_discarded_norm,\n",
    "                'count_discarded_all': count_discarded_all\n",
    "            })\n",
    "\n",
    "# Convert the collected rows into a dataframe\n",
    "df_stats_species_count = pd.DataFrame(data_rows)\n",
    "df_stats_species_count['count_total'] = df_stats_species_count['count_retained'] + df_stats_species_count['count_discarded_norm'] + df_stats_species_count['count_discarded_all']\n",
    "df_stats_species_count"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_stats_species_count[df_stats_species_count['sample'] == 'ALL']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize an empty list to collect data rows\n",
    "data_rows = []\n",
    "\n",
    "# Iterate through modes and S values\n",
    "for mode in [3, 5, 7]:\n",
    "    for S in [0, 1, 2, 3, 4, 5, 6, 7, 10, 15]:\n",
    "        # Read the data\n",
    "        df_pval_HCvsRR = pd.read_csv(f'{RESULTS_DIR}/differential_abundance/mode{mode}_S{S}_HCvsRR.tsv', sep='\\t')\n",
    "        df_pval_HCvsSP = pd.read_csv(f'{RESULTS_DIR}/differential_abundance/mode{mode}_S{S}_HCvsSP.tsv', sep='\\t')\n",
    "        df_pval_RRvsSP = pd.read_csv(f'{RESULTS_DIR}/differential_abundance/mode{mode}_S{S}_RRvsSP.tsv', sep='\\t')\n",
    "        df_pval_sex = pd.read_csv(f'{RESULTS_DIR}/differential_abundance/mode{mode}_S{S}_sex.tsv', sep='\\t')\n",
    "\n",
    "        # Get the list of samples (assuming they are column names)\n",
    "        samples = df_normpipe_retained_species.columns[3:]\n",
    "\n",
    "        data_rows.append({\n",
    "                'mode': mode,\n",
    "                'S': S,\n",
    "                'count_HCvsRR': len(df_pval_HCvsRR[df_pval_HCvsRR['pval_MW'] < 0.05]),\n",
    "                'count_HCvsSP': len(df_pval_HCvsSP[df_pval_HCvsSP['pval_MW'] < 0.05]),\n",
    "                'count_RRvsSP': len(df_pval_RRvsSP[df_pval_RRvsSP['pval_MW'] < 0.05]),\n",
    "                'count_sex': len(df_pval_sex[df_pval_sex['pval_MW'] < 0.05]),\n",
    "                'species_HCvsRR': df_pval_HCvsRR[df_pval_HCvsRR['pval_MW'] < 0.05]['name'].values.tolist(),\n",
    "                'species_HCvsSP': df_pval_HCvsSP[df_pval_HCvsSP['pval_MW'] < 0.05]['name'].values.tolist(),\n",
    "                'species_RRvsSP': df_pval_RRvsSP[df_pval_RRvsSP['pval_MW'] < 0.05]['name'].values.tolist(),\n",
    "                'species_sex': df_pval_sex[df_pval_sex['pval_MW'] < 0.05]['name'].values.tolist()\n",
    "            })\n",
    "        \n",
    "      \n",
    "\n",
    "# Convert the collected rows into a dataframe\n",
    "df_stats_species_diffabundance = pd.DataFrame(data_rows)\n",
    "df_stats_species_diffabundance['count_total'] = df_stats_species_diffabundance['count_HCvsRR'] + df_stats_species_diffabundance['count_HCvsSP'] + \\\n",
    "                                                df_stats_species_diffabundance['count_RRvsSP'] + df_stats_species_diffabundance['count_sex']\n",
    "df_stats_species_diffabundance\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Why normalization with biological controls is relevant"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "mode = 3\n",
    "S = 7"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_normpipe_retained_species = pd.read_csv(f'{RESULTS_DIR}/merged_counts/mode{mode}_S{S}_NORM+_retained.tsv', sep='\\t')\n",
    "df_normpipe_discarded_normplus = pd.read_csv(f'{RESULTS_DIR}/merged_counts/mode{mode}_S{S}_NORM+_discarded_norm+.tsv', sep='\\t')\n",
    "df_normpipe_discarded_common = pd.read_csv(f'{RESULTS_DIR}/merged_counts/mode{mode}_S{S}_NORM+_discarded_common.tsv', sep='\\t')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(df_normpipe_retained_species), len(df_normpipe_discarded_normplus), len(df_normpipe_discarded_common))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "species = 'Janibacter'\n",
    "\n",
    "display(df_normpipe_retained_species[df_normpipe_retained_species['name'] == species])\n",
    "display(df_normpipe_discarded_normplus[df_normpipe_discarded_normplus['name'] == species])\n",
    "display(df_normpipe_discarded_common[df_normpipe_discarded_common['name'] == species])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plotting discarded species"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "def annotation_pval(df, ax, groups, x_col, y_col):\n",
    "    if len(groups) == 2:\n",
    "        pairs = [(groups[0], groups[1])]\n",
    "\n",
    "    annotator = Annotator(ax, pairs, data=df, x=x_col, y=y_col)\n",
    "    annotator.configure(test='Mann-Whitney', text_format='simple', loc='outside', line_width=0.65, fontsize=8)\n",
    "    annotator.apply_and_annotate()\n",
    "\n",
    "def add_medians(df, ax, x_col, y_col, ysync=0.1, xsync=0.1):\n",
    "    groups = df[x_col].unique()\n",
    "    vals = [df[df[x_col] == group][y_col].values for group in groups]\n",
    "\n",
    "    medians = [np.nanmedian(val) for val in vals]\n",
    "    medianslog =  [np.log10(np.nanmedian(val) + 1) for val in vals]\n",
    "    \n",
    "\n",
    "    for i, median in enumerate(medians):\n",
    "            ax.text(i + xsync, medianslog[i] + ysync, f\"{medians[i]:.0f}\", ha='left', va='bottom', fontsize=7)\n",
    "\n",
    "    log2fc = np.log2(medians[1] / medians[0])\n",
    "\n",
    "    if log2fc > 0:\n",
    "        ax.text(-0.4, 5, f\"log$_2$FC: {log2fc:.2f}\", ha='left', va='bottom', fontsize=7)\n",
    "    else:\n",
    "        ax.text(1.5, 5, f\"log$_2$FC: {log2fc:.2f}\", ha='right', va='bottom', fontsize=7)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_ax_2(ax, y1, y2, color1, color2, group1, group2, species):\n",
    "    ax.scatter([0] * len(y1), y1, color=color1, label=group1, alpha=0.8)\n",
    "    ax.scatter([1] * len(y2), y2, color=color2, label=group2, alpha=0.8)\n",
    "\n",
    "    \n",
    "    # Add the means as horizontal lines\n",
    "    ax.plot([0 - 0.2, 0 + 0.2], [np.log10(np.nanmedian(10 ** y1 - 1) + 1), np.log10(np.nanmedian(10 ** y1 - 1) + 1)], color=color1, lw=2)\n",
    "    ax.plot([1 - 0.2, 1 + 0.2], [np.log10(np.nanmedian(10 ** y2 - 1) + 1), np.log10(np.nanmedian(10 ** y2 - 1) + 1)], color=color2, lw=2)\n",
    "\n",
    "    # Add horizontal gridlines at 0.5 intervals\n",
    "    for y in np.arange(0, max(6, max(max(y1) + 0.5, max(y2) + 0.5)), 1):  # Adjust range to match y-axis limits\n",
    "        ax.axhline(y, color='lightgray', linestyle='--', linewidth=0.5, zorder=0)\n",
    "\n",
    "    # Customize the x-axis\n",
    "    ax.set_xticks([0, 1])\n",
    "    ax.set_xticklabels([group1, group2])\n",
    "    ax.set_xlim(-0.5, 1.5)\n",
    "\n",
    "    # Set axis limits and labels\n",
    "    ax.set_ylim(0, max(6, max(max(y1) + 0.5, max(y2) + 0.5)))\n",
    "    ax.set_ylabel('log$_{10}$ counts')\n",
    "\n",
    "    # Remove the x and y axis lines\n",
    "    ax.spines['top'].set_visible(False)\n",
    "    ax.spines['right'].set_visible(False)\n",
    "\n",
    "    # Remove ticks on the y-axis\n",
    "    ax.yaxis.set_ticks_position('none')\n",
    "\n",
    "    ax.set_title(species, fontsize=10, pad=20)\n",
    "\n",
    "    # Set lighter grid aesthetics\n",
    "    ax.grid(False)\n",
    "\n",
    "\n",
    "def plot_ax_4(ax, yHC, ySP, yRR, yCTRL, species):\n",
    "    ax.scatter([0] * len(yHC), yHC, color='#648FFF', label='HC', alpha=0.8)\n",
    "    ax.scatter([1] * len(ySP), ySP, color='#DC267F', label='SP', alpha=0.8)\n",
    "    ax.scatter([2] * len(yRR), yRR, color='#FE6100', label='RR', alpha=0.8)\n",
    "    ax.scatter([3] * len(yCTRL), yCTRL, color='#848484', label='CTRL', alpha=0.8)\n",
    "\n",
    "    # Add the means as horizontal lines\n",
    "    ax.plot([0 - 0.2, 0 + 0.2], [np.nanmean(yHC), np.nanmean(yHC)], color='#648FFF', lw=2)\n",
    "    ax.plot([1 - 0.2, 1 + 0.2], [np.nanmean(ySP), np.nanmean(ySP)], color='#DC267F', lw=2)\n",
    "    ax.plot([2 - 0.2, 2 + 0.2], [np.nanmean(yRR), np.nanmean(yRR)], color='#FE6100', lw=2)\n",
    "    ax.plot([3 - 0.2, 3 + 0.2], [np.nanmean(yCTRL), np.nanmean(yCTRL)], color='#848484', lw=2)\n",
    "\n",
    "    # Add horizontal gridlines at 0.5 intervals\n",
    "    for y in np.arange(0, max(5, max(max(yHC) + 0.5, max(ySP) + 0.5, max(yRR) + 0.5, max(yCTRL) + 0.5)), 1):  # Adjust range to match y-axis limits\n",
    "        ax.axhline(y, color='lightgray', linestyle='--', linewidth=0.5, zorder=0)\n",
    "\n",
    "    # Customize the x-axis\n",
    "    ax.set_xticks([0, 1, 2, 3])\n",
    "    ax.set_xticklabels(['HC', 'SP', 'RR', 'CTRL'])\n",
    "    ax.set_xlim(-0.5, 3.5)\n",
    "\n",
    "    # Set axis limits and labels\n",
    "    ax.set_ylim(0, max(5, max(max(yHC) + 0.5, max(ySP) + 0.5, max(yRR) + 0.5, max(yCTRL) + 0.5)))\n",
    "    ax.set_ylabel('log$_{10}$ counts')\n",
    "\n",
    "    # Remove the x and y axis lines\n",
    "    ax.spines['top'].set_visible(False)\n",
    "    ax.spines['right'].set_visible(False)\n",
    "\n",
    "    # Remove ticks on the y-axis\n",
    "    ax.yaxis.set_ticks_position('none')\n",
    "\n",
    "    ax.set_title(species, fontsize=10)\n",
    "\n",
    "    # Set lighter grid aesthetics\n",
    "    ax.grid(False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_normpipe_discarded_common"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_species(df, ncols_max=6):\n",
    "    nplots = len(df)\n",
    "    nrows = nplots//ncols_max \n",
    "    nrows += (nplots % ncols_max != 0)\n",
    "\n",
    "    if nplots:\n",
    "        # Create the figure and axis\n",
    "        fig, axs = plt.subplots(nrows, ncols_max, figsize=(2 * ncols_max, 2 * nrows))\n",
    "\n",
    "        for i, species in enumerate(df.index):\n",
    "            ax = axs.ravel()[i] if nplots > 1 else axs\n",
    "            yHC = np.log10(df.loc[species][['HC1', 'HC2', 'HC3', 'HC4']].astype(float).values + 1)\n",
    "            ySP = np.log10(df.loc[species][['SP1', 'SP2', 'SP3', 'SP4']].astype(float).values + 1)\n",
    "            yRR = np.log10(df.loc[species][['RR1', 'RR2', 'RR3', 'RR4']].astype(float).values + 1)\n",
    "            yCTRL = np.log10(df.loc[species][['ACIDOLA', 'BLACTIS']].astype(float).values + 1)\n",
    "            \n",
    "            plot_ax_4(ax, yHC, ySP, yRR, yCTRL, species)\n",
    "            df_annot = pd.DataFrame({'group': ['HC'] * 4 + ['SP'] * 4 + ['RR'] * 4 + ['CTRL'] * 2, \n",
    "                                     'exp': df.loc[species, ['HC1', 'HC2', 'HC3', 'HC4', 'SP1', 'SP2', 'SP3', 'SP4', 'RR1', 'RR2', 'RR3', 'RR4', 'ACIDOLA', 'BLACTIS']].astype(float).values})\n",
    "            # add_medians(df_annot, ax, x_col='group', y_col='exp', xsync=0.15, ysync=0)\n",
    "\n",
    "        if nplots % ncols_max:\n",
    "            for i in range(nplots, nrows * ncols_max ):\n",
    "                axs.ravel()[i].axis('off')\n",
    "        plt.tight_layout()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_species(df_normpipe_discarded_common.copy().set_index('name'), ncols_max=6)\n",
    "for format in ['png', 'tiff']: \n",
    "    plt.savefig(f'{RESULTS_DIR}/figures/paper/l2fc_NORM+x_discarded.{format}', dpi=DPI)\n",
    "\n",
    "plot_species(df_normpipe_discarded_normplus.copy().set_index('name'), ncols_max=7)\n",
    "for format in ['png', 'tiff']: \n",
    "    plt.savefig(f'{RESULTS_DIR}/figures/paper/l2fc_NORMx_discarded.{format}', dpi=DPI)\n",
    "\n",
    "plot_species(df_normpipe_retained_species.copy().set_index('name'), ncols_max=8)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plotting differential species"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_diff_sopecies(df_pval, cols_A, cols_B, color1, color2, nrows=2, other_labels = None):\n",
    "    ncols = len(df_pval) // nrows\n",
    "\n",
    "    if other_labels is None:\n",
    "        g1, g2 = cols_A[0][:-1], cols_B[0][:-1]\n",
    "\n",
    "    if ncols:\n",
    "        # Create the figure and axis\n",
    "        fig, axs = plt.subplots(nrows, ncols, figsize=(2 * ncols, nrows * 2))\n",
    "\n",
    "        for i in range(len(df_pval)):\n",
    "            df_i = df_pval.iloc[i]\n",
    "            ax = axs.ravel()[i] if ncols > 1 else axs\n",
    "            y1 = np.log10(df_i[cols_A].astype(float).values + 1)\n",
    "            y2 = np.log10(df_i[cols_B].astype(float).values + 1)\n",
    "            \n",
    "            species = df_i['name']\n",
    "\n",
    "            df_annot = pd.DataFrame({'group': [i[:-1] for i in cols_A + cols_B], 'exp': df_i[cols_A + cols_B].astype(float).values})\n",
    "\n",
    "            plot_ax_2(ax, y1=y1, y2=y2, color1=color1, color2=color2, group1=g1, group2=g2, species=species)\n",
    "            annotation_pval(df_annot, ax, [g1, g2], x_col='group', y_col='exp')\n",
    "            add_medians(df_annot, ax, x_col='group', y_col='exp', xsync=0.15, ysync=0)\n",
    "\n",
    "        plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot HC vs RR\n",
    "df_pval_HCvsRR = pd.read_csv(f'{RESULTS_DIR}/differential_abundance/mode{mode}_S{S}_HCvsRR.tsv', sep='\\t')\n",
    "df_pval_HCvsRR_up = df_pval_HCvsRR[(df_pval_HCvsRR['pval_MW'] < 0.15) & (df_pval_HCvsRR['log2FC'] > 0)].sort_values(by='log2FC').head()\n",
    "df_pval_HCvsRR_down = df_pval_HCvsRR[(df_pval_HCvsRR['pval_MW'] < 0.15) & (df_pval_HCvsRR['log2FC'] <= 0)].sort_values(by='log2FC').tail()\n",
    "\n",
    "df_pval_HCvsRR_updown = pd.concat([df_pval_HCvsRR_up, df_pval_HCvsRR_down])\n",
    "df_pval_HCvsRR_updown = df_pval_HCvsRR_updown.iloc[:len(df_pval_HCvsRR_updown) - len(df_pval_HCvsRR_updown) % 2]\n",
    "\n",
    "\n",
    "ncols = len(df_pval_HCvsRR_updown) \n",
    "\n",
    "plot_diff_sopecies(df_pval_HCvsRR_updown,  ['HC1', 'HC2', 'HC3', 'HC4'], ['RR1', 'RR2', 'RR3', 'RR4'], \"#648FFF\", \"#FE6100\", nrows=1)\n",
    "\n",
    "for format in ['png', 'tiff']: \n",
    "    plt.savefig(f'{RESULTS_DIR}/figures/paper/l2fc_HCvRR.{format}', dpi=DPI)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "df_pval_HCvsRR_up = df_pval_HCvsRR[df_pval_HCvsRR['pval_MW'] < 0.05].sort_values(by='log2FC').head()\n",
    "df_pval_HCvsRR_down = df_pval_HCvsRR[df_pval_HCvsRR['pval_MW'] < 0.05].sort_values(by='log2FC').tail()\n",
    "\n",
    "df_pval_HCvsRR_updown = pd.concat([df_pval_HCvsRR_up, df_pval_HCvsRR_down])\n",
    "\n",
    "ncols = len(df_pval_HCvsRR_updown) \n",
    "\n",
    "plot_diff_sopecies(df_pval_HCvsRR_updown,  ['HC1', 'HC2', 'HC3', 'HC4'], ['RR1', 'RR2', 'RR3', 'RR4'], \"#648FFF\", \"#FE6100\", nrows=1)\n",
    "\n",
    "for format in ['png', 'tiff']: \n",
    "    plt.savefig(f'{RESULTS_DIR}/figures/paper/pval_HCvRR.{format}', dpi=DPI)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot HC vs SP\n",
    "\n",
    "df_pval_HCvsSP = pd.read_csv(f'{RESULTS_DIR}/differential_abundance/mode{mode}_S{S}_HCvsSP.tsv', sep='\\t')\n",
    "\n",
    "df_pval_HCvsSP_up = df_pval_HCvsSP[(df_pval_HCvsSP['pval_MW'] < 0.15) & (df_pval_HCvsSP['log2FC'] > 0)].sort_values(by='log2FC').head()\n",
    "df_pval_HCvsSP_down = df_pval_HCvsSP[(df_pval_HCvsSP['pval_MW'] < 0.15) & (df_pval_HCvsSP['log2FC'] <= 0)].sort_values(by='log2FC').tail()\n",
    "\n",
    "df_pval_HCvsSP_updown = pd.concat([df_pval_HCvsSP_up, df_pval_HCvsSP_down])\n",
    "df_pval_HCvsSP_updown = df_pval_HCvsSP_updown.iloc[:len(df_pval_HCvsSP_updown) - len(df_pval_HCvsSP_updown) % 2]\n",
    "\n",
    "ncols = len(df_pval_HCvsSP_updown)\n",
    "\n",
    "plot_diff_sopecies(df_pval_HCvsSP_updown,  ['HC1', 'HC2', 'HC3', 'HC4'], ['SP1', 'SP2', 'SP3', 'SP4'], \"#648FFF\", \"#DC267F\", nrows=1)\n",
    "\n",
    "for format in ['png', 'tiff']: \n",
    "    plt.savefig(f'{RESULTS_DIR}/figures/paper/l2fc_HCvSP.{format}', dpi=DPI)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "df_pval_HCvsSP_up = df_pval_HCvsSP[df_pval_HCvsSP['pval_MW'] < 0.05].sort_values(by='log2FC').head()\n",
    "\n",
    "ncols = len(df_pval_HCvsSP_up) \n",
    "\n",
    "plot_diff_sopecies(df_pval_HCvsSP_up,  ['HC1', 'HC2', 'HC3', 'HC4'], ['SP1', 'SP2', 'SP3', 'SP4'], \"#648FFF\", \"#DC267F\", nrows=1)                                        \n",
    "\n",
    "for format in ['png', 'tiff']: \n",
    "    plt.savefig(f'{RESULTS_DIR}/figures/paper/pval_HCvsSP.{format}', dpi=DPI)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Plot RR vs SP\n",
    "# df_pval_RRvsSP = pd.read_csv(f'{RESULTS_DIR}/differential_abundance/mode{mode}_S{S}_RRvsSP.tsv', sep='\\t')\n",
    "# df_pval_RRvsSP_up = df_pval_RRvsSP[(df_pval_RRvsSP['pval_MW'] < 0.15) & (df_pval_RRvsSP['log2FC'] > 0)].sort_values(by='log2FC').head(14)\n",
    "# df_pval_RRvsSP_down = df_pval_RRvsSP[(df_pval_RRvsSP['pval_MW'] < 0.15) & (df_pval_RRvsSP['log2FC'] <= 0)].sort_values(by='log2FC').head(14)\n",
    "\n",
    "# df_pval_RRvsSP_updown = pd.concat([df_pval_RRvsSP_up, df_pval_RRvsSP_down])\n",
    "# df_pval_RRvsSP_updown = df_pval_RRvsSP_updown.iloc[:len(df_pval_RRvsSP_updown) - len(df_pval_RRvsSP_updown) % 2]\n",
    "\n",
    "# ncols = len(df_pval_RRvsSP_updown) // 3\n",
    "\n",
    "# plot_diff_sopecies(df_pval_RRvsSP_updown,  ['RR1', 'RR2', 'RR3', 'RR4'], ['SP1', 'SP2', 'SP3', 'SP4'],  \"#FE6100\", \"#DC267F\", nrows=3)\n",
    "\n",
    "# for format in ['png', 'tiff']: \n",
    "#     plt.savefig(f'{RESULTS_DIR}/figures/paper/l2fc_RRvsSP.{format}', dpi=DPI)\n",
    "\n",
    "\n",
    "df_pval_RRvsSP_up = df_pval_RRvsSP[df_pval_RRvsSP['pval_MW'] < 0.05].sort_values(by='log2FC').head(16)\n",
    "\n",
    "ncols = len(df_pval_RRvsSP_up)\n",
    "\n",
    "plot_diff_sopecies(df_pval_RRvsSP_up,  ['RR1', 'RR2', 'RR3', 'RR4'], ['SP1', 'SP2', 'SP3', 'SP4'],  \"#FE6100\", \"#DC267F\", nrows=1) \n",
    "\n",
    "for format in ['png', 'tiff']: \n",
    "    plt.savefig(f'{RESULTS_DIR}/figures/paper/pval_RRvsSP.{format}', dpi=DPI)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_pval_RRvsSP_updown"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot sex\n",
    "\n",
    "df_pval_sex = pd.read_csv(f'{RESULTS_DIR}/differential_abundance/mode{mode}_S{S}_sex.tsv', sep='\\t')\n",
    "df_pval_sex = df_pval_sex.rename(columns={'HC1':'Male1', 'HC2':'Male2', 'RR1':'Male3', 'RR2':'Male4', 'SP1':'Male5', 'SP2':'Male6', 'HC3': 'Female1', 'HC4': 'Female2', 'RR3': 'Female3', 'RR4': 'Female4', 'SP3': 'Female5', 'SP4': 'Female6'})\n",
    "\n",
    "df_pval_sex_up = df_pval_sex[(df_pval_sex['pval_MW'] < 0.15) & (df_pval_sex['log2FC'] > 0)].sort_values(by='log2FC').head(10)\n",
    "df_pval_sex_down = df_pval_sex[(df_pval_sex['pval_MW'] < 0.15) & (df_pval_sex['log2FC'] <= 0)].sort_values(by='log2FC').tail(10)\n",
    "\n",
    "df_pval_sex_updown = pd.concat([df_pval_sex_up, df_pval_sex_down])\n",
    "df_pval_sex_updown = df_pval_sex_updown.iloc[:len(df_pval_sex_updown) - len(df_pval_sex_updown) % 2]\n",
    "\n",
    "ncols = len(df_pval_sex_updown) \n",
    "\n",
    "plot_diff_sopecies(df_pval_sex_updown, [f'Female{i}' for i in range(1,7)], [f'Male{i}' for i in range(1,7)], \"#785EF0\", \"#FFB000\", nrows=1)\n",
    "\n",
    "for format in ['png', 'tiff']: \n",
    "    plt.savefig(f'{RESULTS_DIR}/figures/paper/l2fc_sex.{format}', dpi=DPI)\n",
    "\n",
    "\n",
    "df_pval_sex_up = df_pval_sex[df_pval_sex['pval_MW'] < 0.05].sort_values(by='log2FC').head()\n",
    "\n",
    "ncols = len(df_pval_sex_up) \n",
    "\n",
    "plot_diff_sopecies(df_pval_sex_up,  [f'Male{i}' for i in range(1,7)], [f'Female{i}' for i in range(1,7)], \"#785EF0\", \"#FFB000\", nrows=1) \n",
    "\n",
    "\n",
    "for format in ['png', 'tiff']: \n",
    "    plt.savefig(f'{RESULTS_DIR}/figures/paper/pval_sex.{format}', dpi=DPI)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "species = 'Sutterella'\n",
    "\n",
    "df_all = pd.concat([df_normpipe_retained_species, df_normpipe_discarded_normplus, df_normpipe_discarded_common]).set_index('name')\n",
    "\n",
    "yHC = np.log10(df_all.loc[species][['HC1', 'HC2', 'HC3', 'HC4']].astype(float).values + 1)\n",
    "ySP = np.log10(df_all.loc[species][['SP1', 'SP2', 'SP3', 'SP4']].astype(float).values + 1)\n",
    "yRR = np.log10(df_all.loc[species][['RR1', 'RR2', 'RR3', 'RR4']].astype(float).values + 1)\n",
    "yCTRL = np.log10(df_all.loc[species][['ACIDOLA', 'BLACTIS']].astype(float).values + 1)\n",
    "\n",
    "fig, ax = plt.subplots(1, 1)\n",
    "plot_ax_4(ax, yHC, ySP, yRR, yCTRL, species)\n",
    "\n",
    "\n",
    "display(df_normpipe_discarded_common[df_normpipe_discarded_common['name'] == species])\n",
    "display(df_normpipe_discarded_normplus[df_normpipe_discarded_normplus['name'] == species])\n",
    "display(df_normpipe_retained_species[df_normpipe_retained_species['name'] == species])\n",
    "\n",
    "\n",
    "display(df_pval_HCvsRR[df_pval_HCvsRR['name'] == species])\n",
    "display(df_pval_HCvsSP[df_pval_HCvsSP['name'] == species])\n",
    "display(df_pval_RRvsSP[df_pval_RRvsSP['name'] == species])\n",
    "display(df_pval_sex[df_pval_sex['name'] == species])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "EVs",
   "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.13.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
