如何在散点图中进行自定义着色？

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]。我希望我的动画分别根据person.cat 用[r,g,b] 着色。这些值不断更新，所以我需要保持我的颜色更新。如上所述，我的代码在颜色映射中出现了一些错误。有什么解决办法吗？

下面是 ind 类的 person 对象的构造函数供参考：

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()