04b - KEGGCompletenessProcessing.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "db172e3a-85f0-4df4-b831-1cb4e65ae41e",
   "metadata": {},
   "source": [
    "This part of the pipeline processes the raw output generated by the microbeAnnotator tool and generates a clustermap of the KEGG module completenesses."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "75fdf9bd-eae5-4976-ad01-cf6f8188f970",
   "metadata": {},
   "source": [
    "### Paths and parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb2f0be9-a77b-451a-9f4d-3e55dccb4275",
   "metadata": {},
   "source": [
    "#### Pipeline input folders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8dd70edd-e8e5-4a78-aaa8-73709f4df7cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "output_folder = \"./04-KEGGCompleteness/output\"\n",
    "metadata = \"./genomes_metadata\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a7b6353-25ea-4599-b760-3d0273ba9b4b",
   "metadata": {},
   "source": [
    "#### Pipeline output folders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f19ac486-9a3f-459a-9340-48caaeb2db46",
   "metadata": {},
   "outputs": [],
   "source": [
    "task_root = \"./04-KEGGCompleteness\"\n",
    "results_folder = \"./04-KEGGCompleteness/processed_output\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56b9f8ce-9dce-4f5b-80cd-ad09b5010ffb",
   "metadata": {},
   "source": [
    "#### Tool pointers and parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38b22c85-e252-4039-b3f0-e73ef0823d59",
   "metadata": {},
   "outputs": [],
   "source": [
    "inclusion_threshold = 50"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7cace78f-b691-44b1-a25c-f6b5f29bd69a",
   "metadata": {},
   "source": [
    "### Processing results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ef07831-4b42-4776-acf4-72f0fa9cd695",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.patches as ppt\n",
    "import seaborn as sns\n",
    "import os\n",
    "import pandas as pd\n",
    "from colorcet import glasbey\n",
    "from scipy.cluster import hierarchy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ae3093f-b6be-40c5-9dac-4ef5eedde17f",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.makedirs(results_folder, exist_ok=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d95eb21-1f62-4771-9946-55c213431d83",
   "metadata": {},
   "source": [
    "#### Parsing output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a0a498a-aca2-4f14-ba51-5e41b972f4e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "completeness = pd.read_table(output_folder + '/metabolic_summary__module_completeness.tab')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91dfeed2-9200-458c-aa9e-dc301d67c8ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Filtering out KEGG modules that have a maximum completeness less than 50%.\n",
    "completeness.columns = list(map(lambda x: '.'.join(x.split('.')[:2]), list(completeness.columns)))\n",
    "completeness['maximum_presence'] = completeness.max(axis = 1, numeric_only = True)\n",
    "completeness_toplot = completeness[completeness['maximum_presence'] >= inclusion_threshold].drop(columns = ['maximum_presence'])\n",
    "completeness_toplot"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26109a0b-a8ea-494c-b0e9-10618d8f5de6",
   "metadata": {},
   "source": [
    "#### Load cluster annotations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "388cda79-5103-45ba-bf71-0e981b032bc3",
   "metadata": {},
   "outputs": [],
   "source": [
    "metadata = pd.read_table(metadata, sep = '\\t', usecols = [0,1,2,4,5])\n",
    "metadata.columns = ['name','assemblyID', 'cluster', 'size', 'no. genes']\n",
    "cluster_annotations_0 = metadata[['assemblyID', 'cluster']]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "045d5ecf-f366-485b-8c0e-3541e35212fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "cluster_annotations = cluster_annotations_0.to_dict(orient = 'list')\n",
    "cluster_annotations = dict(zip(*cluster_annotations.values()))\n",
    "cluster_annotations = pd.Series(cluster_annotations, name = \"Cluster\")\n",
    "cluster_annotations"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b14d9f77-c245-408f-ac5f-f30bc6ba579d",
   "metadata": {},
   "source": [
    "#### Defining colour groups for rRNA cluster metadata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6127642a-af6f-439d-9cd9-3af9492c94aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "assembly_IDs = list(filter(lambda x: 'GC' in x, list(completeness.columns)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0385d06-ccf3-407a-870a-3c298c5f6674",
   "metadata": {},
   "outputs": [],
   "source": [
    "clusters = cluster_annotations.loc[assembly_IDs]\n",
    "\n",
    "# define colour mapping\n",
    "groupColour = clusters.map(dict(zip(['1','4','14a','14b'], 'rbgk')))\n",
    "groupColour.name = \"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f93b7ca-522e-4c55-beb9-35fee42116a2",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "groupColour"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff32c115-3500-42a7-a275-557aba45fe97",
   "metadata": {},
   "source": [
    "#### Defining colour groups for KEGG pathway modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "716f7364-bda5-491f-8438-75c3dd5fd85f",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Make a pathway to colour mapping via the pathway group\n",
    "# get all pathway groups\n",
    "pathway_groups = list(completeness_toplot['pathway group'].unique())\n",
    "\n",
    "# mapping pathway groups to colours\n",
    "palette = sns.color_palette(glasbey, n_colors = len(pathway_groups), as_cmap = True) \n",
    "pathway_colours = dict(zip(pathway_groups, palette))\n",
    "\n",
    "# mapping pathways to pathway groups\n",
    "name_pathway = dict(zip(*completeness_toplot[['name', 'pathway group']].to_dict(orient = 'list').values()))\n",
    "\n",
    "# connecting the two mappings for all pathways\n",
    "group_colours = {name: pathway_colours[name_pathway[name]] for name in completeness_toplot['name']}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2772756b-5232-4135-9cd2-663c5ece7725",
   "metadata": {},
   "source": [
    "#### Clustering the completenesses"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1ae6094-16da-465c-adf5-01f29310bd99",
   "metadata": {},
   "outputs": [],
   "source": [
    "completeness_numeric = completeness_toplot.select_dtypes(include = \"number\")\n",
    "linkage = hierarchy.linkage(completeness_numeric.T, method = 'average', metric = 'euclidean', \n",
    "                            optimal_ordering = True) # with optimal ordering for visualising any lower-level pattern"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca536be7-ca1b-4a23-9bea-4535a18ad08c",
   "metadata": {},
   "source": [
    "#### Clustermap"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "510d42dc-0096-4e2d-a186-0c6baffe80ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "h = sns.clustermap(completeness_numeric, col_linkage = linkage, col_cluster = True, row_cluster = False,\n",
    "                   xticklabels = False, yticklabels = True,\n",
    "                   figsize = (9,20), dendrogram_ratio = (0.2, 0.05),\n",
    "                   col_colors = groupColour, cmap = \"magma_r\")\n",
    "h.ax_heatmap.set_yticklabels(completeness_toplot['name'])\n",
    "h.ax_heatmap.set_xlabel('Genome assemblies', size = 18)\n",
    "h.ax_heatmap.set_ylabel('KEGG module', size = 18)\n",
    "h.ax_heatmap.collections[0].colorbar.set_label('Module completeness (%)', size = 16)\n",
    "\n",
    "# The legend has to be composed on the fly. Otherwise, we can't use our pathway group-level label colouring\n",
    "legend_patches = []\n",
    "pathways_covered = []\n",
    "for label in h.ax_heatmap.get_yticklabels():\n",
    "    text = label.get_text()\n",
    "    pathway = name_pathway[text]\n",
    "    color = group_colours[text]\n",
    "    # change the label of the pathway module at the tick label to the appropriate colour\n",
    "    label.set_color(color)\n",
    "    # only add a new entry to the legend if we haven't encountered a pathway from this group yet\n",
    "    if pathway not in pathways_covered:\n",
    "        legend_patches.append(ppt.Patch(color = color, label = pathway))\n",
    "        pathways_covered.append(pathway)\n",
    "        \n",
    "plt.legend(handles = legend_patches, ncol = 1, bbox_to_anchor = (42,1,1,0))\n",
    "plt.savefig(results_folder + '/' + 'module_completeness.svg')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cca4b44d-4365-4e8d-a202-f959bea80782",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The metadata of the assemblies ordered by the clustering\n",
    "clustered_ids = completeness_numeric.columns[h.dendrogram_col.reordered_ind]\n",
    "clustered_ids.name = 'assemblyID'\n",
    "clustered_ids = pd.DataFrame(clustered_ids)\n",
    "clustered = clustered_ids.merge(metadata, on='assemblyID')\n",
    "clustered"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fff0665f-8f72-46f5-9fcf-b51f3aa12654",
   "metadata": {},
   "outputs": [],
   "source": [
    "clustered.to_csv(results_folder + '/' + 'clustered_strains.tsv', sep = \"\\t\", index = False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aaae799f-98d5-4a9f-a92c-a1b071b8ffb5",
   "metadata": {},
   "source": [
    "#### Clustermap without pathway labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5fe1d02d-e4dc-4a5e-a750-246332d76c82",
   "metadata": {},
   "outputs": [],
   "source": [
    "h = sns.clustermap(completeness_numeric, col_linkage = linkage, col_cluster = True, row_cluster = False, \n",
    "                   xticklabels = False, yticklabels = False, cbar_pos = (0.01,0.5,0.05,0.4),\n",
    "                   figsize = (6,6), dendrogram_ratio = (0.2, 0.05),\n",
    "                   col_colors = groupColour, cmap = 'magma_r')\n",
    "h.ax_heatmap.set_xlabel('Genome assemblies', size = 12)\n",
    "h.ax_heatmap.collections[0].colorbar.set_label('Module completeness (%)', size = 11)\n",
    "\n",
    "plt.savefig(results_folder + '/' + 'module_completeness_noTicks.svg')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39a78c8d-80cf-4944-8864-87ffbe17f31e",
   "metadata": {},
   "source": [
    "#### Extract metabolically versatile and non-versatile group members"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c595b16-03e0-4775-9567-6108005d1b66",
   "metadata": {},
   "outputs": [],
   "source": [
    "reordered_inds = h.dendrogram_col.reordered_ind\n",
    "versatile_ids = list(hierarchy.fcluster(linkage, 2, criterion=\"maxclust\")[reordered_inds])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17f2fe5d-891d-4625-b127-6dd08cf01120",
   "metadata": {},
   "outputs": [],
   "source": [
    "edge = versatile_ids.index(2) # ids are grouped so the first one of the other group indicates the boundary between the two groups\n",
    "versatile = completeness_numeric.columns[reordered_inds][edge:].to_series(name = \"versatile\")\n",
    "nonversatile = completeness_numeric.columns[reordered_inds][:edge].to_series(name = \"non-versatile\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28a332b2-6102-48ec-a8f8-3c2b856e5453",
   "metadata": {},
   "outputs": [],
   "source": [
    "versatile.to_csv(task_root + '/versatile', header = False, index = False)\n",
    "nonversatile.to_csv(task_root + '/non-versatile', header = False, index = False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1108465c-f30b-4c32-b604-c4481516fa3d",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}