10b5 - MobilomeAnalysis_ProphagesProcessing.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3dba3548-5d86-4e72-82da-4c0fb6c70927",
   "metadata": {},
   "source": [
    "This part of the pipeline processes the raw PHASTEST output and statistically compares the normalised prophage counts by rRNA cluster."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82bcb274-5cc8-48bc-8674-29c93d6c1be4",
   "metadata": {},
   "source": [
    "### Paths and parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1657f39b-d4bf-42e1-b3e4-3c2ae24032b3",
   "metadata": {},
   "source": [
    "#### Pipeline input folders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57b92397-a77e-4e1e-8ad9-e90f2fc93b21",
   "metadata": {},
   "outputs": [],
   "source": [
    "metadata = \"./genomes_metadata\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3ad9bd4-a3df-49c9-8193-e65f97def25c",
   "metadata": {},
   "source": [
    "#### Pipeline output folders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3758f1b-563c-433a-8d30-e26a14b66c01",
   "metadata": {},
   "outputs": [],
   "source": [
    "task_root = \"./10-MGEs/prophages\"\n",
    "output_folder = task_root+\"/output\"\n",
    "results_folder = task_root+\"/processed_output\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25b879dc-7552-4420-98f8-76851980393d",
   "metadata": {},
   "source": [
    "#### Tool pointers and parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54cd1156-e437-497e-ae1a-976444ee1b1c",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "fa865e3a-313c-4860-828a-e37fe8e20efd",
   "metadata": {},
   "source": [
    "#### Libraries and other setups"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ca60d05-8ad0-44bb-8816-d18851305e15",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import os\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "from supervenn import supervenn\n",
    "import seaborn as sns\n",
    "import itertools as it\n",
    "import scipy.stats as sts\n",
    "import numpy as np\n",
    "from statannotations.Annotator import Annotator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b77af638-7624-4483-9017-76a034ea44f5",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# custom_palette = sns.husl_palette()\n",
    "# custom_palette = [custom_palette[0], custom_palette[2], custom_palette[4], custom_palette[5]]\n",
    "custom_palette = {'1': '#ee6677',\n",
    "                  '4': '#4477aa',\n",
    "                  '14a': '#228833',\n",
    "                  '14b': '#ccbb44'}\n",
    "custom_palette"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f909f76-6bbf-49ac-be4f-47d439817672",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.makedirs(results_folder, exist_ok=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b66b3058-5230-4ba6-abe0-fbded21eb018",
   "metadata": {},
   "source": [
    "## Phage identifications"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9f6c600-d9d5-4abd-b508-bec0c6d38d21",
   "metadata": {},
   "source": [
    "### Parsing PHASTEST phage regions output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc88269f-42f1-493f-9c8a-613f35b5b3af",
   "metadata": {},
   "outputs": [],
   "source": [
    "result_dirs = filter(lambda x: '.log' not in x, os.listdir(output_folder))\n",
    "hits = []\n",
    "# Prophage hit regions can be read directly from the json_input_regions file\n",
    "for dir in result_dirs:\n",
    "    with open(output_folder + '/' + dir + '/' + \"json_input_regions\", \"r\") as handle:\n",
    "        regions = json.load(handle)\n",
    "    for r in regions:\n",
    "        phage_info = r['most_common_phage']\n",
    "        phage_genome_ID = '_'.join(phage_info.split('_')[-2:])\n",
    "        phage_name = '_'.join(phage_info.split('_')[1:-2])\n",
    "        record = {'assembly_ID': dir, 'length': r['stop'] - r['start'], 'completeness': r['completeness'], \n",
    "                  'name': phage_name, 'phage_genome_ID': phage_genome_ID}\n",
    "        hits.append(record)\n",
    "hits = pd.DataFrame(hits)\n",
    "hits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cec9d977-28fe-4c61-a0e8-b5fc2e403105",
   "metadata": {},
   "outputs": [],
   "source": [
    "hits.to_csv(results_folder + \"/all_hits\", sep = \"\\t\", index = False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93ed9b50-7ceb-46ff-ad36-f9ad5ba2ed48",
   "metadata": {},
   "source": [
    "#### Intact prophages only"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3b2847a-7fde-4b8d-a60d-32712f0014e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "hits_intact = hits[hits['completeness'] == 'intact'].reset_index(drop=True)\n",
    "hits_intact"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2180b122-cf16-48c2-9b03-134c0f13405d",
   "metadata": {},
   "outputs": [],
   "source": [
    "hits_intact.to_csv(results_folder + \"/intact_hits\", sep = \"\\t\", index = False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6328bcc-8ed0-4072-a9b6-1d89c9429b0b",
   "metadata": {},
   "source": [
    "#### Decayed prophages only"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ab2244e-96e1-46c1-b2f6-5f6722988466",
   "metadata": {},
   "outputs": [],
   "source": [
    "hits_decayed = hits[hits['completeness'] != 'intact'].reset_index(drop=True)\n",
    "hits_decayed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ef9170e-fc80-45a2-812f-178dbc2f3685",
   "metadata": {},
   "outputs": [],
   "source": [
    "hits_decayed.to_csv(results_folder + \"/decayed_hits\", sep = \"\\t\", index = False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33c7304c-538e-43e3-a30a-1e5eb09ce67e",
   "metadata": {},
   "source": [
    "## Loading the cluster annotations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33c6a71f-d16a-4500-b6ba-048c7778872d",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "cluster_annotations_0 = pd.read_table(metadata, sep = '\\t', usecols = [1,2,11])\n",
    "cluster_annotations_0.columns = ['assemblyID', 'cluster', 'Failed_QC']\n",
    "cluster_annotations_0.drop(cluster_annotations_0[cluster_annotations_0['Failed_QC'] != False].index, inplace = True)\n",
    "cluster_annotations = cluster_annotations_0[['assemblyID', 'cluster']].to_dict(orient = 'list')\n",
    "cluster_annotations = dict(zip(*cluster_annotations.values()))\n",
    "cluster_annotations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba8a6958-3c45-40ec-b989-e26b258e0f78",
   "metadata": {},
   "outputs": [],
   "source": [
    "annots = pd.DataFrame.from_records(list(zip(cluster_annotations.keys(), cluster_annotations.values())), columns = ('assembly_ID', 'cluster'))\n",
    "annots"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c93851ba-0dac-437e-a3fa-4da3245fe406",
   "metadata": {},
   "source": [
    "## Setting up the count pivot tables"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33e9748e-5142-4905-8dff-e9b068c2c90c",
   "metadata": {},
   "source": [
    "### All prophages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40def38d-6cfe-4eab-b2dc-ac14800d19d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "counts = pd.DataFrame(hits.groupby(['assembly_ID', 'name'])['length'].count()\n",
    "                     ).reset_index().rename(columns = {'length': 'counts'})\n",
    "counts = counts.pivot(columns = 'name', index = 'assembly_ID', values = 'counts')\n",
    "counts = counts.merge(annots, on = \"assembly_ID\", how = \"right\"\n",
    "                     ).set_index('assembly_ID').fillna(0).drop('cluster', axis = 1).astype(int)\n",
    "counts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed40ff44-cd6d-4df0-a16c-f7b257773bbd",
   "metadata": {},
   "outputs": [],
   "source": [
    "counts.to_csv(results_folder + \"/all_counts\", sep = \"\\t\", index = False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d72f9d14-34c5-4b71-84ce-985917494c1e",
   "metadata": {},
   "source": [
    "### Intact prophages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "801abed8-c11e-4755-9a4d-d0a1d22005bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "counts_intact = pd.DataFrame(hits_intact.groupby(['assembly_ID', 'name'])['length'].count()\n",
    "                            ).reset_index().rename(columns = {'length': 'counts'})\n",
    "counts_intact = counts_intact.pivot(columns = 'name', index = 'assembly_ID', values = 'counts')\n",
    "counts_intact = counts_intact.merge(annots, on = \"assembly_ID\", how = \"right\"\n",
    "                                   ).set_index('assembly_ID').fillna(0).drop('cluster', axis = 1).astype(int)\n",
    "counts_intact"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b1c18ad-eed2-4fee-b6d3-e3250b55bc16",
   "metadata": {},
   "outputs": [],
   "source": [
    "counts_intact.to_csv(results_folder + \"/intact_counts\", sep = \"\\t\", index = False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "292df749-4f55-41f5-8d45-f2ad34512425",
   "metadata": {},
   "source": [
    "### Decayed prophages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "644f8763-06e6-4525-bcb3-e5f50d35bfa9",
   "metadata": {},
   "outputs": [],
   "source": [
    "counts_decayed = pd.DataFrame(hits_decayed.groupby(['assembly_ID', 'name'])['length'].count()\n",
    "                             ).reset_index().rename(columns = {'length': 'counts'})\n",
    "counts_decayed = counts_decayed.pivot(columns = 'name', index = 'assembly_ID', values = 'counts')\n",
    "counts_decayed = counts_decayed.merge(annots, on = \"assembly_ID\", how = \"right\"\n",
    "                                     ).set_index('assembly_ID').fillna(0).drop('cluster', axis = 1).astype(int)\n",
    "counts_decayed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34a0ddac-9ed4-4130-89a0-cc688c632496",
   "metadata": {},
   "outputs": [],
   "source": [
    "counts_decayed.to_csv(results_folder + \"/decayed_counts\", sep = \"\\t\", index = False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62e6e2a0-d3d7-4dc2-89e9-f8de80dabb2c",
   "metadata": {},
   "source": [
    "## Count plotting"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef742462-d486-4c6e-a2e9-0e41e5b092d0",
   "metadata": {},
   "source": [
    "### all"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3aec689-755d-4f0c-afc2-a1d52f87a63d",
   "metadata": {},
   "source": [
    "Adding rRNA cluster annotations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c798cbf-60a3-4164-b897-78ac171dbc98",
   "metadata": {},
   "outputs": [],
   "source": [
    "counts_av = pd.DataFrame(counts.sum(axis = 1), columns = ['counts'])\n",
    "counts_av['cluster'] = counts_av.index.to_series().apply(lambda x: cluster_annotations[x])\n",
    "counts_av"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5035d175-d47d-4550-97d3-199990f0aef8",
   "metadata": {},
   "source": [
    "#### Barplot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1ea6bb5-3967-4e6b-910e-ef0d2e431980",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize = (5,2))\n",
    "ax = sns.barplot(ax = ax, data = counts_av, estimator = \"mean\", errorbar = \"se\",\n",
    "                 x = \"counts\", y = \"cluster\", palette = custom_palette, \n",
    "                 width = 0.9, orient = \"h\")\n",
    "plt.xlabel('Avg. no. prophages')\n",
    "plt.ylabel('Taxonomic cluster')\n",
    "plt.title('Prophages')\n",
    "plt.savefig(results_folder + \"/\" + \"av_counts_bar.svg\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3f7533e-efdb-44ff-b63b-6c4ce5e95555",
   "metadata": {},
   "source": [
    "#### Violinplot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "602f7f3d-574e-456a-8c88-8b9e81c3c452",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize = (5,3))\n",
    "ax = sns.violinplot(ax = ax, data = counts_av, x = 'counts', y = 'cluster', orient = 'h', palette = custom_palette, cut = 0)\n",
    "plt.xlabel('No. prophages')\n",
    "plt.ylabel('rRNA cluster')\n",
    "plt.title('Prophages')\n",
    "\n",
    "# Adding statistical significance markers\n",
    "pairs = list(it.combinations(counts_av['cluster'].unique(), 2))\n",
    "annotator = Annotator(ax = ax, pairs = pairs, data = counts_av, x = 'counts', y = 'cluster', orient = 'h', cut = 0)\n",
    "annotator.configure(test = 'Mann-Whitney', text_format = 'star', loc = 'inside')\n",
    "annotator.apply_and_annotate()\n",
    "\n",
    "plt.savefig(results_folder + \"/\" + 'av_counts_violin.svg')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8405602-ec45-4168-8a95-b1afd58b5a97",
   "metadata": {},
   "source": [
    "#### Exact stats"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "896a386a-d759-432b-9df0-c121edf9e501",
   "metadata": {},
   "source": [
    "Grouping the counts by rRNA cluster"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21a8e4d7-dda0-4aea-ae65-7a8eaa47d994",
   "metadata": {},
   "outputs": [],
   "source": [
    "counts_av_stats = counts_av[['counts', 'cluster']].to_dict(orient = 'list')\n",
    "counts_av_stats = list(zip(*counts_av_stats.values()))\n",
    "counts_av_stats_dict = {}\n",
    "for record in counts_av_stats:\n",
    "    try:\n",
    "        counts_av_stats_dict[record[1]].append(record[0])\n",
    "    except KeyError:\n",
    "        counts_av_stats_dict[record[1]] = [record[0]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e5c3d4d-3047-4bdd-83e8-2fa0d2c6e328",
   "metadata": {},
   "outputs": [],
   "source": [
    "[(i, [np.mean(j), np.std(j)]) for i,j in counts_av_stats_dict.items()]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a35ba65b-d3f6-4e77-b7b3-6bed5c079abb",
   "metadata": {},
   "outputs": [],
   "source": [
    "tests = it.combinations(counts_av_stats_dict.keys(), 2)\n",
    "for comb in tests:\n",
    "    print(str(comb) + ': ' + \n",
    "          str(sts.mannwhitneyu(counts_av_stats_dict[comb[0]], \n",
    "                               counts_av_stats_dict[comb[1]])\n",
    "              [1])\n",
    "         )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61e1bd61-f976-45bd-a17f-57319ef48461",
   "metadata": {},
   "source": [
    "### intact"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea98353f-db62-4678-9a65-8b79a42007b1",
   "metadata": {},
   "source": [
    "Adding rRNA cluster annotations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8276f803-13d7-4f91-9f23-356afe26d609",
   "metadata": {},
   "outputs": [],
   "source": [
    "counts_intact_av = pd.DataFrame(counts_intact.sum(axis = 1), columns = ['counts'])\n",
    "counts_intact_av['cluster'] = counts_intact_av.index.to_series().apply(lambda x: cluster_annotations[x])\n",
    "counts_intact_av"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e36a8c08-a761-4837-bdeb-dac76c616d51",
   "metadata": {},
   "source": [
    "#### Barplot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e5a2583-0e96-47f6-86da-6a0c1a02d736",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize = (5,2))\n",
    "ax = sns.barplot(ax = ax, data = counts_intact_av, estimator = \"mean\", errorbar = \"se\",\n",
    "                 x = \"counts\", y = \"cluster\", palette = custom_palette,\n",
    "                 width = 0.9, orient = \"h\")\n",
    "plt.xlabel('Avg. no. intact prophages')\n",
    "plt.ylabel('Taxonomic cluster')\n",
    "plt.title('Complete prophages')\n",
    "plt.savefig(results_folder + \"/\" + \"av_counts_intact_bar.svg\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a4b79e9-4926-4dd2-a60a-199fd51a5d6d",
   "metadata": {},
   "source": [
    "#### Violinplot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5f236f3-f61b-4108-87dd-c9f4bb0c4373",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize = (5,3))\n",
    "ax = sns.violinplot(ax = ax, data = counts_intact_av, x = 'counts', y = 'cluster', orient = 'h', palette = custom_palette, cut = 0)\n",
    "plt.xlabel('No. intact prophages')\n",
    "plt.ylabel('rRNA cluster')\n",
    "plt.title('Complete prophages')\n",
    "\n",
    "# Adding statistical significance markers\n",
    "pairs = list(it.combinations(counts_intact_av['cluster'].unique(), 2))\n",
    "annotator = Annotator(ax = ax, pairs = pairs, data = counts_intact_av, x = 'counts', y = 'cluster', orient = 'h', cut = 0)\n",
    "annotator.configure(test = 'Mann-Whitney', text_format = 'star', loc = 'inside')\n",
    "annotator.apply_and_annotate()\n",
    "\n",
    "plt.savefig(results_folder + \"/\" + 'av_counts_intact_violin.svg')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "99b2e924-6bf3-45e8-a68e-6ad375035d9b",
   "metadata": {},
   "source": [
    "#### Exact stats"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f69a894c-b0bd-4a2f-9155-f2d7ea31eb37",
   "metadata": {},
   "source": [
    "Grouping the counts by rRNA cluster"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4261b3e2-0e05-496b-8f57-b4105eaa5e94",
   "metadata": {},
   "outputs": [],
   "source": [
    "counts_intact_av_stats = counts_intact_av[['counts', 'cluster']].to_dict(orient = 'list')\n",
    "counts_intact_av_stats = list(zip(*counts_intact_av_stats.values()))\n",
    "counts_intact_av_stats_dict = {}\n",
    "for record in counts_intact_av_stats:\n",
    "    try:\n",
    "        counts_intact_av_stats_dict[record[1]].append(record[0])\n",
    "    except KeyError:\n",
    "        counts_intact_av_stats_dict[record[1]] = [record[0]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2523f9ec-38f1-43ac-ae02-68867ad08e4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "[(i, [np.mean(j), np.std(j)]) for i,j in counts_intact_av_stats_dict.items()]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ffa5be74-344b-45fc-9300-bf258e45a1ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "tests = it.combinations(counts_intact_av_stats_dict.keys(), 2)\n",
    "for comb in tests:\n",
    "    print(str(comb) + ': ' + \n",
    "          str(sts.mannwhitneyu(counts_intact_av_stats_dict[comb[0]], \n",
    "                               counts_intact_av_stats_dict[comb[1]])\n",
    "              [1])\n",
    "         )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "340356ef-2fcc-40fb-9cf9-56b86c691bec",
   "metadata": {},
   "source": [
    "### decayed"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2aec6bcb-c24e-4e1b-9eb9-209d13363560",
   "metadata": {},
   "source": [
    "Adding rRNA cluster annotations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64e6a662-f016-4fc1-8551-0c0645e90b14",
   "metadata": {},
   "outputs": [],
   "source": [
    "counts_decayed_av = pd.DataFrame(counts_decayed.sum(axis = 1), columns = ['counts'])\n",
    "counts_decayed_av['cluster'] = counts_decayed_av.index.to_series().apply(lambda x: cluster_annotations[x])\n",
    "counts_decayed_av"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb4d443f-f6b7-48a1-83a7-272ba74a2f31",
   "metadata": {},
   "source": [
    "#### Barplot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3fd14fcf-a297-4bfd-bbe2-ad43f2150f84",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize = (5,2))\n",
    "ax = sns.barplot(ax = ax, data = counts_decayed_av, estimator = \"mean\", errorbar = \"se\",\n",
    "                 x = \"counts\", y = \"cluster\", palette = custom_palette,\n",
    "                 width = 0.9, orient = \"h\")\n",
    "plt.xlabel('Avg. no. decayed prophages')\n",
    "plt.ylabel('Taxonomic cluster')\n",
    "plt.title('Decayed prophages')\n",
    "plt.savefig(results_folder + \"/\" + \"av_counts_decayed_bar.svg\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "445edbc4-5340-4854-9e15-4b0a6ac0e7e1",
   "metadata": {},
   "source": [
    "#### Violinplot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13f61218-db25-4d87-9b83-efcc3d95221e",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize = (5,3))\n",
    "ax = sns.violinplot(ax = ax, data = counts_decayed_av, x = 'counts', y = 'cluster', orient = 'h', palette = custom_palette, cut = 0)\n",
    "plt.xlabel('No. decayed prophages')\n",
    "plt.ylabel('rRNA cluster')\n",
    "plt.title('Decayed prophages')\n",
    "\n",
    "pairs = list(it.combinations(counts_decayed_av['cluster'].unique(), 2))\n",
    "annotator = Annotator(ax = ax, pairs = pairs, data = counts_decayed_av, x = 'counts', y = 'cluster', orient = 'h', cut = 0)\n",
    "annotator.configure(test = 'Mann-Whitney', text_format = 'star', loc = 'inside')\n",
    "annotator.apply_and_annotate()\n",
    "\n",
    "plt.savefig(results_folder + \"/\" + 'av_counts_decayed_violin.svg')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3703c819-33f1-41db-bdb2-5f24d59d44ff",
   "metadata": {},
   "source": [
    "#### Exact stats"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98b8a479-3ea4-4828-a85b-0273bb454800",
   "metadata": {},
   "source": [
    "Grouping the counts by rRNA cluster"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cec20840-1cb0-487f-a39d-6a993d05c544",
   "metadata": {},
   "outputs": [],
   "source": [
    "counts_decayed_av_stats = counts_decayed_av[['counts', 'cluster']].to_dict(orient = 'list')\n",
    "counts_decayed_av_stats = list(zip(*counts_decayed_av_stats.values()))\n",
    "counts_decayed_av_stats_dict = {}\n",
    "for record in counts_decayed_av_stats:\n",
    "    try:\n",
    "        counts_decayed_av_stats_dict[record[1]].append(record[0])\n",
    "    except KeyError:\n",
    "        counts_decayed_av_stats_dict[record[1]] = [record[0]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c2b5891-4229-4512-90cc-6e3d09f89999",
   "metadata": {},
   "outputs": [],
   "source": [
    "[(i, [np.mean(j), np.std(j)]) for i,j in counts_decayed_av_stats_dict.items()]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2ca0368-b39b-4c7c-9f97-e63895b60ca8",
   "metadata": {},
   "outputs": [],
   "source": [
    "tests = it.combinations(counts_decayed_av_stats_dict.keys(), 2)\n",
    "for comb in tests:\n",
    "    print(str(comb) + ': ' + \n",
    "          str(sts.mannwhitneyu(counts_decayed_av_stats_dict[comb[0]], \n",
    "                               counts_decayed_av_stats_dict[comb[1]])\n",
    "              [1])\n",
    "         )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a83c98f-b584-4f6a-9076-e1b777a0287b",
   "metadata": {},
   "source": [
    "## Type and infectivity plotting"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce6ab3a5-0c3b-417f-895a-58963cb69ddc",
   "metadata": {},
   "source": [
    "### all"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06fa6031-5e98-469e-896a-fd715971b2fe",
   "metadata": {},
   "source": [
    "Adding rRNA cluster annotations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f428582f-adc4-4b90-86f0-86c16ee579eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "hits_molten = counts.reset_index().melt(id_vars = 'assembly_ID', var_name = \"phage name\", value_name = 'counts')\n",
    "hits_molten['cluster'] = hits_molten['assembly_ID'].apply(lambda x: cluster_annotations[x])\n",
    "hits_molten"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62d62736-8d52-4781-bcf1-723624849a39",
   "metadata": {},
   "source": [
    "#### Prophage types"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb8bd854-a97b-4ebb-8823-df2317b34b16",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize = (6,25))\n",
    "ax = sns.barplot(ax = ax, data = hits_molten, estimator = \"mean\", errorbar = \"se\",\n",
    "                 x = \"counts\", y = \"phage name\", hue = \"cluster\", order = sorted(hits_molten['phage name'].unique()),\n",
    "                 width = 0.9, orient = \"h\", palette = custom_palette)\n",
    "plt.xlabel('Avg. no. prophages')\n",
    "plt.ylabel('Prophage')\n",
    "plt.savefig(results_folder + \"/\" + \"av_counts_prophage.svg\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "47e6cef8-8c0a-44f1-9f75-b213d0f2b32f",
   "metadata": {},
   "source": [
    "#### Prophage infectivity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf7df977-5882-44bb-984c-c0af14669ea0",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "infect_freq = pd.DataFrame(hits_molten.groupby(['phage name', 'cluster'])['counts'].mean()).reset_index().rename(columns = {'counts': 'freq'})\n",
    "infect_freq['infects'] = infect_freq['freq'] > 0\n",
    "infect_freq"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1aea42b4-b2d7-4d4f-89c8-af4381264deb",
   "metadata": {},
   "source": [
    "Get all infecting phages by rRNA cluster of the host"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9de139e5-b1ed-4030-ab4a-c471d077e256",
   "metadata": {},
   "outputs": [],
   "source": [
    "infectivities = []\n",
    "for c in infect_freq['cluster'].unique():\n",
    "    to_add = infect_freq[(infect_freq['cluster'] == c) & (infect_freq['infects'] == True)]['phage name'].squeeze()\n",
    "    to_add.name = 'Cluster ' + c\n",
    "    infectivities.append(to_add)\n",
    "\n",
    "infectivities_sets = [set(i) for i in infectivities]\n",
    "labels = [i.name for i in infectivities]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dce1265c-f2fa-4856-a425-4027dccfdacd",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize = (7.5,4.5))\n",
    "supervenn(infectivities_sets, labels, side_plots = 'right',\n",
    "          chunks_ordering = 'minimize gaps', sets_ordering = 'minimize gaps', \n",
    "          widths_minmax_ratio = 0.15)\n",
    "plt.xlabel('# infecting phages')\n",
    "plt.ylabel('rRNA cluster')\n",
    "plt.savefig(results_folder + \"/\" + \"venn_infectivity.svg\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc8f78ba-f0ec-41ae-a677-1dc17c8fdd8d",
   "metadata": {},
   "source": [
    "### intact"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5dbd12b-c5ff-4a2b-9c9f-2931f4e7e027",
   "metadata": {},
   "source": [
    "Adding rRNA cluster annotations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d621bc0-fffb-49ff-a8e2-060d84d8c181",
   "metadata": {},
   "outputs": [],
   "source": [
    "hits_intact_molten = counts_intact.reset_index().melt(id_vars = 'assembly_ID', var_name = \"phage name\", value_name = 'counts')\n",
    "hits_intact_molten['cluster'] = hits_intact_molten['assembly_ID'].apply(lambda x: cluster_annotations[x])\n",
    "hits_intact_molten"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d95513ff-6c84-4384-be63-ad46863f6813",
   "metadata": {},
   "source": [
    "#### Prophage types"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aeb00db4-0d5c-4abc-9780-955c4df947dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize = (6,15))\n",
    "ax = sns.barplot(ax = ax, data = hits_intact_molten, estimator = \"mean\", errorbar = \"se\",\n",
    "                 x = \"counts\", y = \"phage name\", hue = \"cluster\", order = sorted(hits_intact_molten['phage name'].unique()),\n",
    "                 width = 0.9, orient = \"h\", palette = custom_palette)\n",
    "plt.xlabel('Avg. no. intact prophages')\n",
    "plt.ylabel('Prophage')\n",
    "plt.savefig(results_folder + \"/\" + \"av_counts_prophage_intact.svg\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4caa865-9033-439a-b1c8-b3014599ac26",
   "metadata": {},
   "source": [
    "#### Prophage infectivity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "faaf1963-d8e5-41ae-a969-9a6b70a30f06",
   "metadata": {},
   "outputs": [],
   "source": [
    "infect_intact_freq = pd.DataFrame(hits_intact_molten.groupby(['phage name', 'cluster'])['counts'].mean()\n",
    "                                 ).reset_index().rename(columns = {'counts': 'freq'})\n",
    "infect_intact_freq['infects'] = infect_intact_freq['freq'] > 0\n",
    "infect_intact_freq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45361d22-9e23-4394-8a28-1b7af4d464f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "infectivities_intact = []\n",
    "for c in infect_intact_freq['cluster'].unique():\n",
    "    to_add = infect_intact_freq[(infect_intact_freq['cluster'] == c) & (infect_intact_freq['infects'] == True)]['phage name'].squeeze()\n",
    "    to_add.name = 'Cluster ' + c\n",
    "    infectivities_intact.append(to_add)\n",
    "\n",
    "infectivities_intact_sets = [set(i) for i in infectivities_intact]\n",
    "labels_intact = [i.name for i in infectivities_intact]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42ebb8ef-ccc3-4acc-967e-328595f0c351",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize = (7.5,4.5))\n",
    "supervenn(infectivities_intact_sets, labels_intact, side_plots = 'right',\n",
    "          chunks_ordering = 'minimize gaps', sets_ordering = 'minimize gaps', \n",
    "          widths_minmax_ratio = 0.15)\n",
    "plt.xlabel('# infecting phages')\n",
    "plt.ylabel('rRNA cluster')\n",
    "plt.savefig(results_folder + \"/\" + \"venn_infectivity_intact.svg\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26fd5709-4de2-4c3f-bffd-de2bfdccbf7c",
   "metadata": {},
   "source": [
    "### decayed"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5d19b40-141d-4d79-8387-a2f072d1068a",
   "metadata": {},
   "source": [
    "Adding rRNA cluster annotations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94ccdd57-e259-4107-9f34-c9c500c5c594",
   "metadata": {},
   "outputs": [],
   "source": [
    "hits_decayed_molten = counts_decayed.reset_index().melt(id_vars = 'assembly_ID', var_name = \"phage name\", value_name = 'counts')\n",
    "hits_decayed_molten['cluster'] = hits_decayed_molten['assembly_ID'].apply(lambda x: cluster_annotations[x])\n",
    "hits_decayed_molten"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ecd79385-fc86-4240-9b7d-75b3e1dd3b4d",
   "metadata": {},
   "source": [
    "#### Prophage types"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b04f1c05-b50c-40ef-8ea0-193de1aeda65",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize = (6,15))\n",
    "ax = sns.barplot(ax = ax, data = hits_decayed_molten, estimator = \"mean\", errorbar = \"se\",\n",
    "                 x = \"counts\", y = \"phage name\", hue = \"cluster\", order = sorted(hits_decayed_molten['phage name'].unique()),\n",
    "                 width = 0.9, orient = \"h\", palette = custom_palette)\n",
    "plt.xlabel('Avg. no. decayed prophages')\n",
    "plt.ylabel('Prophage')\n",
    "plt.savefig(results_folder + \"/\" + \"av_counts_prophage_decayed.svg\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14ec371f-470f-4ac6-89de-c39d846efba7",
   "metadata": {},
   "source": [
    "#### Prophage infectivity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c0da6af-f7e8-41f6-a0ea-9bbbea521e7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "infect_decayed_freq = pd.DataFrame(hits_decayed_molten.groupby(['phage name', 'cluster'])['counts'].mean()\n",
    "                                  ).reset_index().rename(columns = {'counts': 'freq'})\n",
    "infect_decayed_freq['infects'] = infect_decayed_freq['freq'] > 0\n",
    "infect_decayed_freq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cbb762ce-74dc-488f-801f-cadeca6f3f67",
   "metadata": {},
   "outputs": [],
   "source": [
    "infectivities_decayed = []\n",
    "for c in infect_decayed_freq['cluster'].unique():\n",
    "    to_add = infect_decayed_freq[(infect_decayed_freq['cluster'] == c) & (infect_decayed_freq['infects'] == True)]['phage name'].squeeze()\n",
    "    to_add.name = 'Cluster ' + c\n",
    "    infectivities_decayed.append(to_add)\n",
    "\n",
    "infectivities_decayed_sets = [set(i) for i in infectivities_decayed]\n",
    "labels_decayed = [i.name for i in infectivities_decayed]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4711cf0d-25a7-4694-a20c-8cc3e43e225f",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize = (7.5,4.5))\n",
    "supervenn(infectivities_decayed_sets, labels_decayed, side_plots = 'right',\n",
    "          chunks_ordering = 'minimize gaps', sets_ordering = 'minimize gaps', \n",
    "          widths_minmax_ratio = 0.15)\n",
    "plt.xlabel('# infecting phages')\n",
    "plt.ylabel('rRNA cluster')\n",
    "plt.savefig(results_folder + \"/\" + \"venn_infectivity_decayed.svg\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "71b2559d-32e1-4c46-a1ed-3f18fde5c5ed",
   "metadata": {},
   "source": [
    "## Prophage protein annotations"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a690c42d-f520-4cc2-8900-faf2b43ef39c",
   "metadata": {},
   "source": [
    "### Parsing PHASTEST prophage region protein annotations"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa272102-73d9-469a-ac48-ff225eb26c22",
   "metadata": {},
   "source": [
    "The protein annotation of the proteins in the prophage regions can be found in the `json_input` file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a50dd0c7-009d-4264-a133-952bd0a43ffe",
   "metadata": {},
   "outputs": [],
   "source": [
    "result_dirs = filter(lambda x: '.log' not in x, os.listdir(output_folder))\n",
    "hits = []\n",
    "for dir in result_dirs:\n",
    "    with open(output_folder + '/' + dir + '/' + \"json_input\", \"r\") as handle:\n",
    "        proteins = json.load(handle)\n",
    "    for p in proteins:\n",
    "        protein_type = p['type']\n",
    "        protein_tags = p['name']\n",
    "        tags = p['name'].split(';')\n",
    "        try:\n",
    "            protein_annotation = tags[1]\n",
    "        except IndexError:\n",
    "            protein_annotation = tags[0]\n",
    "        try:\n",
    "            protein_length = int(p['sequence_length'].split(' ')[0])\n",
    "        except ValueError: # skip non-delineated proteins\n",
    "            continue\n",
    "        protein_start = p['start']\n",
    "        protein_stop = p['stop']\n",
    "        record = {'strain': dir, 'type': protein_type, 'annotation': protein_annotation, 'size': protein_length, \n",
    "                  'start': protein_start, 'stop': protein_stop, 'tags': protein_tags}\n",
    "        hits.append(record)\n",
    "hits = pd.DataFrame(hits)\n",
    "hits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "489371d7-bbfd-4116-94e9-019360aa3706",
   "metadata": {},
   "outputs": [],
   "source": [
    "hits.to_csv(results_folder + '/prophage_proteins', sep = \"\\t\", index = False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab9b70bb-961f-4c0d-bdd8-eaf8ca08f1e6",
   "metadata": {},
   "source": [
    "### Querying the prophage protein annotation dataframe"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "190cb4c2-d8b3-45c5-81a2-e20e3d7cffe6",
   "metadata": {},
   "source": [
    "**Auxiliary function**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12e393d4-a11f-4bd7-8e0b-ba60296d8e8a",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Queries the prophage protein annotation dataframe for a given search query and returns how many hits there are for each assembly\n",
    "##\n",
    "## PARAMS\n",
    "## query      the search query string\n",
    "## hits       the prophage protein annotation dataframe\n",
    "##\n",
    "## OUTPUT\n",
    "## a dataframe of the number of hits for that search query for each genome assembly\n",
    "##\n",
    "def query_hits(query, hits):\n",
    "    strains = pd.Series(sorted(list(hits['strain'].unique())), name = 'strain')\n",
    "    res = hits.loc[hits['annotation'].str.contains(query, regex = True, case = False),]\n",
    "    res = res.groupby('strain')['type'].count()\n",
    "    res = pd.merge(res, strains, on = \"strain\", how = \"right\").fillna(0).set_index('strain').astype(int)\n",
    "    return list(res['type'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1de3b26f-de6b-4e6f-b3f6-ea6ecfec579a",
   "metadata": {},
   "source": [
    "#### Queries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bbb9149f-33b3-46a6-9457-ce1ac6bd1831",
   "metadata": {},
   "outputs": [],
   "source": [
    "sum(query_hits('Type [I|V|X]+ .*secretion|secretion.*Type [I|V|X]+', hits))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d49b2198-ba9d-4789-9adc-36856aa30aea",
   "metadata": {},
   "outputs": [],
   "source": [
    "sum(query_hits('Type II .*secretion|secretion.*Type II', hits))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0da7dbb-d689-4be2-b47b-a28a97135651",
   "metadata": {},
   "outputs": [],
   "source": [
    "sum(query_hits('Type IV .*secretion|secretion.*Type IV', hits))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30a25c60-ade5-4bf4-a0b4-2f1ba4ce7bd0",
   "metadata": {},
   "outputs": [],
   "source": [
    "sum(query_hits('Type VI .*secretion|secretion.*Type VI', hits))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea6d15b1-1cf1-4a92-97ba-5fc9889278ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "sum(query_hits('Type VII .*secretion|secretion.*Type VII', hits))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90b84a0b-a789-4e6d-8204-0ae756e314b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "sum(query_hits('toxin', hits))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c2a25f3-9b8a-4c99-9763-b51d5fb0c51f",
   "metadata": {},
   "outputs": [],
   "source": [
    "sum(query_hits('toxin-antitoxin', hits))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "214853b5-515a-4213-82f5-fdd8b3c1ae14",
   "metadata": {},
   "outputs": [],
   "source": [
    "sum(query_hits('Type II .*toxin-antitoxin|toxin-antitoxin.*Type II', hits))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a60d9c5-6c1d-4f02-b619-f144990dcc6b",
   "metadata": {},
   "outputs": [],
   "source": [
    "sum(query_hits('hemolysin', hits))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eac26aca-74fb-4bb1-99e2-2e795151d7ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "sum(query_hits('phage.*struct', hits))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "deb19493-0299-4aec-8fe6-73dd7d79d4a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "sum(query_hits('head', hits))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "765bc8d0-b421-4ba9-b351-406f3ef8e2af",
   "metadata": {},
   "outputs": [],
   "source": [
    "sum(query_hits('tail', hits))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec7fe0f7-fbcd-4006-bd04-9097146a9ca5",
   "metadata": {},
   "outputs": [],
   "source": [
    "sum(query_hits('capsid', hits))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1a1b0ea-ac15-48cf-9002-5486b84bc6ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "sum(query_hits('methyltransferase', hits))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "815a2a70-ba45-4e1a-8074-38b1853a8741",
   "metadata": {},
   "outputs": [],
   "source": [
    "sum(query_hits('acetyltransferase', hits))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "785ddef2-805d-4a4f-8337-575dd3d693ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "sum(query_hits('type II methyltransferase', hits))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ae456eb-4e4f-466b-b39b-8502cb46facb",
   "metadata": {},
   "outputs": [],
   "source": [
    "sum(query_hits('restriction', hits))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "649814f3-8068-48b0-8d82-9d6bf65b223f",
   "metadata": {},
   "outputs": [],
   "source": [
    "sum(query_hits('type 1 restriction', hits))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ecc8463-a5e6-4dda-9937-30e05503713c",
   "metadata": {},
   "outputs": [],
   "source": [
    "sum(query_hits('type I restriction', hits))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62948ea8-d30c-42ee-a3a9-791a940cb9e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "sum(query_hits('type 2 restriction', hits))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b4a3226-3152-46e2-ab90-c3f24a3236ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "sum(query_hits('type II restriction', hits))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1245a41e-8bf1-4a7a-baaa-6015352c4f0e",
   "metadata": {},
   "outputs": [],
   "source": [
    "sum(query_hits('type 3 restriction', hits))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3fd4f56-c2a3-4585-afb0-6b98288c7264",
   "metadata": {},
   "outputs": [],
   "source": [
    "sum(query_hits('type III restriction', hits))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82508d6c-8a4d-4dad-a2d5-174e64109710",
   "metadata": {},
   "outputs": [],
   "source": [
    "sum(query_hits('type 4 restriction', hits))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea29f335-8091-4892-8dff-d2f0fe3ff9e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "sum(query_hits('type IV restriction', hits))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "828926e6-7e9c-488f-bb7a-dd1acb8b2d55",
   "metadata": {},
   "outputs": [],
   "source": [
    "sum(query_hits('type 5 restriction', hits))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1479924-4654-4e7b-a81f-9bc2759aab61",
   "metadata": {},
   "outputs": [],
   "source": [
    "sum(query_hits('type V restriction', hits))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb493b4e-5cd9-4a92-aeac-4a8190237a2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "sum(query_hits('transporter', hits))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b8f2913-fd1c-465f-be8d-afc29eeb749a",
   "metadata": {},
   "outputs": [],
   "source": [
    "sum(query_hits('ABC transporter', hits))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1fc94c5-398c-4a91-b76d-51f86674a69d",
   "metadata": {},
   "outputs": [],
   "source": [
    "sum(query_hits('beta-lactam', hits))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7739b17-32de-4f2c-a0fb-c9f9a09ecbc3",
   "metadata": {},
   "outputs": [],
   "source": [
    "sum(query_hits('aminoglycoside', hits))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c85f2bc-4680-4602-91a2-cbb06dbb7a46",
   "metadata": {},
   "outputs": [],
   "source": [
    "sum(query_hits('tetracycline', hits))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8435510a-ae46-432f-81a7-480734e16080",
   "metadata": {},
   "outputs": [],
   "source": [
    "sum(query_hits('erythromycin', hits))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4154d9dc-93ac-4da1-8f4d-281591e0ab92",
   "metadata": {},
   "outputs": [],
   "source": [
    "sum(query_hits('chloramphenicol', hits))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ad0fb1d-e96d-4be5-b777-8aa5de188f01",
   "metadata": {},
   "outputs": [],
   "source": [
    "sum(query_hits('resistance', hits))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4778881c-4cf9-4782-9380-29a71e7c20e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "sum(query_hits('arsen', hits))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9390e932-6bd4-4520-9387-f22c8c35d076",
   "metadata": {},
   "outputs": [],
   "source": [
    "sum(query_hits('copper', hits))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38245b57-3ea6-470a-88d4-51b564cc2806",
   "metadata": {},
   "outputs": [],
   "source": [
    "sum(query_hits('cadmium', hits))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b1d97d3-42be-4281-a9a5-856ad2828431",
   "metadata": {},
   "outputs": [],
   "source": [
    "sum(query_hits('zinc .*transport', hits))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea6eb4bc-60af-4b57-97d5-36b001a9ea1c",
   "metadata": {},
   "outputs": [],
   "source": [
    "sum(query_hits('cobalt .*transport', hits))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa29b64a-1daa-4f92-b451-6f3c04f19b75",
   "metadata": {},
   "outputs": [],
   "source": [
    "sum(query_hits('nitrate reductase', hits))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd874fc5-bd43-48e4-beee-adc74506ae99",
   "metadata": {},
   "outputs": [],
   "source": [
    "sum(query_hits('cytochrome', hits))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7e0a2ea-2516-468f-8b31-5f803c7dd692",
   "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
}