Using matplotlib.animate to animate a contour plot in python

Question

I have a 3D array of data (2 spatial dimensions and 1 time dimension) and I'm trying to produce an animated contour plot using matplotlib.animate. I'm using this link as a basis:

http://jakevdp.github.io/blog/2012/08/18/matplotlib-animation-tutorial/

And here's my attempt:

import numpy as np
from matplotlib import pyplot as plt
from matplotlib import animation
from numpy import array, zeros, linspace, meshgrid
from boutdata import collect

# First collect data from files
n = collect("n")   #  This is a routine to collect data
Nx = n.shape[1]
Nz = n.shape[2]
Ny = n.shape[3]
Nt = n.shape[0]

fig = plt.figure()
ax = plt.axes(xlim=(0, 200), ylim=(0, 100))
cont, = ax.contourf([], [], [], 500)

# initialisation function
def init():
    cont.set_data([],[],[])
    return cont,

# animation function
def animate(i): 
    x = linspace(0, 200, Nx)
    y = linspace(0, 100, Ny)
    x,y = meshgrid(x,y)
    z = n[i,:,0,:].T
    cont.set_data(x,y,z)
    return cont, 

anim = animation.FuncAnimation(fig, animate, init_func=init,
                           frames=200, interval=20, blit=True)

plt.show()

But when I do this, I get the following error:

Traceback (most recent call last):
  File "showdata.py", line 16, in <module>
    cont, = ax.contourf([], [], [], 500)
  File "/usr/lib/pymodules/python2.7/matplotlib/axes.py", line 7387, in contourf
    return mcontour.QuadContourSet(self, *args, **kwargs)
  File "/usr/lib/pymodules/python2.7/matplotlib/contour.py", line 1112, in __init__
    ContourSet.__init__(self, ax, *args, **kwargs)
  File "/usr/lib/pymodules/python2.7/matplotlib/contour.py", line 703, in __init__
    self._process_args(*args, **kwargs)
  File "/usr/lib/pymodules/python2.7/matplotlib/contour.py", line 1125, in _process_args
    x, y, z = self._contour_args(args, kwargs)
  File "/usr/lib/pymodules/python2.7/matplotlib/contour.py", line 1172, in _contour_args
    x,y,z = self._check_xyz(args[:3], kwargs)
  File "/usr/lib/pymodules/python2.7/matplotlib/contour.py", line 1204, in _check_xyz
    raise TypeError("Input z must be a 2D array.")
TypeError: Input z must be a 2D array.

So I've tried replacing all the [] by [[],[]] but this then produces:

Traceback (most recent call last):
  File "showdata.py", line 16, in <module>
    cont, = ax.contourf([[],[]], [[],[]], [[],[]],500)
  File "/usr/lib/pymodules/python2.7/matplotlib/axes.py", line 7387, in contourf
    return mcontour.QuadContourSet(self, *args, **kwargs)
  File "/usr/lib/pymodules/python2.7/matplotlib/contour.py", line 1112, in __init__
    ContourSet.__init__(self, ax, *args, **kwargs)
  File "/usr/lib/pymodules/python2.7/matplotlib/contour.py", line 703, in __init__
    self._process_args(*args, **kwargs)
  File "/usr/lib/pymodules/python2.7/matplotlib/contour.py", line 1125, in _process_args
    x, y, z = self._contour_args(args, kwargs)
  File "/usr/lib/pymodules/python2.7/matplotlib/contour.py", line 1177, in _contour_args
    self.zmax = ma.maximum(z)
  File "/usr/lib/python2.7/dist-packages/numpy/ma/core.py", line 5806, in __call__
    return self.reduce(a)
  File "/usr/lib/python2.7/dist-packages/numpy/ma/core.py", line 5824, in reduce
    t = self.ufunc.reduce(target, **kargs)
ValueError: zero-size array to maximum.reduce without identity

Thanks in advance!

Answer 1

I have been looking at this a while ago. I my situation I had a few subplots with contours which I wanted to animate. I did not want to use the plt.clf() solution as Rehman ali suggest as I used some special setup of my axis (with pi symbols etc) which would be cleaned as well, so I preferred the 'remove()' approach suggest be Felix. The thing is that only using 'remove' does not clean up memory and will clog your computer eventually, so you need to explicitly delete of the contours by setting it to an empty list as well.

In order to have a generic remove routine which is able to take away contours as well as text, I wrote the routine 'clean_up_artists' which you should use on every time step on all the axis.

This routine cleans up the artists which are passed in a list called 'artist_list' in a given axis 'axis'. This means that for animating multiple subplots, we need to store the lists of artists for each axis which we need to clean every time step.

Below the full code to animate a number of subplots of random data. It is pretty self-explanatory, so hopefully it becomes clear what happens. Anyhow, I just thought to post it, as it combines several ideas I found on stack overflow which I just to come up with this working example.

Anybody with suggestions to improve the code, please shoot-)

import matplotlib.pyplot as plt
from matplotlib import cm
import matplotlib.animation as animation
import string
import numpy as np


