Basic CSV x-y Plotting with Python

Question

Trying to write a python script that reads a csv file and prints an xy plot. I have a csv file with a few rows and columns worth of data. I'd like to plot an xy plot for the first and second column. Here is what I've got so far...

import csv
def getColumn(filename, column):
    results = csv.reader(open(filename), delimiter="\t")
    return [result[column] for result in results]

x = getColumn("TableOne.csv",0)
y = getColumn("TableOne.csv",1)

plt.figure("X-Y Plot")
plt.xlabel("Site")
plt.ylabel("Average")
plt.plot(x,y)

... but it's reading my csv file by row, not column, which outputs a bunch of mixed data instead of the one specific data I want. Any suggestions?

Answer 1

Look at the zip function and some answers regarding zip splats .

You can easily do this with:

In [1]: data = [['x1', 'y1'], ['x2', 'y2'], ['x3', 'y3']]

In [2]: zip(*data)
Out[2]: [('x1', 'x2', 'x3'), ('y1', 'y2', 'y3')]

Here, the *data unpacks data into 3 lists of [x, y] pairs. Then zip packages these together by taking the first element of each list into one group, the second element of each list into another group, and so on. Since you only have 2 elements per list, zip(*data) returns two groups, the x and y elements separately.

In your case, change

def getColumns(filename):
    results = csv.reader(open(filename), delimiter="\t")
    return zip(*list(results)) # may not need the list() but just in case

x, y  = getColumns("TableOne.csv")

plt.figure("X-Y Plot")
plt.xlabel("Site")
plt.ylabel("Average")
plt.plot(x,y)