Improve task level plot

Following modifications have been added:

• new colormap
• grid
• labels are rotated by 90 degrees now
• fix a small bug where the highest level has not been plotted

For comparison, here's the default output of the current version of the script:

And here is the modified version:

Additionally, I added the -c and -d flags.

• -c: counts on how many levels each task operates, and writes it in brackets on the task name label
• -d: if a task has the same number of counts for multiple levels, displace the counts by a tiny amount to make them visible.

Here's an example output: Now we can for example see that there is more than one level of grav_down_in with 8 tasks calls.

tools/plot_task_level.py

-#!/usr/bin/python
-"""
+#!/usr/bin/env python3
+description = """
+Plot the number of tasks for each depth level and each type of task.
+
 Usage:
   ./plot_task_level.py task_level.txt
 
-Description:
-  Plot the number of tasks for each depth level and each type of task.
 """
 
 
 import pandas as pd
+import numpy as np
 import matplotlib.pyplot as plt
 import sys
+import argparse
+from os import path
+
+
+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 the required filename provided on the command line doesn't exist
+    """
 
-# get filename
-filename = sys.argv[-1]
-
-# Column names
-names = ["type", "subtype", "depth", "count"]
-
-# read file
-data = pd.read_csv(filename, sep=" ", comment="#", names=names)
-
-# generate color map
-cmap = plt.get_cmap("hsv")
-N = data["depth"].max() + 5
-
-# plot data
-for i in range(data["depth"].max()):
-    ind = data["depth"] == i
-    label = "depth = %i" % i
-    c = cmap(i / N)
-    plt.plot(
-        data["type"][ind] + "_" + data["subtype"][ind],
-        data["count"][ind],
-        ".",
-        label=label,
-        color=c,
+    # description is string at the top of this file.
+    parser = argparse.ArgumentParser(description=description)
+
+    parser.add_argument(
+        "-c",
+        "--count",
+        dest="count_levels",
+        help="count on how many different levels the tasks can be found"
+        " and add it to the label",
+        action="store_true",
+    )
+
+    parser.add_argument(
+        "-d",
+        "--displace",
+        dest="displace",
+        help="attempt to displace overlapping point on the plot a bit"
+        " and try to make them visible",
+        action="store_true",
+    )
+
+    parser.add_argument(
+        "file",
+        type=str,
+        help="Required file name of .csv file(s) of the task levels "
+        "generated by swift.",
     )
 
-# modify figure parameters and show it
-plt.gca().set_yscale("log")
-plt.xticks(rotation=45)
-plt.ylabel("Number of Tasks")
-plt.gcf().subplots_adjust(bottom=0.15)
-plt.legend()
-plt.show()
+    args = parser.parse_args()
+    filename = args.file
+    print(filename)
+
+    if not path.exists(filename):
+        raise FileNotFoundError("File not found:'" + filename + "'")
+
+    return args, filename
+
+
+def read_data(filename):
+    """
+    Reads in data from the csv file.
+
+    Parameters
+    ----------
+
+    filename: str
+        filename to be read from
+
+
+    Returns
+    -------
+
+    data: pandas dataset
+        dataset containing read in data
+    """
+
+    # Column names
+    names = ["type", "subtype", "depth", "count"]
+
+    # read file
+    data = pd.read_csv(filename, sep=" ", comment="#", names=names)
+
+    return data
+
+
+def get_discrete_cmap(nentries):
+    """
+    Returns a discrete colormap.
+
+    Parameters
+    ----------
+
+    nentries: int
+        how many entries you want for your colormap
+
+
+    Returns
+    -------
+
+    cmap: list
+        list of colors
+
+
+    Raises
+    ------
+
+    IndexError:
+        When you want more entries than there are available.
+        Current maximum is 21.
+    """
+
+    fullcolorlist = [
+        "red",
+        "green",
+        "blue",
+        "gold",
+        "magenta",
+        "cyan",
+        "lime",
+        "saddlebrown",
+        "darkolivegreen",
+        "cornflowerblue",
+        "orange",
+        "dimgrey",
+        "navajowhite",
+        "darkslategray",
+        "mediumpurple",
+        "lightpink",
+        "mediumseagreen",
+        "maroon",
+        "midnightblue",
+        "silver",
+        "black",
+    ]
+
+    if nentries >= len(fullcolorlist) - 1:
+        raise IndexError(
+            "I can't handle more than 21 different colors."
+            "Add more manually in get_discrete_cmap() function"
+        )
+
+    return fullcolorlist[: nentries + 1]
+
+
+def add_levelcounts(data):
+    """
+    Adds a column to the data with the counts on how many levels a given task
+    is executed on.
+
+    Parameters
+    ----------
+
+    data: pandas dataframe
+        The dataframe to use
+
+    Returns
+    -------
+
+    data: pandas dataframe
+        the modified dataframe
+    """
+
+    # add new column
+    data["nlevels"] = ["1" for _ in range(data.shape[0])]
+
+    # count on how many levels each task exists
+    for i, (ttype, tsubtype) in enumerate(zip(data["type"], data["subtype"])):
+        istype = data["type"] == ttype
+        issubtype = data["subtype"] == tsubtype
+        isthis = np.logical_and(istype, issubtype)
+        count = np.count_nonzero(isthis)
+        data.at[i, "nlevels"] = str(count)
+
+    return data
+
+
+def add_displacement(data):
+    """
+    Add small displacements to the task number counts and try
+    to make them better visible
+
+    Parameters
+    ----------
+
+    data: pandas dataframe
+        the data to be modified
+
+
+    Returns
+    -------
+
+    data: pandas dataframe
+        the modified data
+    """
+    # add new column
+    data["yvals"] = data["count"] * 1.0
+    data["yval_modified"] = False
+    inds = np.arange(0, data.shape[0])
+
+    # count on how many levels each task exists
+    for i, (ttype, tsubtype) in enumerate(zip(data["type"], data["subtype"])):
+        if data["yval_modified"][i]:
+            continue
+        istype = data["type"] == ttype
+        issubtype = data["subtype"] == tsubtype
+        isthis = np.logical_and(istype, issubtype)
+        uniques = np.unique(data["count"][isthis])
+        for u in uniques:
+            occurances = np.count_nonzero(data["count"][isthis] == u)
+            for o in range(occurances):
+                thisind = inds[isthis][o]
+                step = 0.05 * max(int(np.log10(u) + 0.5), 1) * o * (-1) ** o
+                data.at[thisind, "yvals"] += step
+                data.at[thisind, "yval_modified"] = True
+
+    return data
+
+
+if __name__ == "__main__":
+
+    args, filename = parse_args()
+
+    data = read_data(filename)
+    cmap = get_discrete_cmap(data["depth"].max())
+
+    # are we counting the levels?
+    if args.count_levels:
+        data = add_levelcounts(data)
+
+    # are we displacing the particles on the y axis?
+    if args.displace:
+        data = add_displacement(data)
+
+    # plot data
+    for i in range(data["depth"].max() + 1):
+        ind = data["depth"] == i
+        label = "depth = {0:d}".format(i)
+
+        if args.count_levels:
+            xvals = (
+                data["type"][ind]
+                + "_"
+                + data["subtype"][ind]
+                + "["
+                + data["nlevels"][ind]
+                + "]"
+            )
+        else:
+            xvals = data["type"][ind] + "_" + data["subtype"][ind]
+
+        if args.displace:
+            yvals = data["yvals"][ind]
+        else:
+            yvals = data["count"][ind]
+        c = cmap[i]
+        plt.plot(xvals, yvals, "o", label=label, color=c, alpha=0.7)
+
+    # modify figure parameters and show it
+    plt.gca().set_yscale("log")
+    plt.xticks(rotation=90)
+    plt.ylabel("Number of Tasks")
+    plt.gcf().subplots_adjust(bottom=0.225)
+    plt.legend()
+    plt.grid()
+    plt.show()