Plot data using matplotlib

plot.py

#!/usr/bin/env python

""" A script to plot data using matplotlib. 
  Usage: [prog] [options] file1 file2 ..

The files passed are assumed to contain two values (x, y) per line.

There are two ways of using this script -- from the shell, or by calling the
Plot function with the appropriate arguments.

Example Run:

$ python plot.py data1.txt data2.txt --labels="old:new" --title='Performance vs $\alpha$' --xlabel='$\alpha$' --ylabel='Performance'

Or in a python prompt
>>> import plot
>>> plot.Plot(['data1.txt', 'data2.txt']
             ,labels = ['old', 'new']
             ,title = 'Performance vs $\alpha'
             ,xlabel = '$\alpha$'
             ,ylabel = 'Performance
             ,outfile = None
             ,colors = None
             ,linestyles = None)

An outfile argument can also be given, in which case the graph is saved.
matplotlib supports eps, png, etc.. apart from other formats.
"""

import matplotlib.pyplot as p
from optparse import OptionParser

def SplitTokens(value, options, parser):
    s = getattr(options, value)
    if not s:
        return None
    s = s.split(':')
    if len(s) != len(options.filenames):
        parser.error("%s was given arguments, but does not match number of filenames."%value)
    return s
        

def ParseArgs():
    default_string = "[default: %default]"
    usage = "%prog [options] file1 file2 ..."
    parser = OptionParser(usage)

    parser.add_option("-l"
                     ,"--labels"
                     ,dest = "labels"
                     ,type = "string"
                     ,default = None
                     ,help = "Labels as a colon separated string." + default_string)

    parser.add_option("-c"
                     ,"--colors"
                     ,dest = "colors"
                     ,type = "string"
                     ,default = None
                     ,help = "Matplotlib colors as a colon separated string." + default_string)

    parser.add_option("-s"
                     ,"--linestyles"
                     ,dest = "linestyles"
                     ,type = "string"
                     ,default = None
                     ,help = " Linestyles as a colon separated string." + default_string)

    parser.add_option("-o"
                     ,"--outfile"
                     ,dest = "outfile"
                     ,type = "string"
                     ,default = None
                     ,help = "Outfile. If None, display graph." + default_string)

    parser.add_option("-x"
                     ,"--xlabel"
                     ,dest = "xlabel"
                     ,type = "string"
                     ,default = None
                     ,help = "X-axis Label" + default_string)

    parser.add_option("-y"
                     ,"--ylabel"
                     ,dest = "ylabel"
                     ,type = "string"
                     ,default = None
                     ,help = "Y-axis Label" + default_string)

    parser.add_option("-t"
                     ,"--title"
                     ,dest = "title"
                     ,type = "string"
                     ,default = None
                     ,help = "Plot Title." + default_string)

    options, args = parser.parse_args()
    if not len(args):
        parser.error("No files given")

    options.filenames = args

    options.labels = SplitTokens('labels', options, parser)
    options.colors = SplitTokens('colors', options, parser)
    options.linestyles = SplitTokens('linestyles', options, parser)

    return options


def Plot(filenames,
         labels,
         title,
         xlabel,
         ylabel,
         outfile,
         colors,
         linestyles):

    for i in range(len(filenames)):
        args = {}
        s = None
        if colors:
            s = colors[i]
        if linestyles:
            s += linestyles[i]
        if labels:
            args['label'] = labels[i]
        lines = open(filenames[i]).readlines()
        xdata = map(lambda x: float(x.split()[0]), lines)
        ydata = map(lambda y: float(y.split()[1]), lines)
        if s:
            p.plot(xdata, ydata, s, **args)
        else:
            p.plot(xdata, ydata, **args)
    
    
    if xlabel:
        p.xlabel(xlabel)
    if ylabel: 
        p.ylabel(ylabel)
    
    if title:
        p.title(title)
    if labels:
        p.legend(tuple(labels))

    if outfile:
        p.savefig(outfile)
    else:
        p.show()
    return


if __name__ == '__main__':
    options = ParseArgs()
    Plot(options.filenames,
         options.labels,
         options.title,
         options.xlabel,
         options.ylabel,
         options.outfile,
         options.colors,
         options.linestyles)