scripts/lpls/plots_lpls.py

@@ -1,5 +1,6 @@
 import pylab
 import matplotlib
+import networkx as nx
 
 def plot_corrloads(R, pc1=0,pc2=1,s=20, c='b', zorder=5,expvar=None,ax=None,drawback=True, labels=None):
     """ Correlation loading plot."""
@@ -36,3 +37,21 @@ def plot_corrloads(R, pc1=0,pc2=1,s=20, c='b', zorder=5,expvar=None,ax=None,draw
         for name, r in zip(labels, R):
             pylab.text(r[pc1], r[pc2], "  " + name)
     #pylab.show()
+
+def plot_dag(edge_dict, node_color='b', node_size=30,labels=None,nodelist=None,pos=None):
+    clean_edges = {}
+    for head, neigb in edge_dict.items():
+        head = head.replace(":", "_")
+        nei = [i.replace(":", "_") for i in neigb]
+        clean_edges[head] = nei
+    if pos==None:
+        G = nx.from_dict_of_lists(clean_edges, nx.DiGraph(name='GO'))
+        pos = nx.pydot_layout(G, prog='dot')
+    G = nx.from_dict_of_lists(edge_dict, nx.DiGraph(name='GO'))
+    if len(node_color)>1:
+        assert(len(node_color)==len(nodelist))
+    if labels!=None:
+        with_labels=True
+        
+    nx.draw_networkx(G,pos, with_labels=with_labels, node_size=node_size, node_color=node_color, nodelist=nodelist)
+