How to create a scatter plot with color specified for each point individually?

Question

Using matplotlib I create a scatter plot animation that shows a new point after each second and shows all old points partly transparent. Each point is defined by x and y, but also by a category s. I want the color of the points to be tied to its category. Ideally that means that the array s contains values 1, 2 and 3, and the colors belonging to those values are defined seperately. However, I can not get this to work.

What I do get to work is to specify the edgecolors of each point individually in s, the code for this is shown below.

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.image as plti
import matplotlib.animation

s = [[1,0,0],[0,1,0],[0,0,1]];
x = [525,480,260];
y = [215,180,180];

img = plti.imread('myimage.png')
fig, ax = plt.subplots()
plt.imshow(img)
plt.axis('off')

x_vals = []
y_vals = []
intensity = []
iterations = len(x)
colors = []

t_vals = np.linspace(0,iterations-1,iterations,dtype=int)
scatter = ax.scatter(x_vals, y_vals, s=100, c=colors, vmin=0, vmax=1)

def init():
    pass

def update(t):
    global x, y, x_vals, y_vals, intensity
    x_vals.extend([x[t]])
    y_vals.extend([y[t]])
    scatter.set_offsets(np.c_[x_vals,y_vals])
    intensity = np.concatenate((np.array(intensity), np.ones(1)))
    if len(intensity) > 1:
        intensity[-2] = 0.5
    scatter.set_array(intensity)
    colors.extend([s[t]])
    scatter.set_color(colors)

    return ani

ani = matplotlib.animation.FuncAnimation(fig, update, frames=t_vals, interval=1000, repeat=False, init_func=init)
plt.show()

Simply changing c=colors to facecolor=colors does not work. Also I have tried to use colormaps but I cannot get it to work using that either.

The resulting animation from the code above looks as below.

However, the animation should look like this..

So my question is; does someone know how to tie the facecolor of each point to the category that that point belongs to?

Answer 1

The problem occurred because the line scatter.set_array(intensity) was called before scatter.set_color(colors). So instead of defining the intensity by a seperate variable, it is instead integrated into the colors directly. The following code produces the desired result.

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.image as plti
import matplotlib.animation

s = [1,2,3];
x = [525,480,260];
y = [215,180,180];

img = plti.imread('myimage.png')
fig, ax = plt.subplots()
plt.imshow(img)
plt.axis('off')

x_vals = []
y_vals = []
iterations = len(x)
colors = []

t_vals = np.linspace(0,iterations-1,iterations,dtype=int)
scatter = ax.scatter(x_vals, y_vals, s=100, color=colors, vmin=0, vmax=1)

def init():
    pass

def update(t):
    global x, y, x_vals, y_vals
    x_vals.extend([x[t]])
    y_vals.extend([y[t]])
    scatter.set_offsets(np.c_[x_vals,y_vals])

    if t > 0:
        if s[t-1] == 1:
            colors[t-1] = [1,0,0,0.5];
        elif s[t-1] == 2:
            colors[t-1] = [0,1,0,0.5];
        else:
            colors[t-1] = [0,0,1,0.5];

    if s[t] == 1:
        colors.extend([[1,0,0,1]])
    elif s[t] == 2:
        colors.extend([[0,1,0,1]])
    else:
        colors.extend([[0,0,1,1]])
    scatter.set_color(colors);

    return ani

ani = matplotlib.animation.FuncAnimation(fig, update, frames=t_vals, init_func=init, interval=1000, repeat=False)
plt.show()

Answer 2

The normal way to plot plots with points in different colors in matplotlib is to pass a list of colors as a parameter.

E.g.:

import matplotlib.pyplot
matplotlib.pyplot.scatter([1,2,3],[4,5,6],color=['red','green','blue'])

But if for some reason you wanted to do it with just one call, you can make a big list of colors, with a list comprehension and a bit of flooring division:

import matplotlib
import numpy as np

X = [1,2,3,4]
Ys = np.array([[4,8,12,16],
      [1,4,9,16],
      [17, 10, 13, 18],
      [9, 10, 18, 11],
      [4, 15, 17, 6],
      [7, 10, 8, 7],
      [9, 0, 10, 11],
      [14, 1, 15, 5],
      [8, 15, 9, 14],
       [20, 7, 1, 5]])
nCols = len(X)  
nRows = Ys.shape[0]

colors = matplotlib.cm.rainbow(np.linspace(0, 1, len(Ys)))

cs = [colors[i//len(X)] for i in range(len(Ys)*len(X))] #could be done with numpy's repmat
Xs=X*nRows #use list multiplication for repetition
matplotlib.pyplot.scatter(Xs,Ys.flatten(),color=cs)