plotprofile.py

   1 #
   2 #  plotprofile.py
   3 #  ENRAGE
   4 #
   5
   6 try:
   7     import psyco
   8     psyco.full()
   9 except:
  10     pass
  11
  12 import sys
  13 import optparse
  14 from pylab import *
  15 from math import *
  16 import matplotlib
  17
  18
  19 print "%prog: version 2.2"
  20
  21 def main(argv=None):
  22     if not argv:
  23         argv = sys.argv
  24
  25     usage = "usage: python %s infile outfile.png [--scale] [--max weightMax] [--ymin bottom] [--ymax top] [--subtractEvens]"
  26
  27     parser = optparse.OptionParser(usage=usage)
  28     parser.add_option("--scale", action="store_true", dest="doScale")
  29     parser.add_option("--max", type="float", dest="weightMax")
  30     parser.add_option("--ymin", type="float", dest="ymin")
  31     parser.add_option("--ymax", type="float", dest="ymax")
  32     parser.add_option("--subtractEvens", action="store_true", dest="subtractEvens")
  33     parser.set_defaults(doScale=False, weightMax=-1, ymin=None, ymax=None, subtractEvens=False)
  34     (options, args) = parser.parse_args(argv[1:])
  35
  36     if len(args) < 2:
  37         print usage
  38         sys.exit(1)
  39
  40     infile = args[0]
  41     pngfilename = args[1]
  42
  43     plotprofile(infile, pngfilename, options.doScale, options.weightMax, options.ymin, options.ymax, options.subtractEvens)
  44
  45
  46 def plotprofile(inFileName, pngfilename, doScale=False, weightMax=-1, ymin=None, ymax=None, subtractEvens=False):
  47     infile = open(inFileName)
  48     limitYscale = False
  49     if ymax is not None:
  50         limitYscale = True
  51     else:
  52         ymax = 0.
  53
  54     if ymin is not None:
  55         limitYscale = True
  56     else:
  57         ymin = 0.
  58
  59     matplotlib.use("Agg")
  60
  61     labelList = []
  62     dataList = []
  63     plotList = []
  64     xmin = 10**20
  65     xmax = -10**20
  66
  67     xcoordList = []
  68     datapointList = []
  69     weightList = []
  70     line = infile.readline()
  71     fields = line.strip().split()
  72     for data in fields[1:-1]:
  73         datapoint = float(data)
  74         if datapoint < xmin:
  75             xmin = datapoint
  76
  77         if datapoint > xmax:
  78             xmax = datapoint
  79
  80         xcoordList.append(datapoint)
  81
  82     index = 1
  83     for line in infile:
  84         fields = line.strip().split()
  85         datapointList = []
  86         for data in fields[1:-1]:
  87             datapointList.append(float(data))
  88
  89         if subtractEvens and index % 2 == 0:
  90             for dataIndex in range(len(datapointList)):
  91                 dataList[-1][dataIndex] -= datapointList[dataIndex]
  92         else:
  93             dataList.append(datapointList)
  94
  95         weight = float(fields[-1])
  96         if subtractEvens and index % 2 == 0:
  97             pass
  98         else:
  99             labelList.append(fields[0])
 100             if weight > weightMax:
 101                 weightMax = weight
 102
 103             weightList.append(weight)
 104
 105         index += 1
 106
 107     for index in range(len(dataList)):
 108         newList = []
 109         if doScale:
 110             scale = weightList[index] / weightMax
 111             print weightList[index], weightMax, scale
 112             for val in dataList[index]:
 113                 newList.append(val * scale)
 114         else:
 115             newList = dataList[index]
 116
 117         plotList.append(plot(xcoordList, newList, linewidth=3.0))
 118
 119     xticks(xcoordList, rotation="vertical")
 120     xlim(xmin - 0.1, xmax + 0.1)
 121     if limitYscale:
 122         ylim(ymin, ymax)
 123
 124     legend(plotList, labelList)
 125     savefig(pngfilename)
 126
 127
 128 if __name__ == "__main__":
 129     main(sys.argv)