source: misc/tools/plotcsv/plotcsv.py @ 6338

#!/usr/bin/env python

'''A module to draw a graph on the screen given a data file.

The code is written such that it may be run as a program or imported
and used as a module.
'''

import sys
import os.path
import types
import csv
import time
import getopt
import pylab


# Flag strings - keys in the 'options' dictionary
X_DATACOL = 'x_datacol'
Y_DATACOL = 'y_datacol'
X_RANGE = 'x_range'
Y_RANGE = 'y_range'
FILENAME = 'filename'
SIZE = 'size'
TITLE = 'title'
X_LABEL = 'x_label'
Y_LABEL = 'y_label'


##
# @brief Issue an error message.
# @param msg The message.
def error(msg):
    print >>sys.stderr, msg
    sys.exit(10)

 
##
# @brief Plot a sequence of data.
# @param x_data The X data sequence to plot.
# @param y_data A list of one or more Y data sequences to plot.
# @param options A dictionary of plot options.
def plot_data(x_data, y_data, options):
    pylab.plot(x_data, y_data)
    pylab.title(options.get(TITLE, ''))
    pylab.grid(True)

    # if user request a particular Y range
    if not options[Y_RANGE] is None:
        try:
            (minimum, maximum) = options[Y_RANGE].split(',')
            minimum = float(minimum)
            maximum = float(maximum)
        except:
            error('Sorry, got a bad value for Y range: %s' % options[Y_RANGE])
        pylab.ylim(ymin=minimum, ymax=maximum)

    pylab.xlabel(options.get(X_LABEL, ''))
    pylab.ylabel(options.get(Y_LABEL, ''))

    pylab.show()


##
# @brief Plot a data file.
# @param filename Path to the data file to plot.
# @param options A dictionary of options.
def plot_file(filename, options=None):
    # set options defaults
    opts = options
    if opts is None:
        opts = {}

    opts[X_DATACOL] = opts.get(X_DATACOL, '0')
    opts[Y_DATACOL] = opts.get(Y_DATACOL, '1')
    opts[SIZE] = opts.get(SIZE, '800,600')
    opts[Y_RANGE] = opts.get(Y_RANGE, None)

    # get contents of data file
    # after this, 'header' is list of column header strings
    #             'data' is a list of lists of data
    fd = open(filename)
    c = csv.reader(fd)
    data = []
    for row in c:
        data.append(row)
    fd.close()
    header = data[0]
    del data[0]

    # convert column specifiers to 'int' if required
    try:
        index = int(opts[X_DATACOL])
    except:
        try:
            index = header.index(opts[X_DATACOL])
        except ValueError:
            error("Sorry, X column header '%s' isn't in the data file." % opts[X_DATACOL])
    opts[X_DATACOL] = index

    try:
        index = int(opts[Y_DATACOL])
    except:
        try:
            index = header.index(opts[Y_DATACOL])
        except ValueError:
            error("Sorry, Y column header '%s' isn't in the data file." % opts[Y_DATACOL])
    opts[Y_DATACOL] = index
    
    # extract required columns from the data
    x_col = opts[X_DATACOL]
    y_col = opts[Y_DATACOL]

    # get max column number, check requested columns
    max_col = len(header)
    if x_col >= max_col or y_col >= max_col:
        error('Sorry, maximum column number for that file is %d.' % (max_col-1))

    x_label = header[x_col].title()
    x_data = map(lambda x: x[x_col], data)
    if x_label == 'Time':
        x_label = 'Time (hours)'
        x_data = map(lambda x: float(x)/3600., x_data)
    y_data = map(lambda x: x[y_col], data)
    y_label = header[y_col].title()

    opts[TITLE] = 'File: %s' % filename
    opts[X_LABEL] = x_label
    opts[Y_LABEL] = y_label

    plot_data(x_data, y_data, opts)


##
# @brief Help for the befuddled user.
def usage():
    print 'usage: %s <options> <filename>' % (ProgName)
    print 'where <filename> is the path to the data file to plot'
    print '  and <options>  is zero or more of:'
    print '                   -x <datacol> where <datacol> is a column specifier,'
    print '                   -y <datacol>     either a number or a header string'
    print '                   -v <range>   force a range in Y axis values, of the form'
    print '                                    <min>,<max>, eq, "-5,10"'
    print '                   -s <min,max> sets PNG output picture size'
    sys.exit(10)

if __name__ == '__main__':
    # dictionary to convert option string to internal name
    of_dict = { '-x': X_DATACOL,
                '-y': Y_DATACOL,
                '-s': SIZE,
                '-v': Y_RANGE }

    # set name of the program (for error reporting)
    global ProgName
    try:
        ProgName = os.path.basename(sys.argv[0])
    except:
        pass

    # get options
    opt_dict = {}
    params = sys.argv[1:]
    (opts, args) = getopt.gnu_getopt(params, 'x:y:s:v:')
    if len(args) != 1:
        usage()
    for (o, v) in opts:
        try:
            opt_dict[of_dict[o]] = v.lower()
        except:
            usage()

    # get name of data file
    if len(args) != 1:
        usage()
    filename = args[0]
    (filename_minus, _) = filename.split('.', 1)

    # set default option values
    opt_dict[X_DATACOL] = opt_dict.get(X_DATACOL, 0)
    opt_dict[Y_DATACOL] = opt_dict.get(Y_DATACOL, 1)
    opt_dict[SIZE] = opt_dict.get(SIZE, '800,600')
#    opt_dict[X_RANGE] = opt_dict.get(X_RANGE, None)
    opt_dict[Y_RANGE] = opt_dict.get(Y_RANGE, None)
    opt_dict[FILENAME] = filename_minus

    # plot the file
    plot_file(args[0], opt_dict)