{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Prepare tree rendering package"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebooks manipulates a tree and generates rendering packages. It does:\n",
    "\n",
    " - Collapse a tree at given taxonomic rank(s) based on serveral criteria.\n",
    " - Generate color gradient for branch support values.\n",
    " - Generate files that can be directly parsed and rendered using [**iTOL**](https://itol.embl.de/) and [**FigTree**](http://tree.bio.ed.ac.uk/software/figtree/)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Preparation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "import numpy as np\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from skbio import TreeNode"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Input files"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Tree file (with node IDs, and without support values)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "tree_fp = '../trees/release/astral.cons.nwk'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Taxonomic information file (original or tax2tree-curated)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "taxonomy_fp = '../taxonomy/tax2tree/ncbi/astral/filled_ranks.tsv'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Custom node attributes (bootstrap, estimated time range, metadata category, additional name, etc.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "custom_attrs_fps = {\n",
    "    'lpp': '../trees/release/supports/astral.txt'}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Collapse the tree from this rank up. For example, \"class\" will have the tree collapsed at class (if possible) or phylum. Leave empty or None if not needed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# collapse_rank = None\n",
    "collapse_rank = 'genus'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Determine the visual length of a collapsed clade (triangle or line). Options are: mean, std (don't use), min, 25%, 50% (median), 75% and max."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "collapse_length_func = '50%'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Clades with descendants less than this threshold will not be collapsed. Either a fixed number, or a rank-to-number dictionary. Example: phylum = 1, class = 10. Leave 0 if not needed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# min_clade_size = 0\n",
    "# for full-scale (10k-taxon) trees\n",
    "min_clade_size = {'kingdom': 1, 'phylum': 1, 'class': 5, 'order': 50, 'family': 50, 'genus': 50, 'species': 50}\n",
    "# for 1k-taxon trees\n",
    "# min_clade_size = {'kingdom': 1, 'phylum': 1, 'class': 1, 'order': 5, 'family': 10, 'genus': 10, 'species': 10}\n",
    "# for class / phylum only:\n",
    "# min_clade_size = {'kingdom': 1, 'phylum': 1, 'class': 10, 'order': 0, 'family': 0, 'genus': 0, 'species': 0}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Split clades with descendants less than this threshold of *fraction* of the dominant clade of the same taxon will not be collapsed. For example, `Firmicutes_1` has 1000 tips, then if `Firmicutes_10` has 45 tips (< 1000 * 5%), it will not be collapsed.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# min_split_clade_frac = 0\n",
    "min_split_clade_frac = 0.05"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Whether to delete tips not belonging to any collapsed clades."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "delete_uncollapsed = True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Whether to hide uncollapsed tip names. Effective when `delete_uncollapsed` is `False`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "hide_uncollapsed = True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Manipulate node labels using the following regular expressions (pairs of pattern and replacement)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "label_format_regexes = [\n",
    "    (r'^Candidatus ', r'Ca. '),\n",
    "    (r'^candidate division ', r'Ca. ')\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Append rank code to taxon (e.g.,: `Bacteria` => `k__Bacteria`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "append_rank_code = True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Append clade size to taxon"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "append_clade_size = True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Low and high end of color gradient."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "color_range = ('#f0f0f0', '#191919')  # gray\n",
    "# color_range = ('#deebf7', '#3182bd')  # blue"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Helpers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Basic utilities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sort_node_ids(d):\n",
    "    \"\"\"Sort names of tips and internal nodes.\"\"\"\n",
    "    return sorted(d, key=lambda x: (x[0], int(x[1:])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "def digits(num):\n",
    "    \"\"\"Get number digits after decimal point.\"\"\"\n",
    "    if not num.replace('.', '').isdigit() or num.count('.') != 1:\n",
    "        raise ValueError('Not a valid float number: %s' % num)\n",
    "    return len(num.split('.')[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "def de_suffix(taxon, names):\n",
    "    \"\"\"Restore suffixed taxon name.\"\"\"\n",
    "    if '_' not in taxon:\n",
    "        return taxon\n",
    "    res = '_'.join(taxon.split('_')[:-1])\n",
    "    return res if res in names else taxon"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Node dimension calculation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_clade_dimensions(node):\n",
    "    \"\"\"Calculate the dimensions of a clade.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    node : skbio.TreeNode\n",
    "        clade to calculate\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    pd.Series\n",
    "        count, mean, std, min, 25%, 50%, 75%, max\n",
    "    \"\"\"\n",
    "    lengths = pd.Series(x.accumulate_to_ancestor(node) for x in node.tips())\n",
    "    return lengths.describe()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Selective tree shearing and pruning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "def selective_prune(tree, tips_to_keep, nodes_to_keep=[]):\n",
    "    \"\"\"Shear a tree and selectively prune it.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    tree : skbio.TreeNode\n",
    "        tree to shear\n",
    "    tips_to_keep : iterable of str\n",
    "        tip names to keep\n",
    "    nodes_to_keep : iterable of str\n",
    "        internal node names to keep\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    tree : skbio.TreeNode\n",
    "        resulting tree\n",
    "\n",
    "    Notes\n",
    "    -----\n",
    "    Inherited from scikit-bio's `shear` and  `prune` functions, but will\n",
    "    selectively remove internal nodes.\n",
    "    \"\"\"\n",
    "    tcopy = tree.deepcopy()\n",
    "\n",
    "    ids = set(tips_to_keep)\n",
    "    marked = set()\n",
    "    for tip in tcopy.tips():\n",
    "        if tip.name in ids:\n",
    "            marked.add(tip)\n",
    "            for anc in tip.ancestors():\n",
    "                if anc in marked:\n",
    "                    break\n",
    "                else:\n",
    "                    marked.add(anc)\n",
    "\n",
    "    for node in list(tcopy.traverse()):\n",
    "        if node not in marked:\n",
    "            node.parent.remove(node)\n",
    "\n",
    "    ids = set(nodes_to_keep)\n",
    "    nodes_to_remove = []\n",
    "    for node in tcopy.traverse(include_self=False):\n",
    "        if len(node.children) == 1:\n",
    "            if node.name not in ids:\n",
    "                nodes_to_remove.append(node)\n",
    "\n",
    "    for node in nodes_to_remove:\n",
    "        child = node.children[0]\n",
    "\n",
    "        if child.length is None or node.length is None:\n",
    "            child.length = child.length or node.length\n",
    "        else:\n",
    "            child.length += node.length\n",
    "\n",
    "        if node.parent is None:\n",
    "            continue\n",
    "\n",
    "        node.parent.append(child)\n",
    "        node.parent.remove(node)\n",
    "\n",
    "    return tcopy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Newick string formatting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "def format_newick(tree, operators=',:_;()[] ', digits=None):\n",
    "    \"\"\"Generate a Newick string from a tree.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    tree : skbio.TreeNode\n",
    "        tree to convert to a Newick string\n",
    "    operators : str\n",
    "        list of characters that have special meaning in a tree file so that\n",
    "        a node name containing any of them must be quoted\n",
    "    digits : int or tuple of (int, int)\n",
    "        number of digits (float and scientific) to print in a branch length\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    str\n",
    "        formatted Newick string\n",
    "\n",
    "    Notes\n",
    "    -----\n",
    "    Modified from scikit-bio's `_tree_node_to_newick`. In addition to the\n",
    "    prototype, it can do:\n",
    "    \n",
    "    1. Keep spaces without converting them to underscores.\n",
    "    2. Print branch lengths based on given precision.\n",
    "    \"\"\"\n",
    "    res = ''\n",
    "    operators = set(operators or '')\n",
    "    if isinstance(digits, int):\n",
    "        digits = (digits, digits) \n",
    "    current_depth = 0\n",
    "    nodes_left = [(tree, 0)]\n",
    "    while len(nodes_left) > 0:\n",
    "        entry = nodes_left.pop()\n",
    "        node, node_depth = entry\n",
    "        if node.children and node_depth >= current_depth:\n",
    "            res += '('\n",
    "            nodes_left.append(entry)\n",
    "            nodes_left += ((child, node_depth + 1) for child in\n",
    "                           reversed(node.children))\n",
    "            current_depth = node_depth + 1\n",
    "        else:\n",
    "            if node_depth < current_depth:\n",
    "                res += ')'\n",
    "                current_depth -= 1\n",
    "            if node.name:\n",
    "                escaped = \"%s\" % node.name.replace(\"'\", \"''\")\n",
    "                if any(t in operators for t in node.name):\n",
    "                    res += \"'\"\n",
    "                    res += escaped\n",
    "                    res += \"'\"\n",
    "                else:\n",
    "                    res += escaped\n",
    "            if node.length is not None:\n",
    "                res += ':'\n",
    "                length = str(node.length)\n",
    "                if digits:\n",
    "                    length = '%.*g' % ((digits[0] if 'e' in length\n",
    "                                        else digits[1]), node.length)\n",
    "                res += length\n",
    "\n",
    "            if nodes_left and nodes_left[-1][1] == current_depth:\n",
    "                res += ','\n",
    "    return res + ';'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Color gradient generation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "def hex2rgb(h):\n",
    "    return tuple(int(h.lstrip('#')[i: i + 2], 16) for i in (0, 2 ,4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rgb2hex(r):\n",
    "    return '#{:02x}{:02x}{:02x}'.format(r[0], r[1], r[2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_color_palette(start, end, n=101):\n",
    "    \"\"\"Generate a gradient of 101 colors.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    start : str\n",
    "        start color in hex format\n",
    "    end : str\n",
    "        end color in hex format\n",
    "    n : int\n",
    "        number of colors to return\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    list of str\n",
    "        colors in hex format\n",
    "    \"\"\"\n",
    "    start_, end_ = hex2rgb(start), hex2rgb(end)\n",
    "    seqs = [np.linspace(start_[i], end_[i], n).astype(int) for i in range(3)]\n",
    "    rgbs = [[seqs[x][i] for x in range(3)] for i in range(n)]\n",
    "    return [rgb2hex(x) for x in rgbs]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_color_gradient(node2val, colors):\n",
    "    \"\"\"Deal with polytomic taxa.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    node2val : dict of float or int\n",
    "        node ID to value map\n",
    "    colors : list of str\n",
    "        101 colors for values of 0 to 100\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    dict of str\n",
    "        node ID to color map\n",
    "    \"\"\"\n",
    "    for id_, val in node2val.items():\n",
    "        if val is None or np.isnan(val) or val == '':\n",
    "            node2val[id_] = 0\n",
    "        elif not isinstance(val, int) and not isinstance(val, float):\n",
    "            raise ValueError('Invalid number %s.' % val)\n",
    "\n",
    "    # shrink larger integers to 0-100 range\n",
    "    max_val = max(node2val.values())\n",
    "    if max_val > 100:\n",
    "        for id_ in node2val:\n",
    "            node2val[id_] /= (max_val / 100)\n",
    "\n",
    "    # convert fraction into percentage, and percentage to integer\n",
    "    convert = True if max_val <= 1 else False\n",
    "    for id_ in node2val:\n",
    "        try:\n",
    "            node2val[id_] = (\n",
    "                int(node2val[id_] * 100) if convert else int(node2val[id_]))\n",
    "        except ValueError:\n",
    "            print('%s' % id_)\n",
    "\n",
    "    # map support to color\n",
    "    return {k: colors[v] for k, v in node2val.items()}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "iTOL file generation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "def write_itol_label(f, id2label):\n",
    "    \"\"\"Generate iTOL node label file.\"\"\"\n",
    "    f.write('LABELS\\n')\n",
    "    f.write('SEPARATOR TAB\\n')\n",
    "    f.write('DATA\\n')\n",
    "    for id_ in sort_node_ids(id2label):\n",
    "        f.write('%s\\t%s\\n' % (id_, id2label[id_]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "def write_itol_collapse(f, nodes_to_collapse):\n",
    "    \"\"\"Generate an iTOL collapse file.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    nodes_to_collapse : iterable of str\n",
    "        node IDs to collapse\n",
    "    f : file handle\n",
    "        file to write collapse information\n",
    "    \"\"\"\n",
    "    f.write('COLLAPSE\\n')\n",
    "    f.write('DATA\\n')\n",
    "    for id_ in nodes_to_collapse:\n",
    "        f.write('%s\\n' % id_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "def write_itol_tree_colors(f, id2color, target='branch',\n",
    "                           label_or_style='normal', size='1'):\n",
    "    \"\"\"Generate an iTOL tree colors file.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    id2label : dict of str\n",
    "        node ID to text map\n",
    "    f : file handle\n",
    "        file to write node texts\n",
    "    target, label_or_style, size : str or dict of str\n",
    "        iToL flavors, either a fixed value or a node ID to value map\n",
    "        target == \"type\" in iTOL jargon\n",
    "    \"\"\"\n",
    "    f.write('TREE_COLORS\\n')\n",
    "    f.write('SEPARATOR TAB\\n')\n",
    "    f.write('DATA\\n')\n",
    "    # format: ID, target, color, label_or_style, size_factor\n",
    "    for id_ in sort_node_ids(id2color):\n",
    "        f.write('%s\\t%s\\t%s\\t%s\\t%s\\n' % (\n",
    "            id_, target[id_] if isinstance(target, dict) else target,\n",
    "            id2color[id_], label_or_style[id_] if isinstance(\n",
    "                label_or_style, dict) else label_or_style,\n",
    "            size[id_] if isinstance(size, dict) else size))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "def write_itol_dataset_text(f, title, id2text, position='0.5',\n",
    "                            color='#000000', style='normal',\n",
    "                            size='1', rotation='0'):\n",
    "    \"\"\"Generate an iTOL text dataset file.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    id2label : dict of str\n",
    "        node ID to text map\n",
    "    f : file handle\n",
    "        file to write node texts\n",
    "    title : str\n",
    "        title of this dataset\n",
    "    position, color, style, size, rotation : str or dict of str\n",
    "        iToL flavors, either a fixed value or a node ID to value map\n",
    "    \"\"\"\n",
    "    f.write('DATASET_TEXT\\n')\n",
    "    f.write('SEPARATOR TAB\\n')\n",
    "    f.write('DATASET_LABEL\\t%s\\n' % title)\n",
    "    f.write('SHOW_INTERNAL\\t1\\n')\n",
    "    f.write('DATA\\n')\n",
    "    # format: ID, label, position, color, style, size_factor, rotation\n",
    "    for id_ in sort_node_ids(id2text):\n",
    "        text = id2text[id_]\n",
    "        if isinstance(text, float):\n",
    "            text = '%.3g' % text\n",
    "        f.write('%s\\t%s\\t%s\\t%s\\t%s\\t%s\\t%s\\n' % (\n",
    "            id_,\n",
    "            text,\n",
    "            position[id_] if isinstance(position, dict) else position,\n",
    "            color[id_] if isinstance(color, dict) else color,\n",
    "            style[id_] if isinstance(style, dict) else style,\n",
    "            size[id_] if isinstance(size, dict) else size,\n",
    "            rotation[id_] if isinstance(rotation, dict) else rotation))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "def write_itol_dataset_style(f, title, ids, target='branch', what='node',\n",
    "                             color='#000000', factor='normal', style='1',\n",
    "                             bgcolor=None):\n",
    "    \"\"\"Generate an iTOL style dataset file.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    f : file handle\n",
    "        file to write node texts\n",
    "    title : str\n",
    "        title of this dataset\n",
    "    ids : iterable of str\n",
    "        node ID list\n",
    "    target, what, color, factor, style, bgcolor : str or str or dict\n",
    "        iToL flavors, either a fixed value or a node ID to value map\n",
    "    \"\"\"\n",
    "    f.write('DATASET_STYLE\\n')\n",
    "    f.write('SEPARATOR TAB\\n')\n",
    "    f.write('DATASET_LABEL\\t%s\\n' % title)\n",
    "    f.write('COLOR\\t#000000\\n')\n",
    "    f.write('DATA\\n')\n",
    "    # format: ID, target, what, color, factor, style, bgcolor\n",
    "    for id_ in sort_node_ids(ids):\n",
    "        f.write('%s\\t%s\\t%s\\t%s\\t%s\\t%s' % (\n",
    "            id_,\n",
    "            target[id_] if isinstance(target, dict) else target,\n",
    "            what[id_] if isinstance(what, dict) else what,\n",
    "            color[id_] if isinstance(color, dict) else color,\n",
    "            factor[id_] if isinstance(factor, dict) else factor,\n",
    "            style[id_] if isinstance(style, dict) else style))\n",
    "        if bgcolor is not None:\n",
    "            f.write('\\t%s' % bgcolor[id_] if isinstance(bgcolor, dict) else bgcolor)\n",
    "        f.write('\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "FigTree file generation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In a FigTree-compatible Nexus tree file, nodes (tips and internal nodes) and taxa may contain attributes in the following format:\n",
    "```\n",
    "(taxon1,taxon2)[&!name=\"Escherichia coli\",support=90,range={80,95},!color=#ff0000]:1.234,...\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here \"!name\", \"support\", \"range\" and \"!color\" are the attributes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_figtree_attr_str(name, attr_db, attrs={}):\n",
    "    \"\"\"Generate a FigTree-compatible attribute string.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    name : str\n",
    "        name of node or taxon to annotate\n",
    "    attr_db : dict of dict of str\n",
    "        map of names to attributes\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    str\n",
    "        formatted attribute string\n",
    "\n",
    "    Notes\n",
    "    -----\n",
    "    For example, attr_db = {\n",
    "        '!name': {'N1': '\"spA\"', 'N2': '\"spB\"'...},\n",
    "        'support': {'N1': '90', 'N2': '75'...}\n",
    "    }\n",
    "    \n",
    "    For node \"N2\", the result will be `[&!name=\"spB\",support=75]`.\n",
    "    \n",
    "    All values should be str. Strings should be double-quoted. Tuples should be\n",
    "    written like `{\"spA\",0.95,#000000}`. Special FigTree-aware attributes such\n",
    "    as \"name\", \"color\" and \"collapse\" should have a prefix `!`.\n",
    "    \"\"\"\n",
    "    for attr in attr_db:\n",
    "        if name in attr_db[attr]:\n",
    "            val = attr_db[attr][name]\n",
    "            if val:  # omit null or empty string\n",
    "                attrs[attr] = val\n",
    "    attr_strs = []\n",
    "    for attr, val in sorted(attrs.items(), key=lambda x: x[0]):\n",
    "        attr_strs.append('%s=%s' % (attr, val))\n",
    "    return '[&%s]' % ','.join(attr_strs) if len(attr_strs) > 0 else ''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "code_folding": []
   },
   "outputs": [],
   "source": [
    "def add_figtree_node_attrs(tree, node2attrs):\n",
    "    \"\"\"Add FigTree-compatible attributes to nodes of a tree.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    tree : skbio.TreeNode\n",
    "        tree to add node attributes to\n",
    "    node2attrs : dict of dict\n",
    "        map of node names to attributes\n",
    "    \"\"\"\n",
    "    for node in tree.traverse(include_self=True):\n",
    "        if not node.name:\n",
    "            continue\n",
    "        attrs = {} if node.is_tip() else {'id': '\"%s\"' % node.name}\n",
    "        attr_str = make_figtree_attr_str(node.name, node2attrs, attrs)\n",
    "        node.name = ('%s%s' % (node.name, attr_str) if node.is_tip()\n",
    "                     else attr_str)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_figtree_taxon_attrs(tree, taxon2attrs):\n",
    "    \"\"\"Add FigTree-compatible attributes to the taxon labels.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    tree : skbio.TreeNode\n",
    "        tree to add node attributes to\n",
    "    taxon2attrs : dict of dict\n",
    "        map of taxa to attributes\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    list of str\n",
    "        taxon labels with attributes appended\n",
    "    \"\"\"\n",
    "    res = []\n",
    "    for taxon in sorted(tree.subset()):\n",
    "        attr_str = make_figtree_attr_str(taxon, taxon2attrs)\n",
    "        res.append('%s%s' % (taxon, attr_str))\n",
    "    return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "def write_figtree_nexus(tree, f, title='tree1', taxlabels=None):\n",
    "    \"\"\"Generate a FigTree-compatible Nexus tree file.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    tree : skbio.TreeNode\n",
    "        tree to add node attributes to\n",
    "    f : file handle\n",
    "        file to write nexus tree\n",
    "    title : str\n",
    "        title of the tree\n",
    "    taxlabels : list of str\n",
    "        custom taxon labels to write\n",
    "    \"\"\"\n",
    "    f.write('#NEXUS\\n')\n",
    "    f.write('begin taxa;\\n')\n",
    "    f.write('\\tdimensions ntax=%d;\\n' % tree.count(tips=True))\n",
    "    f.write('\\ttaxlabels\\n')\n",
    "    if taxlabels is None:\n",
    "        taxlabels = sorted(tree.subset())\n",
    "    for taxon in taxlabels:\n",
    "        f.write('\\t%s\\n' % taxon)\n",
    "    f.write(';\\n')\n",
    "    f.write('end;\\n')\n",
    "    f.write('\\n')\n",
    "    f.write('begin trees;\\n')\n",
    "    f.write('\\ttree %s = [&%s] ' % (\n",
    "        title, 'R' if len(tree.children) == 2 else 'U'))\n",
    "    f.write(format_newick(tree, operators=None))\n",
    "    f.write('\\n')\n",
    "    f.write('end;\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pre-processing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Read and process tree"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Read tree."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tree has 10575 tips and 10574 internal nodes.\n"
     ]
    }
   ],
   "source": [
    "tree = TreeNode.read(tree_fp)\n",
    "n, m = tree.count(), tree.count(tips=True)\n",
    "print('Tree has %d tips and %d internal nodes.' % (m, n - m))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "tips = tree.subset()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Convert null branch lengths to zero."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "for node in tree.traverse(include_self=False):\n",
    "    node.length = node.length or 0.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Get the precision (maximum number of float or scientific notion digits) of branch lengths. Will be useful in the correct formatting of branch lengths after collapsing the tree."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(10, 5)"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "max_f, max_e = 0, 0\n",
    "for node in tree.traverse():\n",
    "    if node.length is not None:\n",
    "        x = str(float(node.length))\n",
    "        if 'e' in x:\n",
    "            max_e = max(max_e, digits(str(float(x.split('e')[0]))))\n",
    "        else:\n",
    "            max_f = max(max_f, digits(x))\n",
    "max_f, max_e"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Calculate number of descendants of each node."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "node2n = {}\n",
    "for node in tree.postorder(include_self=True):\n",
    "    if node.is_tip():\n",
    "        node2n[node.name] = 1\n",
    "    else:\n",
    "        node2n[node.name] = sum([node2n[x.name] for x in node.children])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Read and process taxonomy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "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>kingdom</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>node</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>G000005825</th>\n",
       "      <td>Bacteria</td>\n",
       "      <td>Firmicutes_1</td>\n",
       "      <td>Bacilli_1</td>\n",
       "      <td>Bacillales_1</td>\n",
       "      <td>Bacillaceae_3</td>\n",
       "      <td>Bacillus_2</td>\n",
       "      <td>Bacillus pseudofirmus</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>G000006175</th>\n",
       "      <td>Archaea</td>\n",
       "      <td>Euryarchaeota_1</td>\n",
       "      <td>Methanococci</td>\n",
       "      <td>Methanococcales</td>\n",
       "      <td>Methanococcaceae</td>\n",
       "      <td>Methanococcus</td>\n",
       "      <td>Methanococcus voltae</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>G000006605</th>\n",
       "      <td>Bacteria</td>\n",
       "      <td>Actinobacteria</td>\n",
       "      <td>Actinobacteria</td>\n",
       "      <td>Corynebacteriales</td>\n",
       "      <td>Corynebacteriaceae</td>\n",
       "      <td>Corynebacterium</td>\n",
       "      <td>Corynebacterium falsenii</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>G000006725</th>\n",
       "      <td>Bacteria</td>\n",
       "      <td>Proteobacteria_1</td>\n",
       "      <td>Gammaproteobacteria</td>\n",
       "      <td>Xanthomonadales</td>\n",
       "      <td>Xanthomonadaceae</td>\n",
       "      <td>Xylella</td>\n",
       "      <td>Xylella fastidiosa</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>G000006785</th>\n",
       "      <td>Bacteria</td>\n",
       "      <td>Firmicutes_1</td>\n",
       "      <td>Bacilli_1</td>\n",
       "      <td>Lactobacillales</td>\n",
       "      <td>Streptococcaceae</td>\n",
       "      <td>Streptococcus</td>\n",
       "      <td>Streptococcus pyogenes</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "             kingdom            phylum                class  \\\n",
       "node                                                          \n",
       "G000005825  Bacteria      Firmicutes_1            Bacilli_1   \n",
       "G000006175   Archaea   Euryarchaeota_1         Methanococci   \n",
       "G000006605  Bacteria    Actinobacteria       Actinobacteria   \n",
       "G000006725  Bacteria  Proteobacteria_1  Gammaproteobacteria   \n",
       "G000006785  Bacteria      Firmicutes_1            Bacilli_1   \n",
       "\n",
       "                        order              family            genus  \\\n",
       "node                                                                 \n",
       "G000005825       Bacillales_1       Bacillaceae_3       Bacillus_2   \n",
       "G000006175    Methanococcales    Methanococcaceae    Methanococcus   \n",
       "G000006605  Corynebacteriales  Corynebacteriaceae  Corynebacterium   \n",
       "G000006725    Xanthomonadales    Xanthomonadaceae          Xylella   \n",
       "G000006785    Lactobacillales    Streptococcaceae    Streptococcus   \n",
       "\n",
       "                             species  \n",
       "node                                  \n",
       "G000005825     Bacillus pseudofirmus  \n",
       "G000006175      Methanococcus voltae  \n",
       "G000006605  Corynebacterium falsenii  \n",
       "G000006725        Xylella fastidiosa  \n",
       "G000006785    Streptococcus pyogenes  "
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dfr = pd.read_csv(taxonomy_fp, sep='\\t', index_col=0)\n",
    "dfr = dfr[dfr.index.isin(tips)]\n",
    "dfr.index.name = 'node'\n",
    "dfr.dropna().head(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['kingdom', 'phylum', 'class', 'order', 'family', 'genus', 'species']"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ranks = dfr.columns.tolist()\n",
    "ranks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tree annotation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generate node labels"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The lowest common ancestor (LCA) of genomes represented by each taxon will receive this taxon as the node label. One node may receive multiple taxa if they all meet this criterium.\n",
    "\n",
    "If this operation is applied to the tax2tree consensus strings (`consensus_ranks.tsv`), the outcome should match the labels decorated to the tree by tax2tree (`decorations_by_rank.tsv`).\n",
    "\n",
    "In the current analysis, the input file should be the tax2tree consensus string filled by taxa representing single genomes (`filled_ranks.tsv`). Therefore the outcome will contain more information. Both tips and internal nodes will be included."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "labels = {}\n",
    "for rank in ranks:\n",
    "    for taxon in dfr[rank].value_counts().index:\n",
    "        indices = dfr[dfr[rank] == taxon].index.tolist()\n",
    "        node = (indices[0] if len(indices) == 1\n",
    "                else tree.lca(list(tips.intersection(indices))).name)\n",
    "        labels.setdefault(node, {})[rank] = taxon"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "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>kingdom</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>node</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>G000005825</th>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>Bacillus pseudofirmus</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>G000006175</th>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>Methanococcus voltae</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>G000006725</th>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>Xylella fastidiosa</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "           kingdom phylum class order family genus                species\n",
       "node                                                                     \n",
       "G000005825     NaN    NaN   NaN   NaN    NaN   NaN  Bacillus pseudofirmus\n",
       "G000006175     NaN    NaN   NaN   NaN    NaN   NaN   Methanococcus voltae\n",
       "G000006725     NaN    NaN   NaN   NaN    NaN   NaN     Xylella fastidiosa"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dfl = pd.DataFrame.from_dict(labels, orient='index')\n",
    "dfl.index.name = 'node'\n",
    "dfl = dfl[ranks]\n",
    "dfl = dfl.loc[sorted(dfl.index, key=lambda x: (x[0], int(x[1:])))]\n",
    "dfl.head(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Get the highest-rank name when multiple ranks have names in a node label."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_highest_taxon(row):\n",
    "    \"\"\"Get the highest taxon in a row.\"\"\"\n",
    "    for rank in row.index:\n",
    "        if pd.notnull(row[rank]):\n",
    "            return rank, row[rank]\n",
    "    return np.nan, np.nan"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {
    "scrolled": true
   },
   "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>hrank</th>\n",
       "      <th>htaxon</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>node</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>G000005825</th>\n",
       "      <td>species</td>\n",
       "      <td>Bacillus pseudofirmus</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>G000006175</th>\n",
       "      <td>species</td>\n",
       "      <td>Methanococcus voltae</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>G000006725</th>\n",
       "      <td>species</td>\n",
       "      <td>Xylella fastidiosa</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>G000006745</th>\n",
       "      <td>species</td>\n",
       "      <td>Vibrio cholerae</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>G000006845</th>\n",
       "      <td>species</td>\n",
       "      <td>Neisseria gonorrhoeae</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "              hrank                 htaxon\n",
       "node                                      \n",
       "G000005825  species  Bacillus pseudofirmus\n",
       "G000006175  species   Methanococcus voltae\n",
       "G000006725  species     Xylella fastidiosa\n",
       "G000006745  species        Vibrio cholerae\n",
       "G000006845  species  Neisseria gonorrhoeae"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dfl['hrank'], dfl['htaxon'] = zip(*dfl.apply(get_highest_taxon, axis=1))\n",
    "dfl[['hrank', 'htaxon']].dropna().head(5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Collapse clades at or above a rank"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Identify the ranks to collapse."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collapse at the following ranks: kingdom, phylum, class, order, family, genus.\n"
     ]
    }
   ],
   "source": [
    "collapse_ranks = []\n",
    "for rank in ranks:\n",
    "    collapse_ranks.append(rank)\n",
    "    if collapse_rank and rank == collapse_rank:\n",
    "        break\n",
    "print('Collapse at the following ranks: %s.' % ', '.join(collapse_ranks))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Generate a list of candidate nodes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "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>hrank</th>\n",
       "      <th>htaxon</th>\n",
       "      <th>size</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>G000007185</th>\n",
       "      <td>class</td>\n",
       "      <td>Methanopyri</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>G000007485</th>\n",
       "      <td>genus</td>\n",
       "      <td>Tropheryma</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>G000008265</th>\n",
       "      <td>family</td>\n",
       "      <td>Picrophilaceae</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>G000008325</th>\n",
       "      <td>genus</td>\n",
       "      <td>Methylococcus</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>G000008625</th>\n",
       "      <td>genus</td>\n",
       "      <td>Aquifex</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "             hrank          htaxon  size\n",
       "node                                    \n",
       "G000007185   class     Methanopyri     1\n",
       "G000007485   genus      Tropheryma     1\n",
       "G000008265  family  Picrophilaceae     1\n",
       "G000008325   genus   Methylococcus     1\n",
       "G000008625   genus         Aquifex     1"
      ]
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_can = dfl[dfl['hrank'].isin(collapse_ranks)][['hrank', 'htaxon']]\n",
    "df_can['size'] = df_can.index.to_series().map(node2n)\n",
    "df_can.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Exclude nodes with number of descendants below threshold."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "288"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "if min_clade_size:\n",
    "    to_keep = []\n",
    "    for row in df_can.itertuples():\n",
    "        th = min_clade_size[row.hrank] if isinstance(min_clade_size, dict) else min_clade_size\n",
    "        if row.size >= th:\n",
    "            to_keep.append(row.Index)\n",
    "    df_can = df_can[df_can.index.isin(to_keep)]\n",
    "df_can.shape[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_ = df_can[df_can['htaxon'].str.contains('_\\d+$', regex=True)].copy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Exclude split clades of the same taxon which has less than a fraction of the dominant clade."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "270"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "if min_split_clade_frac > 0:\n",
    "    df_ = df_can[df_can['htaxon'].str.contains('_\\d+$', regex=True)].copy()\n",
    "    df_['taxon'], df_['idx'] = zip(*df_['htaxon'].apply(lambda x: x[::-1]).str.split(\n",
    "        '_', n=1).apply(lambda x: (x[1][::-1], x[0][::-1])))\n",
    "    top_clade_sizes = dict(df_.query('idx == \"1\"')[['taxon', 'size']].values.tolist())\n",
    "    df_ = df_[df_['size'] >= df_['taxon'].map(top_clade_sizes) * min_split_clade_frac]\n",
    "    df_can = df_can[df_can.index.isin(df_.index) | ~df_can['htaxon'].str.contains('_\\d+$', regex=True)]\n",
    "df_can.shape[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Get the dimensions of clades represented by internal nodes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "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>count</th>\n",
       "      <th>mean</th>\n",
       "      <th>std</th>\n",
       "      <th>min</th>\n",
       "      <th>25%</th>\n",
       "      <th>50%</th>\n",
       "      <th>75%</th>\n",
       "      <th>max</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>N2</th>\n",
       "      <td>669.0</td>\n",
       "      <td>0.923035</td>\n",
       "      <td>0.145750</td>\n",
       "      <td>0.592703</td>\n",
       "      <td>0.796175</td>\n",
       "      <td>0.900001</td>\n",
       "      <td>1.048704</td>\n",
       "      <td>1.226321</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>N3</th>\n",
       "      <td>9906.0</td>\n",
       "      <td>0.911626</td>\n",
       "      <td>0.165300</td>\n",
       "      <td>0.475954</td>\n",
       "      <td>0.786587</td>\n",
       "      <td>0.910457</td>\n",
       "      <td>1.016349</td>\n",
       "      <td>1.715364</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>N17</th>\n",
       "      <td>2.0</td>\n",
       "      <td>0.426866</td>\n",
       "      <td>0.000009</td>\n",
       "      <td>0.426860</td>\n",
       "      <td>0.426863</td>\n",
       "      <td>0.426866</td>\n",
       "      <td>0.426869</td>\n",
       "      <td>0.426872</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>N18</th>\n",
       "      <td>366.0</td>\n",
       "      <td>0.755701</td>\n",
       "      <td>0.107231</td>\n",
       "      <td>0.491407</td>\n",
       "      <td>0.673813</td>\n",
       "      <td>0.730512</td>\n",
       "      <td>0.844795</td>\n",
       "      <td>1.077242</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>N21</th>\n",
       "      <td>3.0</td>\n",
       "      <td>0.410779</td>\n",
       "      <td>0.093041</td>\n",
       "      <td>0.350268</td>\n",
       "      <td>0.357211</td>\n",
       "      <td>0.364154</td>\n",
       "      <td>0.441034</td>\n",
       "      <td>0.517913</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "      count      mean       std       min       25%       50%       75%  \\\n",
       "N2    669.0  0.923035  0.145750  0.592703  0.796175  0.900001  1.048704   \n",
       "N3   9906.0  0.911626  0.165300  0.475954  0.786587  0.910457  1.016349   \n",
       "N17     2.0  0.426866  0.000009  0.426860  0.426863  0.426866  0.426869   \n",
       "N18   366.0  0.755701  0.107231  0.491407  0.673813  0.730512  0.844795   \n",
       "N21     3.0  0.410779  0.093041  0.350268  0.357211  0.364154  0.441034   \n",
       "\n",
       "          max  \n",
       "N2   1.226321  \n",
       "N3   1.715364  \n",
       "N17  0.426872  \n",
       "N18  1.077242  \n",
       "N21  0.517913  "
      ]
     },
     "execution_count": 51,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tips = tree.subset()\n",
    "dimensions = {x: get_clade_dimensions(tree.find(x)) for x in df_can.index if x not in tips}\n",
    "df_dim = pd.DataFrame.from_dict(dimensions, orient='index')\n",
    "df_dim = df_dim.loc[sorted(df_dim.index, key=lambda x: int(x[1:]))]\n",
    "df_dim.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Determine which clades (as represented by nodes) should be collapsed."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The rationale is: Start from the lowest rank, move up the hierarchy. If a node is already marked as \"collapsed\", all its ancestral nodes will be prohibited from being selected."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Nodes to collapse: 207.\n"
     ]
    }
   ],
   "source": [
    "nodes_to_collapse = []\n",
    "nodes_to_skip = set()\n",
    "for rank in collapse_ranks[::-1]:\n",
    "    for node in df_can[df_can['hrank'] == rank].index:\n",
    "        if node not in nodes_to_skip:\n",
    "            nodes_to_collapse.append(node)\n",
    "            for anc in tree.find(node).ancestors():\n",
    "                nodes_to_skip.add(anc.name)\n",
    "print('Nodes to collapse: %d.' % len(nodes_to_collapse))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Calculate how many tips (genomes) are covered by the collapsed clades."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tips covered: 6944. Tips missed: 3631.\n"
     ]
    }
   ],
   "source": [
    "tips_covered = set()\n",
    "for name in nodes_to_collapse:\n",
    "    node = tree.find(name)\n",
    "    tips = set([name]) if node.is_tip() else node.subset()\n",
    "    if len(tips.intersection(tips_covered)) > 0:\n",
    "        raise ValueError('Overlapping clades detected.')\n",
    "    tips_covered.update(tips)\n",
    "tips_missed = tree.subset() - tips_covered\n",
    "print('Tips covered: %d. Tips missed: %d.'\n",
    "      % (len(tips_covered), len(tips_missed)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tree visualization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Tree and labels manipulation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Tree pruning"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Original tree dimensions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Original tree has 10575 tips and 10574 internal nodes.\n"
     ]
    }
   ],
   "source": [
    "tree_tips = tree.subset()\n",
    "tree_nodes = set(x.name for x in tree.non_tips(include_self=True))\n",
    "print('Original tree has %d tips and %d internal nodes.'\n",
    "      % (len(tree_tips), len(tree_nodes)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Prune tree to include collapsed clades only."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output tree has 6944 tips and 6996 internal nodes.\n"
     ]
    }
   ],
   "source": [
    "nodes_w_labels = [x for x in dfl['htaxon'].dropna().index if x in tree_nodes]\n",
    "tree1 = selective_prune(tree, tips_covered, nodes_w_labels) if collapse_rank and delete_uncollapsed else tree.copy()\n",
    "tree1_tips = tree1.subset()\n",
    "tree1_nodes = set(x.name for x in tree1.non_tips(include_self=True))\n",
    "print('Output tree has %d tips and %d internal nodes.'\n",
    "      % (len(tree1_tips), len(tree1_nodes)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Export pruned tree."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('pruned_tree.nwk', 'w') as f:\n",
    "    f.write('%s\\n' % format_newick(tree1, operators=None))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Tree shrinking"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Generate a tree in which the collapsed clades are actually deleted."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collapsed tree has 207 tips.\n"
     ]
    }
   ],
   "source": [
    "tree3 = tree1.copy()\n",
    "nodes_to_remove = []\n",
    "for node in tree3.non_tips():\n",
    "    if node.name in nodes_to_collapse:\n",
    "        node.length += df_dim[collapse_length_func][node.name]\n",
    "        nodes_to_remove.extend(node.children)\n",
    "tree3.remove_deleted(lambda x: x in nodes_to_remove)\n",
    "tree3.prune()\n",
    "print('Collapsed tree has %d tips.' % tree3.count(tips=True))\n",
    "with open('collapsed_tree.nwk', 'w') as f:\n",
    "    f.write(format_newick(tree3, operators=None))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Node label formatting"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Format node label strings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[('G000005825', 's__Bacillus pseudofirmus'),\n",
       " ('G000006175', 's__Methanococcus voltae'),\n",
       " ('G000006725', 's__Xylella fastidiosa'),\n",
       " ('G000006745', 's__Vibrio cholerae'),\n",
       " ('G000006845', 's__Neisseria gonorrhoeae')]"
      ]
     },
     "execution_count": 58,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "name_map = dfl['htaxon'].to_dict()\n",
    "if len(label_format_regexes) > 0:\n",
    "    for id_ in name_map:\n",
    "        for pattern, repl in label_format_regexes:\n",
    "            name_map[id_] = re.sub(pattern, repl, name_map[id_])\n",
    "if append_rank_code is True:\n",
    "    for id_ in name_map:\n",
    "        name_map[id_] = '%s__%s' % (dfl['hrank'][id_][0], name_map[id_])\n",
    "if append_clade_size is True:\n",
    "    for id_ in name_map:\n",
    "        n = node2n[id_]\n",
    "        if n > 1:\n",
    "            name_map[id_] = '%s (%d)' % (name_map[id_], node2n[id_])\n",
    "sorted(name_map.items())[:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [],
   "source": [
    "tip_name_map, node_name_map = {}, {}\n",
    "for id_, name in name_map.items():\n",
    "    if id_ in tree_tips:\n",
    "        tip_name_map[id_] = name\n",
    "    elif id_ in tree_nodes:\n",
    "        node_name_map[id_] = name"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Additional attributes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [],
   "source": [
    "if custom_attrs_fps:\n",
    "    dfa = {}\n",
    "    for name, fp in custom_attrs_fps.items():\n",
    "        dfa[name] = pd.read_table(fp, index_col=0, names=[name])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### FigTree file generation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Generate FigTree tip and node name maps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [],
   "source": [
    "figtree_tip_name_map = {k: '\"%s\"' % v for k, v in tip_name_map.items()}\n",
    "figtree_node_name_map = {k: '\"%s\"' % v for k, v in node_name_map.items()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [],
   "source": [
    "if collapse_rank and hide_uncollapsed:\n",
    "    for name in figtree_tip_name_map:\n",
    "        if name in tips_missed:\n",
    "            figtree_tip_name_map[name] = '\"\"'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let FigTree display internal node labels without displaying labels of tips (including collapsed clades)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [],
   "source": [
    "no_labels = set(nodes_to_collapse).union(tree1_tips)\n",
    "figtree_label_map = {k: v for k, v in figtree_node_name_map.items() if k not in no_labels}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [],
   "source": [
    "tip2attrs = {'!name': figtree_tip_name_map}\n",
    "node2attrs = {'!name': figtree_node_name_map, 'label': figtree_label_map}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Generate a FigTree collapse map."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[('N100', '{\"collapsed\",1.080952001}'),\n",
       " ('N10087', '{\"collapsed\",0.6039728893}'),\n",
       " ('N1016', '{\"collapsed\",0.9363783172}'),\n",
       " ('N1040', '{\"collapsed\",0.7580735353}'),\n",
       " ('N1043', '{\"collapsed\",0.5273656381}')]"
      ]
     },
     "execution_count": 65,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "figtree_collapse_map = {}\n",
    "tree_radius = max(x.accumulate_to_ancestor(tree1) for x in tree1.tips())\n",
    "for name in nodes_to_collapse:\n",
    "    if name not in tree1_tips:\n",
    "        length = df_dim[collapse_length_func][name]\n",
    "        height = tree_radius - tree1.find(name).accumulate_to_ancestor(tree1) - length\n",
    "        figtree_collapse_map[name] = '{\"collapsed\",%.*g}' % (max_f, height)\n",
    "sorted(figtree_collapse_map.items())[:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [],
   "source": [
    "node2attrs['!collapse'] = figtree_collapse_map"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Generate FigTree size map."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [],
   "source": [
    "node2attrs['size'] = node2n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Generate additional attributes for FigTree."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [],
   "source": [
    "if custom_attrs_fps:\n",
    "    for name, df_ in dfa.items():\n",
    "        if np.issubdtype(df_[name], np.number):\n",
    "            map_ = {k: str(v) for k, v in df_[name].iteritems()}\n",
    "        else:\n",
    "            map_ = {k: '\"%s\"' % v for k, v in df_[name].iteritems()}\n",
    "        tip2attrs[name] = node2attrs[name] = map_"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Write FigTree files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [],
   "source": [
    "tree2 = tree1.copy()\n",
    "taxlabels = add_figtree_taxon_attrs(tree2, tip2attrs)\n",
    "add_figtree_node_attrs(tree2, node2attrs)\n",
    "with open('figtree.tre', 'w') as f:\n",
    "    write_figtree_nexus(tree2, f, taxlabels=taxlabels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Task completed.\n"
     ]
    }
   ],
   "source": [
    "print('Task completed.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### iTOL files generation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Step 1: Upload the already exported pruned tree file (Newick format) to iTOL."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Write iTOL node label file. (Applies to both tips and internal nodes, including collapsed triangles.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('label.txt', 'w') as f:\n",
    "    write_itol_label(f, {**node_name_map, **tip_name_map})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Write iTOL branch text file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {},
   "outputs": [],
   "source": [
    "branch_name_map = {k: v for k, v in node_name_map.items() if k not in nodes_to_collapse}\n",
    "with open('branch_text.txt', 'w') as f:\n",
    "    # position = 0.5: at the middle of branch\n",
    "    write_itol_dataset_text(f, 'branch text', branch_name_map, position='0.5', size='1')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Write iTOL collapse file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "metadata": {},
   "outputs": [],
   "source": [
    "if collapse_rank:\n",
    "    with open('collapse.txt', 'w') as f:\n",
    "        write_itol_collapse(f, sorted(\n",
    "            x for x in nodes_to_collapse if x in tree1_nodes))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Write iTOL files for extra node attributes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {},
   "outputs": [],
   "source": [
    "color_gradient = make_color_palette(color_range[0], color_range[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "metadata": {},
   "outputs": [],
   "source": [
    "if custom_attrs_fps:\n",
    "    for name, df_ in dfa.items():\n",
    "        # node text\n",
    "        with open('%s_node_text.txt' % name, 'w') as f:\n",
    "            write_itol_dataset_text(\n",
    "                f, '%s node text' % name, df_[name].to_dict(), position='1', size='1')\n",
    "        # branch color gradient\n",
    "        if np.issubdtype(df_[name], np.number):\n",
    "            branch_color_map = make_color_gradient(df_[name].to_dict(), color_gradient)\n",
    "            with open('%s_branch_color.txt' % name, 'w') as f:\n",
    "                write_itol_dataset_style(\n",
    "                    f, '%s color gradient' % name, branch_color_map, target='branch',\n",
    "                    what='node', color=branch_color_map)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Task completed!\n"
     ]
    }
   ],
   "source": [
    "print('Task completed!')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Misc features"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A funky workaround for rendering collapsed clade with size proportional to number of descendants.\n",
    "\n",
    "In the output tree, collapsed clades are converted into polytomic branches with equal length. FigTree or iTOL will render it like a sector. But actually the sector constitutes numerous tiny branches.\n",
    "\n",
    "Then one may further recover the actual sectors by hacking the resulting SVG file."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```\n",
    "tree3 = tree1.copy()\n",
    "polytomized_labels = {}\n",
    "for node in tree3.levelorder():\n",
    "    if node.name in nodes_to_collapse:\n",
    "        if node.is_tip():\n",
    "            polytomized_labels[node.name] = name_map[node.name]\n",
    "        else:\n",
    "            length = df_dim[collapse_length_func][node.name]\n",
    "            tips = [x for x in node.tips()]\n",
    "            n = len(tips)\n",
    "            for i, tip in enumerate(tips):\n",
    "                tip.length = length\n",
    "                if i == int(n / 2):\n",
    "                    polytomized_labels[tip.name] = name_map[node.name]\n",
    "            node.children = tips\n",
    "with open('polytomized_pruned_tree.nwk', 'w') as f:\n",
    "    f.write(format_newick(tree3, operators=None))\n",
    "with open('polytomized_label.txt', 'w') as f:\n",
    "    write_itol_label(f, polytomized_labels)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The followng code will set branch widths proportional to number of descendants. However, iTOL is frozen when dealing with this number of styles."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```\n",
    "from math import sqrt, ceil\n",
    "node2n_intree = {k: v for k, v in node2n.items() if k in tree_tips or k in tree_nodes}\n",
    "factor_map = {k: sqrt(v) for k, v in node2n_intree.items()}\n",
    "max_val = max(factor_map.values())\n",
    "factor_map = {k: ceil(v / max_val * 100) / 10 for k, v in factor_map.items()}\n",
    "with open('branch_size_factor.txt', 'w') as f:\n",
    "    write_itol_dataset_style(\n",
    "        f, 'branch size factor', factor_map, target='branch',\n",
    "        what='node', factor=factor_map, style='normal')\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}