notebooks/S_gene_demo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f71646ba",
   "metadata": {},
   "source": [
    "# Basic BVAS demo using S-gene data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "341a0926",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import gzip\n",
    "from bvas import BVASSelector, map_inference"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "995d788d",
   "metadata": {},
   "source": [
    "### Load data\n",
    "\n",
    "The analysis in this notebook is meant for demonstration purposes only.\n",
    "Our data only includes alleles from the S gene. As such results are expected to be biased as compared to a genome wide analysis. <br> <br>\n",
    "Powered by <br>\n",
    "<img src=\"https://www.gisaid.org/fileadmin/gisaid/img/schild.png\" alt=\"GISAID\" width=\"80\" align=\"left\"/>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "3521707d",
   "metadata": {},
   "outputs": [],
   "source": [
    "data = torch.load(gzip.GzipFile(\"../data/S_gene.pt.gz\", \"rb\"), map_location='cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "af888137",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Gamma torch.Size([415, 415])\n",
      "Y torch.Size([415])\n",
      "num_alleles 415\n",
      "num_regions 74\n",
      "mutations 415\n"
     ]
    }
   ],
   "source": [
    "# inspect data\n",
    "for k, v in data.items():\n",
    "    if hasattr(v, 'shape'):\n",
    "        print(k, v.shape)\n",
    "    elif isinstance(v, list):\n",
    "        print(k, len(v))\n",
    "    else:\n",
    "        print(k, v)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "887d5318",
   "metadata": {},
   "source": [
    "### Instantiate BVASSelector object"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "1864637f",
   "metadata": {},
   "outputs": [],
   "source": [
    "selector = BVASSelector(data['Y'].double(),  # use 64-bit precision\n",
    "                        data['Gamma'].double(), \n",
    "                        data['mutations'], \n",
    "                        nu_eff_multiplier=0.25,\n",
    "                        S=5.0,\n",
    "                        tau=100.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a22ba98e",
   "metadata": {},
   "source": [
    "### Run BVAS MCMC-based inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "cf64f287",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "937e513eeb644c159ac896165230c608",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/6000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "selector.run(T=5000, T_burnin=1000, seed=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b02aac64",
   "metadata": {},
   "source": [
    "### Inspect results\n",
    "\n",
    "The results can be found in the `selector.summary` Pandas DataFrame."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "3d8c3b4d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "              PIP      Beta   BetaStd  Rank\n",
      "S:L452R  1.000000  0.500928  0.049617     1\n",
      "S:T478K  1.000000  0.507602  0.074132     2\n",
      "S:R346K  1.000000  0.481718  0.059620     3\n",
      "S:T19R   1.000000  0.541961  0.075788     4\n",
      "S:N440K  1.000000  0.482342  0.072035     5\n",
      "S:E484K  1.000000  0.325571  0.033246     6\n",
      "S:P681R  1.000000  0.421493  0.049610     7\n",
      "S:N501Y  0.999877  0.288173  0.049724     8\n",
      "S:T95I   0.999205  0.336165  0.078261     9\n",
      "S:N969K  0.996847  0.432681  0.102006    10\n",
      "S:Q954H  0.996489  0.432219  0.106664    11\n",
      "S:G339D  0.996196  0.430900  0.087655    12\n",
      "S:N679K  0.994718  0.399761  0.095743    13\n",
      "S:N764K  0.985545  0.386436  0.102249    14\n",
      "S:S375F  0.979299  0.374481  0.107152    15\n",
      "S:S373P  0.978723  0.369183  0.113199    16\n",
      "S:T859N  0.964351  0.255766  0.080352    17\n",
      "S:S477N  0.934212  0.137568  0.050168    18\n",
      "S:T716I  0.912222  0.289920  0.119026    19\n",
      "S:Y145H  0.872734  0.186306  0.094216    20\n",
      "S:H655Y  0.779720  0.224054  0.135953    21\n",
      "S:D405N  0.698377  0.220314  0.166125    22\n",
      "S:T19I   0.692335  0.212242  0.168075    23\n",
      "S:S371F  0.689166  0.213947  0.171222    24\n",
      "S:V213G  0.689152  0.218743  0.165125    25\n",
      "S:R408S  0.680633  0.203690  0.169732    26\n",
      "S:T376A  0.680403  0.200668  0.170551    27\n",
      "S:K77T   0.607354 -0.115821  0.112460    28\n",
      "S:D796Y  0.549913  0.157896  0.143488    29\n",
      "S:P681H  0.541714  0.078355  0.076483    30\n"
     ]
    }
   ],
   "source": [
    "print(selector.summary.iloc[:30][['PIP', 'Beta', 'BetaStd', 'Rank']])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2fc49c9",
   "metadata": {},
   "source": [
    "## Let's compare to MAP inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "6eb9d959",
   "metadata": {},
   "outputs": [],
   "source": [
    "map_results = map_inference(data['Y'].double(), data['Gamma'].double(), data['mutations'], tau_reg=2048.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "98636892",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Beta</th>\n",
       "      <th>BetaStd</th>\n",
       "      <th>Rank</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>S:T478K</th>\n",
       "      <td>0.479279</td>\n",
       "      <td>0.016720</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>S:L452R</th>\n",
       "      <td>0.403808</td>\n",
       "      <td>0.015633</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>S:T19R</th>\n",
       "      <td>0.401982</td>\n",
       "      <td>0.018214</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>S:P681R</th>\n",
       "      <td>0.385613</td>\n",
       "      <td>0.017929</td>\n",
       "      <td>4</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>S:N440K</th>\n",
       "      <td>0.255411</td>\n",
       "      <td>0.019505</td>\n",
       "      <td>5</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>S:R346K</th>\n",
       "      <td>0.232116</td>\n",
       "      <td>0.018944</td>\n",
       "      <td>6</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>S:T95I</th>\n",
       "      <td>0.226314</td>\n",
       "      <td>0.016852</td>\n",
       "      <td>7</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>S:N969K</th>\n",
       "      <td>0.223328</td>\n",
       "      <td>0.021557</td>\n",
       "      <td>8</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>S:Q954H</th>\n",
       "      <td>0.222891</td>\n",
       "      <td>0.021557</td>\n",
       "      <td>9</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>S:G339D</th>\n",
       "      <td>0.222524</td>\n",
       "      <td>0.021557</td>\n",
       "      <td>10</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>S:N679K</th>\n",
       "      <td>0.220374</td>\n",
       "      <td>0.021292</td>\n",
       "      <td>11</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>S:H655Y</th>\n",
       "      <td>0.214477</td>\n",
       "      <td>0.020056</td>\n",
       "      <td>12</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>S:S373P</th>\n",
       "      <td>0.213697</td>\n",
       "      <td>0.021375</td>\n",
       "      <td>13</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>S:S375F</th>\n",
       "      <td>0.213697</td>\n",
       "      <td>0.021375</td>\n",
       "      <td>14</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>S:N764K</th>\n",
       "      <td>0.213502</td>\n",
       "      <td>0.021403</td>\n",
       "      <td>15</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>S:N501Y</th>\n",
       "      <td>0.212104</td>\n",
       "      <td>0.016656</td>\n",
       "      <td>16</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>S:D796Y</th>\n",
       "      <td>0.195433</td>\n",
       "      <td>0.021086</td>\n",
       "      <td>17</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>S:Q493R</th>\n",
       "      <td>0.175371</td>\n",
       "      <td>0.021422</td>\n",
       "      <td>18</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>S:Y505H</th>\n",
       "      <td>0.175360</td>\n",
       "      <td>0.021422</td>\n",
       "      <td>19</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>S:E484A</th>\n",
       "      <td>0.175350</td>\n",
       "      <td>0.021422</td>\n",
       "      <td>20</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>S:E484K</th>\n",
       "      <td>0.172017</td>\n",
       "      <td>0.015439</td>\n",
       "      <td>21</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>S:P681H</th>\n",
       "      <td>0.164017</td>\n",
       "      <td>0.014328</td>\n",
       "      <td>22</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>S:K417N</th>\n",
       "      <td>0.157449</td>\n",
       "      <td>0.018647</td>\n",
       "      <td>23</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>S:Q498R</th>\n",
       "      <td>0.151624</td>\n",
       "      <td>0.021281</td>\n",
       "      <td>24</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>S:D405N</th>\n",
       "      <td>0.125156</td>\n",
       "      <td>0.021687</td>\n",
       "      <td>25</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>S:S371F</th>\n",
       "      <td>0.125145</td>\n",
       "      <td>0.021686</td>\n",
       "      <td>26</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>S:T19I</th>\n",
       "      <td>0.125106</td>\n",
       "      <td>0.021674</td>\n",
       "      <td>27</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>S:V213G</th>\n",
       "      <td>0.125066</td>\n",
       "      <td>0.021670</td>\n",
       "      <td>28</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>S:T376A</th>\n",
       "      <td>0.125014</td>\n",
       "      <td>0.021687</td>\n",
       "      <td>29</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>S:R408S</th>\n",
       "      <td>0.125014</td>\n",
       "      <td>0.021687</td>\n",
       "      <td>30</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "             Beta   BetaStd  Rank\n",
       "S:T478K  0.479279  0.016720     1\n",
       "S:L452R  0.403808  0.015633     2\n",
       "S:T19R   0.401982  0.018214     3\n",
       "S:P681R  0.385613  0.017929     4\n",
       "S:N440K  0.255411  0.019505     5\n",
       "S:R346K  0.232116  0.018944     6\n",
       "S:T95I   0.226314  0.016852     7\n",
       "S:N969K  0.223328  0.021557     8\n",
       "S:Q954H  0.222891  0.021557     9\n",
       "S:G339D  0.222524  0.021557    10\n",
       "S:N679K  0.220374  0.021292    11\n",
       "S:H655Y  0.214477  0.020056    12\n",
       "S:S373P  0.213697  0.021375    13\n",
       "S:S375F  0.213697  0.021375    14\n",
       "S:N764K  0.213502  0.021403    15\n",
       "S:N501Y  0.212104  0.016656    16\n",
       "S:D796Y  0.195433  0.021086    17\n",
       "S:Q493R  0.175371  0.021422    18\n",
       "S:Y505H  0.175360  0.021422    19\n",
       "S:E484A  0.175350  0.021422    20\n",
       "S:E484K  0.172017  0.015439    21\n",
       "S:P681H  0.164017  0.014328    22\n",
       "S:K417N  0.157449  0.018647    23\n",
       "S:Q498R  0.151624  0.021281    24\n",
       "S:D405N  0.125156  0.021687    25\n",
       "S:S371F  0.125145  0.021686    26\n",
       "S:T19I   0.125106  0.021674    27\n",
       "S:V213G  0.125066  0.021670    28\n",
       "S:T376A  0.125014  0.021687    29\n",
       "S:R408S  0.125014  0.021687    30"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "map_results.iloc[:30]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea560727",
   "metadata": {},
   "source": [
    "### Compare uncertainty estimates\n",
    "\n",
    "We note that the MAP uncertainty estimates are much narrower than the \n",
    "corresponding BVAS uncertainty estimates. This is ultimatedly due to\n",
    "the fact that BVAS considers multiple hypotheses about which alleles\n",
    "are neutral and which are not."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "00c6da28",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.07413167570495986"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# BVAS posterior standard deviation of selection coefficient for S:T478K \n",
    "selector.summary.loc['S:T478K'].BetaStd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "2f5bfec5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.016719752782775"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# MAP posterior standard deviation of selection coefficient for S:T478K \n",
    "map_results.loc['S:T478K'].BetaStd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "74219a2d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "4.433778218379622"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# compute ratio\n",
    "selector.summary.loc['S:T478K'].BetaStd / map_results.loc['S:T478K'].BetaStd"
   ]
  }
 ],
 "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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}