Added DAG plot to gobrowser module and smokers workflow.

workflows/gobrowser.py

@@ -2,10 +2,14 @@
 import gtk
 from fluents import dataset, logger, plots, workflow, fluents, project
 import geneontology
+from matplotlib.nxutils import points_inside_poly
+import matplotlib
 #from scipy import array, randn, log, ones, zeros
 from scipy import *
+from numpy import matlib
 import networkx
 import re
+import rpy
 
 EVIDENCE_CODES=[('IMP', 'Inferred from mutant phenotype'),
                 ('IGI', 'Inferred from genetic interaction'),
@@ -449,6 +453,47 @@ class TTestFunction(workflow.Function):
             return options
 
 
+class PlotDagFunction(workflow.Function):
+    def __init__(self):
+        workflow.Function.__init__(self, 'go-dag', 'Build DAG')
+
+    def run(self, selection):
+        g = self.get_network(list(selection['go-terms']))
+#        print g.edges()
+        ds = dataset.GraphDataset(networkx.adj_matrix(g),
+                                  [('go-terms', g.nodes()), ('_go-terms', g.nodes())],
+                                  name="DAG")
+        return [DagPlot(g)]
+
+    def get_network(self, terms, subtree='bp'):
+        """Returns a DAG connecting the given terms by including their parents
+        up to the level needed to connect them. The subtree parameter is one of
+           mf - molecular function
+           bp - biological process
+           cc - cellular component"""
+
+        rpy.r.library("GOstats")
+    
+        if subtree == 'mf':
+            subtree_r = rpy.r.GOMFPARENTS
+        elif subtree == 'bp':
+            subtree_r = rpy.r.GOBPPARENTS
+        elif subtree == 'cc':
+            subtree_r = rpy.r.GOCCPARENTS
+        else:
+            raise Exception("Unknown subtree. Use mf, bp or cc.")
+        
+        g = rpy.r.GOGraph(terms, subtree_r)
+        edges = rpy.r.edges(g)
+    
+        nxgraph = networkx.DiGraph()
+        for child, d in edges.items():
+            for parent in d.keys():
+                nxgraph.add_edge(parent, child)
+
+        return nxgraph
+    
+
 class TTestOptions(workflow.Options):
 
     def __init__(self):
@@ -473,7 +518,6 @@ class GOWeightOptions(workflow.Options):
         self['similarity_threshold'] = 0.0
         self['rank_threshold'] = 0.0
 
-
 class ProbabilityHistogramPlot(plots.HistogramPlot):
     def __init__(self, ds):
         plots.HistogramPlot.__init__(self, ds, name="Confidence", bins=50)
@@ -486,3 +530,129 @@ class VolcanoPlot(plots.ScatterPlot):
                                    name="Volcano plot", 
                                    sel_dim_2='_p', **kw)
 
