first try on matsengrp#250; using ecgtheow script to color pruned nodes

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):

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)