Animate time-dependent 3D-function using matplotlib in Python

Question

Hope someone can help me. My goal is to animate a function U(x,z,t) and in particular the time evolution of the 3D-function U(x,z). Firstly I wrote some code that discretize the time stepping and saved each "snapshot" in an array U[t] where t is an integer in {0,1,...,Nt} for some Nt number of timesteps. U[0] is the initial function and is given. I want to remark that U[t] (t in {0,1,...,Nt} are all matrices containg numbers and are already suited to be plotted. I now want to animate a plot such that:

1. The first image displayed in the window is U[0], i.e. the initial function with an appropriate label.
2. The window then refreshes and displays U1 with an appropriate label.
3. It then repeats step (2) up to timestep Nt.

My first approach was to use only a "for cycle". Unfortunately, instead of a refreshing image, I created a sequence of different figures. The code is:

for i in range(0,Nt+1):
        fig = plt.figure()
        ax = plt.axes(projection='3d')
        ax.plot_surface(gridx, gridz, U[i],cmap='viridis', edgecolor='none')
        #gridx,gridz are created via np.meshgrid(...)
        ax.set_title("Function U at time-step " + str(i)+ ". Time "+ str(i*dt)+ "." )
        ax.set_xlabel('X')
        ax.set_ylabel('Z')
        ax.set_zlabel('u(x,z)')
plt.show()

Now I tried searching online and found many different examples using matplotlib.animation with FuncAnimation or just Animation. The problem is that I do not understand how to adapt the examples to my cases. Could someone help me in this case?

Update: For the case where the function can also be a variable. I wrote the following code:

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.animation import FuncAnimation
def Animate(gridx,gridz,Nt,Name, func,dt):
    
    fig = plt.figure()
    ax = fig.gca(projection= "3d")
    ax.set_xlabel('X')
    ax.set_ylabel('Z')
    ax.set_zlabel(Name)

    def update(frame,fig):
        if len(fig.axes[0].collections) != 0:
            fig.axes[0].collections = []
            surf = fig.axes[0].plot_surface(gridx, gridz, func[frame], cmap= "viridis")
            ax.set_title("Function "+ Name + " at time-step " + str(frame)+ ". Time "+ str(frame*dt)+ ".", y = 1  )
        else:
            surf = fig.axes[0].plot_surface(gridx, gridz, func[frame], cmap= "viridis")
            ax.set_title("Function " + Name + " at time-step " + str(frame)+ ". Time "+ str(frame*dt)+ "." )
        fig.canvas.draw()
        return surf,
    ani =FuncAnimation(fig,update,fargs=[fig],frames = Nt+1, blit = True)
    ani.save("Gif for the function " + Name + ".gif")
    return

This code can be called with different functions to ease the treatment in case you need to visualize more functions.

The result is: Example of Gif for a function P

Answer 1

I used following code to do animation by adding random z-direction displacement to surface, hope it can solve your issue. Pls read carefully the animation tutorial, and a 3d line example 3D animation lines. My solution might not be efficeint as I redraw the surface, but it should be easy to understand.

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



fig = plt.figure()
ax = fig.gca(projection='3d')


def update(frame, fig):
    X = np.arange(-5, 5, 0.25)
    Y = np.arange(-5, 5, 0.25)
    X, Y = np.meshgrid(X, Y)
    R = np.sqrt(X**2 + Y**2)
    # add random shift in z directions
    Z = np.sin(R)+np.random.random_sample()
    if len(fig.axes[0].collections) != 0:
        fig.axes[0].collections = []
        surf = fig.axes[0].plot_surface(X, Y, Z, cmap=cm.coolwarm, linewidth=0, antialiased=False)
    else:
        surf = fig.axes[0].plot_surface(X, Y, Z, cmap=cm.coolwarm, linewidth=0, antialiased=False)
    ax.set_zlim(-1.5, 1.5)

    fig.canvas.draw()
    return surf,
    
ani = FuncAnimation(fig, update, fargs=[fig], frames=5, blit=True)