How do I assign different dots different colors in a scatter plot made with MatPlotLib?

Question

Using MatPlotLib and Python, I am making a scatter plot of laptops that has cost on the x-axis and rating on the y-axis. The data is now stored in a tab-separated file.

This link describes MatPlotLib scatter plots. http://matplotlib.sourceforge.net/examples/api/unicode_minus.html

Each laptop's brand is Lenova, HP, or Dell. I want to assign different colors for dots representing laptops of each brand. For instance, Dell laptops are marked by a blue point, while Lenova laptops are represented by a red point.

How can I do that?

Answer 1

A better example for you would be the scatter plot example.

In your case, in place of the close array, you want some array of strings that are one of 'r', 'g', or 'b'. So if you can read an array of the brands, you can then translate that into an array of those characters. (And I guess you don't want the equivalent of the volume array, which makes the circles different sizes.)

You could translate your brand array to colors with something like

brandcolors = [brand[i].replace('HP','g').replace('Lenovo','r').replace('Dell','b') for i in range(len(brand))]

Then, the plotting function would look like

scatter(cost, rating, c=brandcolors)

Adding a legend gets a little tricky. The matplotlib docs suggest using a "proxy artist", where you create something, but don't put it on the plot. For example:

pr = Rectangle((0, 0), 1, 1, fc='r')
pg = Rectangle((0, 0), 1, 1, fc='g')
pb = Rectangle((0, 0), 1, 1, fc='b')
legend([pr, pg, pb], ['Lenovo', 'HP', 'Dell'])

That's a little ugly because it comes out as a rectangle regardless of the dot shape. The other possible hack is to actually plot a single dot, but make sure it stays off the plot area.

plot([-100], 'ro', label='Lenovo')
...
ax.set_ylim(bottom=0)
legend()