2_find_SNPs.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Important\n",
    "\n",
    "`run_raxml.sh` in `Genome_tree` to build tree this notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "from collections import Counter, defaultdict\n",
    "\n",
    "import pandas as pd\n",
    "from Bio import SeqIO, AlignIO, Phylo\n",
    "from Bio.Seq import Seq\n",
    "from Bio.SeqRecord import SeqRecord\n",
    "from treetime import TreeAnc\n",
    "\n",
    "genomeTreeDir = \"Genome_tree\"\n",
    "snpDir = \"SNPs\"\n",
    "dataDir = \"Data\"\n",
    "\n",
    "if not os.path.exists(snpDir):\n",
    "    os.mkdir(snpDir)\n",
    "\n",
    "ambiguity = {\n",
    "    \"Y\": [\"C\", \"T\"], \"R\": [\"A\", \"G\"], \"W\": [\"A\", \"T\"],\n",
    "    \"S\": [\"G\", \"C\"], \"K\": [\"T\", \"G\"], \"M\": [\"C\", \"A\"],\n",
    "    \"D\": [\"A\", \"G\", \"T\"], \"V\": [\"A\", \"C\", \"G\"],\n",
    "    \"H\": [\"A\", \"C\", \"T\"], \"B\": [\"C\", \"G\", \"T\"]\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(os.path.join(genomeTreeDir, \"outgroup.txt\")) as f:\n",
    "    for row in f:\n",
    "        exclude = row.strip(\"\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "tree = Phylo.read(os.path.join(genomeTreeDir, \"RAxML_bestTree.GTRGAMMA\"), \"newick\")\n",
    "alignedCDS = AlignIO.read(os.path.join(genomeTreeDir, \"aligned_cds.fasta\"), \"fasta\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "6.54\tWARNING: Previous versions of TreeTime (<0.7.0) RECONSTRUCTED sequences of\n",
      "    \ttips at positions with AMBIGUOUS bases. This resulted in unexpected\n",
      "    \tbehavior is some cases and is no longer done by default. If you want to\n",
      "    \treplace those ambiguous sites with their most likely state, rerun with\n",
      "    \t`reconstruct_tip_states=True` or `--reconstruct-tip-states`.\n"
     ]
    }
   ],
   "source": [
    "treeAnc = TreeAnc(tree = tree, aln = alignedCDS, gtr=\"Jukes-Cantor\", verbose=False)\n",
    "treeAnc.infer_ancestral_sequences(infer_gtr=True, marginal=True)\n",
    "alignedAnc = SeqRecord(Seq(''.join(treeAnc.tree.root.sequence)), \"ancestral\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "alignedCDS.add_sequence(\"ancestral\", str(alignedAnc.seq))\n",
    "AlignIO.write(alignedCDS, os.path.join(snpDir, \"ancestral.fasta\"), \"fasta\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "info = pd.read_csv(os.path.join(dataDir, \"info.csv\"))\n",
    "info.columns.values[list(info.columns.values).index(\"Location\")] = \"City\"\n",
    "info.columns.values[list(info.columns.values).index(\"Area\")] = \"Province\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "meta = pd.read_csv(os.path.join(dataDir, \"original.csv\"))\n",
    "meta.columns.values[list(meta.columns.values).index(\"Accession ID\")] = \"Accession ID\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(os.path.join(snpDir, \"cds_product.json\")) as f:\n",
    "    cdsAnno = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "snpPosIndex = []\n",
    "for i in range(alignedCDS.get_alignment_length()):\n",
    "    aaSum = Counter(alignedCDS[:, i])\n",
    "    if len(aaSum) > 1:\n",
    "        snpPosIndex.append(i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "allSNP = pd.DataFrame(columns=(\n",
    "    \"Pos\", \"SNP\", \"Ref\", \n",
    "    \"Genome pos\", \"Start index\",\n",
    "    \"Product\", \"Product ID\", \"Accession ID\"\n",
    "))\n",
    "\n",
    "for record in alignedCDS:\n",
    "    if record.id != exclude:\n",
    "        for index in snpPosIndex:\n",
    "            ref = alignedAnc[index].upper()\n",
    "            n = record[index].upper()\n",
    "            if n != 'N' and n != '-' and ref != n:\n",
    "                if n in ambiguity:\n",
    "                    if ref in ambiguity[n]:\n",
    "                        continue\n",
    "                pos = index + 1\n",
    "                posAnno = cdsAnno[str(pos)]\n",
    "                row = pd.Series(data={\n",
    "                    \"Pos\": pos, \n",
    "                    \"SNP\": n,\n",
    "                    \"Ref\": ref,\n",
    "                    \"Genome pos\": posAnno[\"genome_pos\"],\n",
    "                    \"Start index\": posAnno[\"start_index\"],\n",
    "                    \"Product\": posAnno[\"product_name\"],\n",
    "                    \"Product ID\": posAnno[\"product_id\"],\n",
    "                    \"Accession ID\": record.id\n",
    "                })\n",
    "                allSNP = allSNP.append(row, ignore_index=True)\n",
    "                \n",
    "allSNP = pd.merge(allSNP, info, on=\"Accession ID\", how=\"left\", suffixes=(\"\", \"_y\"))\n",
    "allSNP = pd.merge(allSNP, meta, on=\"Accession ID\", how=\"left\", suffixes=(\"\", \"_y\"))\n",
    "allSNP = allSNP.sort_values([\"Pos\", \"SNP\"])\n",
    "allSNP.to_csv(os.path.join(snpDir, \"all.csv\"), index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "hubeiSNP = allSNP[allSNP[\"Province\"] == \"Hubei\"][[\"Pos\", \"SNP\"]].drop_duplicates()\n",
    "\n",
    "provinceSNP = pd.DataFrame(columns=allSNP.columns)\n",
    "citySNP = pd.DataFrame(columns=allSNP.columns)\n",
    "\n",
    "for i in allSNP[allSNP[\"Province\"] != \"Hubei\"].groupby([\"Pos\", \"Ref\"]):\n",
    "    inHubei = any((hubeiSNP[\"Pos\"] == i[0][0]) & (hubeiSNP[\"SNP\"] == i[0][1]))\n",
    "    isProvince = len(i[1][\"Province\"].drop_duplicates()) >= 2\n",
    "    isCity = len(i[1][\"City\"].drop_duplicates()) >= 2\n",
    "    if not inHubei:\n",
    "        if isProvince:\n",
    "            provinceSNP = provinceSNP.append(i[1], ignore_index=True)\n",
    "        if isCity:\n",
    "            citySNP = citySNP.append(i[1], ignore_index=True)\n",
    "\n",
    "provinceSNP.to_csv(os.path.join(snpDir, \"province.csv\"), index=False)\n",
    "citySNP.to_csv(os.path.join(snpDir, \"city.csv\"), index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "nonSNP = defaultdict(list)\n",
    "\n",
    "for _, row in allSNP[[\"Ref\", \"Pos\", \"SNP\"]].drop_duplicates().iterrows():\n",
    "    mut = ''.join((row[\"Ref\"], str(row[\"Pos\"]), row[\"SNP\"]))\n",
    "    nonSNP[mut].extend(allSNP.loc[(allSNP[\"Pos\"] == row[\"Pos\"]) & (allSNP[\"SNP\"] == row[\"SNP\"]), \"Accession ID\"].values)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "nonSNPdir = os.path.join(snpDir, \"nonSNP\")\n",
    "if not os.path.exists(nonSNPdir):\n",
    "    os.mkdir(nonSNPdir)\n",
    "\n",
    "for mut, accession in nonSNP.items():\n",
    "    meta[meta[\"Accession ID\"].isin(accession)].to_csv(os.path.join(snpDir, \"nonSNP\", mut + \".csv\"), index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "seqs = SeqIO.index(os.path.join(genomeTreeDir, \"aligned_cds.fasta\"), \"fasta\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "ancIdentical = pd.read_csv(os.path.join(dataDir, \"original.csv\"))\n",
    "ancIdentical = ancIdentical[ancIdentical[\"Accession ID\"].isin(list(seqs.keys()))]\n",
    "ancIdentical = ancIdentical[~ancIdentical[\"Accession ID\"].isin(allSNP[\"Accession ID\"])]\n",
    "ancIdentical.to_csv(os.path.join(snpDir, \"identical.csv\"), index=False)"
   ]
  }
 ],
 "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.5.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}