{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Subsample taxa in phylogeny"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Overview"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Goal**: Generate a smaller group of taxa from the original 10,575 genomes in the phylogeny, which can then be used for phylogenetic analyses using more robust methods. Because the original phylogeny is already obtained, the subsampling process will take advantage of the distance and diversity indicated by the phylogeny.\n",
    "\n",
    "Two strategies are used here:\n",
    "\n",
    "I. Use **prototype selection** based on the phylogenetic distance (sum of lengths of branches connecting each pair of tips).\n",
    "\n",
    "- Pro: Maximizes the sum of phylogenetic distance (hence the largest diversity).\n",
    "- Con: Favors long branches, and multiple closely related taxa at the far end of a long branch.\n",
    "\n",
    "II. Select clades with the smallest relative evolutionary divergence (**RED**) ([Parks, et al., 2018](https://www.nature.com/articles/nbt.4229)). Then select one taxon within each clade.\n",
    "\n",
    "- Pro: Ensures that the basal phylogeny is sufficiently represented.\n",
    "- Con: Misses long branches, such as the single-taxon phyla.\n",
    "\n",
    "With strategy II, the following criteria are applied sequentially to further select one taxon within each clade:\n",
    "\n",
    "1. Contains the most marker genes.\n",
    "2. Contamination level is the lowest.\n",
    "3. DNA quality score is the highest.\n",
    "\n",
    "The process terminates whenever there is one candidate left. If not after all three criteria are applied, the program will randomly choose one taxon.\n",
    "\n",
    "The following two rules are optional:\n",
    "\n",
    "1. Must be represented in the r-protein tree.\n",
    "2. Must have a standard Latin species name.\n",
    "\n",
    "Finally, for specific applications, manual curation is involved.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from random import seed, choice"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import scipy as sp\n",
    "import pandas as pd\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "from skbio.tree import TreeNode\n",
    "from skbio.stats.distance import DistanceMatrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "seed(42)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Number of taxa to keep"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "ntaxa = 92"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Input files"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Original, full-scale phylogenetic tree (10,575 taxa)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "tree_fp = 'astral.cons.nid.nwk'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "10575"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tree = TreeNode.read('astral.cons.nid.nwk')\n",
    "tree.count(tips=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Branch support values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "supports_fp = 'astral.supports.tsv'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": false
   },
   "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>EN</th>\n",
       "      <th>LPP</th>\n",
       "      <th>QT</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>#node</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>N2</th>\n",
       "      <td>196.0</td>\n",
       "      <td>0.998406</td>\n",
       "      <td>0.450953</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>N3</th>\n",
       "      <td>196.0</td>\n",
       "      <td>0.998406</td>\n",
       "      <td>0.450953</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>N4</th>\n",
       "      <td>124.0</td>\n",
       "      <td>0.999993</td>\n",
       "      <td>0.535120</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "          EN       LPP        QT\n",
       "#node                           \n",
       "N2     196.0  0.998406  0.450953\n",
       "N3     196.0  0.998406  0.450953\n",
       "N4     124.0  0.999993  0.535120"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dfs = pd.read_table(supports_fp, index_col=0)\n",
    "dfs.head(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Ignore low-support branches (EN <= 5 or LPP <= 0.5 (abbr.: e5p50))."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "en_th = 5\n",
    "lpp_th = 0.5"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Genome metadata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "genomes_fp = '../../genomes/metadata.tsv.xz'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Index(['asm_name', 'assembly_accession', 'bioproject', 'biosample',\n",
       "       'wgs_master', 'seq_rel_date', 'submitter', 'ftp_path', 'img_id',\n",
       "       'gtdb_id', 'scope', 'assembly_level', 'genome_rep', 'refseq_category',\n",
       "       'release_type', 'taxid', 'species_taxid', 'organism_name',\n",
       "       'infraspecific_name', 'isolate', 'superkingdom', 'phylum', 'class',\n",
       "       'order', 'family', 'genus', 'species', 'classified', 'unique_name',\n",
       "       'lv1_group', 'lv2_group', 'score_faa', 'score_fna', 'score_rrna',\n",
       "       'score_trna', 'total_length', 'contigs', 'gc', 'n50', 'l50', 'proteins',\n",
       "       'protein_length', 'coding_density', 'completeness', 'contamination',\n",
       "       'strain_heterogeneity', 'markers', '5s_rrna', '16s_rrna', '23s_rrna',\n",
       "       'trnas', 'draft_quality'],\n",
       "      dtype='object')"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dfg = pd.read_table(genomes_fp, index_col=0)\n",
    "dfg.columns"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Taxa in the r-protein tree."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "rpls_fp = 'rpls.txt'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "9814"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "with open(rpls_fp, 'r') as f:\n",
    "    rpls = set(f.read().splitlines())\n",
    "len(rpls)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "melai_fp = 'dte/Melainabacteria.txt'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "28"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "with open(melai_fp, 'r') as f:\n",
    "    melai = set(f.read().splitlines())\n",
    "len(melai)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Helper functions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Prototype selection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def prototype_selection_destructive_maxdist(dm, num_prototypes, seedset=None):\n",
    "    \"\"\"Prototype selection function (minified).\"\"\"\n",
    "    numRemain = len(dm.ids)\n",
    "    currDists = dm.data.sum(axis=1)\n",
    "    maxVal = currDists.max()\n",
    "    if seedset is not None:\n",
    "        for e in seedset:\n",
    "            currDists[dm.index(e)] = maxVal * 2\n",
    "    minElmIdx = currDists.argmin()\n",
    "    currDists[minElmIdx], numRemain = np.infty, numRemain - 1\n",
    "    while (numRemain > num_prototypes):\n",
    "        currDists -= dm.data[minElmIdx]\n",
    "        minElmIdx = currDists.argmin()\n",
    "        currDists[minElmIdx], numRemain = np.infty, numRemain - 1\n",
    "    return [dm.ids[idx]\n",
    "            for idx, dist in enumerate(currDists)\n",
    "            if dist != np.infty]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "RED calculation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def calc_brlen_metrics(tree):\n",
    "    \"\"\"Calculate branch length-related metrics.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    tree : skbio.TreeNode\n",
    "\n",
    "    Notes\n",
    "    -----\n",
    "    The following metrics are calculated:\n",
    "    \n",
    "        - height: Sum of branch lengths from the root to the node.\n",
    "\n",
    "        - depths: Sums of branch lengths from all descendants to current node.\n",
    "\n",
    "        - red: Relative evolutionary divergence (RED), introduced by Parks,\n",
    "          et al., 2018, Nat Biotechnol.\n",
    "\n",
    "              RED = p + (d / u) * (1 - p)\n",
    "\n",
    "          where p = RED of parent, d = length, u = mean depth of parent\n",
    "    \n",
    "    Metrics will be appended to each node of tree in place.\n",
    "    \"\"\"\n",
    "    # calculate depths\n",
    "    for node in tree.postorder(include_self=True):\n",
    "        if node.name is None:\n",
    "            raise ValueError('Error: Found an unnamed node.')\n",
    "        if node.length is None:\n",
    "            node.length = 0.0\n",
    "        if node.is_tip():\n",
    "            node.depths = [0.0]\n",
    "            node.taxa = [node.name]\n",
    "        else:\n",
    "            node.depths = [\n",
    "                y + x.length for x in node.children for y in x.depths]\n",
    "            node.taxa = sorted(set().union(*[x.taxa for x in node.children]))\n",
    "\n",
    "    # calculate heights and REDs\n",
    "    for node in tree.preorder(include_self=True):\n",
    "        if node.is_root():\n",
    "            node.height = 0.0\n",
    "            node.red = 0.0\n",
    "        else:\n",
    "            node.height = node.parent.height + node.length\n",
    "            if node.is_tip():\n",
    "                node.red = 1.0\n",
    "            else:\n",
    "                node.red = node.parent.red + node.length \\\n",
    "                    / (node.length + sum(node.depths) / len(node.depths)) \\\n",
    "                    * (1 - node.parent.red)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Latin species name checker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def is_latin(str):\n",
    "    if str == '':\n",
    "        return False\n",
    "    elif str.count(' ') != 1:\n",
    "        return False\n",
    "    str_ = str.replace(' ', '')\n",
    "    if not str_.istitle():\n",
    "        return False\n",
    "    elif not str_.isalpha():\n",
    "        return False\n",
    "    return True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Analysis"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Optional: Filter tree to keep taxa with Latin species name."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "g2species = dfg['species'].to_dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "%%script false\n",
    "to_keep = [x for x in tree.subset() if is_latin(g2species[x])]\n",
    "tree = tree.shear(to_keep)\n",
    "print('Tree has %d taxa with Latin species names.' % tree.count(tips=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Optional: Filter tree to keep taxa represented in the r-protein tree."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tree has 9814 taxa represented in the r-protein tree.\n"
     ]
    }
   ],
   "source": [
    "to_keep = [x for x in tree.subset() if x in rpls]\n",
    "tree = tree.shear(to_keep)\n",
    "print('Tree has %d taxa represented in the r-protein tree.' % tree.count(tips=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Solution I: Prototype selection based on phylogenetic distance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "%%script false\n",
    "print('Tree has %d taxa.' % tree.count(tips=True))\n",
    "\n",
    "print('Calculating tip-to-tip distances...')\n",
    "dm = tree.tip_tip_distances()\n",
    "print('Sum of distances: %d.' % np.tril(dm.data).sum())\n",
    "\n",
    "print('Performing prototype selection...')\n",
    "prototypes = prototype_selection_destructive_maxdist(dm, ntaxa)\n",
    "print('Downsampled to %d taxa.' % len(prototypes))\n",
    "print('Sum of distances: %d.' % np.tril(dm.filter(prototypes).data).sum())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "%%script false\n",
    "tout = tree.shear(prototypes)\n",
    "tout.write('proto.%d.nwk' % ntaxa)\n",
    "\n",
    "with open('proto.%d.txt' % ntaxa, 'w') as f:\n",
    "    for g in sorted(prototypes):\n",
    "        print(g, file=f)\n",
    "\n",
    "gs = prototypes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Solution II: Choose clades by minimizing RED."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Calculate branch length-related metrics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "calc_brlen_metrics(tree)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "19627"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data = {}\n",
    "for node in tree.postorder(include_self=True):\n",
    "    data[node.name] = {\n",
    "        'parent': None if node.is_root() else node.parent.name,\n",
    "        'taxa': node.taxa,\n",
    "        'red': node.red,\n",
    "    }\n",
    "len(data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Filter internal nodes by branch support value."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "10027"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "valid = set(dfs.query('EN > %d and LPP > %f' % (en_th, lpp_th)).index.tolist())\n",
    "len(valid)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A helper function to remove any ancestral node of current node from the pool (so that the chosen clades do not overlap)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def remove_ancestor(nid, chosen, data):\n",
    "    cid = nid\n",
    "    while True:\n",
    "        if cid in chosen:\n",
    "            chosen.remove(cid)\n",
    "            return cid\n",
    "        else:\n",
    "            cid = data[cid]['parent']\n",
    "            if cid is None:\n",
    "                return None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Iterate from the low end of the RED list, adding clades to the pool, until the desired number of taxa is reached."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "chosen_clades = set()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "92"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "for nid in sorted(data, key=lambda x: data[x]['red']):\n",
    "    if nid.startswith('G') or nid in valid:\n",
    "        remove_ancestor(nid, chosen_clades, data)\n",
    "        chosen_clades.add(nid)\n",
    "        if len(chosen_clades) == ntaxa:\n",
    "            break\n",
    "len(chosen_clades)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Within each clade, select one taxon by sequentially applying the following criteria:\n",
    "1. Contains the most marker genes.\n",
    "2. Contamination level is the lowest.\n",
    "3. DNA quality score is the highest.\n",
    "4. Randomly choose one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "chosen_taxa = {}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Criterium 1: Most marker genes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "g2markers = dfg['markers'].to_dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def most_markers(gs, g2markers):\n",
    "    max_gs = []\n",
    "    max_markers = 0\n",
    "    for g in sorted(gs, key=lambda x: g2markers[x], reverse=True):\n",
    "        if max_markers == 0:\n",
    "            max_markers = g2markers[g]\n",
    "            max_gs.append(g)\n",
    "        elif max_markers == g2markers[g]:\n",
    "            max_gs.append(g)\n",
    "        else:\n",
    "            break\n",
    "    return max_gs, max_markers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Criterium 2: Lowest contamination."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "g2contam = dfg['contamination'].to_dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def least_contaminated(gs, g2contam):\n",
    "    min_gs = []\n",
    "    min_contam = None\n",
    "    for g in sorted(gs, key=lambda x: g2contam[x]):\n",
    "        if min_contam is None:\n",
    "            min_contam = g2contam[g]\n",
    "            min_gs.append(g)\n",
    "        elif min_contam == g2contam[g]:\n",
    "            min_gs.append(g)\n",
    "        else:\n",
    "            break\n",
    "    return min_gs, min_contam"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Criterium 3: Highest DNA quality score."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "g2dnaqty = dfg['score_fna'].to_dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def best_dna(gs, g2dnaqty):\n",
    "    max_gs = []\n",
    "    max_dnaqty = 0\n",
    "    for g in sorted(gs, key=lambda x: g2dnaqty[x], reverse=True):\n",
    "        if max_dnaqty == 0:\n",
    "            max_dnaqty = g2dnaqty[g]\n",
    "            max_gs.append(g)\n",
    "        elif max_dnaqty == g2dnaqty[g]:\n",
    "            max_gs.append(g)\n",
    "        else:\n",
    "            break\n",
    "    return max_gs, max_dnaqty"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Perform taxon selection."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "chosen_taxa = {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Clade N250: Equally good: G000007005, G900079115.\n",
      "Clade N1637: Equally good: G000283655, G001305595.\n"
     ]
    }
   ],
   "source": [
    "for nid in chosen_clades:\n",
    "    gs = data[nid]['taxa']\n",
    "\n",
    "    # optional: r-protein\n",
    "#     gs = rpls.intersection(gs)\n",
    "#     if len(gs) == 0:\n",
    "#         raise ValueError('%s: No taxon is represented in the r-protein tree.' % nid)\n",
    "\n",
    "    # optional: Latin name\n",
    "#     gs_ = [x for x in gs if is_latin(g2species[x])]\n",
    "#     if len(gs_) == 0:\n",
    "#         print('%s: No taxon has a Latin species name.' % nid)\n",
    "#     else:\n",
    "#         gs = gs_\n",
    "\n",
    "    # maximize marker count\n",
    "    gs, max_markers = most_markers(gs, g2markers)\n",
    "\n",
    "    if len(gs) > 1:\n",
    "        # minimize contamination\n",
    "        gs, min_contam = least_contaminated(gs, g2contam)\n",
    "\n",
    "    if len(gs) > 1:\n",
    "        # maximize DNA quality\n",
    "        gs, max_dnaqty = best_dna(gs, g2dnaqty)\n",
    "\n",
    "    if len(gs) > 1:\n",
    "        # random choice\n",
    "        print('Clade %s: Equally good: %s.' % (nid, ', '.join(gs)))\n",
    "        g = choice(gs)\n",
    "    else:\n",
    "        g = max(gs)\n",
    "\n",
    "    chosen_taxa[nid] = g"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "N741 - G001768645: Actinobacteria bacterium RBG_13_55_18\n",
      "N1122 - G001399675: Alicyclobacillus ferrooxydans\n",
      "N118 - G002011215: Archaeoglobus sp. JdFR-24\n",
      "N966 - G000092245: Arcobacter nitrofigilis\n",
      "N1637 - G000283655: Azospirillum lipoferum\n",
      "N1334 - G000311725: Bacillus massiliosenegalensis\n",
      "N1898 - G001769855: Bdellovibrionales bacterium GWA2_49_15\n",
      "N1304 - G000281175: Caldilinea aerophila\n",
      "N3441 - G001886815: Caldithrix abyssi\n",
      "N94 - G002010385: Candidatus Acetothermia bacterium JdFR-52\n",
      "N242 - G001593935: Candidatus Bathyarchaeota archaeon B26-2\n",
      "N198 - G001774395: Candidatus Berkelbacteria bacterium RIFCSPLOWO2_01_FULL_50_28\n",
      "N109 - G000270325: Candidatus Caldiarchaeum subterraneum\n",
      "N190 - G001777435: Candidatus Daviesbacteria bacterium RIFCSPLOWO2_02_FULL_38_18\n",
      "N1361 - G001777945: Candidatus Edwardsbacteria bacterium RIFOXYD12_FULL_50_11\n",
      "N275 - G001779195: Candidatus Gottesmanbacteria bacterium RIFCSPHIGHO2_02_FULL_39_11\n",
      "N15 - G001940645: Candidatus Heimdallarchaeota archaeon LC_3\n",
      "N1607 - G002011615: Candidatus Hydrothermae bacterium JdFR-72\n",
      "N543 - G001818155: Candidatus Jacksonbacteria bacterium RIFCSPLOWO2_02_FULL_44_20\n",
      "N2149 - G000315095: Candidatus Kuenenia stuttgartiensis\n",
      "N879 - G000998635: Candidatus Magasanikbacteria bacterium GW2011_GWC2_42_27\n",
      "N145 - G001784485: Candidatus Margulisbacteria bacterium GWE2_39_32\n",
      "N3087 - G001872685: Candidatus Marinimicrobia bacterium CG1_02_48_14\n",
      "N330 - G001749745: Candidatus Marispirochaeta associata\n",
      "N72 - G001717035: Candidatus Methanomethylicus mesodigestum\n",
      "N314 - G001899315: Candidatus Obscuribacter phosphatis\n",
      "N56 - G001788755: Candidatus Peregrinibacteria bacterium RIFOXYA2_FULL_33_21\n",
      "N956 - G000406005: Candidatus Poribacteria bacterium WGA-3G\n",
      "N1876 - G001789275: Candidatus Raymondbacteria bacterium RIFOXYD12_FULL_49_13\n",
      "N1894 - G000014905: Candidatus Solibacter usitatus\n",
      "N710 - G001819985: Candidatus Spechtbacteria bacterium RIFCSPHIGHO2_02_FULL_43_15b\n",
      "N38 - G001791795: Candidatus Wallbacteria bacterium GWC2_49_35\n",
      "N9 - G001871415: Candidatus Woesearchaeota archaeon CG1_02_57_44\n",
      "N125 - G001816645: Candidatus Woykebacteria bacterium GWB1_45_5\n",
      "N144 - G000585335: Cloacibacillus evryensis\n",
      "N1330 - G001042715: Clostridium aceticum\n",
      "N36 - G000020945: Coprothermobacter proteolyticus\n",
      "N1896 - G000335475: Cystobacter fuscus\n",
      "N1551 - G000716715: Dactylosporangium aurantiacum\n",
      "N1538 - G002011495: Dehalococcoides sp. JdFR-56\n",
      "N4915 - G001595385: Deltaproteobacteria bacterium GWC2_55_46\n",
      "N2193 - G001799095: Deltaproteobacteria bacterium RIFCSPLOWO2_02_FULL_44_34\n",
      "N778 - G000025725: Denitrovibrio acetiphilus\n",
      "N3144 - G900143695: Desulfopila aestuarii\n",
      "N1556 - G000235605: Desulfosporosinus orientis\n",
      "N1555 - G900113335: Desulfotomaculum arcticum\n",
      "N1121 - G000174415: Dethiobacter alkaliphilus\n",
      "N923 - G000024265: Eggerthella lenta\n",
      "N773 - G001800475: Elusimicrobia bacterium RIFOXYB2_FULL_46_23\n",
      "N575 - G000724625: Fimbriimonas ginsengisoli\n",
      "N3447 - G000379765: Flexithrix dorotheae\n",
      "N1875 - G000522985: Gemmatirosa kalamazoonesis\n",
      "N437 - G900156285: Halanaerobium kushneri\n",
      "N1303 - G000018565: Herpetosiphon aurantiacus\n",
      "N1380 - G000242615: Holophaga foetida\n",
      "N1864 - G001803205: Lentisphaerae bacterium GWF2_38_69\n",
      "N232 - G000017685: Leptospira biflexa\n",
      "N746 - G001544015: Limnochorda pilosa\n",
      "N3088 - G001803345: Melioribacter sp. GWF2_38_21\n",
      "N112 - G000745485: Methanobacterium veterum\n",
      "N258 - G000308215: Methanomassiliicoccus luminyensis\n",
      "N267 - G000969945: Methanosarcina sp. Kolksee\n",
      "N276 - G001803415: Microgenomates group bacterium RBG_16_45_19\n",
      "N2498 - G000341545: Nitrospina gracilis\n",
      "N2499 - G000196815: Nitrospira defluvii\n",
      "N2731 - G001804395: Omnitrophica WOR_2 bacterium GWA2_45_18\n",
      "N411 - G000992445: Parcubacteria group bacterium GW2011_GWA2_38_13b\n",
      "N700 - G000999255: Parcubacteria group bacterium GW2011_GWA2_43_17\n",
      "N1080 - G001824395: Parcubacteria group bacterium RIFCSPHIGHO2_02_FULL_48_10b\n",
      "N1865 - G000172555: Pedosphaera parvula\n",
      "N752 - G000271665: Pelosinus fermentans\n",
      "N1372 - G000703085: Persephonella sp. KM09-Lau-8\n",
      "N37 - G000018605: Petrotoga mobilis\n",
      "N1641 - G001708485: Pseudomonas fluorescens\n",
      "N45 - G000263735: Pyrococcus sp. ST04\n",
      "N605 - G001469005: SAR324 cluster bacterium lautmerah10\n",
      "N1863 - G000255655: Schlesneria paludicola\n",
      "N1108 - G000423665: Solirubrobacter soli\n",
      "N445 - G001829415: Spirochaetes bacterium GWF1_51_8\n",
      "N250 - G000007005: Sulfolobus solfataricus\n",
      "N581 - G000009905: Symbiobacterium thermophilum\n",
      "N1318 - G000305935: Thermacetogenium phaeum\n",
      "N1329 - G000806225: Thermoanaerobacter sp. YS13\n",
      "N933 - G002010535: Thermoanaerobacterales bacterium JdFR-68\n",
      "N21 - G000212395: Thermodesulfobium narugense\n",
      "N315 - G001858525: Vampirovibrio chlorellavorus\n",
      "N133 - G000997185: candidate division CPR2 bacterium GW2011_GWC1_41_48\n",
      "N31 - G001567355: candidate division WS6 bacterium OLB20\n",
      "N82 - G001773295: candidate division WWE3 bacterium RIFOXYD1_FULL_39_9\n",
      "N2468 - G001304015: candidate division Zixibacteria bacterium SM23_81\n",
      "N582 - G900016865: uncultured Clostridia bacterium\n",
      "N115 - G000965745: uncultured archaeon\n"
     ]
    }
   ],
   "source": [
    "for nid, g in sorted(chosen_taxa.items(), key=lambda x: g2species[x[1]]):\n",
    "    print('%s - %s: %s' % (nid, g, g2species[g]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "%%script false\n",
    "for g in data['N2193']['taxa']:\n",
    "    print('%s: %s' % (g, dfg.loc['G000992445']['species']))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Output results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "with open('red.%d.txt' % ntaxa, 'w') as f:\n",
    "    for nid in sorted(chosen_taxa, key=lambda x: int(x[1:])):\n",
    "        f.write('%s\\t%s\\n' % (chosen_taxa[nid], nid))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "gs = sorted(chosen_taxa.values())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'red.92.nwk'"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tout = tree.shear(gs)\n",
    "tout.write('red.%d.nwk' % ntaxa)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "%%script false\n",
    "print('Sum of distances: %d.' % np.tril(dm.filter(gs).data).sum())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Explore results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def check_keyword(gs, df, col, word):\n",
    "    res = []\n",
    "    d_ = df[col].dropna().to_dict()\n",
    "    for g in gs:\n",
    "        if g in d_ and word in d_[g]:\n",
    "            res.append('%s: %s' % (g, word))\n",
    "    return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[]"
      ]
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "check_keyword(gs, dfg, 'phylum', 'Cyanobacteria')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'G001858525', 'G001899315'}"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "melai.intersection(gs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "12"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(check_keyword(gs, dfg, 'lv1_group', 'Archaea'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "15"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(check_keyword(gs, dfg, 'lv1_group', 'CPR'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Firmicutes: 16\n",
      "Euryarchaeota: 5\n",
      "Actinobacteria: 4\n",
      "Proteobacteria: 9\n",
      "FCB: 4\n",
      "Crenarchaeota: 1\n",
      "Chlamydiae: 0\n",
      "Bacteria: 17\n",
      "Cyanobacteria: 0\n",
      "Spirochaetes: 3\n",
      "Terrabacteria: 3\n",
      "Chloroflexi: 3\n",
      "Bacteroidetes: 1\n",
      "TACK: 3\n",
      "PVC: 5\n",
      "Archaea: 1\n",
      "CPR: 5\n",
      "Parcubacteria: 6\n",
      "DPANN: 1\n",
      "Microgenomates: 4\n",
      "Asgard: 1\n"
     ]
    }
   ],
   "source": [
    "for lv2 in dfg['lv2_group'].unique():\n",
    "    print('%s: %d' % (lv2, len(check_keyword(gs, dfg, 'lv2_group', lv2))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}