added 'colored' graph to graph2use

Author: miykael

nipype/pipeline/engine.py

@@ -412,10 +412,14 @@ def write_graph(self, dotfilename='graph.dot', graph2use='hierarchical',
         Parameters
         ----------
 
-        graph2use: 'orig', 'hierarchical' (default), 'flat', 'exec'
+        graph2use: 'orig', 'hierarchical' (default), 'flat', 'exec', 'colored'
             orig - creates a top level graph without expanding internal
             workflow nodes;
             flat - expands workflow nodes recursively;
+            hierarchical - expands workflow nodes recursively with a
+            notion on hierarchy; 
+            colored - expands workflow nodes recursively with a
+            notion on hierarchy in color;
             exec - expands workflows to depict iterables
 
         format: 'png', 'svg'
@@ -426,7 +430,7 @@ def write_graph(self, dotfilename='graph.dot', graph2use='hierarchical',
             False.
 
         """
-        graphtypes = ['orig', 'flat', 'hierarchical', 'exec']
+        graphtypes = ['orig', 'flat', 'hierarchical', 'exec', 'colored']
         if graph2use not in graphtypes:
             raise ValueError('Unknown graph2use keyword. Must be one of: ' +
                              str(graphtypes))
@@ -439,11 +443,17 @@ def write_graph(self, dotfilename='graph.dot', graph2use='hierarchical',
             else:
                 base_dir = os.getcwd()
         base_dir = make_output_dir(base_dir)
-        if graph2use == 'hierarchical':
+        if graph2use == 'hierarchical' or graph2use == 'colored':
             dotfilename = os.path.join(base_dir, dotfilename)
-            self.write_hierarchical_dotfile(dotfilename=dotfilename,
-                                            colored=False,
-                                            simple_form=simple_form)
+            if graph2use == 'colored':
+                self.write_hierarchical_dotfile(dotfilename=dotfilename,
+                                                colored=True,
+                                                simple_form=simple_form)
+
+            else:
+                self.write_hierarchical_dotfile(dotfilename=dotfilename,
+                                                colored=False,
+                                                simple_form=simple_form)
             format_dot(dotfilename, format=format)
         else:
             graph = self._graph
@@ -454,11 +464,9 @@ def write_graph(self, dotfilename='graph.dot', graph2use='hierarchical',
             export_graph(graph, base_dir, dotfilename=dotfilename,
                          format=format, simple_form=simple_form)
 
-    def write_hierarchical_dotfile(self, dotfilename=None, colored=True,
+    def write_hierarchical_dotfile(self, dotfilename=None, colored=False,
                                    simple_form=True):
         dotlist = ['digraph %s{' % self.name]
-        if colored:
-            dotlist.append('  ' + 'colorscheme=pastel28;')
         dotlist.append(self._get_dot(prefix='  ', colored=colored,
                                      simple_form=simple_form))
         dotlist.append('}')
@@ -827,18 +835,18 @@ def _generate_flatgraph(self):
             self._graph.remove_nodes_from(nodes2remove)
         logger.debug('finished expanding workflow: %s', self)
 
-    def _get_dot(self, prefix=None, hierarchy=None, colored=True,
-                 simple_form=True):
+    def _get_dot(self, prefix=None, hierarchy=None, colored=False,
+                 simple_form=True, level=0):
         """Create a dot file with connection info
         """
         if prefix is None:
             prefix = '  '
         if hierarchy is None:
             hierarchy = []
-        level = (len(prefix) / 2) + 1
+        colorset = ['#FFFFC8','#0000FF','#B4B4FF','#E6E6FF','#FF0000',
+                    '#FFB4B4','#FFE6E6','#00A300','#B4FFB4','#E6FFE6']
+
         dotlist = ['%slabel="%s";' % (prefix, self.name)]
-        if colored:
-            dotlist.append('%scolor=%d;' % (prefix, level))
         for node in nx.topological_sort(self._graph):
             fullname = '.'.join(hierarchy + [node.fullname])
             nodename = fullname.replace('.', '_')
@@ -847,24 +855,35 @@ def _get_dot(self, prefix=None, hierarchy=None, colored=True,
                 if not simple_form:
                     node_class_name = '.'.join(node_class_name.split('.')[1:])
                 if hasattr(node, 'iterables') and node.iterables:
-                    dotlist.append(('%s[label="%s", style=filled, colorscheme'
-                                    '=greys7 color=2];') % (nodename,
+                    dotlist.append(('%s[label="%s", shape=box3d,'
+                                    'style=filled, color=black, colorscheme'
+                                    '=greys7 fillcolor=2];') % (nodename,
                                                             node_class_name))
                 else:
-                    dotlist.append('%s[label="%s"];' % (nodename,
-                                                        node_class_name))
+                    if colored:
+                        dotlist.append(('%s[label="%s", style=filled,'
+                                        ' fillcolor="%s"];')
+                                        % (nodename,node_class_name,
+                                           colorset[level]))
+                    else:
+                        dotlist.append(('%s[label="%s"];')
+                                        % (nodename,node_class_name))
+
         for node in nx.topological_sort(self._graph):
             if isinstance(node, Workflow):
                 fullname = '.'.join(hierarchy + [node.fullname])
                 nodename = fullname.replace('.', '_')
                 dotlist.append('subgraph cluster_%s {' % nodename)
                 if colored:
+                    dotlist.append(prefix + prefix + 'edge [color="%s"];' % (colorset[level+1]))
                     dotlist.append(prefix + prefix + 'style=filled;')
+                    dotlist.append(prefix + prefix + 'fillcolor="%s";' % (colorset[level+2]))
                 dotlist.append(node._get_dot(prefix=prefix + prefix,
                                              hierarchy=hierarchy + [self.name],
                                              colored=colored,
-                                             simple_form=simple_form))
+                                             simple_form=simple_form, level=level+3))
                 dotlist.append('}')
+                if level==6:level=2
             else:
                 for subnode in self._graph.successors_iter(node):
                     if node._hierarchy != subnode._hierarchy: