Back to index

python-biopython  1.60
NexusIO.py
Go to the documentation of this file.
00001 # Copyright (C) 2009 by Eric Talevich (eric.talevich@gmail.com)
00002 # This code is part of the Biopython distribution and governed by its
00003 # license. Please see the LICENSE file that should have been included
00004 # as part of this package.
00005 
00006 """I/O function wrappers for `Bio.Nexus` trees."""
00007 __docformat__ = "restructuredtext en"
00008 
00009 from itertools import chain
00010 
00011 from Bio.Nexus import Nexus
00012 from Bio.Phylo import Newick, NewickIO
00013 
00014 
00015 # Structure of a Nexus tree-only file
00016 NEX_TEMPLATE = """\
00017 #NEXUS
00018 Begin Taxa;
00019  Dimensions NTax=%(count)d;
00020  TaxLabels %(labels)s;
00021 End;
00022 Begin Trees;
00023  %(trees)s
00024 End;
00025 """
00026 
00027 # 'index' starts from 1; 'tree' is the Newick tree string
00028 TREE_TEMPLATE = "Tree tree%(index)d=[&U]%(tree)s;"
00029 
00030 
00031 def parse(handle):
00032     """Parse the trees in a Nexus file.
00033 
00034     Uses the old Nexus.Trees parser to extract the trees, converts them back to
00035     plain Newick trees, and feeds those strings through the new Newick parser.
00036     This way we don't have to modify the Nexus module yet. (Perhaps we'll
00037     eventually change Nexus to use the new NewickIO parser directly.)
00038     """
00039     nex = Nexus.Nexus(handle)
00040     # NB: Once Nexus.Trees is modified to use Tree.Newick objects, do this:
00041     # return iter(nex.trees)
00042     # Until then, convert the Nexus.Trees.Tree object hierarchy:
00043     def node2clade(nxtree, node):
00044         subclades = [node2clade(nxtree, nxtree.node(n)) for n in node.succ]
00045         return Newick.Clade(
00046                 branch_length=node.data.branchlength,
00047                 name=node.data.taxon,
00048                 clades=subclades,
00049                 confidence=node.data.support,
00050                 comment=node.data.comment)
00051 
00052     for nxtree in nex.trees:
00053         newroot = node2clade(nxtree, nxtree.node(nxtree.root))
00054         yield Newick.Tree(root=newroot, rooted=nxtree.rooted, name=nxtree.name,
00055                           weight=nxtree.weight)
00056 
00057 def write(obj, handle, **kwargs):
00058     """Write a new Nexus file containing the given trees.
00059 
00060     Uses a simple Nexus template and the NewickIO writer to serialize just the
00061     trees and minimal supporting info needed for a valid Nexus file.
00062     """
00063     trees = list(obj)
00064     writer = NewickIO.Writer(trees)
00065     nexus_trees = [TREE_TEMPLATE % {'index': idx+1, 'tree': nwk}
00066                    for idx, nwk in enumerate(
00067                         writer.to_strings(plain=False, plain_newick=True,
00068                                           **kwargs))]
00069     tax_labels = map(str, chain(*(t.get_terminals() for t in trees)))
00070     text = NEX_TEMPLATE % {
00071             'count':    len(tax_labels),
00072             'labels':   ' '.join(tax_labels),
00073             'trees':    '\n'.join(nexus_trees),
00074             }
00075     handle.write(text)
00076     return len(nexus_trees)