Generating scatter plot from a CSV file

Question

Assume you have the following data in the form of a csv-file. The content looks something like this:

,Action,Comedy,Horror 1,650,819, ,76,63, 2,,462,19 ,,18,96 3,652,457,18 ,75,36,89 

which can be interpreted as a table of the form:

 Action Comedy Horror 1 650 819 76 63 2 462 19 18 96 3 652 457 18 75 36 89 

The goal was to write a function that takes a lst with genre names as elements in form of a str and returns a scatter plot of the data, where the data that should appear on the scatter plot is in the second row of every index (76, 63 , and , 18, 96 and 75, 36, 89). The function should be able to distinguish between two-dimensional and three-dimensional scatter plots depending on the input.

from pandas import DataFrame from csv import reader import matplotlib.pyplot as plt from mpl_toolkits.mplot3d import Axes3D def genre_scatter(lst): """ Creates an scatter plot using the data from genre_scores.csv. :param lst: a list with names of the genres considered :return: saves a pdf-file to the folder Fig with the name gen_1_ge_2.pdf """ # First we need to determine the right columns of genre_scores. first_row = [row for row in reader(open('genre_scores.csv', 'r'))][0] index = [first_row.index(x) for x in lst] # Get the relevant data in the form of a DataFrame. # Please note that the first row of data for every index is not necessary for this task. data = DataFrame.from_csv('genre_scores.csv') gen_scores = [data.dropna().iloc[1::2, ind - 1].transpose() for ind in index] # rewrite the values in an flattened array for plotting coordinates = [gen.as_matrix().flatten() for gen in gen_scores] # Plot the results fig = plt.figure() if len(coordinates) == 2: plt.scatter(*coordinates) plt.text(70, 110, "pearson={}".format(round(pearson_coeff(coordinates[0], coordinates[1]), 3))) plt.xlabel(lst[0]) plt.ylabel(lst[1]) plt.savefig("Fig/{}_{}.pdf".format(*lst)) else: ax = fig.add_subplot(111, projection='3d') ax.scatter(*coordinates) ax.update({'xlabel': lst[0], 'ylabel': lst[1], 'zlabel': lst[2]}) plt.savefig("Fig/{}_{}_{}.pdf".format(*lst)) plt.show() plt.close("all") if __name__ == "__main__": genre_scatter(['Action', 'Horror', 'Comedy']) 

The code works and I'm happy with the output but there are a few things that bug me and I'm not sure if I used them right.

1. I'm not incredibly familiar with list comprehension (I think that is what you call expressions of the form [x for x in list], please correct me if I'm wrong) and haven't used them very often, so I'm not quite sure if this here was the right approach for the problem. My biggest concern is the first use of this kind of expression, where I basically need the first row of the CSV file but create a list with all the rows only to use the first. Is there a smarter way to do this?
2. Is there a better way to label the axes? Ideally some function where I just could pass the *lst argument?

Please forget the pearson_coeff() part in the code, it's not really relevant for this.

Answer 1

This really isn't bad, in terms of base Python. The only thing that stands out to me is this:

first_row = [row for row in reader(open('genre_scores.csv', 'r'))][0] 

Firstly, you aren't closing the file. Always close the file after you're done.

'r' is implicit, so you don't need to write it in the arguments to open.

Also, you're building up an entire list in memory from the CSV file, and then throwing it all away only to use the first row. Instead, you should use something like:

with open('genre_scores.csv') as f: csv_reader = reader(f) first_row = next(csv_reader) 

You also ask:

I'd like to implement something that makes sure that lst isn't longer than three elements (since four dimensional plots aren't really a thing). The only way I know to do this is assert len(lst) <=3, which gets the job done but it would be nice if it also could raise a useful error message.

Fairly straightforward, and I'll also assume that the minimum is 2:

if not (2 <= len(lst) <= 3): raise ValueError(f'Invalid lst length of {len(lst)}')