Provide ASTRAL as an alternative to greedy consensus tree

stag/stag.py

@@ -222,8 +222,21 @@ def ProcessTrees(dir_in, dir_matrices, dir_trees_out, GeneToSpecies, qVerbose=Tr
     if qVerbose: print("\nExamined %d trees" % (nSuccess + nNotAllPresent + nFail))
     print("%d trees had all species present and will be used by STAG to infer the species tree\n" % nSuccess)
 
-def InferSpeciesTree(tree_dir, species, outputFN):
-    t = cons.ConsensusTree(tree_dir)
+def Astral(tree_dir, astral_jar_file):
+    treesFN = tree_dir + "../TreesFile.txt"
+    with open(treesFN, 'wb') as outfile:
+        for fn in glob.glob(tree_dir + "/*"):
+            t = tree.Tree(fn)    
+            outfile.write(t.write(format=9) + "\n")
+    speciesTreeFN = tree_dir + "../SpeciesTree_ids.txt"
+    subprocess.call(" ".join(["java", "-Xmx6000M", "-jar", astral_jar_file, "-i", treesFN, "-o", speciesTreeFN]), shell=True)
+    return tree.Tree(speciesTreeFN)     
+
+def InferSpeciesTree(tree_dir, species, outputFN, astral_jar=None):
+    if astral_jar == None:
+        t = cons.ConsensusTree(tree_dir)
+    else:
+        t = Astral(tree_dir, astral_jar)
     for n in t:
         n.name = species[int(n.name)]
     t.write(outfile=outputFN)
@@ -250,7 +263,10 @@ def main(args):
     os.mkdir(dir_trees_out)
     ProcessTrees(dir_in, dir_matrices, dir_trees_out, gene_to_species, qVerbose=(not args.quiet))
     outputFN = dir_out + "SpeciesTree.tre"
-    InferSpeciesTree(dir_trees_out, gene_to_species.species, outputFN)
+    if args.astral_jar == None:
+        InferSpeciesTree(dir_trees_out, gene_to_species.species, outputFN)
+    else:
+        InferSpeciesTree(dir_trees_out, gene_to_species.species, outputFN, astral_jar=args.astral_jar)
     print("STAG species tree: " + os.path.abspath(outputFN) + "\n")
 
 if __name__ == "__main__":
@@ -265,5 +281,6 @@ def main(args):
     parser.add_argument("species_map", help = "Map file from gene names to species names, or SpeciesIDs.txt file from OrthoFinder")
     parser.add_argument("gene_trees", help = "Directory conaining gene trees")
     parser.add_argument("-q", "--quiet", help = "Only print sparse output", action="store_true")
+    parser.add_argument("-a", "--astral_jar", help = "ASTRAL jar file. Use ASTRAL to combine STAG species tree estimates instead of greedy consensus tree.")
     args = parser.parse_args()
     main(args)