Source code for schrodinger.ui.sequencealignment.tree_node

"""
Implementation of multiple sequence viewer TreeNode class.

Copyright Schrodinger, LLC. All rights reserved.
"""

# Contributors: Piotr Rotkiewicz


class TreeNode:
    """
    The TreeNode class is a simple implementation of tree structure
    used to represent and display phylogenetic trees in the sequence viewer.
    """

    def __init__(self, level):

        self.branches = []
        self.lengths = []  # of branches
        self.sequences = []  # Associated sequences
        self.name = None  # of the node
        self.level = level  # node depth level
        self.coordinates = None  #Physical screen coords corresponding to node.

        # When drawing the tree, we use a copy (visible subset) of the original
        # tree. The node coordinates are stored in the original tree. Therefore
        # we need to remember a pointer to the original tree node in order
        # to retrieve the coordinates for mouse interactions.

        #: Corresponding node in the original tree.
        self.original_node = self

    def maxLevel(self, level=0):
        """
        Returns a maximum tree depth.

        :type level: int (optional)
        :param level: optional initial level

        :rtype: int
        :return: Maximum tree depth level.
        """
        if self.level > level:
            level = self.level
        for branch in self.branches:
            level = branch.maxLevel(level=level)
        return level

    def findNode(self, x, y, radius):
        """
        Finds a tree node at a specified position within a specified radius.

        :type x: int
        :param x: x coordinate

        :type y: int
        :param y: y coordinate

        :type radius: int
        :param radius: node search radius (used as a Manhattan distance)

        :rtype: `TreeNode`
        :return: Tree node at (x, y) or None if node was not found there.
        """
        if self.coordinates:
            bx, by = self.coordinates
            if x >= bx - radius and x <= bx + radius and \
               y >= by - radius and y <= by + radius:
                return self

        found = None
        for branch in self.branches:
            if branch.coordinates:
                bx, by = branch.coordinates
                if x >= bx - radius and x <= bx + radius and \
                   y >= by - radius and y <= by + radius:
                    found = branch
                    break
            if not found:
                found = branch.findNode(x, y, radius)
        return found

    def selectSequences(self):
        """
        Recursively select all sequences within self.
        """
        for branch in self.branches:
            for seq in branch.sequences:
                seq.selected = True
            branch.selectSequences()

    def hideSequences(self):
        """
        Recursively hide all sequences within self.
        """
        for branch in self.branches:
            for seq in branch.sequences:
                seq.selected = False
                seq.visible = False
            branch.hideSequences()

    def sequencesVisible(self):
        """
        Return True if any sequence within self is visible.

        :return: True if any sequence is visible
        :rtype: boolean
        """
        for seq in self.sequences:
            if seq.visible:
                return True
        for branch in self.branches:
            if branch.sequencesVisible():
                return True
        return False

    def swap(self):
        """
        Swaps tree branches.
        """
        if self.branches:
            for branch in self.branches:
                branch.swap()
        self.branches.reverse()
        self.lengths.reverse()
        self.sequences.reverse()

    def getSequenceList(self, sequence_list):
        """
        Builds a list of sequences by traversing the tree starting from self
        and appending the sequences to the sequence list.

        :type sequence_list: list of `Sequence`
        :param sequence_list: list of sequences to be extended

        :rtype: list of `Sequence`
        :return: built list of sequences
        """
        for seq in self.sequences:
            if seq.visible:
                sequence_list.append(seq)
        for branch in self.branches:
            branch.getSequenceList(sequence_list)
        return sequence_list

    def getVisibleTree(self):
        """
        Return a subset of the tree that includes only the visible sequences.
        This method is used to actually draw the tree with branches that are
        partially hidden.

        :rtype: `TreeNode`
        :return: visible tree
        """
        new_node = TreeNode(self.level)
        for index, branch in enumerate(self.branches):
            if branch.sequencesVisible():
                new_node.branches.append(branch.getVisibleTree())
                new_node.name = self.name
                new_node.original_node = self
        for seq in self.sequences:
            if seq.visible:
                new_node.sequences.append(seq)
        return new_node

    def buildSequenceList(self, sequence_group):
        """
        Build a list of sequences using a given sequence group as a source.

        Initially, the tree has no sequences associated with it. Every node
        has a name corresposing to a particular sequence index in the sequence
        group. This method finds these correspoding sequences in the sequence
        group and adds them to the node's sequence list.

        This method expects the sequence names to be in format used by
        fileio.load_clustal_file function, i.e. in format SEQxxxx where xxxx
        corresponds to a sequence index in the sequence group.

        :type sequence_group: `SequenceGroup`
        :param sequence_group: source sequence group
        """
        self.sequences = []
        if len(self.branches) > 0:
            for branch in self.branches:
                branch.buildSequenceList(sequence_group)
        elif self.name:
            index = int(self.name[3:])
            if index >= 0 and index < len(sequence_group.sequences):
                self.sequences.append(sequence_group.sequences[index])