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

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...