Animate 3D surface over an initial 3D plot with matplotlib

Question

I am trying to animate a 3D surface over an initial 3D plot. But am struggling with keeping the initial 3D plot on the animation. I have to call the clafunction to be able to plot the new surface and this will delete my initial state.

As an example, I have the animation of a surface while keeping the x,y,z axis. The end goal would be to plot the initial state once, and animate the surface.

Is it possible to create plot layers? Send the initial state to one layer and the animation to another?

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

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

def initialState(ax):
    s=1.15 
    x_axis = np.array([s*5, 0, 0])
    y_axis = np.array([0, s*5, 0])
    z_axis = np.array([0, 0, s*5])
    
    ax.quiver(0,0,0,*x_axis,color = 'k',  alpha =1, arrow_length_ratio=0.05)
    ax.text(*(s*x_axis),'x',fontsize=10) 
    ax.quiver(0,0,0,*y_axis, color = 'k', alpha =1, arrow_length_ratio=0.05)
    ax.text(*(s*y_axis),'y',fontsize=10) 
    ax.quiver(0,0,0,*z_axis,color = 'k',  alpha =1, arrow_length_ratio=0.05)
    ax.text(*(s*z_axis),'z',fontsize=10)
    ax.set_xlim3d(-10, 10)
    ax.set_ylim3d(-5, 5)
    ax.set_zlim3d(-5, 5)

def plot(i):
    X = np.arange(-5+i, 5+i, 0.25)
    Y = np.arange(-5, 5, 0.25)
    X, Y = np.meshgrid(X, Y)
    R = np.sqrt(X**2 + Y**2)
    Z = np.sin(R)
    
    return [X,Y,Z]


def anime(i, ax, stepFactor):
    ax.cla()
    points = plot(i*stepFactor)
    ax.plot_surface(*points, alpha=0.5)
    ax.set_xlim3d(-10, 10)
    ax.set_ylim3d(-5, 5)
    ax.set_zlim3d(-5, 5)

        
animation = FuncAnimation(
                              fig,
                              anime,
                              init_func=initialState(ax),
                              frames=range(100),
                              fargs=(ax, 0.1)
                          )

Thanks for the help!

Answer 1

This is how I would do it: add two identical surfaces (at t=0) at the end of the init method. This will add two Poly3DCollection in the last two positions of ax.collections. Then, ax.collections[-2] represents the surface at t=0, and ax.collections[-1] will be removed and readded in the animation function: it will represent the surface at t=i.

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

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

def initialState(ax):
    s=1.15 
    x_axis = np.array([s*5, 0, 0])
    y_axis = np.array([0, s*5, 0])
    z_axis = np.array([0, 0, s*5])
    
    ax.quiver(0,0,0,*x_axis,color = 'k',  alpha =1, arrow_length_ratio=0.05)
    ax.text(*(s*x_axis),'x',fontsize=10) 
    ax.quiver(0,0,0,*y_axis, color = 'k', alpha =1, arrow_length_ratio=0.05)
    ax.text(*(s*y_axis),'y',fontsize=10) 
    ax.quiver(0,0,0,*z_axis,color = 'k',  alpha =1, arrow_length_ratio=0.05)
    ax.text(*(s*z_axis),'z',fontsize=10)
    # add the initial surface
    surf0 = ax.plot_surface(*plot(0), color="tab:blue")
    # add another copy of the initial surface: this will be removed (and later
    # added again) by the anim function.
    surf1 = ax.plot_surface(*plot(0), color="tab:blue")
    ax.set_xlim3d(-10, 10)
    ax.set_ylim3d(-5, 5)
    ax.set_zlim3d(-5, 5)

def plot(i):
    X = np.arange(-5+i, 5+i, 0.25)
    Y = np.arange(-5, 5, 0.25)
    X, Y = np.meshgrid(X, Y)
    R = np.sqrt(X**2 + Y**2)
    Z = np.sin(R)
    return [X,Y,Z]


def anime(i, stepFactor):
    points = plot(i*stepFactor)
    # remove the last surface added in the previous iteration
    # (or in the init function)
    ax.collections[-1].remove()
    anim_surf = ax.plot_surface(*points, color="tab:blue", alpha=0.5)
  
animation = FuncAnimation(
    fig,
    anime,
    init_func=initialState(ax),
    frames=range(100),
    fargs=(0.1, )
)
plt.show()

EDIT: Second solution to accommodate comment: you can also add the surface in the global namespace, like i did below with surf_to_update. Then, inside anime you would have to use global surf_to_update. It is not the nicest solution, but it is definitely easier than keeping track of all the things you are adding/updating.

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

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

surf_to_update = ax.plot_surface(*plot(0), color="tab:blue")

def initialState(ax):
    s=1.15 
    x_axis = np.array([s*5, 0, 0])
    y_axis = np.array([0, s*5, 0])
    z_axis = np.array([0, 0, s*5])
    
    ax.quiver(0,0,0,*x_axis,color = 'k',  alpha =1, arrow_length_ratio=0.05)
    ax.text(*(s*x_axis),'x',fontsize=10) 
    ax.quiver(0,0,0,*y_axis, color = 'k', alpha =1, arrow_length_ratio=0.05)
    ax.text(*(s*y_axis),'y',fontsize=10) 
    ax.quiver(0,0,0,*z_axis,color = 'k',  alpha =1, arrow_length_ratio=0.05)
    ax.text(*(s*z_axis),'z',fontsize=10)
    # add the initial surface
    surf0 = ax.plot_surface(*plot(0), color="tab:blue")
    
    ax.set_xlim3d(-10, 10)
    ax.set_ylim3d(-5, 5)
    ax.set_zlim3d(-5, 5)

def plot(i):
    X = np.arange(-5+i, 5+i, 0.25)
    Y = np.arange(-5, 5, 0.25)
    X, Y = np.meshgrid(X, Y)
    R = np.sqrt(X**2 + Y**2)
    Z = np.sin(R)
    return [X,Y,Z]


def anime(i, stepFactor):
    global surf_to_update
    points = plot(i*stepFactor)
    # remove the last surface added in the previous iteration
    # (or in the init function)
    surf_to_update.remove()
    surf_to_update = ax.plot_surface(*points, color="tab:blue", alpha=0.5)
  
animation = FuncAnimation(
    fig,
    anime,
    init_func=initialState(ax),
    frames=range(100),
    fargs=(0.1, )
)
plt.show()