matplotlib_csv_example

#!/usr/bin/python

# 
# Copyright (c) 2011 Mateusz Jasinski                                                 #
# required: python-matplotlib python-tk python-argparse                               # 
#                              texlive-latex-extra                                    #
#                                                                                     #
# This program is free software; you can redistribute it and/or modify it under       #
# the terms of the GNU General Public License as published by the Free Software       #
# Foundation; either version 2 of the License, or (at your option) any later          #
# version.                                                                            #
#                                                                                     #
# This program is distributed in the hope that it will be useful, but WITHOUT ANY     #
# WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A     #
# PARTICULAR PURPOSE. See the GNU General Public License for more details.            #
#                                                                                     #
# You should have received a copy of the GNU General Public License along with        #
# this program.  If not, see <http://www.gnu.org/licenses/>.                          #


from matplotlib import mlab
import matplotlib.pyplot as plt
import numpy as np
from numpy import sqrt, mod, linspace
#from scipy import linspace, polyval, polyfit, poly1d, sqrt, stats, randn
import sys
import argparse

def main(argv):
    parser = argparse.ArgumentParser()
    parser.add_argument('-i', type=str, dest='ifile_name', default="./input.csv", help='CSV file name')
    parser.add_argument('-o', type=str, dest='ofile_name', default="./output.eps", help='output file name')
    parser.add_argument('-x', type=str, dest='x', default="0", help='column to use as x axis')
    parser.add_argument('-c', type=str, dest='cols', default="", help='columns to plot separated by commas, default = all')
    parser.add_argument('-n', type=int, dest='curve_number', default=0, help='number of curves to plot')
    parser.add_argument('-d', type=int, dest='dashed', default=False, help='draw plots as dashed lines')
    parser.add_argument('-l', type=int, dest='line', default=True, help='if true draw plots as lines, else draw points')
    parser.add_argument('-log', type=int, dest='log', default=True, help='if true draw plots in log scale on y axis')
    parser.add_argument('-loc', type=int, dest='loc', default=1, help='show legend in corner')
    parser.add_argument('-s', type=int, dest='show', default=1, help='show plot')
    args = parser.parse_args()
#    print args.file_name
#    print args.data_type

    # ustawienia rozmiaru wykresow:
    fig_width_pt = 420.0  		    # z LaTeX \showthe\columnwidth
    inches_per_pt = 1.0/72.27               # konwersja na cale
    golden_mean = (sqrt(5)-1.0)/2.0         # zloty podzial
    fig_width = fig_width_pt*inches_per_pt  # szerokosc w calach
    fig_height = fig_width*golden_mean      # wysokosc w calach
    fig_size =  [fig_width,fig_height]
    params = {'backend': 'ps',
	      'axes.labelsize': 10,
	      'text.fontsize': 10,
	      'legend.fontsize': 10,
	      'xtick.labelsize': 8,
	      'ytick.labelsize': 8,
	      'font.family': 'serif',       # czcionki jak w tekscie pracy
	      'font.serif': 'fourier',
	      'font.sans-serif': 'helvet',
	      'text.usetex': True,
	      'figure.figsize': fig_size}
    plt.rcParams.update(params)
    plt.figure(1)
    plt.clf()
    plt.axes([0.12,0.16,0.96-0.12,0.92-0.16]) # ustawienie marginesow tak aby opisy osi byly widoczne

    dashes = ['--', #    : dashed line
	      #'-', #     : solid line
	      '-.', #   : dash-dot line
	      ':', #    : dotted line
	      '-']
	      
    markers = [
               #'.', # point
	       #',', # pixel
	       'D', #  diamond
	       '+', # plus
	       '*', # star
	       #'*', # star
	       'x', # x
	       '^', # triangle_up
	       #'v', # triangle_down
	       '.', # point
	       'o', # circle
	       #'s', # square
	       'p', # pentagon
	       '<', # triangle_left
	       '>', # triangle_right
	       '1', # tri_down
	       '2', # tri_up
	       '3', # tri_left
	       '4', # tri_right
	       'h', # hexagon1
	       'H', # hexagon2
	       'D']#  diamond

    # wczytanie danych z pliku tekstowego
    f = open(args.ifile_name)
    #names = list(a.dtype.names) # ta funkcja zwraca opisy ale nazwy zwieraja wylacznie male litery
    names = map(str, f.readline().split(';'))
    names = [n.replace('\n','') for n in names]
    #print type(names)
    #print names

    
    a = mlab.csv2rec(args.ifile_name, delimiter=';', skiprows=1) # omin pierwszy wiersz - opisy osi

    if args.cols == "":
      cols = range(1, len(list(a.dtype.names)))
    else:
      cols = map(int, args.cols.split(','))
    print cols

#    print type(a)
#    print a.dtype.fields
#    print a[list(a.dtype.names)].view((float,len(a.dtype)))
#    print a[list(a.dtype.names)]
    
    if args.curve_number == 0 or args.curve_number > len(list(a.dtype.names)):
	args.curve_number = len(list(a.dtype.names))
    
    x_col = map(int, args.x)
    #x = [row[x_col] for row in a[list(a.dtype.names)].view((float,len(a.dtype)))] # os x
    for col in cols:
      y = [row[col] for row in a[list(a.dtype.names)].view((float,len(a.dtype)))]

      x = [row[x_col] for row in a[list(a.dtype.names)].view((float,len(a.dtype)))] # w tym wypadku trzeba odswiezyc wartosci osi x	
      for j in x: # upewnij sie czy usuniete sa wszystkie wartosci 0.0 :P
	for i in y:
	  if i == 0.0:
	    pos = y.index(i)
	    del x[pos]
	    del y[pos]    
	  
      if args.line == True:
	if args.dashed == True:
	  if args.log == True:
	    plt.semilogy(x, y, label=names[col + 2], linestyle = dashes[mod(col,len(dashes))])
	  else:
	    plt.plot(x, y, label=names[col + 2], linestyle = dashes[mod(col,len(dashes))])
	else:
	  if args.log == True:
	    plt.semilogy(x, y, '-', label=names[col + 2], marker = markers[mod(col,len(markers))])
	    #plt.semilogy(x, y, '.')
	  else:
	    plt.plot(x, y, label=names[col + 2], marker = markers[mod(col,len(markers))])
      else:
	if args.log == True:
	  plt.semilogy(x, y, '.', label=names[col + 2])
	else:
	  plt.plot(x, y, '.', label=names[col + 2])

#    plt.xlim(-6.0)
   
    if args.x == "":
        plt.xlabel(names[x_col[0]])
        plt.ylabel(names[cols[0]])
    else:
	#kjjkkplt.suptitle(names[0])
        plt.xlabel(names[1])
        plt.ylabel(names[2])

    plt.grid(True)
    plt.legend(loc=args.loc)
    #plt.legend(loc=2)
    plt.savefig(args.ofile_name)
    #plt.savefig(''.join([args.file_name, '.eps']))
    if args.show == True:
        plt.show()

if __name__ == "__main__":
    main(sys.argv)