Как задать индивидуальную раскраску в точечной диаграмме?

Question

def main():
    numframes = 30*simulation_time
    numpoints = n
    color_data = np.random.random((numframes, numpoints))
    #x, y = [person.posx[0] for person in persons],[person.posy[0] for person in persons]
    #color_data = [person.cat for person in persons]
    fig = plt.figure()
    axis = plt.axes(xlim = (0,sz), ylim = (0,sz))
    scat = plt.scatter([person.posx for person in persons],[person.posy for person in persons],
                      c = np.asarray([person.cat for person in persons]))
    #scat.set_array([person.cat for person in persons])

    anim = animation.FuncAnimation(fig, update_plot, frames=numframes,fargs=(color_data,scat))
    anim.save('p1.mp4', writer = 'ffmpeg', fps = 30)
    plt.show()

def update_plot(i,color_data, scat):
    
    for person in persons:
        person.update_pos(1)
    
    for person in persons:
        if person.quar == True or person.cat != 1:
            continue
        for neighbour in persons:
            if neighbour.cat!=2 or person.ID == neighbour.ID or np.random.rand()>trans_prob:
                continue
            if(np.sqrt((person.posx-neighbour.posx)**2+(person.posy-neighbour.posy)**2)<inf_rad):
                person.cat=2
                person.tm = 0
                
    for person in persons:
        if person.cat==2 and person.tm>=recov_tm:
            person.cat = 3
            person.quar = False
        elif person.cat==2 and person.quar==False and np.random.rand()<quar_prob:
            person.quar = True
            
    
    data = np.zeros((n,3))
    data[:,0] = [person.posx for person in persons]
    data[:,1] = [person.posy for person in persons]
    data[:,2] = [person.cat for person in persons]
    scat.set_offsets(data[:, :2])
    scat.set_array(data[:, 2])
    return scat,

main()

Я пытаюсь оживить график рассеивания. Person - это мой объект с несколькими атрибутами. Один из атрибутов - person.cat, который принимает дискретные значения [1,2,3]. Я хочу, чтобы моя анимация была окрашена в [r,g,b] соответственно в соответствии с person.cat. Эти значения постоянно обновляются, поэтому мне нужно постоянно обновлять мои цвета. Мой код дает некоторую ошибку в отображении цветов, как упоминалось выше. Есть ли какое-то решение для этого?

Для справки ниже приведен конструктор для person объекта класса ind:

class ind:
        
    def __init__(self, ID):
        self.ID = ID
        #self.init_posx = np.random.rand(1)*sz
        #self.init_posy = np.random.rand(1)*sz
        self.posx = np.random.rand(1)*sz
        self.posy = np.random.rand(1)*sz
        self.tm = 0
        self.spd = np.random.rand(1)*max_spd
        self.theta = np.random.rand(1)*2*np.pi
        self.cat = 1 #this can be either of 1,2,3.
        self.rec_tm = recov_tm
        self.quar = False

Answer 1

Я обнаружил, что устанавливать цвета в виде списков значений RGBA довольно мрачно, так как он мешает внутренним атрибутам диаграммы рассеяния. Но вот оно:

import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
import matplotlib.colors as c

fig, ax = plt.subplots()

scatt = ax.scatter([1, 2, 3], [1, 2, 3], c = [1, 2, 3])
scatt._A = None #To ensure that the colormap won't override our RGBA values
color = c.to_rgba_array(["black", "white", "blue"], 1.) #Just to get the RGBA values
scatt._facecolors = c.to_rgba_array(["black", "white", "blue"], 1.0) #This is where the magic happens
'''
print(scatt._facecolors)
outputs it's facecolors as RGBA's in order:
[[0. 0. 0. 1.]
 [1. 1. 1. 1.]
 [0. 0. 1. 1.]]
 '''
#Note that the scale  here is 0. to 1..
scatt._edgecolors = [color[0]]*3 #This is just the same as above, but for edgecolors.
plt.show()