python - Updating z data on a surface_plot in Matplotlib animation

Question

I am looking to create an animation in a surface plot. The animation has fixed x and y data (1 to 64 in each dimension), and reads through an np array for the z information. An outline of the code is like so:

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

def update_plot(frame_number, zarray, plot):
    #plot.set_3d_properties(zarray[:,:,frame_number])
    ax.collections.clear()
    plot = ax.plot_surface(x, y, zarray[:,:,frame_number], color='0.75')

fig = plt.figure()
ax = plt.add_subplot(111, projection='3d')

N = 64
x = np.arange(N+1)
y = np.arange(N+1)
x, y = np.meshgrid(x, y)
zarray = np.zeros((N+1, N+1, nmax+1))

for i in range(nmax):
  #Generate the data in array z
  #store data into zarray
  #zarray[:,:,i] = np.copy(z)

plot = ax.plot_surface(x, y, zarray[:,:,0], color='0.75')

animate = animation.FuncAnimation(fig, update_plot, 25, fargs=(zarray, plot))
plt.show()

So the code generates the z data and updates the plot in FuncAnimation. This is very slow however, I suspect it is due to the plot being redrawn every loop.

I tried the function

ax.set_3d_properties(zarray[:,:,frame_number])

but it comes up with an error

AttributeError: 'Axes3DSubplot' object has no attribute 'set_3d_properties'

How can I update the data in only the z direction without redrawing the whole plot? (Or otherwise increase the framerate of the graphing procedure)

Answer 1

There is a lot going on under the surface when calling plot_surface. You would need to replicate all of it when trying to set new data to the Poly3DCollection.

This might actually be possible and there might also be a way to do that slightly more efficient than the matplotlib code does it. The idea would then be to calculate all the vertices from the gridpoints and directly supply them to Poly3DCollection._vec.

However, the speed of the animation is mainly determined by the time it takes to perform the 3D->2D projection and the time to draw the actual plot. Hence the above will not help much, when it comes to drawing speed.

At the end, you might simply stick to the current way of animating the surface, which is to remove the previous plot and plot a new one. Using less points on the surface will significantly increase speed though.

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D 
import matplotlib.animation as animation

def update_plot(frame_number, zarray, plot):
    plot[0].remove()
    plot[0] = ax.plot_surface(x, y, zarray[:,:,frame_number], cmap="magma")

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')

N = 14
nmax=20
x = np.linspace(-4,4,N+1)
x, y = np.meshgrid(x, x)
zarray = np.zeros((N+1, N+1, nmax))

f = lambda x,y,sig : 1/np.sqrt(sig)*np.exp(-(x**2+y**2)/sig**2)

for i in range(nmax):
    zarray[:,:,i] = f(x,y,1.5+np.sin(i*2*np.pi/nmax))

plot = [ax.plot_surface(x, y, zarray[:,:,0], color='0.75', rstride=1, cstride=1)]
ax.set_zlim(0,1.5)
animate = animation.FuncAnimation(fig, update_plot, nmax, fargs=(zarray, plot))
plt.show()

Note that the speed of the animation itself is determined by the interval argument to FuncAnimation. In the above it is not specified and hence the default of 200 milliseconds. Depending on the data, you can still decrease this value before running into issues of lagging frames, e.g. try 40 milliseconds and adapt it depending on your needs.

animate = animation.FuncAnimation(fig, update_plot, ..., interval=40,  ...)