Enumerate-descendants-in-divergence-trees-with-nextstrain-jsons.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Parse divergence tree and output tips with evidence of onward transmission using nextstrain jsons\n",
    "\n",
    "October 14, 2019 \n",
    "\n",
    "Using the same method as used for mumps, I would like to try parsing the H5N1 nextstrain tree (json format) and quantifying tips with evidence of onward transmission vs. tips without evidence of onward transmission. My thought here is that this could be used to replace my current aim 2. This way, I could actually associated metadata from tips with a metric for transmission. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The rpy2.ipython extension is already loaded. To reload it, use:\n",
      "  %reload_ext rpy2.ipython\n"
     ]
    }
   ],
   "source": [
    "import sys, subprocess, glob, os, shutil, re, importlib\n",
    "from subprocess import call\n",
    "import imp\n",
    "bt = imp.load_source('baltic', '/Users/lmoncla/src/baltic/baltic-iqtree.py')\n",
    "\n",
    "%matplotlib inline\n",
    "import matplotlib as mpl\n",
    "from matplotlib import pyplot as plt\n",
    "import matplotlib.patheffects as path_effects\n",
    "import matplotlib.lines as mlines\n",
    "from matplotlib.font_manager import FontProperties\n",
    "import matplotlib.colors as clr\n",
    "import textwrap as textwrap\n",
    "from textwrap import wrap\n",
    "\n",
    "import pandas as pd\n",
    "\n",
    "import numpy as np\n",
    "from scipy.special import binom\n",
    "\n",
    "import rpy2\n",
    "import seaborn as sns\n",
    "%load_ext rpy2.ipython\n",
    "\n",
    "import json\n",
    "import collections\n",
    "from collections import Counter\n",
    "from Bio import SeqIO\n",
    "from Bio import Seq\n",
    "import Bio.Phylo"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Functions from augur to read in nextstrain jsons and convert to trees"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# function to use the json module to read in a json file and store it as \"data\"                \n",
    "def read_json(file_name):\n",
    "    try:\n",
    "        handle = open(file_name, 'r')\n",
    "    except IOError:\n",
    "        pass\n",
    "    else:\n",
    "        data = json.load(handle)\n",
    "        handle.close()\n",
    "    return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# original code that Trevor gave me for parsing through tree jsons and returning descendents\n",
    "def all_descendants(node):\n",
    "    \"\"\"Take node, ie. dict, and return a flattened list of all nodes descending from this node\"\"\"\n",
    "    yield node\n",
    "    \n",
    "    # this will recursively return all internal nodes (nodes with children)\n",
    "    if 'children' in node:\n",
    "        for child in node['children']:\n",
    "            for desc in all_descendants(child):\n",
    "                yield desc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Biopython's trees don't store links to node parents, so we need to build\n",
    "# a map of each node to its parent.\n",
    "# Code from the Bio.Phylo cookbook: http://biopython.org/wiki/Phylo_cookbook\n",
    "def all_parents(tree):\n",
    "    parents = {}\n",
    "    for clade in tree.find_clades(order='level'):\n",
    "        for child in clade:\n",
    "            parents[child] = clade\n",
    "    return parents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def annotate_parents(tree):\n",
    "    # Get all parent nodes by node.\n",
    "    parents_by_node = all_parents(tree)\n",
    "\n",
    "    # Next, annotate each node with its parent.\n",
    "    for node in tree.find_clades():\n",
    "        if node == tree.root:\n",
    "            node.up = None\n",
    "        else:\n",
    "            node.up = parents_by_node[node]\n",
    "\n",
    "    # Return the tree.\n",
    "    return tree"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def json_to_tree(json_dict, root=True):\n",
    "    \"\"\"Returns a Bio.Phylo tree corresponding to the given JSON dictionary exported\n",
    "    by `tree_to_json`.\n",
    "\n",
    "    Assigns links back to parent nodes for the root of the tree.\n",
    "\n",
    "    >>> import json\n",
    "    >>> json_fh = open(\"tests/data/json_tree_to_nexus/flu_h3n2_ha_3y_tree.json\", \"r\")\n",
    "    >>> json_dict = json.load(json_fh)\n",
    "    >>> tree = json_to_tree(json_dict)\n",
    "    >>> tree.name\n",
    "    u'NODE_0002020'\n",
    "    >>> len(tree.clades)\n",
    "    2\n",
    "    >>> tree.clades[0].name\n",
    "    u'NODE_0001489'\n",
    "    >>> hasattr(tree, \"attr\")\n",
    "    True\n",
    "    >>> \"dTiter\" in tree.attr\n",
    "    True\n",
    "    \"\"\"\n",
    "    node = Bio.Phylo.Newick.Clade()\n",
    "    node.name = json_dict[\"strain\"]\n",
    "\n",
    "    if \"children\" in json_dict:\n",
    "        # Recursively add children to the current node.\n",
    "        node.clades = [json_to_tree(child, root=False) for child in json_dict[\"children\"]]\n",
    "\n",
    "    # Assign all non-children attributes.\n",
    "    for attr, value in json_dict.items():\n",
    "        if attr != \"children\":\n",
    "            setattr(node, attr, value)\n",
    "\n",
    "    node.numdate = node.attr.get(\"num_date\")\n",
    "    node.divergence = node.attr.get(\"div\")\n",
    "\n",
    "    if \"translations\" in node.attr:\n",
    "        node.translations = node.attr[\"translations\"]\n",
    "\n",
    "    if root:\n",
    "        node = annotate_parents(node)\n",
    "\n",
    "    return node"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Read in metadata, find proper parent node, and add branch lengths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_metadata(metadata_path):\n",
    "    metadata = {}\n",
    "    \n",
    "    with open(metadata_path, \"r\") as infile:\n",
    "        metadata_categories = {}\n",
    "        \n",
    "        for line in infile: \n",
    "            if \"isolate_id\" in line: # skip first line\n",
    "                for i in range(len(line.split(\"\\t\"))):\n",
    "                    metadata_categories[i] = line.split(\"\\t\")[i]\n",
    "                    \n",
    "            elif \"isolate_id\" not in line:\n",
    "                strain = line.split(\"\\t\")[0]\n",
    "                metadata[strain] = {}\n",
    "                \n",
    "                for i in range(len(line.split(\"\\t\"))):\n",
    "                    ID = metadata_categories[i]\n",
    "                    metadata[strain][ID] = line.split(\"\\t\")[i]\n",
    "\n",
    "    return(metadata)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "def return_country(strain_name, metadata):\n",
    "        \n",
    "    country = metadata[strain_name][\"country\"]\n",
    "    \n",
    "    return(country)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {},
   "outputs": [],
   "source": [
    "def return_region(strain_name, metadata):\n",
    "        \n",
    "    region = metadata[strain_name][\"region\"]\n",
    "    \n",
    "    return(region)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "metadata": {},
   "outputs": [],
   "source": [
    "def return_host(strain_name, metadata):\n",
    "        \n",
    "    host = metadata[strain_name][\"host\"]\n",
    "    \n",
    "    return(host)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 97,
   "metadata": {},
   "outputs": [],
   "source": [
    "def return_proper_parent_node(node):\n",
    "    \"\"\"given an internal node, traverse back up the tree to find a parental node with a \n",
    "    real branch length (basically, collapse the polytomy)\"\"\"\n",
    "    \n",
    "    #print(node, node.length)\n",
    "    if abs(node.divergence - node.up.divergence) < (1/length_alignment): \n",
    "        \n",
    "        #print(\"going up 1 node\")\n",
    "        if node.up !=None:\n",
    "            parent_node = return_proper_parent_node(node.up)\n",
    "        \n",
    "        else:\n",
    "            #print(\"root is proper parent\")\n",
    "            parent_node = node\n",
    "    \n",
    "    else: \n",
    "        #print(\"current node has proper length\")\n",
    "        parent_node = node\n",
    "    \n",
    "    return(parent_node)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 113,
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_nodes(node):\n",
    "    \"\"\"Take node and add up branch lengths for total subtending tree from that node\"\"\"\n",
    "    total_lengths = 0\n",
    "    \n",
    "    if node.is_terminal() == True: \n",
    "        total_lengths += node.branch_length\n",
    "\n",
    "    \n",
    "    else:\n",
    "        total_lengths += node.branch_length\n",
    "        for child in node.clades:\n",
    "            \n",
    "            total_lengths += add_nodes(child)\n",
    "                            \n",
    "    return(total_lengths)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 192,
   "metadata": {},
   "outputs": [],
   "source": [
    "def return_dates_of_children(proper_parent,current_tip,long_spillover_threshold):\n",
    "    current_tip_date = float(current_tip.attr['num_date'])\n",
    "    number_long = 0\n",
    "    number_short = 0\n",
    "        \n",
    "    for i in proper_parent.get_terminals():\n",
    "        date = float(i.attr['num_date'])\n",
    "\n",
    "        if date - current_tip_date >= long_spillover_threshold:\n",
    "            number_long += 1\n",
    "        else:\n",
    "            number_short += 1\n",
    "            \n",
    "#         if current_tip.strain == 'A/chicken/Vietnam/DT-171/2004':\n",
    "#             print(current_tip_date, date, number_long, number_short)\n",
    "\n",
    "            \n",
    "    return(number_short, number_long)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 193,
   "metadata": {},
   "outputs": [],
   "source": [
    "def return_descendants_dict(tree, metadata, length_alignment):\n",
    "    \n",
    "    output_dict = {}\n",
    "    not_polytomies = []\n",
    "    polytomies = []\n",
    "    \n",
    "    for i in tree.find_clades(): ## iterate over objects in tree            \n",
    "        if i.is_terminal() == True:\n",
    "\n",
    "            country = return_country(i.strain, metadata)\n",
    "            region = return_region(i.strain,metadata)\n",
    "            host = return_host(i.strain, metadata)\n",
    "            date = i.attr['num_date']\n",
    "\n",
    "            if abs(i.divergence - i.up.divergence < (1/length_alignment)): \n",
    "                polytomies.append(i.strain)\n",
    "\n",
    "                proper_parent = return_proper_parent_node(i.up)\n",
    "                branch_length = add_nodes(proper_parent) - i.branch_length\n",
    "                number_children = len(proper_parent.get_terminals()) - 1\n",
    "                short_spillovers,long_spillovers = return_dates_of_children(proper_parent,i,long_spillover_threshold)\n",
    "\n",
    "                output_dict[i.strain] = {}\n",
    "                output_dict[i.strain]['branch_lengths'] = branch_length\n",
    "                output_dict[i.strain]['number_children'] = number_children\n",
    "                output_dict[i.strain]['country'] = country\n",
    "                output_dict[i.strain]['region'] = region\n",
    "                output_dict[i.strain]['host'] = host\n",
    "                output_dict[i.strain]['long_spillovers'] = long_spillovers\n",
    "                output_dict[i.strain]['short_spillovers'] = short_spillovers\n",
    "\n",
    "            else:\n",
    "                not_polytomies.append(i.strain)\n",
    "                output_dict[i.strain] = {}\n",
    "                output_dict[i.strain]['branch_lengths'] = 0\n",
    "                output_dict[i.strain]['number_children'] = 0\n",
    "                output_dict[i.strain]['country'] = country\n",
    "                output_dict[i.strain]['region'] = region\n",
    "                output_dict[i.strain]['host'] = host\n",
    "                output_dict[i.strain]['long_spillovers'] = 0\n",
    "                output_dict[i.strain]['short_spillovers'] = 0\n",
    "    \n",
    "    return(polytomies,not_polytomies,output_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Set paths, run it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 198,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set paths\n",
    "tree_path = \"/Users/lmoncla/src/avian-flu/auspice/flu_avian_h5n1_ha_tree.json\"\n",
    "metadata_path = \"/Users/lmoncla/src/avian-flu/results/metadata_h5n1_ha.tsv\"\n",
    "length_alignment = 1762\n",
    "long_spillover_threshold = 1  #0.083 # this is in terms of years\n",
    "\n",
    "# run\n",
    "metadata = read_metadata(metadata_path)\n",
    "tree = read_json(tree_path)\n",
    "tree = json_to_tree(tree)\n",
    "polytomies,not_polytomies,output_dict = return_descendants_dict(tree, metadata, length_alignment)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Convert to dataframe "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 199,
   "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>index</th>\n",
       "      <th>branch_lengths</th>\n",
       "      <th>number_children</th>\n",
       "      <th>country</th>\n",
       "      <th>region</th>\n",
       "      <th>host</th>\n",
       "      <th>long_spillovers</th>\n",
       "      <th>short_spillovers</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>A/Alberta/1/2014</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0</td>\n",
       "      <td>canada</td>\n",
       "      <td>north_america</td>\n",
       "      <td>human</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>A/American_green_winged_teal/Washington/195750...</td>\n",
       "      <td>0.000608</td>\n",
       "      <td>1</td>\n",
       "      <td>usa</td>\n",
       "      <td>north_america</td>\n",
       "      <td>avian</td>\n",
       "      <td>0</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>A/American_wigeon/Washington/196336/2015</td>\n",
       "      <td>0.020912</td>\n",
       "      <td>7</td>\n",
       "      <td>usa</td>\n",
       "      <td>north_america</td>\n",
       "      <td>avian</td>\n",
       "      <td>0</td>\n",
       "      <td>8</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>A/American_wigeon/Washington/196340/2015</td>\n",
       "      <td>0.000653</td>\n",
       "      <td>1</td>\n",
       "      <td>usa</td>\n",
       "      <td>north_america</td>\n",
       "      <td>avian</td>\n",
       "      <td>0</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>A/Americanblackduck/Alberta/118/2016</td>\n",
       "      <td>0.005564</td>\n",
       "      <td>2</td>\n",
       "      <td>canada</td>\n",
       "      <td>north_america</td>\n",
       "      <td>avian</td>\n",
       "      <td>0</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                               index  branch_lengths  \\\n",
       "0                                   A/Alberta/1/2014        0.000000   \n",
       "1  A/American_green_winged_teal/Washington/195750...        0.000608   \n",
       "2           A/American_wigeon/Washington/196336/2015        0.020912   \n",
       "3           A/American_wigeon/Washington/196340/2015        0.000653   \n",
       "4               A/Americanblackduck/Alberta/118/2016        0.005564   \n",
       "\n",
       "   number_children country         region   host  long_spillovers  \\\n",
       "0                0  canada  north_america  human                0   \n",
       "1                1     usa  north_america  avian                0   \n",
       "2                7     usa  north_america  avian                0   \n",
       "3                1     usa  north_america  avian                0   \n",
       "4                2  canada  north_america  avian                0   \n",
       "\n",
       "   short_spillovers  \n",
       "0                 0  \n",
       "1                 2  \n",
       "2                 8  \n",
       "3                 2  \n",
       "4                 3  "
      ]
     },
     "execution_count": 199,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df = pd.DataFrame.from_dict(output_dict, orient=\"index\")\n",
    "df.reset_index(inplace=True)\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 180,
   "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>index</th>\n",
       "      <th>branch_lengths</th>\n",
       "      <th>number_children</th>\n",
       "      <th>country</th>\n",
       "      <th>region</th>\n",
       "      <th>host</th>\n",
       "      <th>long_spillovers</th>\n",
       "      <th>short_spillovers</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>915</th>\n",
       "      <td>A/chicken/Vietnam/DT-171/2004</td>\n",
       "      <td>0.306898</td>\n",
       "      <td>109</td>\n",
       "      <td>vietnam</td>\n",
       "      <td>southeast_asia</td>\n",
       "      <td>avian</td>\n",
       "      <td>0</td>\n",
       "      <td>110</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                             index  branch_lengths  number_children  country  \\\n",
       "915  A/chicken/Vietnam/DT-171/2004        0.306898              109  vietnam   \n",
       "\n",
       "             region   host  long_spillovers  short_spillovers  \n",
       "915  southeast_asia  avian                0               110  "
      ]
     },
     "execution_count": 180,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "slothing = df[df['index'] == 'A/chicken/Vietnam/DT-171/2004']\n",
    "slothing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Given a category to compare and 2 members of that category, perform a Fisher's exact test and return the results "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 206,
   "metadata": {},
   "outputs": [],
   "source": [
    "def return_fishers_exact_test(df,variable,group1,group2):\n",
    "    group1_with_desc = len(df[(df['number_children'] > 0) & (df[variable] == group1)])\n",
    "    group1_no_desc = len(df[df[variable] == group1]) - group1_with_desc\n",
    "\n",
    "    group2_with_desc = len(df[(df['number_children'] > 0) & (df[variable] == group2)])\n",
    "    group2_no_desc = len(df[df[variable] == group2]) - group2_with_desc\n",
    "    \n",
    "    table = [[group1_with_desc,group1_no_desc],[group2_with_desc,group2_no_desc]]\n",
    "    result = scipy.stats.fisher_exact(table)\n",
    "    return(table, result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 207,
   "metadata": {},
   "outputs": [],
   "source": [
    "def return_dateframe_of_Fishers_exact_results(variable,group1,group2,table, result):\n",
    "    \n",
    "    group1_with_desc = table[0][0]\n",
    "    group1_no_desc = table[0][1]\n",
    "    group2_with_desc = table[1][0]\n",
    "    group2_no_desc = table[1][1]\n",
    "    p_value = result[1]\n",
    "    \n",
    "    x = {'variable':[variable],'group1':[group1],'group2':[group2],\n",
    "         'group1_with_desc':[group1_with_desc],'group1_no_desc':[group1_no_desc],\n",
    "        'group2_with_desc':[group2_with_desc],'group2_no_desc':[group2_no_desc],'p-value':[p_value]}\n",
    "    \n",
    "    d = pd.DataFrame.from_dict(x)\n",
    "    return(d)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 208,
   "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>variable</th>\n",
       "      <th>group1</th>\n",
       "      <th>group2</th>\n",
       "      <th>group1_with_desc</th>\n",
       "      <th>group1_no_desc</th>\n",
       "      <th>group2_with_desc</th>\n",
       "      <th>group2_no_desc</th>\n",
       "      <th>p-value</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>country</td>\n",
       "      <td>china</td>\n",
       "      <td>vietnam</td>\n",
       "      <td>23</td>\n",
       "      <td>142</td>\n",
       "      <td>44</td>\n",
       "      <td>108</td>\n",
       "      <td>0.001436</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  variable group1   group2  group1_with_desc  group1_no_desc  \\\n",
       "0  country  china  vietnam                23             142   \n",
       "\n",
       "   group2_with_desc  group2_no_desc   p-value  \n",
       "0                44             108  0.001436  "
      ]
     },
     "execution_count": 208,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "variable = \"country\"\n",
    "group1 = \"china\"\n",
    "group2 = \"vietnam\"\n",
    "table,result = return_fishers_exact_test(df,variable,group1,group2)\n",
    "Fishers_df = return_dateframe_of_Fishers_exact_results(variable,group1,group2,table, result)\n",
    "Fishers_df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run on multiple variables \n",
    "\n",
    "Now, given a list of variables to test, run them all. Here, I will first subset to include only avian tips, and then compare regions or countries to each other. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 209,
   "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>variable</th>\n",
       "      <th>group1</th>\n",
       "      <th>group2</th>\n",
       "      <th>group1_with_desc</th>\n",
       "      <th>group1_no_desc</th>\n",
       "      <th>group2_with_desc</th>\n",
       "      <th>group2_no_desc</th>\n",
       "      <th>p-value</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>country</td>\n",
       "      <td>china</td>\n",
       "      <td>vietnam</td>\n",
       "      <td>15</td>\n",
       "      <td>110</td>\n",
       "      <td>42</td>\n",
       "      <td>99</td>\n",
       "      <td>0.000510</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>region</td>\n",
       "      <td>china</td>\n",
       "      <td>southeast_asia</td>\n",
       "      <td>48</td>\n",
       "      <td>174</td>\n",
       "      <td>122</td>\n",
       "      <td>257</td>\n",
       "      <td>0.006412</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>region</td>\n",
       "      <td>china</td>\n",
       "      <td>south_asia</td>\n",
       "      <td>48</td>\n",
       "      <td>174</td>\n",
       "      <td>56</td>\n",
       "      <td>159</td>\n",
       "      <td>0.312290</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>region</td>\n",
       "      <td>africa</td>\n",
       "      <td>south_asia</td>\n",
       "      <td>98</td>\n",
       "      <td>141</td>\n",
       "      <td>56</td>\n",
       "      <td>159</td>\n",
       "      <td>0.001022</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  variable  group1          group2  group1_with_desc  group1_no_desc  \\\n",
       "0  country   china         vietnam                15             110   \n",
       "0   region   china  southeast_asia                48             174   \n",
       "0   region   china      south_asia                48             174   \n",
       "0   region  africa      south_asia                98             141   \n",
       "\n",
       "   group2_with_desc  group2_no_desc   p-value  \n",
       "0                42              99  0.000510  \n",
       "0               122             257  0.006412  \n",
       "0                56             159  0.312290  \n",
       "0                56             159  0.001022  "
      ]
     },
     "execution_count": 209,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Fishers_df1 = pd.DataFrame()\n",
    "\n",
    "variables = [{\"country\":['china','vietnam']},{'region':['china','southeast_asia']},\n",
    "            {'region':['china','south_asia']},{'region':['africa','south_asia']}]\n",
    "df_test = df[df['host'] == 'avian']\n",
    "\n",
    "for i in variables: \n",
    "    for key, value in i.items() :\n",
    "        variable = key\n",
    "        group1 = value[0]\n",
    "        group2 = value[1]\n",
    "\n",
    "        table,result = return_fishers_exact_test(df_test,variable,group1,group2)\n",
    "        Fishers_df = return_dateframe_of_Fishers_exact_results(variable,group1,group2,table, result)\n",
    "        Fishers_df1 = Fishers_df1.append(Fishers_df)\n",
    "    \n",
    "Fishers_df1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This suggests that Vietnamese birds are more likely to transmit than Chinese birds, which is a little surprising. This also suggests that African birds are more likely to transmit than South Asian (Middle Eastern) birds, which makes sense. \n",
    "\n",
    "I think that there are a few things I will need to really think about/test. \n",
    "1. How does uneven sampling impact this? I suppose that I can just test in subsampled trees to see whether this is robust to sampling or not. The current nextstrain jsons I'm using here are subsampled for H5N1 by country to 10 sequences per country per year. \n",
    "2. How do I do power testing on this? \n",
    "3. How will I validate this as a method? It is a little strange to me that China is showing up as less likely to be an ancestor than Vietnam, especially when more internal nodes on the tree are in China than Vietnam. This makes me a little concerned that this isn't going to work super well. \n",
    "4. How do I do multiple comparisons? "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 202,
   "metadata": {},
   "outputs": [],
   "source": [
    "def return_fishers_exact_test_longevity(df,variable,group1,group2):\n",
    "    group1_with_desc = len(df[(df['long_spillovers'] > 0) & (df[variable] == group1)])\n",
    "    group1_no_desc = len(df[df[variable] == group1]) - group1_with_desc\n",
    "\n",
    "    group2_with_desc = len(df[(df['long_spillovers'] > 0) & (df[variable] == group2)])\n",
    "    group2_no_desc = len(df[df[variable] == group2]) - group2_with_desc\n",
    "    \n",
    "    table = [[group1_with_desc,group1_no_desc],[group2_with_desc,group2_no_desc]]\n",
    "    result = scipy.stats.fisher_exact(table)\n",
    "    return(table, result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 203,
   "metadata": {},
   "outputs": [],
   "source": [
    "def return_dateframe_of_Fishers_exact_results_longevity(variable,group1,group2,table, result):\n",
    "    \n",
    "    group1_with_long = table[0][0]\n",
    "    group1_no_long = table[0][1]\n",
    "    group2_with_long = table[1][0]\n",
    "    group2_no_long = table[1][1]\n",
    "    p_value = result[1]\n",
    "    \n",
    "    x = {'variable':[variable],'group1':[group1],'group2':[group2],\n",
    "         'group1_with_long':[group1_with_long],'group1_no_long':[group1_no_long],\n",
    "        'group2_with_long':[group2_with_long],'group2_no_long':[group2_no_long],'p-value':[p_value]}\n",
    "    \n",
    "    d = pd.DataFrame.from_dict(x)\n",
    "    return(d)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 204,
   "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>variable</th>\n",
       "      <th>group1</th>\n",
       "      <th>group2</th>\n",
       "      <th>group1_with_long</th>\n",
       "      <th>group1_no_long</th>\n",
       "      <th>group2_with_long</th>\n",
       "      <th>group2_no_long</th>\n",
       "      <th>p-value</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>country</td>\n",
       "      <td>china</td>\n",
       "      <td>vietnam</td>\n",
       "      <td>4</td>\n",
       "      <td>121</td>\n",
       "      <td>5</td>\n",
       "      <td>136</td>\n",
       "      <td>1.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>region</td>\n",
       "      <td>china</td>\n",
       "      <td>southeast_asia</td>\n",
       "      <td>8</td>\n",
       "      <td>214</td>\n",
       "      <td>12</td>\n",
       "      <td>367</td>\n",
       "      <td>0.815650</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>region</td>\n",
       "      <td>china</td>\n",
       "      <td>south_asia</td>\n",
       "      <td>8</td>\n",
       "      <td>214</td>\n",
       "      <td>9</td>\n",
       "      <td>206</td>\n",
       "      <td>0.808307</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>region</td>\n",
       "      <td>africa</td>\n",
       "      <td>south_asia</td>\n",
       "      <td>13</td>\n",
       "      <td>226</td>\n",
       "      <td>9</td>\n",
       "      <td>206</td>\n",
       "      <td>0.662846</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  variable  group1          group2  group1_with_long  group1_no_long  \\\n",
       "0  country   china         vietnam                 4             121   \n",
       "0   region   china  southeast_asia                 8             214   \n",
       "0   region   china      south_asia                 8             214   \n",
       "0   region  africa      south_asia                13             226   \n",
       "\n",
       "   group2_with_long  group2_no_long   p-value  \n",
       "0                 5             136  1.000000  \n",
       "0                12             367  0.815650  \n",
       "0                 9             206  0.808307  \n",
       "0                 9             206  0.662846  "
      ]
     },
     "execution_count": 204,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Fishers_df2 = pd.DataFrame()\n",
    "\n",
    "variables = [{\"country\":['china','vietnam']},{'region':['china','southeast_asia']},\n",
    "            {'region':['china','south_asia']},{'region':['africa','south_asia']}]\n",
    "df_test = df[df['host'] == 'avian']\n",
    "\n",
    "for i in variables: \n",
    "    for key, value in i.items() :\n",
    "        variable = key\n",
    "        group1 = value[0]\n",
    "        group2 = value[1]\n",
    "\n",
    "        table,result = return_fishers_exact_test_longevity(df_test,variable,group1,group2)\n",
    "        Fishers_df = return_dateframe_of_Fishers_exact_results_longevity(variable,group1,group2,table, result)\n",
    "        Fishers_df2 = Fishers_df2.append(Fishers_df)\n",
    "    \n",
    "Fishers_df2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 200,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "653"
      ]
     },
     "execution_count": 200,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sloth = df[df['number_children'] > 0]\n",
    "len(sloth)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 201,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "98"
      ]
     },
     "execution_count": 201,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sloth = df[df['long_spillovers'] > 0]\n",
    "len(sloth)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If I define a long spillover as 6 months or greater, there are 142 tips. If I define it as at least 1 year, then I have 98 tips. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"Using Fisher.test in R, I tried this, but it did not work: \n",
    "\n",
    "sloth <- matrix(c(8,214,12,367,13,226),nrow=3)\n",
    "fisher.test(sloth, simulate.p.value = TRUE, B=10000)\"\"\""
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "LHM-basics (python3)",
   "language": "python",
   "name": "lhm-basics"
  },
  "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}