diff --git a/examples/plot_tasks_MPI.py b/examples/plot_tasks_MPI.py
index 769d969b5411c831f48d0e763b77b7468567656c..01649aec0edf5e2efd564bbbf9e60dbea27731fa 100755
--- a/examples/plot_tasks_MPI.py
+++ b/examples/plot_tasks_MPI.py
@@ -58,49 +58,37 @@ params = {"axes.labelsize": 10,
 }
 pl.rcParams.update(params)
 
-types = {"0": "none",
-         "1": "sort",
-         "2": "self",
-         "3": "pair",
-         "4": "sub",
-         "5": "ghost",
-         "6": "kick1",
-         "7": "kick2",
-         "8": "send",
-         "9": "recv",
-         "10": "grav_pp",
-         "11": "grav_mm",
-         "12": "grav_up",
-         "13": "grav_down",
-         "14": "psort",
-         "15": "split_cell",
-         "16": "rewait",
-         "17": "count"}
-
-subtypes = {"0": "none",
-            "1": "density",
-            "2": "force",
-            "3": "grav",
-            "4": "count"}
-
-
-# Assign colours for all types.
-#colours = ["blue","green","red","cyan","magenta","yellow","black"]
-#colours = colours + list(matplotlib.colors.cnames)
-colours = list(matplotlib.colors.cnames)
-
-index = 0
-subtypecolours = {}
-for key in subtypes:
-    subtypecolours[subtypes[key]] = colours[index]
-    print subtypes[key], " = ", colours[index]
-    index = index + 1
-
-taskcolours = {}
-for key in types:
-    taskcolours[types[key]] = colours[index]
-    print types[key], " = ", colours[index]
-    index = index + 1
+#  Tasks and subtypes. Indexed as in tasks.h.
+types = ["none", "sort", "self", "pair", "sub", "ghost", "kick1", "kick2",
+         "send", "recv", "grav_pp", "grav_mm", "grav_up", "grav_down",
+         "psort", "split_cell", "rewait", "count"]
+
+taskcolours = {"none": "black",
+               "sort": "lightblue",
+               "self": "greenyellow",
+               "pair": "navy",
+               "sub": "hotpink",
+               "ghost": "cyan",
+               "kick1": "maroon",
+               "kick2": "green",
+               "send": "yellow",
+               "recv": "magenta",
+               "grav_pp": "mediumorchid",
+               "grav_mm": "mediumturquoise",
+               "grav_up": "mediumvioletred",
+               "grav_down": "mediumnightblue",
+               "psort": "steelblue",
+               "split_cell": "seagreen",
+               "rewait": "olive",
+               "count": "powerblue"}
+
+subtypes = ["none", "density", "force", "grav", "count"]
+
+subtypecolours = {"none": "black",
+                  "density": "red",
+                  "force": "blue",
+                  "grav": "indigo",
+                  "count": "purple"}
 
 #  Read input.
 data = pl.loadtxt( infile )
@@ -110,6 +98,7 @@ print "Number of ranks:", nranks
 if rank < 0 or rank >= nranks:
     print "Cannot read rank: ", rank, ", out of range"
     sys.exit(1)
+print "Processing rank:", rank
 
 nthread = int(max(data[:,1])) + 1
 print "Number of threads:", nthread
@@ -135,14 +124,16 @@ num_lines = pl.size(sdata) / 10
 for line in range(num_lines):
     thread = int(sdata[line,1])
     tasks[thread].append({})
-    tasks[thread][-1]["type"] = types[str(int(sdata[line,2]))]
-    tasks[thread][-1]["subtype"] = subtypes[str(int(sdata[line,3]))]
+    tasks[thread][-1]["type"] = types[int(sdata[line,2])]
+    tasks[thread][-1]["subtype"] = subtypes[int(sdata[line,3])]
     tic = int(sdata[line,5]) / CPU_CLOCK * 1000
     toc = int(sdata[line,6]) / CPU_CLOCK * 1000
     tasks[thread][-1]["tic"] = tic
     tasks[thread][-1]["toc"] = toc
     tasks[thread][-1]["t"] = (toc + tic)/ 2
 
+print "Collection done..."
+
 combtasks = {}
 combtasks[-1] = []
 for i in range(nthread):
@@ -169,6 +160,8 @@ for thread in range(nthread):
         else:
             combtasks[thread][-1]["toc"] = task["toc"]
 
+print "Combination done..."
+
 typesseen = []
 for i in range(nthread):
     for task in combtasks[i]:
@@ -193,3 +186,6 @@ pl.yticks(pl.array(range(nthread)) + 0.5, pl.array(range(nthread)))
 
 pl.show()
 pl.savefig(outpng)
+print "Graphics done, output written to", outpng
+
+sys.exit(0)