diff --git a/examples/plot_tasks.py b/examples/plot_tasks.py
index 832de37e2e7c2a071d94c79d4bd34df3d6f994ce..84a09524665a6c0935648eb728b7d9447982343a 100644
--- a/examples/plot_tasks.py
+++ b/examples/plot_tasks.py
@@ -1,260 +1,155 @@
-#!/usr/bin/env python
-"""
-Usage:
- plot_tasks.py input.dat output.png [time-range-ms]
-
-where input.dat is a thread info file for a step. Use the '-y interval'
-flag of the swift MPI commands to create these. The output plot will be
-called 'output.png'. Use the time-range-ms in millisecs to produce
-plots with the same time span.
-
-This file is part of SWIFT.
-Copyright (c) 2015 Pedro Gonnet (pedro.gonnet@durham.ac.uk),
- Bert Vandenbroucke (bert.vandenbroucke@ugent.be)
- Matthieu Schaller (matthieu.schaller@durham.ac.uk)
- (c) 2016 Peter W. Draper (p.w.draper@durham.ac.uk)
-
-This program is free software: you can redistribute it and/or modify
-it under the terms of the GNU Lesser General Public License as published
-by the Free Software Foundation, either version 3 of the License, or
-(at your option) any later version.
-
-This program is distributed in the hope that it will be useful,
-but WITHOUT ANY WARRANTY; without even the implied warranty of
-MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
-GNU General Public License for more details.
-
-You should have received a copy of the GNU Lesser General Public License
-along with this program. If not, see <http://www.gnu.org/licenses/>.
-"""
-
-import matplotlib
-import matplotlib.pyplot as plt
-import matplotlib.colors as colors
-import matplotlib.collections as collections
-import matplotlib.cm as cmx
-from matplotlib import gridspec
-matplotlib.use('Agg')
-import pylab as pl
-import numpy as np
-import sys
-import math
-
-# Basic plot configuration.
-PLOT_PARAMS = {"axes.labelsize": 10,
- "axes.titlesize": 10,
- "font.size": 12,
- "legend.fontsize": 12,
- "xtick.labelsize": 10,
- "ytick.labelsize": 10,
- "figure.figsize" : (16., 4.),
- "figure.subplot.left" : 0.03,
- "figure.subplot.right" : 0.995,
- "figure.subplot.bottom" : 0.1,
- "figure.subplot.top" : 0.99,
- "figure.subplot.wspace" : 0.,
- "figure.subplot.hspace" : 0.,
- "lines.markersize" : 6,
- "lines.linewidth" : 3.
- }
-pl.rcParams.update(PLOT_PARAMS)
-
-# Tasks and subtypes. Indexed as in tasks.h.
-TASKTYPES = ["none", "sort", "self", "pair", "sub_self", "sub_pair", "init", "ghost",
- "extra_ghost", "kick", "send", "recv",
- "grav_gather_m", "grav_fft", "grav_mm", "grav_up",
- "grav_external", "cooling", "count"]
-
-TASKCOLOURS = {"none": "black",
- "self": "lightblue",
- "pair": "greenyellow",
- "pair": "navy",
- "sub_self": "greenyellow",
- "sub_pair": "navy",
- "init": "indigo",
- "ghost": "cyan",
- "extra_ghost": "cyan",
- "kick": "green",
- "send": "yellow",
- "recv": "magenta",
- "grav_gather_m": "mediumorchid",
- "grav_fft": "mediumnightblue",
- "grav_mm": "mediumturquoise",
- "grav_up": "mediumvioletred",
- "grav_external": "darkred",
- "cooling": "darkblue",
- "count": "powerblue"}
-
-SUBTYPES = ["none", "density", "gradient", "force", "grav", "tend", "count"]
-
-SUBCOLOURS = {"none": "black",
- "density": "red",
- "gradient": "powerblue",
- "force": "blue",
- "grav": "indigo",
- "tend": "grey",
- "count": "black"}
-
-# Show docs if help is requested.
-if len( sys.argv ) == 2 and ( sys.argv[1][0:2] == "-h" or sys.argv[1][0:3] == "--h" ):
- from pydoc import help
- help( "__main__" )
- sys.exit( 0 )
-
-# Handle command-line.
-if len( sys.argv ) != 3 and len( sys.argv ) != 4:
- print "Usage: ", sys.argv[0], "input.dat output.png [time-range-ms]"
- sys.exit(1)
-
-infile = sys.argv[1]
-outpng = sys.argv[2]
-delta_t = 0
-if len( sys.argv ) == 4:
- delta_t = int(sys.argv[3])
-
-# Read input.
-data = pl.loadtxt( infile )
-
-nthread = int(max(data[:,1])) + 1
-print "Number of threads:", nthread
-
-jet = cm = plt.get_cmap('jet')
-cNorm = colors.Normalize(vmin=0, vmax=max(data[:,0]))
-scalarMap = cmx.ScalarMappable(norm=cNorm, cmap=jet)
-print scalarMap.get_clim()
-
-# Recover the start and end time
-#full_step = data[0,:]
-#tic_step = int(full_step[4])
-#toc_step = int(full_step[5])
-#CPU_CLOCK = float(full_step[-1])
-#data = data[1:,:]
-
-#print "CPU frequency:", CPU_CLOCK
-
-# Avoid start and end times of zero.
-#data = data[data[:,4] != 0]
-#data = data[data[:,5] != 0]
-
-# Calculate the time range, if not given.
-delta_t = max(data[:,3]) - min(data[:,2])
-print "Data range: ",delta_t, "ms"
-
-# Once more doing the real gather and plots this time.
-start_t = min(data[:,2])
-data[:,2] -= start_t
-data[:,3] -= start_t
-end_t = (max(data[:,3]))
-print "Start_t", start_t, "end_t", end_t
-
-tasks = {}
-tasks[-1] = []
-for i in range(nthread):
- tasks[i] = []
-
-num_lines = len(data[:,0])
-for line in range(num_lines):
- thread = int(data[line,1])
- tasks[thread].append({})
- tasks[thread][-1]["tic"] = data[line,2]
- tasks[thread][-1]["toc"] = data[line,3]
- tasks[thread][-1]["t"] = (data[line,1]+data[line,1])/2
- tasks[thread][-1]["colour"] = data[line,0]
-
-combtasks = {}
-combtasks[-1] = []
-for i in range(nthread):
- combtasks[i] = []
-
-for thread in range(nthread):
- tasks[thread] = sorted(tasks[thread], key=lambda l: l["t"])
- lastcolour = ""
- colours = []
- for task in tasks[thread]:
- if task["colour"] not in colours:
- colours.append(task["colour"])
- #if(lastcolour == "" or not lastcolour == task["colour"]):
- combtasks[thread].append({})
- combtasks[thread][-1]["colour"] = task["colour"]
- combtasks[thread][-1]["tic"] = task["tic"]
- combtasks[thread][-1]["toc"] = task["toc"]
- lastcolour = task["colour"]
- #else:
- # combtasks[thread][-1]["toc"] = task["toc"]
-
-
-fig = pl.figure(figsize=(14, 6))
-gs = gridspec.GridSpec(1, 3, width_ratios=[12,1, 1])
-ax = pl.subplot(gs[0])
-ax2 = pl.subplot(gs[2])
-ax.set_xlim(-delta_t * 0.03, delta_t * 1.03)
-ax.set_ylim(0, nthread)
-tictoc = np.zeros(2)
-for i in range(nthread):
-
- # Collect ranges and colours into arrays.
- tictocs = np.zeros(len(combtasks[i])*2)
- colours = np.empty(len(combtasks[i])*2, dtype='object')
- coloursseen = []
- j = 0
- for task in combtasks[i]:
- tictocs[j] = task["tic"]
- tictocs[j+1] = task["toc"]
- colours[j] = task["colour"]
- colours[j+1] = task["colour"]
- j = j + 2
- if task["colour"] not in coloursseen:
- coloursseen.append(task["colour"])
-
- # Legend support, collections don't add to this.
-# if task["subtype"] != "none":
-# qtask = task["type"] + "/" + task["subtype"]
- #else:
- #qtask = task["type"]
- #if qtask not in typesseen:
- #pl.plot([], [], color=task["colour"], label=qtask)
- #typesseen.append(qtask)
-
- # Now plot each colour, faster to use a mask to select colour ranges.
- for colour in coloursseen:
-# for x in np.where(colours==colour):
-# if i == 0:
-# print tictocs[x], "\n: ",min(tictocs[x]), tictocs[np.where(tictocs == min(tictocs[x]))[0]+1]
- collection = collections.BrokenBarHCollection.span_where(tictocs, ymin=i+0.05, ymax=i+0.95,
- where=colours == colour,
- facecolor=scalarMap.to_rgba(colour+1),
- linewidths=0)
- ax.add_collection(collection)
-
-# Legend and room for it.
-#nrow = len(typesseen) / 5
-#if len(typesseen) * 5 < nrow:
- #nrow = nrow + 1
-#ax.fill_between([0, 0], nthread+0.5, nthread + nrow + 0.5, facecolor="white")
-#ax.set_ylim(0, nthread + nrow + 1)
-#ax.legend(loc=1, shadow=True, mode="expand", ncol=5)
-
-# Start and end of time-step
-ax.plot([0, 0], [0, nthread], 'k--', linewidth=1)
-ax.plot([end_t, end_t], [0, nthread], 'k--', linewidth=1)
-
-ax.set_xlabel("Wall clock time [s]")
-ax.set_ylabel("Thread ID" )
-ax.set_yticks(pl.array(range(nthread)), True)
-
-
-dfge= matplotlib.colorbar.ColorbarBase(ax2, jet,norm=cNorm, ticks=[0, int(math.ceil(max(data[:,3])/2)), int(max(data[:,3]))] ,ticklocation='left', label='Number of taskyields')
-#dfge.set_ticks([0, max(data[:,6])/2, max(data[:,6])])
-#dfge.set_label("Number of taskyields")
-#ax2.set_ylabel("Number of taskyields")
-#ax2.set_yticks([0, max(data[:,6]) ])
-
-pl.savefig(outpng, format = 'png', dpi = 1000)
-pl.savefig('tasks.eps', format = 'eps', dpi = 1000)
-
-pl.show()
-
-print "Graphics done, output written to", outpng
-
-sys.exit(0)
-
+#!/usr/bin/env python
+
+"""
+Usage:
+ plot_tasks.py TASKTIMERS output.png [Ticks per ms]
+
+ where TASKTIMERS is a task timers file generated by quicksched.
+
+ This file is part of QuickSched.
+ Copyright (c) 2017 Aidan Chalk (aidan.chalk@stfc.ac.uk)
+
+This program is free software: you can redistribute it and/or modify
+it under the terms of the GNU Lesser General Public License as published
+by the Free Software Foundation, either version 3 of the License, or
+(at your option) any later version.
+
+This program is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+GNU General Public License for more details.
+
+You should have received a copy of the GNU Lesser General Public License
+along with this program. If not, see <http://www.gnu.org/licenses/>.
+"""
+
+import matplotlib
+import matplotlib.pyplot as plt
+import matplotlib.colors as colors
+import matplotlib.collections as collections
+import matplotlib.cm as cmx
+import matplotlib.gridspec as gridspec
+matplotlib.use('Agg')
+import pylab as pl
+import numpy as np
+import sys
+import math
+
+# Basic plot configuration.
+PLOT_PARAMS = {"axes.labelsize": 10,
+ "axes.titlesize": 10,
+ "font.size": 14,
+ "legend.fontsize": 12,
+ "xtick.labelsize": 11,
+ "ytick.labelsize": 11,
+ "figure.figsize" : (16., 4.),
+ "figure.subplot.left" : 0.03,
+ "figure.subplot.right" : 0.995,
+ "figure.subplot.bottom" : 0.1,
+ "figure.subplot.top" : 0.99,
+ "figure.subplot.wspace" : 0.,
+ "figure.subplot.hspace" : 0.,
+ "lines.markersize" : 6,
+ "lines.linewidth" : 3.
+ }
+
+TASKCOLOURS = ["black",
+ "lightblue",
+ "greenyellow",
+ "navy",
+ "greenyellow",
+ "navy",
+ "indigo",
+ "cyan",
+ "cyan",
+ "green",
+ "yellow",
+ "magenta",
+ "mediumorchid",
+ #"mediumnightblue",
+ "mediumturquoise",
+ "mediumvioletred",
+ "darkred",
+ "darkblue"
+ #,"powerblue"
+ ]
+
+pl.rcParams.update(PLOT_PARAMS)
+
+if len(sys.argv) == 2 and (sys.argv[1][0:2] == "-h" or sys.argv[1][0:3] == "--h"):
+ from pydoc import help
+ help( "__main__")
+ sys.exit(0)
+
+#Input from command line
+if len(sys.argv) != 3 and len(sys.argv) != 4:
+ print "Usage: ", sys.argv[0], "TASKTIMERS output.png [ticks-per-ms]"
+ sys.exit(1)
+
+
+
+infile = sys.argv[1]
+outpng = sys.argv[2]
+tpms = 0
+if len(sys.argv) == 4:
+ tpms = int(sys.argv[3])
+
+#read input.
+data = pl.loadtxt(infile)
+nthread=int(max(data[:,1])) + 1
+
+#Recompute timings in MS if possible.
+if(tpms != 0):
+ data[:,2] = data[:,2] / tpms
+ data[:,3] = data[:,3] / tpms
+
+start_t = min(data[:,2])
+data[:,2] -= start_t
+data[:,3] -= start_t
+end_t = (max(data[:,3]))
+start_t = min(data[:,2])
+delta_t = end_t - start_t
+
+
+
+fig = pl.figure()
+ax = fig.add_subplot(1,1,1)
+ax.set_xlim(-delta_t * 0.03, delta_t * 1.03)
+ax.set_ylim(0, nthread)
+
+print len(data[:,0])
+
+for j in range(nthread):
+ tictocs = np.zeros(len(data[:,0])*2)
+ types = np.zeros(len(data[:,0])*2)
+ for i in range(len(data[:,0])):
+ if data[i,1] == j:
+ tictocs[i*2+0] = data[i,2]
+ tictocs[i*2+1] = data[i,3]
+ types[i*2+0] = int(data[i,0])
+ types[i*2+1] = int(data[i,0])
+ #print tictocs
+ #print data[i,1]+0.05
+
+ for i in range(len(TASKCOLOURS)):
+ print i
+ print types==i
+ collection = collections.BrokenBarHCollection.span_where(tictocs, ymin=j+0.05, ymax=j+0.95, where=types==i, facecolor=TASKCOLOURS[i], linewidths=1, edgecolor="white")
+ ax.add_collection(collection)
+
+# Start and end of time-step
+ax.plot([start_t, start_t], [0, nthread], 'k--', linewidth=1)
+ax.plot([end_t, end_t], [0, nthread], 'k--', linewidth=1)
+
+ax.set_xlabel("time [ms or ticks]")
+ax.set_ylabel("Thread ID" )
+ax.set_yticks(pl.array(range(nthread)), True)
+
+pl.savefig(outpng, format = 'png', dpi = 1000)
+pl.savefig('tasks.eps', format = 'eps', dpi = 1000)
+
+pl.show()
+
+print "Graphics done, output written to", outpng
+
+sys.exit(0)