+        
+class DagPlot(plots.Plot):
+    def __init__(self, graph, dim='go-terms', pos=None, nodecolor='b', nodesize=40,
+                 with_labels=False, name='DAG Plot'):
+
+        plots.Plot.__init__(self, name)
+        self.nodes = graph.nodes()
+        self.graph = graph
+        self._pos = pos
+        self._nodesize = nodesize
+        self._nodecolor = nodecolor
+        self._with_labels = with_labels
+        
+        self.current_dim = dim
+
+        if not self._pos:
+            self._pos = self._calc_pos(graph)
+        self._xy = asarray([self._pos[node] for node in self.nodes])
+        self.xaxis_data = self._xy[:,0]
+        self.yaxis_data = self._xy[:,1]
+
+        # Initial draw
+        self.default_props = {'nodesize' : 50,
+                              'nodecolor' : 'blue',
+                              'edge_color' : 'gray',
+                              'edge_color_selected' : 'red'}
+        self.node_collection = None
+        self.edge_collection = None
+        self.node_labels = None
+        lw = zeros(self.xaxis_data.shape)
+        self.node_collection = self.axes.scatter(self.xaxis_data, self.yaxis_data,
+                                                 s=self._nodesize,
+                                                 c=self._nodecolor,
+                                                 linewidth=lw,
+                                                 zorder=3)
+        self._mappable = self.node_collection
+
+        # selected nodes is a transparent graph that adjust node-edge visibility
+        # according to the current selection needed to get get the selected
+        # nodes 'on top' as zorder may not be defined individually
+        self.selected_nodes = self.axes.scatter(self.xaxis_data,
+                                                self.yaxis_data,
+                                                s=self._nodesize,
+                                                c=self._nodecolor,
+                                                edgecolor='r',
+                                                linewidth=lw,
+                                                zorder=4,
+                                                alpha=0)
+
+        edge_color = self.default_props['edge_color']
+        self.edge_collection = networkx.draw_networkx_edges(self.graph,
+                                                            self._pos,
+                                                            ax=self.axes,
+                                                            edge_color=edge_color)
+        # edge color rgba-arrays
+        self._edge_color_rgba = matlib.repmat(plots.ColorConverter().to_rgba(edge_color),
+                                       self.graph.number_of_edges(),1)
+        self._edge_color_selected = plots.ColorConverter().to_rgba(self.default_props['edge_color_selected'])
+        if self._with_labels:
+            self.node_labels = networkx.draw_networkx_labels(self.graph,
+                                                            self._pos,
+                                                            ax=self.axes)
+        
+        # remove axes, frame and grid
+        self.axes.set_xticks([])
+        self.axes.set_yticks([])
+        self.axes.grid(False)
+        self.axes.set_frame_on(False)
+        self.fig.subplots_adjust(left=0, right=1, bottom=0, top=1)
+
+    def _calc_pos(self, graph):
+        """Calculates position for graph nodes."""
+        gv_graph = networkx.DiGraph()
+        for start, end in graph.edges():
+            gv_graph.add_edge(start.replace('GO:', ''), end.replace('GO:', ''))
+
+        pos_gv = networkx.pygraphviz_layout(gv_graph, prog="dot")
+        pos = {}
+        for k, v in pos_gv.items():
+            if k != "all":
+                pos["GO:%s" % k] = v
+            else:
+                pos[k] = v
+        return pos
+
+    def rectangle_select_callback(self, x1, y1, x2, y2, key):
+        ydata = self.yaxis_data
+        xdata = self.xaxis_data
+        
+        # find indices of selected area
+        if x1>x2:
+            x1, x2 = x2, x1
+        if y1>y2:
+            y1, y2 = y2, y1
+        assert x1<=x2
+        assert y1<=y2
+        
+        index = nonzero((xdata>x1) & (xdata<x2) & (ydata>y1) & (ydata<y2))[0]
+        ids = [self.nodes[i] for i in index]
+        ids = self.update_selection(ids, key)
+        self.selection_listener(self.current_dim, ids)
+
+    def lasso_select_callback(self, verts, key=None):
+        xys = c_[self.xaxis_data[:,newaxis], self.yaxis_data[:,newaxis]]
+        index = nonzero(points_inside_poly(xys, verts))[0]
+        ids = [self.nodes[i] for i in index]
+        ids = self.update_selection(ids, key)
+        self.selection_listener(self.current_dim, ids)
+        
+    def set_current_selection(self, selection):
+        linewidth = zeros(self.xaxis_data.shape)
+        edge_color_rgba = self._edge_color_rgba.copy()
+        index = [i for i in range(len(self.nodes)) if self.nodes[i] in selection[self.current_dim]]
+        if len(index) > 0:
+            linewidth[index] = 2
+            idents = selection[self.current_dim]
+            edge_index = [i for i,edge in enumerate(self.graph.edges()) if (edge[0] in idents and  edge[1] in idents)]
+            if len(edge_index)>0:
+                for i in edge_index:
+                    edge_color_rgba[i,:] = self._edge_color_selected
+                self._A = None
+        
+        self.edge_collection._colors = edge_color_rgba
+        self.selected_nodes.set_linewidth(linewidth)
+        self.canvas.draw()
+

workflows/smokers.py

@@ -55,6 +55,7 @@ class SmallTestWorkflow(workflow.Workflow):
         go.add_function(gobrowser.GOWeightFunction())
         go.add_function(gobrowser.DistanceToSelectionFunction())
         go.add_function(gobrowser.TTestFunction())
+        go.add_function(gobrowser.PlotDagFunction())
         self.add_stage(go)
 
         # EXTRA PLOTS