def clean_up_artists(axis, artist_list):
    """
    try to remove the artists stored in the artist list belonging to the 'axis'.
    :param axis: clean artists belonging to these axis
    :param artist_list: list of artist to remove
    :return: nothing
    """
    for artist in artist_list:
        try:
            # fist attempt: try to remove collection of contours for instance
            while artist.collections:
                for col in artist.collections:
                    artist.collections.remove(col)
                    try:
                        axis.collections.remove(col)
                    except ValueError:
                        pass

                artist.collections = []
                axis.collections = []
        except AttributeError:
            pass

        # second attempt, try to remove the text
        try:
            artist.remove()
        except (AttributeError, ValueError):
            pass


def update_plot(frame_index, data_list, fig, axis, n_cols, n_rows, number_of_contour_levels, v_min, v_max,
                changed_artists):
    """
    Update the the contour plots of the time step 'frame_index'

    :param frame_index: integer required by animation running from 0 to n_frames -1. For initialisation of the plot,
    call 'update_plot' with frame_index = -1
    :param data_list: list with the 3D data (time x 2D data) per subplot
    :param fig: reference to the figure
    :param axis: reference to the list of axis with the axes per subplot
    :param n_cols: number of subplot in horizontal direction
    :param n_rows: number of subplot in vertical direction
    :param number_of_contour_levels: number of contour levels
    :param v_min: minimum global data value. If None, take the smallest data value in the 2d data set
    :param v_max: maximum global data value. If None, take the largest value in the 2d data set
    :param changed_artists: list of lists of artists which need to be updated between the time steps
    :return: the changed_artists list
    """

    nr_subplot = 0  # keep the index of the current subplot  (nr_subplot = 0,1,  n_cols x n_rows -1)
    # loop over the subplots
    for j_col in range(n_cols):
        for i_row in range(n_rows):

            # set a short reference to the current axis
            ax = axis[i_row][j_col]

            # for the first setup call, add and empty list which can hold the artists belonging to the current axis
            if frame_index < 0:
                # initialise the changed artist list
                changed_artists.append(list())
            else:
                # for the next calls of update_plot, remove all artists in the list stored in changed_artists[nr_subplot]
                clean_up_artists(ax, changed_artists[nr_subplot])

            # get a reference to 2d data of the current time and subplot
            data_2d = data_list[nr_subplot][frame_index]

            # manually set the levels for better contour range control
            if v_min is None:
                data_min = np.nanmin(data_2d)
            else:
                data_min = v_min
            if v_max is None:
                data_max = np.nanmax(data_2d)
            else:
                data_max = v_max

            # set the contour levels belonging to this subplot
            levels = np.linspace(data_min, data_max, number_of_contour_levels + 1, endpoint=True)

            # create the contour plot
            cs = ax.contourf(data_2d, levels=levels, cmap=cm.rainbow, zorder=0)
            cs.cmap.set_under("k")
            cs.cmap.set_over("k")
            cs.set_clim(v_min, v_max)

            # store the contours artists to the list of artists belonging to the current axis
            changed_artists[nr_subplot].append(cs)

            # set some grid lines on top of the contours
            ax.xaxis.grid(True, zorder=0, color="black", linewidth=0.5, linestyle='--')
            ax.yaxis.grid(True, zorder=0, color="black", linewidth=0.5, linestyle='--')

            # set the x and y label on the bottom row and left column respectively
            if i_row == n_rows - 1:
                ax.set_xlabel(r"Index i ")
            if j_col == 0:
                ax.set_ylabel(r"Index j")

            # set the changing time counter in the top left subplot
            if i_row == 0 and j_col == 1:
                # set a label to show the current time
                time_text = ax.text(0.6, 1.15, "{}".format("Time index : {:4d}".format(frame_index)),
                                    transform=ax.transAxes, fontdict=dict(color="black", size=14))

                # store the artist of this label in the changed artist list
                changed_artists[nr_subplot].append(time_text)

            # for the initialisation call only, set of a contour bar
            if frame_index < 0:
                # the first time we add this  (make sure to pass -1 for the frame_index
                cbar = fig.colorbar(cs, ax=ax)
                cbar.ax.set_ylabel("Random number {}".format(nr_subplot))
                ax.text(0.0, 1.02, "{}) {}".format(string.ascii_lowercase[nr_subplot],
                                                   "Random noise {}/{}".format(i_row, j_col)),
                                         transform=ax.transAxes, fontdict=dict(color="blue", size=12))

            nr_subplot += 1

    return changed_artists


def main():
    n_pixels_x = 50
    n_pixels_y = 30
    number_of_time_steps = 100
    number_of_contour_levels = 10
    delay_of_frames = 1000
    n_rows = 3  # number of subplot rows
    n_cols = 2  # number of subplot columns

    min_data_value = 0.0
    max_data_value = 1.0

    # list containing the random plot per sub plot. Insert you own data here
    data_list = list()
    for j_col in range(n_cols):
        for i_row in range(n_rows):
            data_list.append(np.random.random_sample((number_of_time_steps, n_pixels_x, n_pixels_y)))

    # set up the figure with the axis
    fig, axis = plt.subplots(nrows=n_rows, ncols=n_cols, sharex=True, sharey=True, figsize=(12,8))
    fig.subplots_adjust(wspace=0.05, left=0.08, right=0.98)

    # a list used to store the reference to the axis of each subplot with a list of artists which belong to this subplot
    # this list will be returned and will be updated every time plot which new artists
    changed_artists = list()

    # create first image by calling update_plot with frame_index = -1
    changed_artists = update_plot(-1, data_list, fig, axis, n_cols, n_rows, number_of_contour_levels,
                                                 min_data_value, max_data_value, changed_artists)

    # call the animation function. The fargs argument equals the parameter list of update_plot, except the
    # 'frame_index' parameter.
    ani = animation.FuncAnimation(fig, update_plot, frames=number_of_time_steps,
                                  fargs=(data_list, fig, axis, n_cols, n_rows, number_of_contour_levels, min_data_value,
                                         max_data_value, changed_artists),
                                  interval=delay_of_frames, blit=False, repeat=True)

    plt.show()

if __name__ == "__main__":
    main()

Answer 2

Felix Schneider is correct about the animation becoming very slow. His solution of setting ax.collections = [] removes all old (and superseded) "artist"s. A more surgical approach is to only remove the artists involved in the drawing the contours:

for c in cont.collections:
    c.remove()

which is useful in more complicated cases, in lieu of reconstructing the entire figure for each frame. This also works in Rehman Ali's example; instead of clearing the entire figure with clf() the value returned by contourf() is saved and used in the next iteration.

Here is an example code similar to Luke's from Jun 7 '13, demonstrating removing the contours only:

import pylab as plt
import numpy
import matplotlib.animation as animation
#plt.rcParams['animation.ffmpeg_path'] = r"C:\some_path\ffmpeg.exe"   # if necessary

# Generate data for plotting
Lx = Ly = 3
Nx = Ny = 11
Nt = 20
x = numpy.linspace(0, Lx, Nx)
y = numpy.linspace(0, Ly, Ny)
x,y = numpy.meshgrid(x,y)
z0 = numpy.exp(-(x-Lx/2)**2-(y-Ly/2)**2)   # 2 dimensional Gaussian

def some_data(i):   # function returns a 2D data array
    return z0 * (i/Nt)

fig = plt.figure()
ax = plt.axes(xlim=(0, Lx), ylim=(0, Ly), xlabel='x', ylabel='y')

cvals = numpy.linspace(0,1,Nt+1)      # set contour values 
cont = plt.contourf(x, y, some_data(0), cvals)    # first image on screen
plt.colorbar()

# animation function
def animate(i):
    global cont
    z = some_data(i)
    for c in cont.collections:
        c.remove()  # removes only the contours, leaves the rest intact
    cont = plt.contourf(x, y, z, cvals)
    plt.title('t = %i:  %.2f' % (i,z[5,5]))
    return cont

anim = animation.FuncAnimation(fig, animate, frames=Nt, repeat=False)
anim.save('animation.mp4', writer=animation.FFMpegWriter())

Answer 3

I used Lukes approach (from Jun 7 '13 at 8:08 ), but added

ax.collections = [] 

right before

cont = plt.contourf(x, y, z, 25).

Otherwise I experienced that creating the animation will become very slow for large frame numbers.

Answer 4

Removing the blit=0 or blit = True argument in the FuncAnimation call also helped is important!!!

Answer 5

Here is another way of doing the same thing if matplotlib.animation don't work for you. If you want to continuously update the colorbar and everything else in the figure, use plt.ion() at the very beginning to enable interactive plotting and use a combo of plt.draw() and plt.clf() to continuously update the plot.

import matplotlib.pyplot as plt
import numpy as np

plt.ion(); plt.figure(1);
for k in range(10):
    plt.clf(); plt.subplot(121);
    plt.contourf(np.random.randn(10,10)); plt.colorbar();
    plt.subplot(122,polar=True)
    plt.contourf(np.random.randn(10,10)); plt.colorbar();
    plt.draw();

Note that this works with figures containing different subplots and various types of plots (i.e. polar or cartesian)

Answer 6

This is what I got to work:

# Generate grid for plotting
x = linspace(0, Lx, Nx)
y = linspace(0, Ly, Ny)
x,y = meshgrid(x,y)

fig = plt.figure()
ax = plt.axes(xlim=(0, Lx), ylim=(0, Ly))  
plt.xlabel(r'x')
plt.ylabel(r'y')

# animation function
def animate(i): 
    z = var[i,:,0,:].T
    cont = plt.contourf(x, y, z, 25)
    if (tslice == 0):
        plt.title(r't = %1.2e' % t[i] )
    else:
        plt.title(r't = %i' % i)

    return cont  

anim = animation.FuncAnimation(fig, animate, frames=Nt)

anim.save('animation.mp4')

I found that removing the blit=0 argument in the FuncAnimation call also helped...

Answer 7

This is the line:

cont, = ax.contourf([], [], [], 500)

change to:

 x = linspace(0, 200, Nx)
 y = linspace(0, 100, Ny)
 x, y = meshgrid(x, y)
 z = n[i,:,0,:].T
 cont, = ax.contourf(x, y, z, 500)

You need to intilize with sized arrays.