Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- #!/usr/bin/env python3
- import sys
- class Node:
- def __init__(self, name):
- self.name = name
- def __str__(self):
- if self.name is not None:
- return self.name
- class Edge:
- def __init__(self, nodeA, nodeB):
- self.nodeA = nodeA
- self.nodeB = nodeB
- self.nodes = [nodeA, nodeB]
- class Tree:
- def __init__(self, nodes, edges):
- self.nodes = nodes
- self.edges = edges
- def copy(self):
- node_mem = {node: Node(node.name) for node in self.nodes}
- nodes = list(node_mem.values())
- edges = [
- Edge(node_mem[edge.nodeA], node_mem[edge.nodeB])
- for edge in self.edges
- ]
- tree = Tree(
- nodes=nodes,
- edges=edges
- )
- return tree
- def enumerate_trees(leaves):
- if len(leaves) == 2:
- nodeA = Node(leaves[0])
- nodeB = Node(leaves[1])
- tree = Tree(
- nodes=[nodeA, nodeB],
- edges=[Edge(nodeA, nodeB)]
- )
- return [tree]
- old_trees = enumerate_trees(leaves[:-1])
- new_leaf_name = leaves[-1]
- new_trees = []
- for old_tree in old_trees:
- for index in range(len(old_tree.edges)):
- new_tree = old_tree.copy()
- edge = new_tree.edges[index]
- new_tree.edges.remove(edge)
- internal = Node(None)
- new_tree.nodes.append(internal)
- new_leaf = Node(new_leaf_name)
- new_tree.nodes.append(new_leaf)
- new_tree.edges.append(Edge(edge.nodeA, internal))
- new_tree.edges.append(Edge(edge.nodeB, internal))
- new_tree.edges.append(Edge(new_leaf, internal))
- new_trees.append(new_tree)
- return new_trees
- def newick_format(tree_in):
- tree = tree_in.copy()
- if len(tree.nodes) == 1:
- return "{};".format(tree.nodes[0])
- if len(tree.nodes) == 2:
- return "({},{});".format(*tree.nodes)
- for candidate_node in tree.nodes:
- if candidate_node.name is not None:
- continue
- adjacent_edges = [
- edge for edge in tree.edges if candidate_node in edge.nodes
- ]
- adjacent_nodes = [
- node
- for edge in adjacent_edges
- for node in edge.nodes
- if node in edge.nodes
- and node is not candidate_node
- and node.name is not None
- ]
- if len(adjacent_nodes) == 2 or len(adjacent_nodes) == 3:
- leafA, leafB = adjacent_nodes[0], adjacent_nodes[1]
- edges_to_cut = list(filter(lambda edge: leafA in edge.nodes or leafB in edge.nodes, adjacent_edges))
- candidate_node.name = "({},{})".format(leafA, leafB)
- tree.nodes.remove(leafA)
- tree.nodes.remove(leafB)
- for edge in edges_to_cut:
- tree.edges.remove(edge)
- return newick_format(tree)
- if __name__ == "__main__":
- with open(sys.argv[1]) as file_handle:
- leaves = file_handle.read().split()
- trees = enumerate_trees(leaves)
- print("\n".join([newick_format(tree) for tree in trees]))
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement