Advertisement
Guest User

Untitled

a guest
Jan 29th, 2020
108
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 3.17 KB | None | 0 0
  1. #!/usr/bin/env python3
  2. import sys
  3.  
  4.  
  5. class Node:
  6.     def __init__(self, name):
  7.         self.name = name
  8.  
  9.     def __str__(self):
  10.         if self.name is not None:
  11.             return self.name
  12.  
  13.  
  14. class Edge:
  15.     def __init__(self, nodeA, nodeB):
  16.         self.nodeA = nodeA
  17.         self.nodeB = nodeB
  18.         self.nodes = [nodeA, nodeB]
  19.  
  20.  
  21. class Tree:
  22.     def __init__(self, nodes, edges):
  23.         self.nodes = nodes
  24.         self.edges = edges
  25.  
  26.     def copy(self):
  27.         node_mem = {node: Node(node.name) for node in self.nodes}
  28.         nodes = list(node_mem.values())
  29.         edges = [
  30.             Edge(node_mem[edge.nodeA], node_mem[edge.nodeB])
  31.             for edge in self.edges
  32.         ]
  33.  
  34.         tree = Tree(
  35.             nodes=nodes,
  36.             edges=edges
  37.         )
  38.         return tree
  39.  
  40.  
  41. def enumerate_trees(leaves):
  42.     if len(leaves) == 2:
  43.         nodeA = Node(leaves[0])
  44.         nodeB = Node(leaves[1])
  45.         tree = Tree(
  46.             nodes=[nodeA, nodeB],
  47.             edges=[Edge(nodeA, nodeB)]
  48.         )
  49.         return [tree]
  50.  
  51.     old_trees = enumerate_trees(leaves[:-1])
  52.     new_leaf_name = leaves[-1]
  53.     new_trees = []
  54.  
  55.     for old_tree in old_trees:
  56.         for index in range(len(old_tree.edges)):
  57.             new_tree = old_tree.copy()
  58.             edge = new_tree.edges[index]
  59.  
  60.             new_tree.edges.remove(edge)
  61.  
  62.             internal = Node(None)
  63.             new_tree.nodes.append(internal)
  64.  
  65.             new_leaf = Node(new_leaf_name)
  66.             new_tree.nodes.append(new_leaf)
  67.  
  68.             new_tree.edges.append(Edge(edge.nodeA, internal))
  69.             new_tree.edges.append(Edge(edge.nodeB, internal))
  70.             new_tree.edges.append(Edge(new_leaf, internal))
  71.  
  72.             new_trees.append(new_tree)
  73.  
  74.     return new_trees
  75.  
  76.  
  77. def newick_format(tree_in):
  78.     tree = tree_in.copy()
  79.  
  80.     if len(tree.nodes) == 1:
  81.         return "{};".format(tree.nodes[0])
  82.     if len(tree.nodes) == 2:
  83.         return "({},{});".format(*tree.nodes)
  84.  
  85.     for candidate_node in tree.nodes:
  86.         if candidate_node.name is not None:
  87.             continue
  88.  
  89.         adjacent_edges = [
  90.             edge for edge in tree.edges if candidate_node in edge.nodes
  91.         ]
  92.         adjacent_nodes = [
  93.             node
  94.             for edge in adjacent_edges
  95.             for node in edge.nodes
  96.             if node in edge.nodes
  97.             and node is not candidate_node
  98.             and node.name is not None
  99.         ]
  100.  
  101.         if len(adjacent_nodes) == 2 or len(adjacent_nodes) == 3:
  102.             leafA, leafB = adjacent_nodes[0], adjacent_nodes[1]
  103.             edges_to_cut = list(filter(lambda edge: leafA in edge.nodes or leafB in edge.nodes, adjacent_edges))
  104.  
  105.             candidate_node.name = "({},{})".format(leafA, leafB)
  106.  
  107.             tree.nodes.remove(leafA)
  108.             tree.nodes.remove(leafB)
  109.             for edge in edges_to_cut:
  110.                 tree.edges.remove(edge)
  111.  
  112.             return newick_format(tree)
  113.  
  114.  
  115. if __name__ == "__main__":
  116.     with open(sys.argv[1]) as file_handle:
  117.         leaves = file_handle.read().split()
  118.     trees = enumerate_trees(leaves)
  119.     print("\n".join([newick_format(tree) for tree in trees]))
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement