treeplot.py

#!/usr/bin/env python

"""Script to plot trees with ASR using ete3."""

import argparse
import sys

import pandas
import ete3

def main(args):
    """Run the CLI."""
    parser = argparse.ArgumentParser(
        description="Squish covarion ancestral state reconstructions")
    parser.add_argument(
        "tree",
        nargs="?",
        type=argparse.FileType("r"),
        default=sys.stdin,
        help="The tree file with covarion ancestral states.")
    parser.add_argument(
        "--plot",
        action='append',
        default=[],
        help="Which feature(s) to plot")
    parser.add_argument(
        "--component-string",
        default="recon_lexicon{feature:}_component{component:}",
        help="""Python format string to translate the feature name and component
        number into a compontent string.""")
    parser.add_argument(
        "--data",
        type=argparse.FileType("r"),
        help="Original data file for highlighting tips as “given”")
    parser.add_argument(
        "--strip",
        default=",:()[] ",
        help="Strip these characters from feature names for the tree")
    args = parser.parse_args(args)

    if args.data:
        data = pandas.read_csv(
            args.data,
            sep="\t")
        data = {feature: set(values["Language_ID"])
                for feature, values in data.groupby("Feature_ID")}
    else:
        data = {}
    tree = ete3.Tree(args.tree.read())

    def formatter(feature, component, strip=args.strip):
        for s in args.strip:
            feature = feature.replace(s, "")
        return args.component_string.format(feature=feature,
                                            component=component)
    for meaning in args.plot:
        plot_asr(
            meaning, tree.copy(),
            data.get(meaning, []),
            formatter)


colors = ['#a6cee3', '#1f78b4', '#b2df8a', '#33a02c', '#fb9a99',
          '#e31a1c', '#fdbf6f', '#ff7f00', '#cab2d6', '#6a3d9a',
          '#ffff99', '#b15928', '#000000', '#dddddd', '#ff2222',
          '#22ff22', '#2222ff', '#444444', '#888888', "#ff0000",
          "#ffff00", "#00ff00", "#00ffff", "#0000ff", "#ff00ff",
          "#770000", "#777700", "#007700", "#007777", "#000077",
          "#770077",]


def plot_face(sequence, scale=10, colors=colors, line=False):
    """Generate a StackedBarFace for the binarized features."""
    x = sum(sequence)
    return ete3.StackedBarFace(
        [s*100/x for s in sequence],
        scale*x,
        scale,
        colors=colors,
        line_color="#000000" if line else None)


def plot_asr(item, tree,
             sure_tips=[],
             formatter="recon_lexicon{feature:}_component{component:}",
             colors=colors):
    """Plot the ancestral state reconstruction of item according to tree."""
    ts = ete3.TreeStyle()
    ts.title.add_face(ete3.TextFace(item, fsize=20), column=0)

    c_max = 0
    for node in tree.traverse('preorder'):
        if c_max == 0:
            frequencies = []
            # Obtain the size of the feature
            while hasattr(node, formatter(
                    feature=item, component=c_max)):
                frequencies.append(float(getattr(node, formatter(
                    feature=item, component=c_max))))
                c_max += 1
        else:
            try:
                frequencies = [
                    float(getattr(node, formatter(
                        feature=item, component=c)))
                    for c in range(c_max)]
            except AttributeError:
                node.set_style(ete3.NodeStyle(size=0))
                continue
        if node.is_leaf() and node.name in sure_tips:
            node.add_face(
                plot_face(frequencies, colors=colors, line=True),
                column=0, position='float')
        else:
            node.add_face(
                plot_face(frequencies, colors=colors),
                column=0, position='float')

        node.set_style(ete3.NodeStyle(size=0))

    ts.title.add_face(plot_face([1 for _ in frequencies], scale=20), column=0)

    tree.render("tree_{:}.pdf".format(item.format("")), tree_style=ts)


if __name__ == "__main__":
    main(sys.argv[1:])