Uses dendropy to traverse a tree and modifies the length of each edge

alter_edge_lengths.py

#!/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, ';'