explore-nextclade-counts.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a774cc9c-99e0-4366-a219-d3e07263fafe",
   "metadata": {},
   "source": [
    "# Explore Nextclade Mutations"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a11433c1-e4a0-42cc-9a32-9a94436a33f3",
   "metadata": {},
   "source": [
    "In this notebook we explore enrichment in the input and output mutations to the model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f5a4fd8-78ce-4af1-8920-e10b8b0ee85d",
   "metadata": {},
   "source": [
    "# Init"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf385db8-d4a0-4d84-8159-ce6a2f398f3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import seaborn as sns\n",
    "import numpy as np\n",
    "import matplotlib as plt\n",
    "import pickle\n",
    "import scipy.stats as stats"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7789887-0b54-48ff-91ec-faec48ef4853",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.rcParams['figure.figsize'] = [7, 7]\n",
    "plt.rcParams['figure.dpi'] = 200"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "875ca21e-a421-4305-9082-8feb6130254e",
   "metadata": {},
   "source": [
    "## Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33f7d919-ce95-4dde-9807-4d3ccf19cfae",
   "metadata": {},
   "outputs": [],
   "source": [
    "def preprocess_nextclade_counts(nextclade_counts_df):\n",
    "    # convert index to columns and split, extract aa changes\n",
    "    idx = nextclade_counts_df.index\n",
    "    nextclade_counts_df[['protein','aa_change']] = nextclade_counts_df.index.to_frame()[0].str.split(':',expand=True)\n",
    "    nextclade_counts_df[\"aa_from\"] = nextclade_counts_df['aa_change'].astype(str).str[0]\n",
    "    nextclade_counts_df[\"aa_to\"] = nextclade_counts_df['aa_change'].astype(str).str[-1]\n",
    "    nextclade_counts_df[\"aa_change\"] = nextclade_counts_df[\"aa_from\"] + '_' + nextclade_counts_df[\"aa_to\"]\n",
    "    nextclade_counts_df = nextclade_counts_df.rename(columns={0:'count'})\n",
    "    \n",
    "    return nextclade_counts_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "297063c7-2be0-4d34-86f9-2b0cea60469e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def preprocess_mutations(mutations_tsv):\n",
    "    mutations_tsv[['gene','aa_change']] = mutations_tsv['mutation'].str.split(':',expand=True)\n",
    "    mutations_tsv['aa_from'] = mutations_tsv['aa_change'].str[0:1]\n",
    "    mutations_tsv['aa_to'] = mutations_tsv['aa_change'].str[-1]\n",
    "    mutations_tsv['aa_position'] = mutations_tsv['aa_change'].str[1:-1]\n",
    "    return mutations_tsv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15b031b7-86f6-4bd4-ad78-5da71e0558bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_event_frequencies(mutations_proc):\n",
    "    mutations_1 = mutations_proc[['aa_from','aa_to']].copy()\n",
    "    mutations_1['count'] = 1\n",
    "    mutations_1 = mutations_1.groupby(['aa_from','aa_to']).aggregate({'count':'sum'})\n",
    "    mutations_1 = mutations_1.reset_index()\n",
    "    pre_pivot = mutations_1.copy()\n",
    "    mutations_1 = mutations_1.pivot(index='aa_from',columns='aa_to',values='count')\n",
    "    cols = mutations_1.columns\n",
    "    rows = mutations_1.index\n",
    "    \n",
    "    mutations_1 = mutations_1.values\n",
    "    mutations_1 = np.nan_to_num(mutations_1, 0)\n",
    "    \n",
    "    return {\n",
    "        'vals': mutations_1,\n",
    "        'cols': cols,\n",
    "        'rows': rows,     \n",
    "        'prepivot': pre_pivot,\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58652978-4913-4628-baad-1c687a48be58",
   "metadata": {},
   "outputs": [],
   "source": [
    "def score_mutations(test_data, background_data, exclude_mutations = set(('*')), verbose=True):\n",
    "    # from\n",
    "    test_mutations_from = set(test_data['aa_from'].unique())\n",
    "    bg_mutations_from = set(background_data['aa_from'].unique())\n",
    "    common_mutations_from = test_mutations_from.intersection(bg_mutations_from)\n",
    "    used_mutations_from = common_mutations_from - exclude_mutations\n",
    "    \n",
    "    # to\n",
    "    test_mutations_to = set(test_data['aa_to'].unique())\n",
    "    bg_mutations_to = set(background_data['aa_to'].unique())\n",
    "    common_mutations_to = test_mutations_to.intersection(bg_mutations_to)\n",
    "    used_mutations_to = common_mutations_to - exclude_mutations\n",
    "    \n",
    "    # subset\n",
    "    test_data_subset = test_data[\n",
    "        (test_data['aa_from'].isin(used_mutations_from)) & \n",
    "        (test_data['aa_to'].isin(used_mutations_to))\n",
    "    ]\n",
    "    bg_data_subset = background_data[\n",
    "        (background_data['aa_from'].isin(used_mutations_from)) &\n",
    "        (background_data['aa_to'].isin(used_mutations_to))\n",
    "    ]\n",
    "    \n",
    "    test_sum = test_data_subset['count'].sum()\n",
    "    bg_sum = bg_data_subset['count'].sum()\n",
    "    \n",
    "    # n_tests ignores the diagonal\n",
    "    n_tests = ((len(used_mutations_from) * len(used_mutations_to)) - max(len(used_mutations_from), len(used_mutations_to)))\n",
    "\n",
    "    alpha = 0.05\n",
    "    threshold = alpha / n_tests\n",
    "    \n",
    "    ret = {'from': [], 'to': [], 'pval': [], 'qval': [], 'obs': [], 'exp': [], 'significant': []}\n",
    "    \n",
    "    # Calculate background \n",
    "    for from_mut in used_mutations_from:\n",
    "        for to_mut in used_mutations_to:\n",
    "            obs_count = test_data_subset.loc[\n",
    "                (test_data_subset['aa_from'] == from_mut) &\n",
    "                (test_data_subset['aa_to'] == to_mut)\n",
    "            ]['count'].values\n",
    "            if len(obs_count) ==1:\n",
    "                obs_count = obs_count[0]\n",
    "            else:\n",
    "                obs_count = 0\n",
    "                \n",
    "            bg_count = bg_data_subset.loc[\n",
    "                (bg_data_subset['aa_from'] == from_mut) &\n",
    "                (bg_data_subset['aa_to'] == to_mut)\n",
    "            ]['count'].values\n",
    "            if len(bg_count) == 1:\n",
    "                bg_count = bg_count[0]\n",
    "            else:\n",
    "                bg_count = 0\n",
    "            \n",
    "            bg_prob = bg_count / bg_sum\n",
    "            \n",
    "            pval = stats.binom_test(obs_count, test_sum, bg_prob, alternative='greater')\n",
    "            sign = pval < threshold\n",
    "            \n",
    "            qval = min(1, pval * n_tests)\n",
    "             \n",
    "            if (verbose):\n",
    "                print(f\"{from_mut} -> {to_mut}, obs: {obs_count}/{test_sum} bg_prob: {bg_prob}, pval: {pval}\")\n",
    "    \n",
    "            ret['from'].append(from_mut)\n",
    "            ret['to'].append(to_mut)\n",
    "            ret['obs'].append(obs_count)\n",
    "            ret['exp'].append((bg_prob * test_sum))\n",
    "            ret['pval'].append(pval)\n",
    "            ret['qval'].append(qval)\n",
    "            ret['significant'].append(sign)\n",
    "            \n",
    "    results = pd.DataFrame(ret)\n",
    "    \n",
    "    results = results.sort_values('pval')\n",
    "    \n",
    "    return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48538323-a0ff-4b86-be41-512a47c59e03",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_top55_spike(mutations_proc):\n",
    "    return mutations_proc[mutations_proc['gene'] == 'S'].sort_values('Δ log R', ascending=False).head(55)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79366663-0444-4cea-9f20-6135d631dbe6",
   "metadata": {},
   "source": [
    "# Load and prepare data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be34f278-90aa-4e2e-87b8-287981adfdb2",
   "metadata": {},
   "source": [
    "## Nextclade counts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8066f38-d123-4829-a865-340738219c5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load nextclade counts\n",
    "with open(\"results/nextclade.counts.pkl\", 'rb') as f:\n",
    "    nextclade_counts = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92e36311-1e66-42fc-866e-aa1b47b031ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "# get data frame form index\n",
    "nextclade_counts_df_1 = pd.DataFrame.from_dict(nextclade_counts, orient='index')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e906146-9a61-411c-83f6-512ed22524f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "nextclade_counts_df = preprocess_nextclade_counts(nextclade_counts_df_1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f907d69-ece2-4dc6-b7da-413297f8612b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate a version with all counts set to 1 so that we can just look at frequency of types of mutations\n",
    "nextclade_counts_no_freq = nextclade_counts_df.copy()\n",
    "nextclade_counts_no_freq['count'] = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6fb9d6f4-b78e-4c25-88bd-800425f22b41",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get total summed counts\n",
    "summarized_counts = nextclade_counts_df.groupby(['aa_from','aa_to']).agg({'count':'sum'})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a820e33-2e8b-474e-b2d6-11cd1befb50b",
   "metadata": {},
   "outputs": [],
   "source": [
    "summarized_counts_1 = summarized_counts['count']\n",
    "summarized_counts_1 = summarized_counts_1.reset_index()\n",
    "nextclade_comparison_table = summarized_counts_1.copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24cc9172-60f3-4d54-8be0-1616a29d4331",
   "metadata": {},
   "outputs": [],
   "source": [
    "aa_changes = summarized_counts_1.pivot(index='aa_from',columns='aa_to',values='count')\n",
    "aa_changes_np = aa_changes.to_numpy()\n",
    "aa_changes_np = np.nan_to_num(aa_changes_np, nan=0.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b370448d-df43-47a8-9460-b77b40ef7427",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate Heatmap\n",
    "fig = sns.heatmap(aa_changes_np,\n",
    "            xticklabels=aa_changes.columns,\n",
    "            yticklabels=aa_changes.index,)\n",
    "fig.set(xlabel='aa_to',ylabel='aa_from',title='Mutations in Nextclade')\n",
    "plt.pyplot.savefig('paper/mutation_scoring/nextclade.png')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9688be6-f0e5-4642-a109-97c9c17d7d4c",
   "metadata": {},
   "source": [
    "## Input to model (mutations.tsv)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f78a6dd-ea1d-46bd-975f-f830de75a25f",
   "metadata": {},
   "outputs": [],
   "source": [
    "mutations_tsv = pd.read_table('paper/mutations.tsv')\n",
    "mutations_proc = preprocess_mutations(mutations_tsv)\n",
    "mutations_freq = get_event_frequencies(mutations_proc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39e422c5-ee6d-4048-9c16-788ccbead8d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# heatmap\n",
    "fig = sns.heatmap(mutations_freq['vals'],\n",
    "            xticklabels=mutations_freq['cols'],\n",
    "            yticklabels=mutations_freq['rows'],)\n",
    "fig.set(xlabel='aa_to',ylabel='aa_from',\n",
    "        title='Input to Model Mutation Frequencies')\n",
    "plt.pyplot.savefig('paper/mutation_scoring/input_to_model.png')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50df4cb2-2ac4-4e47-864a-e696442a3532",
   "metadata": {},
   "source": [
    "## Top 255 global amino acids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa7d1f6c-b774-47d7-a417-d96e58798f2f",
   "metadata": {},
   "outputs": [],
   "source": [
    "top255_global = mutations_proc.sort_values('Δ log R',ascending =False).head(255)\n",
    "top255_freq = get_event_frequencies(top255_global)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea613905-28da-4225-9673-4fa4b5ed2de5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# heatmap\n",
    "fig = sns.heatmap(top255_freq['vals'],\n",
    "            xticklabels=top255_freq['cols'],\n",
    "            yticklabels=top255_freq['rows'],)\n",
    "fig.set(xlabel='aa_to',ylabel='aa_from',\n",
    "        title='Top 255 mutation frequencies')\n",
    "plt.pyplot.savefig('paper/mutation_scoring/top_255_global.png')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cc6e88e8-11f6-4527-a5b2-cba831dcd2c0",
   "metadata": {},
   "source": [
    "## Top 55 Spike Vs Features (mutations.tsv)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2602f5d-6554-4f8f-8685-9242bc7b466d",
   "metadata": {},
   "outputs": [],
   "source": [
    "top55_spike = mutations_proc.loc[mutations_proc['gene'] == 'S'].sort_values('Δ log R',ascending =False).head(55)\n",
    "top55_spike_freq = get_event_frequencies(top55_spike)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c985b15c-afc4-4ca9-9e3c-df9d323dda4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# heatmap\n",
    "fig = sns.heatmap(top55_spike_freq['vals'],\n",
    "            xticklabels=top55_spike_freq['cols'],\n",
    "            yticklabels=top55_spike_freq['rows'],)\n",
    "fig.set(xlabel='aa_to',ylabel='aa_from',\n",
    "        title='Top 55 in Spike')\n",
    "plt.pyplot.savefig('paper/mutation_scoring/top_55_spike.png')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "147ece5a-7298-4a51-bbf1-6e26af84f7b8",
   "metadata": {},
   "source": [
    "# Score sets of mutations"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3fd93ebb-e0d1-419e-bd96-323c9b0f5a0e",
   "metadata": {},
   "source": [
    "## Input to model vs Nextclade"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d54a0778-1ab0-4a82-afc3-58ee0b84e2d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "mutations_tsv_vs_nextclade = score_mutations(mutations_freq['prepivot'], nextclade_comparison_table, verbose = False)\n",
    "mutations_tsv_vs_nextclade.to_csv('paper/mutation_scoring/input_vs_nextclade.csv')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "176b1bfb-c42e-4c9f-8003-ca953ed7d730",
   "metadata": {},
   "source": [
    "## Top 255 Vs Input to Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e486077-1496-42a1-ab96-98db9c89225a",
   "metadata": {},
   "outputs": [],
   "source": [
    "top255_vs_mutation_tsv = score_mutations(top255_freq['prepivot'], mutations_freq['prepivot'], verbose=False)\n",
    "top255_vs_mutation_tsv.to_csv('paper/mutation_scoring/top255_VS_input.csv')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3442765-792f-418c-a96f-a291e9cbc1e1",
   "metadata": {},
   "source": [
    "## Top 55 VS Input to Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1354164f-1dd9-4362-984b-b1068849fa59",
   "metadata": {},
   "outputs": [],
   "source": [
    "top55_spike_VS_mutations_tsv = score_mutations(top55_spike_freq['prepivot'], mutations_freq['prepivot'], verbose = False)\n",
    "top55_spike_VS_mutations_tsv.to_csv('paper/mutation_scoring/top55_VS_input.csv')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "78a0fda6-ccf5-4ef9-bae9-181a859303b2",
   "metadata": {},
   "source": [
    "## Top 55 VS nextclade"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb846453-01a2-462c-af8c-235857f69082",
   "metadata": {},
   "outputs": [],
   "source": [
    "top55_spike_VS_nextclade = score_mutations(top55_spike_freq['prepivot'], nextclade_comparison_table, verbose = False)\n",
    "top55_spike_VS_nextclade.to_csv('paper/mutation_scoring/top55_VS_nextclade.csv')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a00817e-f388-4107-af06-646a7842ea0e",
   "metadata": {},
   "source": [
    "## Input to Model vs Nextclade"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a20ed14a-aee8-453b-8a02-712be0e2bf92",
   "metadata": {},
   "outputs": [],
   "source": [
    "mutations_tsv_VS_nextclade = score_mutations(top55_spike_freq['prepivot'], nextclade_comparison_table, verbose = False)\n",
    "mutations_tsv_VS_nextclade.to_csv('paper/mutation_scoring/input_VS_nextclade.csv')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f66fbb9-a5d7-4372-b36f-c73ac9c54e8b",
   "metadata": {},
   "source": [
    "# Combine results into table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48a94453-740c-4dff-aa69-4c1a7e4ef675",
   "metadata": {},
   "outputs": [],
   "source": [
    "mutations_tsv_vs_nextclade = mutations_tsv_vs_nextclade.set_index(['from','to'])\n",
    "top255_vs_mutation_tsv = top255_vs_mutation_tsv.set_index(['from','to'])\n",
    "top55_spike_VS_mutations_tsv = top55_spike_VS_mutations_tsv.set_index(['from','to'])\n",
    "top55_spike_VS_nextclade = top55_spike_VS_nextclade.set_index(['from','to'])\n",
    "mutations_tsv_VS_nextclade = mutations_tsv_VS_nextclade.set_index(['from','to'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c03880ee-1dad-4b5a-87e5-18a850fa4e5c",
   "metadata": {},
   "outputs": [],
   "source": [
    "mutations_tsv_vs_nextclade_premerge = mutations_tsv_vs_nextclade[['qval']]\n",
    "mutations_tsv_vs_nextclade_premerge.columns = ['qval_input_vs_nextclade']\n",
    "\n",
    "top255_vs_mutation_tsv_premerge = top255_vs_mutation_tsv[['qval']]\n",
    "top255_vs_mutation_tsv_premerge.columns = ['qval_top255_vs_input']\n",
    "\n",
    "top55_spike_VS_mutations_tsv_premerge = top55_spike_VS_mutations_tsv[['qval']]\n",
    "top55_spike_VS_mutations_tsv_premerge.columns = ['qval_top55_spike_vs_input']\n",
    "\n",
    "top55_spike_VS_nextclade_premerge = top55_spike_VS_nextclade[['qval']]\n",
    "top55_spike_VS_nextclade_premerge.columns = ['qval_top55_vs_nextclade']\n",
    "\n",
    "mutations_tsv_VS_nextclade_premerge = mutations_tsv_VS_nextclade[['qval']]\n",
    "mutations_tsv_VS_nextclade_premerge.columns = ['qval_mutations_tsv_VS_nextclade']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4ee6dfd-323a-474c-b9a1-eb5cc446f1e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "aggregated_qvals = mutations_tsv_vs_nextclade_premerge.merge(\n",
    "    top255_vs_mutation_tsv_premerge,left_index=True, right_index=True).merge(\n",
    "    top55_spike_VS_mutations_tsv_premerge, left_index=True, right_index=True).merge(\n",
    "    top55_spike_VS_nextclade_premerge, left_index=True, right_index=True).merge(\n",
    "    mutations_tsv_VS_nextclade_premerge, left_index=True, right_index=True).reset_index()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29bbd29b-a300-48e5-8c67-a2d8f98c3432",
   "metadata": {},
   "outputs": [],
   "source": [
    "aggregated_qvals.to_csv('paper/mutation_scoring/aggregated_qvalues.csv')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}