better draw

RNAmigos/ged.py

@@ -351,6 +351,7 @@ def graph_align(g1, g2, ops, title="", save=None, rna=False):
     """
     Draw aligned graphs.
     """
+    from matplotlib import cm
 
     f = nx.compose(g1, g2)
     path = reversed(list(ops.path_iter()))
@@ -359,6 +360,8 @@ def graph_align(g1, g2, ops, title="", save=None, rna=False):
     edit_edges = []
     edge_list = []
     color_list = []
+    nucs = ['A', 'U', 'C', 'G', 'L', 'N']
+    nt_color = {n:i for i,n in enumerate(nucs)}
     for i,o in enumerate(path):
         print(o.op)
         e1, e2 = o.op
@@ -370,16 +373,19 @@ def graph_align(g1, g2, ops, title="", save=None, rna=False):
     c = plt.get_cmap('Paired')
     pos = nx.spring_layout(f)
     # pos = rna_layout.circular_layout(f)
-    # nx.draw_networkx(f, pos, color='red')
-    print(f.edges)
+    nx.draw_networkx_nodes(f, pos, nodelist=['NILL'], color='red')
     config = {'node_size': 500, 'alpha':0.9, 'font_size':25}
-    nx.draw(f, pos, nodelist=['NILL'], color='red', **config)
+    # nx.draw(f, pos, nodelist=['NILL'], color='red', **config)
     nx.draw_networkx_edges(f, pos,
-        edgelist=edit_edges, edge_color='grey', width=3)
-    nx.draw_networkx_nodes(f, pos, nodelist=g1.nodes, node_color='blue', **config)
-    nx.draw_networkx_edges(f, pos, edgelist=g1.edges)
-    nx.draw_networkx_nodes(f, pos, nodelist=g2.nodes, node_color='green', **config)
-    nx.draw_networkx_edges(f, pos, edgelist=g2.edges)
+        edgelist=edit_edges, edge_color='grey', style='dotted', width=3)
+    nx.draw_networkx_nodes(f, pos, nodelist=g1.nodes,
+        node_color=[cm.Set3(nt_color[g1.nodes.data()[n]['nt']]) for n in g1.nodes], **config)
+    nx.draw_networkx_edges(f, pos, edgelist=g1.edges, width=3)
+    nx.draw_networkx_labels(f, pos)
+    # nx.draw_networkx_nodes(f, pos, nodelist=g2.nodes, node_color='green', **config)
+    nx.draw_networkx_nodes(f, pos, nodelist=g2.nodes,
+        node_color=[cm.Set3(nt_color[g2.nodes.data()[n]['nt']]) for n in g2.nodes], **config)
+    nx.draw_networkx_edges(f, pos, edgelist=g2.edges, width=3)
 
     make_label = lambda s: labels[s[:2]] + labels[s[0::2]] if len(set(s[1:])) == 2\
         else labels[s[:2]]
@@ -413,6 +419,7 @@ def graph_align(g1, g2, ops, title="", save=None, rna=False):
     else:
         nx.draw_networkx_edge_labels(f, pos, edge_labels=edge_labels, font_size=config['font_size'])
 
+    plt.axis('off')
     plt.title(f"GED: {cost} " + title)
     if save:
         plt.savefig(save, format="pdf")