Как задать условную окраску в matplotlib.pyplot?

Question

Я пытаюсь сделать анимационный сюжет с помощью функции FuncAnimation. У меня есть объект person, который имеет атрибут с именем person.cat и может принимать значения в [1,2,3]. В настоящее время мой код предоставляет график, на котором все точки данных окрашены в красный цвет.

fig = plt.figure()
axis = plt.axes(xlim = (0,sz), ylim = (0,sz))
d, = axis.plot([person.posx for person in persons],[person.posy for person in persons],
               color = 'red' ,linestyle = 'None', marker = 'o')

Я хочу, чтобы мои точки данных на графике были условно окрашены, т.е. красным, если person.cat = 1, зеленым, если person.cat = 2, и синим, если person.cat = 3.

В любом случае, как этого добиться?

Полный код ниже:

fig = plt.figure()
axis = plt.axes(xlim = (0,sz), ylim = (0,sz))
d, = axis.plot([person.posx for person in persons],
                  [person.posy for person in persons],
                  color = 'red' ,
               linestyle = 'None', marker = 'o')

def init():
    d.set_data([],[])
    return d,

def animate(i):
    for person in persons:
        person.update_pos(1)
    d.set_data([person.posx for person in persons],[person.posy for person in persons])
    return d,

anim = animation.FuncAnimation(fig, animate, init_func = init, frames = 150, interval = 20, blit = True)

anim.save('p1.mp4', writer = 'ffmpeg', fps = 30)