Using multiple colors in matplotlib plot

colors, matplotlib, plot, python

Solution

I assume you want to plot distinct points. In that case, if you define a numpy array:

colormap = np.array(['r', 'g', 'b'])

then you can generate the array of colors with `colormap[categories]`:

In [18]: colormap[categories]
Out[18]: 
array(['r', 'b', 'g', 'g', 'g', 'b', 'r', 'r'], 
      dtype='|S1')

import matplotlib.pyplot as plt
import numpy as np

a = np.array([[ 1, 2, 3, 4, 5, 6, 7, 8 ],
              [ 9, 8, 7, 6, 5, 4, 3, 2 ]])

categories = np.array([0, 2, 1, 1, 1, 2, 0, 0])

colormap = np.array(['r', 'g', 'b'])

plt.scatter(a[0], a[1], s=50, c=colormap[categories])
plt.show()

yields

Problem

I have a numpy array of 2D data points (x,y), which are classified into three categories (0,1,2). ``` a = array([[ 1, 2, 3, 4, 5, 6, 7, 8 ], [ 9, 8, 7, 6, 5, 4, 3, 2 ]]) class = array([0, 2, 1, 1, 1, 2, 0, 0]) ``` My question is if I can plot these points with multiple colors. I would like to do something like this: ``` colors = list() for i in class: if i == 0: colors.append('r') elif i == 1: colors.append('g') else: colors.append('b') print colors ['r', 'b', 'g', 'g', 'g', 'b', 'r', 'r'] pp.plot(a[0], a[1], color = colors) ```

Original source