text: #!/usr/bin/env python3
  
  import argparse
  import collections
  import glob
  import math
  import shelve
  import os.path
  
  import numpy as np
  import matplotlib
  import matplotlib.pyplot as plt
  from matplotlib.ticker import FuncFormatter, LogLocator
  from matplotlib import cm
  from matplotlib import colors
  from mpl_toolkits.mplot3d import Axes3D
  
  import plot_helpers
  
  names = ['FIFO', 'PS', 'SRPT', 'FSP', 'LAS', 'SRPTE', 'SRPTE+PS', 'SRPTE+LAS',
           'FSPE', 'FSPE+PS', 'FSPE+LAS']
  axes = 'shape sigma load timeshape njobs est_factor'.split()
  
  parser = argparse.ArgumentParser(description="3d plot of mean sojourn time")
  parser.add_argument('scheduler', help="scheduler for which to plot results",
                      choices=names)
  parser.add_argument('dirname', help="directory in which results are stored")
  parser.add_argument('--xaxis', default='shape', choices=axes,
                      help='what to put in the x-axis; default: shape')
  parser.add_argument('--yaxis', default='sigma', choices=axes,
                      help='what to put in the y-axis; default: sigma')
  parser.add_argument('--linx', default=False, action='store_true',
                      help='linear (instead of logarithmic) x axis')
  parser.add_argument('--liny', default=False, action='store_true',
                      help='linear (instead of logarithmic) y axis')
  parser.add_argument('--linz', default=False, action='store_true',
                      help='linear (instead of logarithmic) z axis')
  parser.add_argument('--normalize', choices=names,
                      help="normalize against another scheduler")
  parser.add_argument('--shape', type=float, default=0.5,
                      help="shape for job size distribution "
                      "(if not on one of the axes); default: 0.5")
  parser.add_argument('--sigma', type=float, default=0.5,
                      help="sigma for size estimation error log-normal "
                      "distribution (if not on one of the axes); default: 0.5")
  parser.add_argument('--load', type=float, default=0.9,
                      help="load for the generated workload; default: 0.99")
  parser.add_argument('--timeshape', type=float, default=1,
                      help="shape for the Weibull distribution of job "
                      "inter-arrival times; default: 1 (i.e. exponential)")
  parser.add_argument('--njobs', type=int, default=10000,
                      help="number of jobs in the workload; default: 10000")
  parser.add_argument('--est_factor', type=float,
                      help="multiply estimated size by this value")
  parser.add_argument('--notitle', default=False, action='store_true',
                      help="do not render title")
  parser.add_argument('--zmin', type=float,
                      help="minimum value on z axis")
  parser.add_argument('--zmax', type=float,
                      help="maximum value on z axis")
  parser.add_argument('--save', help="don't show but save in target filename")
  args = parser.parse_args()
  
  if not args.est_factor and 'est_factor' not in [args.xaxis, args.yaxis]:
      axes.pop()
  
  xaxis_idx = axes.index(args.xaxis)
  yaxis_idx = axes.index(args.yaxis)
  
  fname_regex = [str(getattr(args, ax)) for ax in axes]
  fname_regex[xaxis_idx] = fname_regex[yaxis_idx] = '[0-9.]*'
  glob_str = os.path.join(args.dirname,
                          'res_{}_[0-9.]*.s'.format('_'.join(fname_regex)))
  fnames = glob.glob(glob_str)
  
  cache = shelve.open(os.path.join(args.dirname, 'cache.s'))
  def getmean(fname, scheduler):
      key = '{}_mean_{}'.format(os.path.split(fname)[1], scheduler)
      try:
          return cache[key]
      except KeyError:
          shelve_ = shelve.open(fname, 'r')
          mean = np.array(shelve_[scheduler]).mean()
          shelve_.close()
          cache[key] = mean
          return mean
  
  results = collections.defaultdict(list)
  xvals, yvals = set(), set()
  for fname in fnames:
      print('.', end='', flush=True)
      split = os.path.splitext(os.path.split(fname)[1])[0].split('_')[1:-1]
      xval = float(split[xaxis_idx])
      if args.xaxis == 'load':
          xval = 1 - xval
      yval = float(split[yaxis_idx])
      if args.yaxis == 'load':
          yval = 1 - yval
      xvals.add(xval)
      yvals.add(yval)
      try:
          mst = getmean(fname, args.scheduler)
          if args.normalize:
              mst = mst / getmean(fname, args.normalize)
      except:
          # the file is being written now
          continue
      results[xval, yval].append(mst)
  cache.close()
  
  print()
  xvals = sorted(xvals)
  yvals = sorted(yvals)
  X, Y = np.meshgrid(xvals, yvals)
  if not args.linx:
      X = np.log2(X)
  if not args.liny:
      Y = np.log2(Y)
  Z = np.zeros_like(X)
  
  for i, xval in enumerate(xvals):
      for j, yval in enumerate(yvals):
          z = np.array(results[xval, yval]).mean()
          Z[j, i] = z if args.linz else np.log2(z)
  
  def format_func(x, pos):
      return '{:.3g}'.format(2 ** x)
  formatter = FuncFormatter(format_func)
  
  def load_format(x, pos):
      return '{:.3g}'.format(1 - 2 ** x)
  load_formatter = FuncFormatter(load_format)
  def load_linformat(x, pos):
      return str(1 - x)
  load_linformatter = FuncFormatter(load_linformat)
  load_ticks = np.log2(1 - np.array([0.5, 0.9, 0.99, 0.999]))
  njobs_ticks = np.log2([100, 1000, 10000, 100000])
          
  fig = plt.figure()
  ax = fig.add_subplot(111, projection='3d')
  ax.set_xlabel(args.xaxis)
  ax.set_ylabel(args.yaxis)
  if args.xaxis == 'njobs':
      ax.xaxis.set_ticks(njobs_ticks)
  if args.yaxis == 'njobs':
      ax.yaxis.set_ticks(njobs_ticks)
  if args.xaxis == 'load':
      ax.xaxis.set_ticks(load_ticks)
      if args.linx:
          ax.xaxis.set_major_formatter(load_linformatter)
      else:
          ax.xaxis.set_major_formatter(load_formatter)
  elif not args.linx:
      ax.xaxis.set_major_formatter(formatter)
  if args.yaxis == 'load':
      ax.yaxis.set_ticks(load_ticks)
      if args.linx:
          ax.yaxis.set_major_formatter(load_linformatter)
      else:
          ax.yaxis.set_major_formatter(load_formatter)
  elif not args.linx:
      ax.yaxis.set_major_formatter(formatter)
  if args.normalize:
      zlabel = "MST / MST({})".format(args.normalize)
  else:
      zlabel = "Mean sojourn time"
  ax.set_zlabel(zlabel)
  if not args.linz:
      ax.zaxis.set_major_formatter(formatter)
  if not args.notitle:
      plt.title(args.scheduler)
  if args.normalize:
      if args.linz:
          surf = ax.plot_surface(X, Y, Z, rstride=1, cstride=1, cmap=cm.Greens)
      else:
          surf = ax.plot_surface(X, Y, Z, rstride=1, cstride=1, cmap=cm.bwr,
                                 vmin=-6, vmax=6)
      cont = ax.contour(X, Y, Z, levels=[0], colors='k', linewidths=5,
                        linestyles='dashed')
  
      # Horrible hack to get the last line always rendered
      # http://stackoverflow.com/questions/20781859/drawing-a-line-on-a-3d-plot-in-matplotlib
      from mpl_toolkits.mplot3d.art3d import Line3DCollection
      class FixZorderCollection(Line3DCollection):
          _zorder = 1000
  
          @property
          def zorder(self):
              return self._zorder
  
          @zorder.setter
          def zorder(self, value):
              pass
      ax.collections[-1].__class__ = FixZorderCollection
  else:
      surf = ax.plot_surface(X, Y, Z, rstride=1, cstride=1, cmap=cm.Greens)
  if args.zmin:
      zmin = args.zmin if args.linz else np.log2(args.zmin)
      ax.set_zlim(bottom=zmin)
  if args.zmax:
      zmax = args.zmax if args.linz else np.log2(args.zmax)
      ax.set_zlim(top=zmax)
  if not args.linz:
      minz, maxz = ax.zaxis.get_view_interval()
      ax.zaxis.set_ticks(range(math.ceil(minz), math.floor(maxz) + 1))
  
  plot_helpers.config_paper(20)
  
  plt.tight_layout(1)
  
  if args.save is not None:
      plt.savefig(args.save)
  else:
      plt.show()
  
  