Attempt to make dependency plots more readable

tools/plot_task_dependencies.py

 #!/usr/bin/env python3
-"""
+
+description = """
 This file generates a graphviz file that represents the SWIFT tasks
 dependencies.
 
@@ -8,28 +9,69 @@ Example: ./plot_task_dependencies.py dependency_graph_*.csv
 from pandas import read_csv
 import numpy as np
 from subprocess import call
-from optparse import OptionParser
+import argparse
+from os import path
+
+# define task colours in the dict here:
+task_colours = {
+    "black_holes": "forestgreen",
+    "stars": "darkorange1",
+    "hydro": "blue3",
+    "gravity": "red3",
+    "RT": "springgreen",
+}
+
+
+def parse_args():
+    """
+    Parses command line arguments.
+
+    Returns
+    -------
+
+    args: Namespace
+        Namespace returned by argparse.ArgumentParser.parse_args()
+        containing all arguments
+
+    files:
+        List of files parsed from the command line.
+
+    Raises
+    ------
 
+    FileNotFoundError
+        If any of the files provided on the command line doesn't exist
+    """
 
-def parseOption():
-    parser = OptionParser()
+    # description is string at the top of this file.
+    parser = argparse.ArgumentParser(description=description)
 
-    parser.add_option(
+    parser.add_argument(
         "-c",
         "--with-calls",
         dest="with_calls",
-        help="Add the function calls in the graph",
+        help="Add the function calls of the tasks to the graph",
         action="store_true",
     )
 
-    opt, files = parser.parse_args()
-    if len(files) != 1:
-        raise Exception("You need to provide one file")
+    parser.add_argument(
+        "files",
+        nargs="+",
+        type=str,
+        help="Required file name(s) of .csv file(s) of the task dependencies generated by swift.",
+    )
+
+    args = parser.parse_args()
+    files = args.files
 
-    return opt, files
+    for f in files:
+        if not path.exists(f):
+            raise FileNotFoundError("You need to provide one file")
 
+    return args, files
 
-def getGitVersion(f, git):
+
+def get_git_version(f, git):
     """
     Read the git version from the file
 
@@ -66,7 +108,7 @@ def getGitVersion(f, git):
     return new_git
 
 
-def appendSingleData(data0, datai):
+def append_single_data(data0, datai):
     """
     Append two DataFrame together
 
@@ -109,9 +151,9 @@ def appendSingleData(data0, datai):
     return data0
 
 
-def appendData(data):
+def append_data(data):
     """
-    Append all the dataframe together
+    Append all the dataframe together, and add a column for colour type
 
     Parameters
     ----------
@@ -127,17 +169,55 @@ def appendData(data):
     """
     N = len(data)
     if N == 1:
+        data[0]["task_colour"] = "black"
         return data[0]
 
     # add number link to data[0]
     for i in range(N - 1):
         i += 1
-        data[0] = appendSingleData(data[0], data[i])
+        data[0] = append_single_data(data[0], data[i])
+
+    data[0]["task_colour"] = "black"
 
     return data[0]
 
 
-def taskIsBlackHoles(name):
+def get_task_colour(taskname):
+    """
+    Get the task colour based on its name.
+    Default colour is black.
+
+
+    Parameters
+    ----------
+
+    taskname: str
+        name of the task at hand
+
+
+    Returns
+    -------
+
+    colour: str
+        colour string directly applicable in the dot file.
+    """
+
+    colour = "black"
+    if task_is_black_holes(taskname):
+        colour = task_colours["black_holes"]
+    elif task_is_stars(taskname):
+        colour = task_colours["stars"]
+    elif task_is_hydro(taskname):
+        colour = task_colours["hydro"]
+    elif task_is_gravity(taskname):
+        colour = task_colours["gravity"]
+    elif task_is_RT(taskname):
+        colour = task_colours["RT"]
+
+    return colour
+
+
+def task_is_black_holes(name):
     """
     Does the task concern black holes?
 
@@ -152,7 +232,7 @@ def taskIsBlackHoles(name):
     return False
 
 
-def taskIsStars(name):
+def task_is_stars(name):
     """
     Does the task concern stars?
 
@@ -170,7 +250,7 @@ def taskIsStars(name):
     return False
 
 
-def taskIsHydro(name):
+def task_is_hydro(name):
     """
     Does the task concern the hydro?
 
@@ -207,7 +287,7 @@ def taskIsHydro(name):
     return False
 
 
-def taskIsGravity(name):
+def task_is_gravity(name):
     """
     Does the task concern the gravity?
 
@@ -224,7 +304,7 @@ def taskIsGravity(name):
     return False
 
 
-def taskIsRT(name):
+def task_is_RT(name):
     """
     Does the task concern Radiative Transfer?
 
@@ -239,7 +319,7 @@ def taskIsRT(name):
     return False
 
 
-def getFunctionCalls(name):
+def get_function_calls(name):
     txt = None
     if name == "ghost":
         txt = """hydro_end_density, chemistry_end_density,<br/>
@@ -295,7 +375,7 @@ def getFunctionCalls(name):
         return pre + txt + app
 
 
-def writeTask(f, name, implicit, mpi, with_calls):
+def write_task(f, name, implicit, mpi, with_calls):
     """
     Write the special task (e.g. implicit and mpi)
 
@@ -321,27 +401,16 @@ def writeTask(f, name, implicit, mpi, with_calls):
     txt = "\t " + name + "["
 
     if implicit:
-        txt += "style=filled,fillcolor=lightgrey,"
+        txt += "style=filled,fillcolor=grey90,"
     if mpi:
-        txt += "shape=diamond,"
-
-    if taskIsBlackHoles(name):
-        txt += "color=forestgreen,"
+        txt += "shape=diamond,style=filled,fillcolor=azure,"
 
-    if taskIsStars(name):
-        txt += "color=darkorange1,"
+    col = get_task_colour(name)
 
-    if taskIsHydro(name):
-        txt += "color=blue3,"
-
-    if taskIsGravity(name):
-        txt += "color=red3,"
-
-    if taskIsRT(name):
-        txt += 'color="springgreen"'
+    txt += "color=%s," % col
 
     if with_calls:
-        func = getFunctionCalls(name)
+        func = get_function_calls(name)
         if func is not None:
             txt += "label=" + func + ","
 
@@ -354,7 +423,7 @@ def writeTask(f, name, implicit, mpi, with_calls):
     f.write(txt)
 
 
-def writeHeader(f, data, git, opt):
+def write_header(f, data, git, opt):
     """
     Write the header and the special tasks
 
@@ -379,8 +448,8 @@ def writeHeader(f, data, git, opt):
     f.write('\t label="Task dependencies for SWIFT %s";\n' % git)
     f.write("\t compound=true;\n")
     f.write("\t ratio=0.66;\n")
-    f.write("\t node[nodesep=0.15];\n")
-
+    f.write("\t node[nodesep=0.15, fontsize=30, penwidth=5.];\n")
+    f.write("\t ranksep=1.2;\n")
     f.write("\n")
 
     # write the special task
@@ -394,7 +463,7 @@ def writeHeader(f, data, git, opt):
             continue
 
         written.append(ta)
-        writeTask(f, ta, data["implicit_in"][i], data["mpi_in"][i], opt.with_calls)
+        write_task(f, ta, data["implicit_in"][i], data["mpi_in"][i], opt.with_calls)
 
     # do task out
     for i in range(N):
@@ -403,12 +472,12 @@ def writeHeader(f, data, git, opt):
             continue
 
         written.append(tb)
-        writeTask(f, tb, data["implicit_out"][i], data["mpi_out"][i], opt.with_calls)
+        write_task(f, tb, data["implicit_out"][i], data["mpi_out"][i], opt.with_calls)
 
     f.write("\n")
 
 
-def writeCluster(f, tasks, cluster):
+def write_cluster(f, tasks, cluster):
     """
     Write a single cluster
 
@@ -423,15 +492,18 @@ def writeCluster(f, tasks, cluster):
 
     cluster: str
         Cluster name
+
     """
+
     f.write("\t subgraph cluster%s {\n" % cluster)
     f.write('\t\t label="";\n')
+    f.write('\t\t bgcolor="grey99";\n')
     for t in tasks:
         f.write("\t\t %s;\n" % t)
     f.write("\t };\n\n")
 
 
-def writeClusters(f, data):
+def write_clusters(f, data):
     """
     Write all the clusters
 
@@ -466,12 +538,12 @@ def writeClusters(f, data):
         tasks = np.unique(tasks)
 
         # write current cluster
-        writeCluster(f, tasks, cluster)
+        write_cluster(f, tasks, cluster)
 
     f.write("\n")
 
 
-def writeDependencies(f, data):
+def write_dependencies(f, data):
     """
     Write all the dependencies between tasks
 
@@ -503,13 +575,16 @@ def writeDependencies(f, data):
         written.append(name)
 
         # write relation
-        arrow = ""
+        arrow = ",color=%s" % data["task_colour"][i]
         if data["number_rank"][i] != max_rank:
-            arrow = ",style=dashed"
-        f.write("\t %s->%s[label=%i%s]\n" % (ta, tb, number_link, arrow))
+            arrow += ",style=dashed"
+        f.write(
+            "\t %s->%s[label=%i%s,fontcolor=%s]\n"
+            % (ta, tb, number_link, arrow, data["task_colour"][i])
+        )
 
 
-def writeFooter(f):
+def write_footer(f):
     """
     Write the footer
 
@@ -522,9 +597,41 @@ def writeFooter(f):
     f.write("}")
 
 
+def set_task_colours(data):
+    """
+    Set the value of the task colour for the plot if you
+    want it non-black. the `task_colours` dict is defined
+    at the top of this script.
+
+    Parameters
+    ----------
+
+    data: DataFrame 
+        DataFrame of all the tasks 
+
+
+    Returns
+    -------
+
+    data: DataFrame 
+        modified DataFrame of all the tasks, now with more
+        colour
+    """
+
+    N = len(data)
+
+    for i in range(N):
+        taskname = data["task_in"][i]
+        col = get_task_colour(taskname)
+        # set the colour
+        data.loc[i, "task_colour"] = col
+
+    return data
+
+
 if __name__ == "__main__":
 
-    opt, files = parseOption()
+    args, files = parse_args()
 
     # output
     dot_output = "dependency_graph.dot"
@@ -535,25 +642,26 @@ if __name__ == "__main__":
     git = None
     for f in files:
         tmp = read_csv(f, delimiter=",", comment="#")
-        git = getGitVersion(f, git)
+        git = get_git_version(f, git)
         data.append(tmp)
 
-    data = appendData(data)
+    data = append_data(data)
+    data = set_task_colours(data)
 
     # write output
     with open(dot_output, "w") as f:
-        writeHeader(f, data, git, opt)
+        write_header(f, data, git, args)
 
-        writeClusters(f, data)
+        write_clusters(f, data)
 
-        writeDependencies(f, data)
+        write_dependencies(f, data)
 
-        writeFooter(f)
+        write_footer(f)
 
     call(["dot", "-Tpng", dot_output, "-o", png_output])
 
     print("You will find the graph in %s" % png_output)
 
-    if opt.with_calls:
+    if args.with_calls:
         print("We recommand to use the python package xdot available on pypi:")
         print("  python -m xdot %s" % dot_output)