Created
June 6, 2011 15:52
-
-
Save mtholder/1010511 to your computer and use it in GitHub Desktop.
Uses dendropy to traverse a tree and modifies the length of each edge
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
#!/usr/bin/env python | |
import sys | |
import dendropy | |
if __name__ == '__main__': | |
from optparse import OptionParser | |
from optparse import OptionGroup | |
import sys, os | |
_script_name = os.path.split(sys.argv[0])[1] | |
_program_version = 0.1 | |
description = '''A script for altering the edge lengths of a tree in predictable ways. | |
In 'epsilon' weighting: | |
epsilon should be a float. | |
If `orig` is the original edge length then the resulting | |
edge length will be orig/(orig + epsilon) | |
In 'scale' weighting: | |
scale should be a float | |
If `orig` is the original edge length then the resulting | |
edge length will be orig*scale | |
''' | |
parser = OptionParser(usage="%(script_name)s [options] focal-tree-filepath" % {'script_name' : _script_name}, | |
add_help_option=True, | |
version=_program_version, | |
description=description) | |
parser.add_option("-f","--format", dest="format", default="newick", help="newick or NEXUS") | |
parser.add_option("-i","--internal", dest="internals", default=False, action="store_true", help="modify internal edges") | |
parser.add_option("-t","--terminal", dest="terminals", default=False, action="store_true", help="modify terminal edges") | |
parser.add_option("--epsilon", dest="epsilon", default=None, type='float', help="The amount to add to the denominator (to reduce the branch length)") | |
parser.add_option("--scale", dest="scale", default=None, type='float', help="The multiplier for each branch") | |
parser.add_option("--cap", dest="cap", default=None, type='float', help="The maximum value for each branch") | |
parser.add_option("--min", dest="min", default=None, type='float', help="The minimum value for each branch") | |
(opts, args) = parser.parse_args() | |
mod_terminals = opts.terminals | |
mod_internals = opts.internals | |
if not (mod_terminals or mod_internals): | |
sys.exit('Modification of internals (-i) and/or terminals (-t) must requested') | |
focal_tree_format = opts.format | |
if len(args) == 0: | |
stream = sys.stdin | |
else: | |
focal_tree_file = args[0] | |
stream=open(focal_tree_file, 'rU') | |
if opts.epsilon is None: | |
if opts.scale is None: | |
sys.exit('--epsilon or --scale option must be used') | |
else: | |
transform_func = lambda x: x*opts.scale | |
else: | |
if opts.scale is not None: | |
sys.exit('--epsilon and --scale options cannot both be used') | |
transform_func = lambda x: x/(x + opts.epsilon) | |
if opts.cap is not None: | |
if opts.min is not None: | |
bound_func = lambda x: max(opts.min, min(opts.cap, transform_func(x))) | |
else: | |
bound_func = lambda x: min(opts.cap, transform_func(x)) | |
else: | |
if opts.min is not None: | |
bound_func = lambda x: max(opts.min, transform_func(x)) | |
else: | |
bound_func = transform_func | |
tree_set = dendropy.TreeList(stream=stream, schema=focal_tree_format) | |
for tree in tree_set: | |
for edge in tree.preorder_edge_iter(): | |
if edge.tail_node: | |
is_term = edge.is_terminal() | |
modify = ((not is_term) and mod_internals) or (is_term and mod_terminals) | |
if modify: | |
edge.length = bound_func(edge.length) | |
print tree, ';' |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment