анимация matplotlib: обновить цвет маркера - PullRequest
Новая
       20

анимация matplotlib: обновить цвет маркера

0 голосов
/ 25 апреля 2020

Я использую FuncAnimation и хочу обновить цвета маркеров во время анимации.

Это мой текущий код. Получающаяся анимация не меняет цвета или маркеры вообще. Я пробовал set_color и многие другие подходы. 

class AnimatedScatter(object):
    def __init__(self, x, y, sgn, scale=1e6, title="", ANNCRE=False, Sources=False, SourcePosX=False, SourcePosY=False, SourceT=False):
        self.x = iter(x)
        self.y = iter(y)
        if ANNCRE:
            self.sgn = iter(sgn)
        else:
            self.sgn = sgn
        if Sources:
            norm = matplotlib.colors.Normalize(vmin=0,vmax=np.max(SourceT), clip=False)
            self.mapper = matplotlib.cm.ScalarMappable(norm=norm, cmap=matplotlib.cm.Reds)

            self.SourceT = iter(SourceT)
            self.SourcePosX = SourcePosX
            self.SourcePosY = SourcePosY


        self.Sources = Sources
        self.scale = scale
        self.n_t = x.__len__()
        self.title = "Simulation Domain\n" + title
        self.ANNCRE = ANNCRE

        # Setup the figure and axes...
        self.fig, self.ax = plt.subplots()

        if Sources:
            self.fig.colorbar(matplotlib.cm.ScalarMappable(norm=norm, cmap=matplotlib.cm.Reds))

    def setup_plot(self):
        x = next(self.x)
        y = next(self.y)
        if self.ANNCRE:
            sgn = next(self.sgn)
        else:
            sgn = self.sgn

        if self.Sources:
            SourceT = self.mapper.to_rgba(np.squeeze(next(self.SourceT)))


            self.scatSource = self.ax.scatter(self.SourcePosX * self.scale, self.SourcePosY * self.scale, marker="x",
                             c=SourceT)

        self.scatP = self.ax.scatter(x[(sgn > 0)] * self.scale, y[(sgn > 0)] * self.scale, marker=r'$\perp$',
                                     c='blue')
        self.scatN = self.ax.scatter(x[(sgn < 0)] * self.scale, y[(sgn < 0)] * self.scale, marker=r'$\top$',
                                     c='red')

        self.ax.axis([0, dim * 1e6, 0, dim * 1e6])

        plt.xlabel('Direction $x$ $(\mu m)$')
        plt.ylabel('Dircetion $y$ $(\mu m)$')
        plt.title(self.title)
        return self.scatP, self.scatN, self.scatSource,

    def update(self, i):
        """Update the scatter plot."""
        x = next(self.x)
        y = next(self.y)
        if self.ANNCRE:
            sgn = next(self.sgn)
        else:
            sgn = self.sgn

        if self.Sources:
            SourceT = self.mapper.to_rgba(np.squeeze(next(self.SourceT)))
            # self.scatSource.edgecolors(SourceT)
            self.scatSource.set_array(SourceT)


        dataP = np.vstack((x[(sgn > 0)], y[(sgn > 0)]))
        dataN = np.vstack((x[(sgn < 0)], y[(sgn < 0)]))

        self.scatP.set_offsets(dataP.T * self.scale)
        self.scatN.set_offsets(dataN.T * self.scale)
        return self.scatP, self.scatN, self.scatSource,

    def save(self, Frames=False, title="animation_thingy"):
        if ~Frames:
            Frames=(self.n_t - 2)
        Writer = animation.FFMpegWriter(fps=15)

        self.ani = animation.FuncAnimation(self.fig,
                                           self.update,
                                           init_func=self.setup_plot,
                                           frames=tqdm(range(Frames)),
                                           interval=1/15,
                                           blit=True)
        self.ani.save("{}.mp4".format(title), writer=Writer)

Я полностью застрял на этом этапе, ваша помощь очень ценится.

...