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