{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# GTDB processing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Background"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[GTDB](http://gtdb.ecogenomic.org/) (Genome Taxonomy Database) is a large-scale effort of curating microbial taxonomy based on phylogeny. The dataset includes NCBI RefSeq (GCF) and GenBank (GCA) genomes, as well as genomes they assembled from publicly available metagenomes (UBA). The current release - r86.1 - contains 125,243 bacterial genomes and 2,075 archaeal genomes, of which 27,372 bacterial genomes and 1,569 archaeal genomes are represented in two phylogenetic trees, respectively. The data are available at the [website](http://gtdb.ecogenomic.org/downloads)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Preparation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "from skbio import TreeNode"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Input files"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "GTDB r86.1 release"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "bac_metadata_fp = 'bac_metadata_r86.tsv.xz'\n",
    "bac_taxonomy_fp = 'bac_taxonomy_r86.tsv.xz'\n",
    "bac_tree_fp = 'bac120_r86.1.tree'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "arc_metadata_fp = 'arc_metadata_r86.tsv.xz'\n",
    "arc_taxonomy_fp = 'arc_taxonomy_r86.tsv.xz'\n",
    "arc_tree_fp = 'arc122_r86.1.tree'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "WoL data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "wgs_fp = '../../../genomes/glists/fna_prok.txt'\n",
    "wgis_fp = '../../../genomes/glists/in.txt'\n",
    "wtree_fp = '../../../trees/astral/newick/astral.nid.nwk'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Examination of GTDB files"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Three files are relevant to this study (exemplified here are `bac`, while there are also `arc`):"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Metadata: `bac_metadata_r86.tsv`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "accession\tscaffold_count\tgc_count\tlongest_scaffold\tgc_percentage\r\n",
      "RS_GCF_001245025.1\t157\t2457823\t329958\t52.2585149986\r\n",
      "RS_GCF_000678935.1\t4\t2880443\t3928455\t65.5995847821\r\n",
      "RS_GCF_000020485.1\t1\t976684\t2578146\t37.8831920302\r\n",
      "RS_GCF_001206855.1\t75\t837175\t119381\t39.5136727904\r\n"
     ]
    }
   ],
   "source": [
    "!xzcat $bac_metadata_fp | cut -f1-5 | head -n5"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Taxonomy: `bac_taxonomy_r86.tsv`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "RS_GCF_001300075.1\td__Bacteria;p__Firmicutes;c__Bacilli;o__Lactobacillales;f__Streptococcaceae;g__St\r\n",
      "RS_GCF_001245025.1\td__Bacteria;p__Proteobacteria;c__Gammaproteobacteria;o__Enterobacterales;f__Enter\r\n",
      "RS_GCF_000678935.1\td__Bacteria;p__Actinobacteriota;c__Actinobacteria;o__Corynebacteriales;f__Coryneb\r\n",
      "RS_GCF_000020485.1\td__Bacteria;p__Firmicutes_F;c__Halanaerobiia;o__Halanaerobiales;f__Halothermotric\r\n",
      "RS_GCF_001206855.1\td__Bacteria;p__Firmicutes;c__Bacilli;o__Lactobacillales;f__Streptococcaceae;g__St\r\n"
     ]
    }
   ],
   "source": [
    "!xzcat $bac_taxonomy_fp | cut -c1-100 | head -n5"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Tree: `bac120_r86.1.tree`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(((((((((((((((((((((((((((((((((((((((((((((((((((((((((((((((((((((((RS_GCF_900110195.1:1e-05,RS_GCF_900187485.1:1e-05)100.0:1e-05,RS_GCF_000783395.1:2e-05)64.0:1e-05,RS_GCF_000802985.1:2e-05)92.0:1e-05,(RS_GCF_002003425.1:1e-05,RS_GCF_900187515.1:2e-05)60.0:1e-05)39.0:1e-05,RS_GCF_000276585.1:2e-\r\n"
     ]
    }
   ],
   "source": [
    "!cat $bac_tree_fp | cut -c1-300"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Processing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### GTDB identifiers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Read GTDB genome IDs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "127318"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "arc_ = !xzcat $arc_metadata_fp | cut -f1 | tail -n+2\n",
    "bac_ = !xzcat $bac_metadata_fp | cut -f1 | tail -n+2\n",
    "ggs = arc_ + bac_\n",
    "len(ggs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Extract UBA genomes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3130"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "uba = sorted([x for x in ggs if x.startswith('UBA')], key=lambda x: int(x[3:]))\n",
    "len(uba)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "UBA7904, UBA7905, UBA7906, UBA7907, UBA7908\n"
     ]
    }
   ],
   "source": [
    "print(', '.join(uba[:5]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Generate a GTDB to WoL genome ID map. e.g.: `GB_GCA_000006155.2` => `GC000006155`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "124188"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "g2w = {x: 'G%s' % x.split('_')[-1].split('.')[0]\n",
    "       for x in ggs if x.startswith('GB_GCA_')\n",
    "       or x.startswith('RS_GCF_')}\n",
    "len(g2w)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check for duplicate genomes in GTDB (same assembly ID represented by both RefSeq and GenBank versions)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'G000009205', 'G000821245', 'G001645235'}"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gdups = set()\n",
    "used = set()\n",
    "for g, w in g2w.items():\n",
    "    if w in used:\n",
    "        gdups.add(w)\n",
    "    else:\n",
    "        used.add(w)\n",
    "gdups"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Read WoL genome IDs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "86200"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "wgs = !cat ../../../genomes/glists/fna_prok.txt\n",
    "len(wgs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Identify shared genome IDs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "83466"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sgs = set(wgs).intersection(set(g2w.values()))\n",
    "len(sgs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Read genome IDs in WoL tree."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "10575"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "wigs = !cat ../../../genomes/glists/in.txt\n",
    "len(wigs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Identify shared genome IDs in the tree."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "9732"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sigs = sgs.intersection(wigs)\n",
    "len(sigs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Export WoL-to-GTDB translation table. If there are duplicates (GCF (RefSeq) + GCA (GenBank)), select the GCF."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "G000009205: RS_GCF_000009205.1\n",
      "G000009205: GB_GCA_000009205.2\n",
      "G000821245: RS_GCF_000821245.1\n",
      "G000821245: GB_GCA_000821245.2\n",
      "G001645235: GB_GCA_001645235.2\n",
      "G001645235: RS_GCF_001645235.1\n"
     ]
    }
   ],
   "source": [
    "w2g = {}\n",
    "with open('glists/wol2gtdb.txt', 'w') as f:\n",
    "    for g, w in sorted(g2w.items(), key=lambda x: x[1]):\n",
    "        if w in gdups:\n",
    "            print('%s: %s' % (w, g))\n",
    "            if 'GCF' not in g:\n",
    "                continue\n",
    "        w2g[w] = g\n",
    "        f.write('%s\\t%s\\n' % (w, g))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Taxonomy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "ranks = ['domain', 'phylum', 'class', 'order', 'family', 'genus', 'species']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load GTDB taxonomy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "taxfile = !xzcat $arc_taxonomy_fp $bac_taxonomy_fp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "127318"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data = {}\n",
    "for line in taxfile:\n",
    "    g, lineage = line.split('\\t')\n",
    "    data[g] = {}\n",
    "    taxa = lineage.split(';')\n",
    "    if len(taxa) != len(ranks):\n",
    "        raise ValueError(taxa)\n",
    "    for i, taxon in enumerate(taxa):\n",
    "        if len(taxon) < 3:\n",
    "            raise ValueError(taxon)\n",
    "        if taxon[1:3] != '__':\n",
    "            raise ValueError(taxon)\n",
    "        if taxon[0] != ranks[i][0]:\n",
    "            raise ValueError(taxon)\n",
    "        data[g][ranks[i]] = taxon[3:]\n",
    "len(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "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>domain</th>\n",
       "      <th>phylum</th>\n",
       "      <th>class</th>\n",
       "      <th>order</th>\n",
       "      <th>family</th>\n",
       "      <th>genus</th>\n",
       "      <th>species</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>genome</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>GB_GCA_000006155.2</th>\n",
       "      <td>Bacteria</td>\n",
       "      <td>Firmicutes</td>\n",
       "      <td>Bacilli</td>\n",
       "      <td>Bacillales</td>\n",
       "      <td>Bacillaceae_G</td>\n",
       "      <td>Bacillus_A</td>\n",
       "      <td>Bacillus_A anthracis</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>GB_GCA_000007185.1</th>\n",
       "      <td>Archaea</td>\n",
       "      <td>Euryarchaeota</td>\n",
       "      <td>Methanopyri</td>\n",
       "      <td>Methanopyrales</td>\n",
       "      <td>Methanopyraceae</td>\n",
       "      <td>Methanopyrus</td>\n",
       "      <td>Methanopyrus kandleri</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>GB_GCA_000007225.1</th>\n",
       "      <td>Archaea</td>\n",
       "      <td>Crenarchaeota</td>\n",
       "      <td>Thermoprotei</td>\n",
       "      <td>Thermoproteales</td>\n",
       "      <td>Thermoproteaceae</td>\n",
       "      <td>Pyrobaculum</td>\n",
       "      <td>Pyrobaculum aerophilum</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>GB_GCA_000007385.1</th>\n",
       "      <td>Bacteria</td>\n",
       "      <td>Proteobacteria</td>\n",
       "      <td>Gammaproteobacteria</td>\n",
       "      <td>Xanthomonadales</td>\n",
       "      <td>Xanthomonadaceae</td>\n",
       "      <td>Xanthomonas</td>\n",
       "      <td>Xanthomonas oryzae</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>GB_GCA_000007405.1</th>\n",
       "      <td>Bacteria</td>\n",
       "      <td>Proteobacteria</td>\n",
       "      <td>Gammaproteobacteria</td>\n",
       "      <td>Enterobacterales</td>\n",
       "      <td>Enterobacteriaceae</td>\n",
       "      <td>Escherichia</td>\n",
       "      <td>Escherichia coli</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                      domain          phylum                class  \\\n",
       "genome                                                              \n",
       "GB_GCA_000006155.2  Bacteria      Firmicutes              Bacilli   \n",
       "GB_GCA_000007185.1   Archaea   Euryarchaeota          Methanopyri   \n",
       "GB_GCA_000007225.1   Archaea   Crenarchaeota         Thermoprotei   \n",
       "GB_GCA_000007385.1  Bacteria  Proteobacteria  Gammaproteobacteria   \n",
       "GB_GCA_000007405.1  Bacteria  Proteobacteria  Gammaproteobacteria   \n",
       "\n",
       "                               order              family         genus  \\\n",
       "genome                                                                   \n",
       "GB_GCA_000006155.2        Bacillales       Bacillaceae_G    Bacillus_A   \n",
       "GB_GCA_000007185.1    Methanopyrales     Methanopyraceae  Methanopyrus   \n",
       "GB_GCA_000007225.1   Thermoproteales    Thermoproteaceae   Pyrobaculum   \n",
       "GB_GCA_000007385.1   Xanthomonadales    Xanthomonadaceae   Xanthomonas   \n",
       "GB_GCA_000007405.1  Enterobacterales  Enterobacteriaceae   Escherichia   \n",
       "\n",
       "                                   species  \n",
       "genome                                      \n",
       "GB_GCA_000006155.2    Bacillus_A anthracis  \n",
       "GB_GCA_000007185.1   Methanopyrus kandleri  \n",
       "GB_GCA_000007225.1  Pyrobaculum aerophilum  \n",
       "GB_GCA_000007385.1      Xanthomonas oryzae  \n",
       "GB_GCA_000007405.1        Escherichia coli  "
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gtax = pd.DataFrame().from_dict(data, orient='index')\n",
    "gtax.index.name = 'genome'\n",
    "gtax.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "gtax.to_csv('ranks/gtdb.tsv', sep='\\t')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note: The GTDB taxonomy is filled in all cells (i.e., there is no unclassified ranks)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "domain: 2\n",
      "phylum: 125\n",
      "class: 324\n",
      "order: 873\n",
      "family: 2025\n",
      "genus: 6918\n",
      "species: 11551\n"
     ]
    }
   ],
   "source": [
    "for rank in gtax.columns:\n",
    "    print('%s: %d' % (rank, gtax[rank].nunique()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Downsample to genomes with WoL-style IDs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "124188"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gwtax = gtax.loc[sorted(g2w)]\n",
    "gwtax['gid'] = gwtax.index.to_series().map(g2w)\n",
    "gwtax.shape[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Make sure duplicated genome IDs have identical taxonomy annotations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['G001645235', 'G000009205', 'G000821245']"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vc = gwtax['gid'].value_counts()\n",
    "dups = vc[vc > 1].index.tolist()\n",
    "dups"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "for dup in dups:\n",
    "    rows = gwtax[gwtax.gid == dup]\n",
    "    for rank in ranks:\n",
    "        if len(set(rows[rank].tolist())) > 1:\n",
    "            raise ValueError('%s has conflicting taxonomy annotations.' % dup)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Subset and translate the genome IDs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "124185"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gwtax = gwtax.sort_values('gid').drop_duplicates('gid', keep='first').set_index('gid')\n",
    "gwtax.shape[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "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>domain</th>\n",
       "      <th>phylum</th>\n",
       "      <th>class</th>\n",
       "      <th>order</th>\n",
       "      <th>family</th>\n",
       "      <th>genus</th>\n",
       "      <th>species</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>gid</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>G000003135</th>\n",
       "      <td>Bacteria</td>\n",
       "      <td>Actinobacteriota</td>\n",
       "      <td>Actinobacteria</td>\n",
       "      <td>Actinomycetales</td>\n",
       "      <td>Bifidobacteriaceae</td>\n",
       "      <td>Bifidobacterium</td>\n",
       "      <td>Bifidobacterium longum</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>G000003215</th>\n",
       "      <td>Bacteria</td>\n",
       "      <td>Firmicutes_A</td>\n",
       "      <td>Clostridia</td>\n",
       "      <td>Peptostreptococcales</td>\n",
       "      <td>Peptostreptococcaceae</td>\n",
       "      <td>Clostridioides</td>\n",
       "      <td>Clostridioides difficile</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>G000003645</th>\n",
       "      <td>Bacteria</td>\n",
       "      <td>Firmicutes</td>\n",
       "      <td>Bacilli</td>\n",
       "      <td>Bacillales</td>\n",
       "      <td>Bacillaceae_G</td>\n",
       "      <td>Bacillus_A</td>\n",
       "      <td>Bacillus_A cereus_AE</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>G000003925</th>\n",
       "      <td>Bacteria</td>\n",
       "      <td>Firmicutes</td>\n",
       "      <td>Bacilli</td>\n",
       "      <td>Bacillales</td>\n",
       "      <td>Bacillaceae_G</td>\n",
       "      <td>Bacillus_A</td>\n",
       "      <td>Bacillus_A mycoides</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>G000003955</th>\n",
       "      <td>Bacteria</td>\n",
       "      <td>Firmicutes</td>\n",
       "      <td>Bacilli</td>\n",
       "      <td>Bacillales</td>\n",
       "      <td>Bacillaceae_G</td>\n",
       "      <td>Bacillus_A</td>\n",
       "      <td>Bacillus_A cereus_P</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "              domain            phylum           class                 order  \\\n",
       "gid                                                                            \n",
       "G000003135  Bacteria  Actinobacteriota  Actinobacteria       Actinomycetales   \n",
       "G000003215  Bacteria      Firmicutes_A      Clostridia  Peptostreptococcales   \n",
       "G000003645  Bacteria        Firmicutes         Bacilli            Bacillales   \n",
       "G000003925  Bacteria        Firmicutes         Bacilli            Bacillales   \n",
       "G000003955  Bacteria        Firmicutes         Bacilli            Bacillales   \n",
       "\n",
       "                           family            genus                   species  \n",
       "gid                                                                           \n",
       "G000003135     Bifidobacteriaceae  Bifidobacterium    Bifidobacterium longum  \n",
       "G000003215  Peptostreptococcaceae   Clostridioides  Clostridioides difficile  \n",
       "G000003645          Bacillaceae_G       Bacillus_A      Bacillus_A cereus_AE  \n",
       "G000003925          Bacillaceae_G       Bacillus_A       Bacillus_A mycoides  \n",
       "G000003955          Bacillaceae_G       Bacillus_A       Bacillus_A cereus_P  "
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gwtax.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "domain: 2\n",
      "phylum: 125\n",
      "class: 317\n",
      "order: 848\n",
      "family: 1943\n",
      "genus: 6655\n",
      "species: 11463\n"
     ]
    }
   ],
   "source": [
    "for rank in gwtax.columns:\n",
    "    print('%s: %d' % (rank, gwtax[rank].nunique()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "gwtax.to_csv('ranks/gtdb.wid.tsv', sep='\\t')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Downsample to genomes shared between GTDB and WoL."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "domain: 2\n",
      "phylum: 103\n",
      "class: 241\n",
      "order: 600\n",
      "family: 1297\n",
      "genus: 4380\n",
      "species: 9615\n"
     ]
    }
   ],
   "source": [
    "gwtax = gwtax.loc[sorted(sgs)]\n",
    "for rank in gwtax.columns:\n",
    "    print('%s: %d' % (rank, gwtax[rank].nunique()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "gwtax.to_csv('ranks/gtdb.wid.prok.tsv', sep='\\t')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Downsample to genomes in the WoL tree."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "domain: 2\n",
      "phylum: 94\n",
      "class: 215\n",
      "order: 523\n",
      "family: 1111\n",
      "genus: 3619\n",
      "species: 5503\n"
     ]
    }
   ],
   "source": [
    "gwtax = gwtax.loc[sorted(sigs)]\n",
    "for rank in gwtax.columns:\n",
    "    print('%s: %d' % (rank, gwtax[rank].nunique()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "gwtax.to_csv('ranks/gtdb.wid.in.tsv', sep='\\t')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "List additional phyla in GTDB but not in the WoL tree."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'4572-55',\n",
       " 'AABM5-125-24',\n",
       " 'BRC1',\n",
       " 'CG03',\n",
       " 'CG2-30-53-67',\n",
       " 'CG2-30-70-394',\n",
       " 'Desantisbacteria',\n",
       " 'Entotheonellota',\n",
       " 'Eremiobacterota',\n",
       " 'Fermentibacterota',\n",
       " 'Firestonebacteria',\n",
       " 'Goldbacteria-1',\n",
       " 'OLB16',\n",
       " 'RBG-13-61-14',\n",
       " 'UAP2',\n",
       " 'UBA2233',\n",
       " 'UBA3054',\n",
       " 'UBA5359',\n",
       " 'UBA6262',\n",
       " 'UBA8481',\n",
       " 'UBP1',\n",
       " 'UBP10',\n",
       " 'UBP12',\n",
       " 'UBP13',\n",
       " 'UBP14',\n",
       " 'UBP17',\n",
       " 'UBP18',\n",
       " 'UBP3',\n",
       " 'UBP4',\n",
       " 'UBP6',\n",
       " 'UBP7'}"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "set(gtax['phylum'].unique()) - set(gwtax['phylum'].unique())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Phylogeny"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Read GTDB trees"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "10575"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "wtree = TreeNode.read(wtree_fp)\n",
    "wtree.count(tips=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1569"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gatree = TreeNode.read(arc_tree_fp, convert_underscores=False)\n",
    "gatree.count(tips=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "27372"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gbtree = TreeNode.read(bac_tree_fp, convert_underscores=False)\n",
    "gbtree.count(tips=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Export trees without internal node labels but only support values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'trees/arc122_r86.1.bs_only.nwk'"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gatree_ = gatree.copy()\n",
    "for node in gatree_.non_tips(include_self=True):\n",
    "    if node.name:\n",
    "        node.name = node.name.split(':')[0].split('.')[0]\n",
    "gatree_.write('trees/arc122_r86.1.bs_only.nwk')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'trees/bac120_r86.1.bs_only.nwk'"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gbtree_ = gbtree.copy()\n",
    "for node in gbtree_.non_tips(include_self=True):\n",
    "    if node.name:\n",
    "        node.name = node.name.split(':')[0].split('.')[0]\n",
    "gbtree_.write('trees/bac120_r86.1.bs_only.nwk')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Overlap taxon sets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Get genome IDs in GTDB trees."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "gatgs = gatree.subset()\n",
    "gbtgs = gbtree.subset()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "28941"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gtgs = sorted(list(gatgs) + list(gbtgs))\n",
    "len(gtgs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GB_GCA_000007185.1, GB_GCA_000007225.1, GB_GCA_000008085.1, GB_GCA_000010565.1, GB_GCA_000011005.1\n"
     ]
    }
   ],
   "source": [
    "print(', '.join(gtgs[:5]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Export GTDB tree taxon list."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('glists/gid_in_gtrees.txt', 'w') as f:\n",
    "    for g in gtgs:\n",
    "        f.write('%s\\n' % g)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Get WoL-style genome IDs of GTDB tree taxa."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "27222"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gtwgs = sorted([g2w[x] for x in gtgs if x in g2w])\n",
    "len(gtwgs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Get GTDB tree taxa that are in the WoL genome pool."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "17514"
      ]
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sgtwgs = sgs.intersection(gtwgs)\n",
    "len(sgtwgs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('glists/wid_in_gtrees.prok.txt', 'w') as f:\n",
    "    for g in sorted(sgtwgs):\n",
    "        f.write('%s\\n' % g)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Get GTDB tree taxa that are in the WoL tree."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "8042"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sgtwigs = sorted(set(wigs).intersection(gtwgs))\n",
    "len(sgtwigs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('glists/wid_in_gtrees.in.txt', 'w') as f:\n",
    "    for g in sorted(sgtwigs):\n",
    "        f.write('%s\\n' % g)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Shrink the two trees to contain common taxa only."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "486"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gatgs_to_keep = [x for x in gatgs if x in g2w and g2w[x] in sgtwigs]\n",
    "gatree_ = gatree_.shear(gatgs_to_keep)\n",
    "gatree_.count(tips=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'trees/arc122_r86.1.lap.nwk'"
      ]
     },
     "execution_count": 51,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gatree_.write('trees/arc122_r86.1.lap.nwk')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'trees/arc122_r86.1.lap.wid.nwk'"
      ]
     },
     "execution_count": 52,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "for tip in gatree_.tips():\n",
    "    tip.name = g2w[tip.name]\n",
    "gatree_.write('trees/arc122_r86.1.lap.wid.nwk')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "7556"
      ]
     },
     "execution_count": 53,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gbtgs_to_keep = [x for x in gbtgs if x in g2w and g2w[x] in sgtwigs]\n",
    "gbtree_ = gbtree_.shear(gbtgs_to_keep)\n",
    "gbtree_.count(tips=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'trees/bac120_r86.1.lap.nwk'"
      ]
     },
     "execution_count": 54,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gbtree_.write('trees/bac120_r86.1.lap.nwk')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'trees/bac120_r86.1.lap.wid.nwk'"
      ]
     },
     "execution_count": 55,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "for tip in gbtree_.tips():\n",
    "    tip.name = g2w[tip.name]\n",
    "gbtree_.write('trees/bac120_r86.1.lap.wid.nwk')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Arbitrarily connect two trees at roots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "8042"
      ]
     },
     "execution_count": 56,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gtree_ = TreeNode(children=[gatree_, gbtree_])\n",
    "gtree_.count(tips=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'trees/joint.lap.wid.nwk'"
      ]
     },
     "execution_count": 57,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gtree_.write('trees/joint.lap.wid.nwk')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "8042"
      ]
     },
     "execution_count": 58,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "wtree_ = wtree.shear(sgtwigs)\n",
    "wtree_.count(tips=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'trees/astral.lap.nwk'"
      ]
     },
     "execution_count": 59,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "wtree_.write('trees/astral.lap.nwk')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'trees/astral.lap.gid.nwk'"
      ]
     },
     "execution_count": 60,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "t_ = wtree_.copy()\n",
    "for tip in t_.tips():\n",
    "    tip.name = w2g[tip.name]\n",
    "t_.write('trees/astral.lap.gid.nwk')"
   ]
  }
 ],
 "metadata": {
  "hide_input": false,
  "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.7.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}