diff --git a/SConstruct b/SConstruct index 3205024..aa0f509 100644 --- a/SConstruct +++ b/SConstruct @@ -446,7 +446,7 @@ def add_cluster_analysis(w): 'prune_strategy': prune_strategy, 'asr_prog': asr_prog, # Just 100 for everyone now - 'prune_count': 100} + 'prune_count': 10} for prune_strategy, asr_prog in itertools.product( ['min_adcl', 'seed_lineage'] if 'seed' in c else ['min_adcl'], @@ -471,6 +471,15 @@ def add_cluster_analysis(w): + (" --seed " + c['seed']['id'] if 'seed' in c else '') + " $SOURCE $TARGET") + # create png showing included seqs (kept in pruning) as red + @w.add_target() + def pruned_cluster_fasttree_png(outdir, c): + pruned_cluster_fasttree_png = env.Command( + path.join(outdir, "pruned_cluster_fasttree.png"), + [c["fasttree"], c["pruned_ids"]], + "xvfb-run -a bin/annotate_fasttree_tree.py $SOURCES " + " --naive %s" % options['inferred_naive_name'] + (" --seed " + c['seed']['id'] if 'seed' in c else '') + " --output-path $TARGET") + env.Depends(pruned_cluster_fasttree_png, "bin/annotate_fasttree_tree.py") + return pruned_cluster_fasttree_png @w.add_target() def cluster_mapping(outdir, c): diff --git a/bin/annotate_fasttree_tree.py b/bin/annotate_fasttree_tree.py new file mode 100755 index 0000000..126f93e --- /dev/null +++ b/bin/annotate_fasttree_tree.py @@ -0,0 +1,51 @@ +#!/usr/bin/env python + +import argparse +import os +from ete3 import Tree, TextFace, TreeStyle + + +if __name__ == '__main__': + parser = argparse.ArgumentParser(description="Annotate FastTree tree with prune ids.") + parser.add_argument( + 'tree_path', type=str, + help="Path to FastTree tree file.") + parser.add_argument( + 'ids_path', type=str, + help="Path to prune id file.") + parser.add_argument( + '--naive', type=str, required=True, + help="The name of the naive sequence.") + parser.add_argument( + '--seed', type=str, + help="The name of the seed sequence.") + parser.add_argument( + '--output-path', type=str, required=True, + help="The PNG output file path.") + + args = parser.parse_args() + + tree = Tree(args.tree_path, format=1) + tree.set_outgroup(tree & args.naive) + + with open(args.ids_path) as f: + ids = f.readlines() + ids = [id.rstrip("\n") for id in ids] + + for leaf_node in tree.get_leaves(): + if leaf_node.name in ([args.naive, args.seed] if args.seed else [args.naive]): + color = "blue" + elif leaf_node.name in ids: + color = "red" + else: + color = "black" + + node_face = TextFace(leaf_node.name, fsize=6, fgcolor=color) + leaf_node.add_face(node_face, column=0, position="float") + + ts = TreeStyle() + ts.mode="c" + ts.show_leaf_name = False + ts.show_scale = False + + tree.render(args.output_path, tree_style=ts)