Update matplotlib's colorbar during animation

Question

I have managed to get matplotlib to update a contour plot using FuncAnimation, but I can't get the colorbar to update the contour levels. When I replace the contour plot, the colorbar is unchanged. If I create the colorbar with each frame of the animation, then I get a whole bunch of colorbars.

How can I update the existing colorbar?

Here's an example that plots an elliptical function that changes over time. The range of values is steadily decreasing, so I need to update the colorbar.

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


class ContourAnimation(object):
    def __init__(self, count):
        self.count = count
        self.xs = np.arange(count)
        self.ys = np.arange(count)
        self.x_coords, self.y_coords = np.meshgrid(self.xs, self.ys)
        self.z = np.zeros(self.x_coords.shape)
        self.contour = None
        self.colorbar = None

    def update(self, n):
        t = (n + 1) / 100
        for x in self.xs:
            for y in self.ys:
                self.z[y][x] = ((x - self.count/2) ** 2 +
                                ((y - self.count/2)/t) ** 2)
        if self.contour:
            for artist in self.contour.collections:
                artist.remove()

        self.contour = plt.contourf(self.x_coords, self.y_coords, self.z)
        # Without this conditional, I get many colorbars.
        if self.colorbar is None:
            self.colorbar = plt.colorbar()
        return self.contour,


def main():
    fig = plt.figure()
    fib = ContourAnimation(30)
    plt.title('Contour Animation')
    fib_ani = animation.FuncAnimation(fig,
                                      fib.update,
                                      interval=500)
    plt.show()


if __name__ == '__main__':
    main()

score 1 · Accepted Answer · answered Nov 17 '18 at 01:33

After looking closer at the colorbar() function, I found that you can tell it which axes to put the colorbar in. If you pass the default of None, it will add a new subplot and use that. To update the colorbar, I record the axes from the colorbar subplot after adding it, and then reuse it next time. Now the colorbar updates nicely.

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


class ContourAnimation(object):
    def __init__(self, count):
        self.count = count
        self.xs = np.arange(count)
        self.ys = np.arange(count)
        self.x_coords, self.y_coords = np.meshgrid(self.xs, self.ys)
        self.z = np.zeros(self.x_coords.shape)
        self.contour = None
        self.colorbar_axes = None

    def update(self, n):
        t = (n + 1) / 100
        for x in self.xs:
            for y in self.ys:
                self.z[y][x] = ((x - self.count/2) ** 2 +
                                ((y - self.count/2)/t) ** 2)
        if self.contour:
            for artist in self.contour.collections:
                artist.remove()

        self.contour = plt.contourf(self.x_coords, self.y_coords, self.z)
        self.colorbar = plt.colorbar(cax=self.colorbar_axes)
        _, self.colorbar_axes = plt.gcf().get_axes()
        return self.contour,


def main():
    fig = plt.figure()
    fib = ContourAnimation(30)
    plt.title('Contour Animation')
    fib_ani = animation.FuncAnimation(fig,
                                      fib.update,
                                      interval=500)
    plt.show()


if __name__ == '__main__':
    main()

Update matplotlib's colorbar during animation

1 Answers1