How to create a 3D animation of a surface defined by 2 spatial coordinates using matplotlib/Mayavi?

Question

Here is a numerical code of a 1D diffusion equation using for discretizing a finite difference scheme. The velocity is obtained for each time step and I would like to animate this solution in order to visualize the evolution of velocity with respect to time under diffusion. Any help would be appreciated. Thank you!

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D ##library for 3d projection plots
from matplotlib import cm ##cm = "colormap" for changing the 3d plot color palette

###variable declarations
nx = 31
ny = 31
nt = 17
nu=.05
dx = 2.0/(nx-1)
dy = 2.0/(ny-1)
sigma = .25
dt = sigma*dx*dy/nu

x = np.linspace(0,2,nx)
y = np.linspace(0,2,ny)

u = np.ones((ny,nx)) ##create a 1xn vector of 1's
un = np.ones((ny,nx)) ##

###Assign initial conditions

u[.5/dy:1/dy+1,.5/dx:1/dx+1]=2 ##set hat function I.C. : u(.5<=x<=1 && .5<=y<=1 ) is 2

fig = plt.figure()
ax = fig.gca(projection='3d')
X,Y = np.meshgrid(x,y)
surf = ax.plot_surface(X,Y,u[:], rstride=1, cstride=1, cmap=cm.coolwarm,
    linewidth=0, antialiased=False)
plt.show()
ax.set_xlim(0,2)
ax.set_ylim(0,2)
ax.set_zlim(1,2.5)
#ax.zaxis.set_major_locator(LinearLocator(5))

###Run through nt timesteps

u[.5/dy:1/dy+1,.5/dx:1/dx+1]=2

for n in range(nt+1): 
    un[:] = u[:]
    u[1:-1,1:-1]=un[1:-1,1:-1]+nu*dt/dx**2*(un[2:,1:-1]-2*un[1:-1,1:-1]+un[0:-2,1:-1])+nu*dt/dy**2*   (un[1:-1,2:]-2*un[1:-1,1:-1]+un[1:-1,0:-2])    

    u[0,:]=1
    u[-1,:]=1

    u[:,0]=1
    u[:,-1]=1

fig = plt.figure()
ax = fig.gca(projection='3d')
surf = ax.plot_surface(X,Y,u[:], rstride=1, cstride=1, cmap=cm.coolwarm,
linewidth=0, antialiased=True)
ax.set_zlim(1,2.5)
plt.show()

Answer 1

Here's one approach using Mayavi - read the comments for some explanation:

import numpy as np
import time

# import mayavi's mlab API for scripting
from mayavi import mlab

###variable declarations
nx = 31
ny = 31
nt = 17
nu=.05
dx = 2.0/(nx-1)
dy = 2.0/(ny-1)
sigma = .25
dt = sigma*dx*dy/nu

x = np.linspace(0,2,nx)
y = np.linspace(0,2,ny)

u = np.ones((ny,nx)) ##create a 1xn vector of 1's
un = np.ones((ny,nx)) ##

###Assign initial conditions

u[.5/dy:1/dy+1,.5/dx:1/dx+1]=2 ##set hat function I.C. : u(.5<=x<=1 && .5<=y<=1 ) is 2
X,Y = np.meshgrid(x,y)

###Run through nt timesteps

u[.5/dy:1/dy+1,.5/dx:1/dx+1]=2

# create a surface from grid-shaped data
surf = mlab.mesh(X,Y,u[:])

t = time.time()
max_framerate = 10
for n in range(nt+1): 
    un[:] = u[:]
    u[1:-1,1:-1]=un[1:-1,1:-1]+nu*dt/dx**2*(un[2:,1:-1]-2*un[1:-1,1:-1]+un[0:-2,1:-1])+nu*dt/dy**2*   (un[1:-1,2:]-2*un[1:-1,1:-1]+un[1:-1,0:-2])    

    u[0,:]=1
    u[-1,:]=1

    u[:,0]=1
    u[:,-1]=1

    # the mlab_source attribute of surf represents the data we're plotting.
    # it has x, y and z attributes as you'd expect. here we only need to
    # update the z attribute
    surf.mlab_source.z = u

    # there's no need to call any equivalent to matplotlib's draw() or show() 
    # functions - another draw event gets triggered automatically whenever
    # surf's data source gets modified

    # put a pause in here to control the maximum framerate
    while time.time() - t < (1./max_framerate):
        pass
    t = time.time()

You can find out more about animating data in Mayavi here.