Maintaining one colorbar for maptlotlib FuncAnimation

Question

I've made a script which uses matplotlib's FuncAnimation function to animate a series of contour plots for paraboloid surface functions. I'd like to add a colorbar for which the range does not change throughout the entire animation. I really have no idea how to do this. The script is shown below:

import numpy as np
import itertools
import matplotlib.pyplot as plt
import matplotlib.mlab as ml
import matplotlib.animation as animation

#Generate some lists

def f(x,y,a):
    return a*(x**2+y**2)

avals = list(np.linspace(0,1,10))
xaxis = list(np.linspace(-2,2,9))
yaxis = list(np.linspace(-2,2,9))

xy = list(itertools.product(xaxis,yaxis))
xy = list(map(list,xy))
xy = np.array(xy)

x = xy[:,0]
y = xy[:,1]
x = list(x)
y = list(y)

zlist = []

for a in avals:
    z = []
    for i, xval in enumerate(x):
        z.append(f(x[i],y[i],a))
    zlist.append(z)

xi = np.linspace(min(x),max(x),len(x))
yi = np.linspace(min(y), max(y), len(y))

fig,ax = plt.subplots()

def animate(index):
    zi = ml.griddata(x, y, zlist[index], xi, yi, interp='linear')
    ax.clear()
    contourplot = ax.contourf(xi, yi, zi, cmap=plt.cm.hsv,origin='lower')
    #cbar = plt.colorbar(contourplot)
    ax.set_title('%03d'%(index))
    return ax

ani = animation.FuncAnimation(fig,animate,np.array([0,1,2,3,4,5,6,7,8,9]),interval=200,blit=False)
plt.show()

Line 42 was my attempt at including said colorbar. The issue here is that because FuncAnimation calls the plotting function multiple times (once for each frame), the colorbar gets plotted multiple times thus messing up the animation. I also can't think of any way to move the colorbar instantiation outside of the animate function since the ax object appears to be local to it.

How can I put one colorbar for the whole animation?

Please note the above is fully working code. It should work on the appropriate python interpreter.

Answer 1

I guess the idea would be to create a contour plot outside the updating function once and give it a colorbar. The contour plot would then need to have defined levels and the colorrange needs to be defined.

ax.contourf(..., levels=levels, vmin=zmin, vmax=zmax)

where zmin and zmax are the minimum and maximum data to be shown, and levels is the list or array of levels to use.

Then, inside the animating function, you would only create a new contour plot with those same parameters without touching the colorbar at all.

import numpy as np
import itertools
import matplotlib.pyplot as plt
import matplotlib.mlab as ml
import matplotlib.animation as animation

def f(x,y,a):
    return a*(x**2+y**2)

avals = list(np.linspace(0,1,10))
xaxis = list(np.linspace(-2,2,9))
yaxis = list(np.linspace(-2,2,9))

xy = list(itertools.product(xaxis,yaxis))
xy = np.array(list(map(list,xy)))

x = xy[:,0]
y = xy[:,1]

zlist = []

for a in avals:
    z = []
    for i, xval in enumerate(x):
        z.append(f(x[i],y[i],a))
    zlist.append(z)

xi = np.linspace(min(x),max(x),len(x))
yi = np.linspace(min(y), max(y), len(y))

zmin = min([min(zl) for zl in zlist])
zmax = max([max(zl) for zl in zlist])
levels = np.linspace(zmin, zmax,41)
kw = dict(levels=levels, cmap=plt.cm.hsv, vmin=zmin, vmax=zmax, origin='lower')

fig,ax = plt.subplots()
zi = ml.griddata(x, y, zlist[0], xi, yi, interp='linear')
contourplot = ax.contourf(xi, yi, zi, **kw)
cbar = plt.colorbar(contourplot)

def animate(index):
    zi = ml.griddata(x, y, zlist[index], xi, yi, interp='linear')
    ax.clear()
    ax.contourf(xi, yi, zi, **kw)
    ax.set_title('%03d'%(index))


ani = animation.FuncAnimation(fig,animate,10,interval=200,blit=False)
plt.show()

Answer 2

As usual, I got beaten to the punch by @ImportanceOfBeingErnest, but I have a slightly different approach, which I thinks works as well.

I created a separate axe for the color bar, and I created a standalone color bar using the example from matplotlib's documentation. This requires to know the extend of the color scale before hand though.

Then I just plot the contourf in the animation using the same colorbar and normalization.

#Generate some lists
def f(x,y,a):
    return a*(x**2+y**2)

avals = list(np.linspace(0,1,10))
xaxis = list(np.linspace(-2,2,9))
yaxis = list(np.linspace(-2,2,9))

xy = list(itertools.product(xaxis,yaxis))
xy = list(map(list,xy))
xy = np.array(xy)

x = xy[:,0]
y = xy[:,1]
x = list(x)
y = list(y)

zlist = []

for a in avals:
    z = []
    for i, xval in enumerate(x):
        z.append(f(x[i],y[i],a))
    zlist.append(z)

xi = np.linspace(min(x),max(x),len(x))
yi = np.linspace(min(y), max(y), len(y))

fig,[ax,cax] = plt.subplots(1,2, gridspec_kw={"width_ratios":[10,1]})


# Set the colormap and norm to correspond to the data for which
# the colorbar will be used.
cmap = mpl.cm.hsv
norm = mpl.colors.Normalize(vmin=0, vmax=10)

cb1 = mpl.colorbar.ColorbarBase(cax, cmap=cmap,
                                norm=norm,
                                orientation='vertical')

def animate(index):
    zi = ml.griddata(x, y, zlist[index], xi, yi, interp='linear')
    ax.clear()
    contourplot = ax.contourf(xi, yi, zi, cmap=cmap, norm=norm, origin='lower')
    #cbar = plt.colorbar(contourplot)
    ax.set_title('%03d'%(index))
    return ax

ani = animation.FuncAnimation(fig,animate,np.array([0,1,2,3,4,5,6,7,8,9]),interval=200,blit=False)

Answer 3

Here is a lazy way to add colorbar. Instead of updating colorbar object, this code delete and create all objects in fig.

N = 10 # number of color steps
vmin, vmax = 0, 10 # this should be min and max of z
V = np.linspace(vmin, vmax, N) 

fig = plt.figure()
def animate(index):
    fig.clear()
    ax = plt.subplot(1,1,1)
    zi = ml.griddata(x, y, zlist[index], xi, yi, interp='linear')
    contourplot = ax.contourf(xi, yi, zi, V, cmap=plt.cm.hsv,origin='lower')
    cbar = plt.colorbar(contourplot)
    ax.set_title('%03d'%(index))
    return ax