etetoolkit / ete

Python package for building, comparing, annotating, manipulating and visualising trees. It provides a comprehensive API and a collection of command line tools, including utilities to work with the NCBI taxonomy tree.
http://etetoolkit.org
GNU General Public License v3.0
797 stars 214 forks source link

Fix issue #721 ultrametric attribute malfunction when setting from TreeStyle() #729

Closed dengzq1234 closed 1 year ago

dengzq1234 commented 1 year ago

Fix issue #721, this fix focus on simplify the complex conditions for initial ultrametric setting from TreeStyle()

the old code is

...
        ultrametric = args.get('ultrametric') == '1' # asked for ultrametric?
        if ultrametric and not tree.style.ultrametric:  # change to on
            tree.tree.to_ultrametric()
            ops.update_sizes_all(tree.tree)
            initialize_tree_style(tree, ultrametric=True)
        elif not ultrametric and tree.style.ultrametric:  # change to off
            app.trees.pop(tid, None)  # delete from memory
            # Forces it to be reloaded from disk next time it is accessed.
...

tree.style.ultrametric in here is unnecessary because regardless ultrametric is set True/False in TreeStyle() or swtich-on/off in control panel or not. ultrametric = args.get('ultrametric') == '1' calling will be consistent as ultrametric on/off.

So we only need

        if ultrametric:
            # If we want ultrametric and the tree is not already ultrametric, change to on
            tree.tree.to_ultrametric()
            ops.update_sizes_all(tree.tree)
            initialize_tree_style(tree, ultrametric=True)
        else:
            # If we do not want ultrametric, regardless of the current state of the tree, delete from memory.
            app.trees.pop(tid, None)  # delete from memory
            # If the tree was ultrametric, this also forces it to be reloaded from disk next time it is accessed.

which can be ultrametric attribute can be fully functional no matter in TreeStyle() setting or control panel

demo code

import argparse
from ete4 import Tree
from ete4.smartview import TreeStyle, TreeLayout

# Set up argument parser
parser = argparse.ArgumentParser(description='Modify tree style based on command line flag.')
parser.add_argument('--ultrametric', action='store_true', help='Set the tree to be ultrametric.')

# Parse arguments
args = parser.parse_args()

# Create your tree
t = Tree()
t.populate(100, random_branches=True)

# Function to modify the tree style
def modify_tree_style(tree_style):
    # tree_style.aligned_grid = False
    tree_style.collapse_size = 10
    tree_style.ultrametric = args.ultrametric  # Set based on command line flag

# Create a TreeLayout object, passing in the function
tree_layout = TreeLayout(name="MyTreeLayout", ts=modify_tree_style)

# Add the layout to a list
layouts = [tree_layout]

# Explore the tree with the specified layout
t.explore(keep_server=True, layouts=layouts)

image