Surface animation and saving with matplotlib

Question

I tried to do a very simple animation and save it with matplotlib, but without success. I want for example to see something oscillating: here the best I could do

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

#Define x,y vectors and meshgrid with function u on it

x = np.arange(0,10,0.1)
y = np.arange(0,10,0.1)
X,Y = np.meshgrid(x,y)
u = np.sin(X + Y)

#Create a figure and an axis object for the surface
#(which by the way is not initialized, because I don't know were to)

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

#Define a kind-of animation function, imitating what I saw many times

def animate(n):
    global u
    for n in (1,20):
        u = np.sin((X + Y)*n)
    return fig,     

#I'm missing many commands, I'm just putting what I know

anim = animation.FuncAnimation(fig,animate)
anim.save('A.mp4',fps=10)

I read tons of online documentation, including the official one, but I cannot find a clear example in which a surface evolve with time. I also find very difficult to grasp the logic behind the construction of plots with matplotlib (I read about figures, objects associated, axes, other objects... I get confused), but I'm also self-learning it for an exam and I'm not really a tech guy, so maybe that's why I'm having so much troubles; fundamentally, if you answer and can spend two more minutes describing in some more detail what you're doing, you'll make me very happy. Many thanks for any help.

Answer 1

First of all, you have to set up the mathematical domain of the surface, which is constant for each animation frame, therefore you can do it at the beginning:

x = np.arange(0,10,0.1)
y = np.arange(0,10,0.1)
X,Y = np.meshgrid(x,y)

Then you can initialize the figure:

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

Here come the animation stuff, you have to define a function which describes what changes between one animation frame and the following one. First of all you have to clear from the plot what you have already drawn in the previous frame, so you can use ax.cla().
Then it comes the math part: you have to define your function and how it changes with time. Pay attention that the animate function takes a parameter, n in you case, which increase by 1 in each frame, you can use it to describe what changes between consecutive frames. For example if I write u = np.sin((X + Y) + n) then the synusoidal surface will increase its phase in each iteration, so it will "move forward"; you can re-define this equation in order to achieve your target by properly use the n parameter.
Then it's time to draw the surface with ax.plot_surface(X, Y, u). It is not mandatory, but I suggest to fix the axis limits, in order to improve the animation, with ax.set_zlim(-2, 2).
In this way the animate function is defined as:

def animate(n):
    ax.cla()

    u = np.sin((X + Y) + n)

    ax.plot_surface(X, Y, u)
    ax.set_zlim(-2, 2)

    return fig,

Finally you have to set up the FuncAnimation object with the parameter you want:

anim = FuncAnimation(fig = fig, func = animate, frames = 10, interval = 1, repeat = False)

Complete code

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


x = np.arange(0,10,0.1)
y = np.arange(0,10,0.1)
X,Y = np.meshgrid(x,y)


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


def animate(n):
    ax.cla()

    u = np.sin((X + Y) + n)

    ax.plot_surface(X, Y, u)
    ax.set_zlim(-2, 2)

    return fig,


anim = FuncAnimation(fig = fig, func = animate, frames = 10, interval = 1, repeat = False)
anim.save('A.mp4',fps=10)

---

Another example changing the surface formula, if you use:

u = np.sin(5/2/np.pi*n)*np.sin((X + Y))

within animate function, then the surface phase will be kept constant, but the synusoidal surface amplitude will change, following itself a synusoidal function, so you will get:

Answer 2

Here is another example. I borrowed it from the gallery of the JavaScript library vis3d, and I turned it into Python thanks to @Zephyr's answer.

from math import pi, cos, sin
import numpy as np
import matplotlib.colors as mcolors
from matplotlib import cm
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

x_ = np.linspace(0, 314, 50)
y_ = np.linspace(0, 314, 50)
X, Y = np.meshgrid(x_, y_)

t_ = np.linspace(0, 2*pi, 90)

def f(x, y):
  return np.sin(x/50) * np.cos(y/50) * 50 + 50

def g(x, y, t):
  return f(x*cos(t) - y*sin(t), x*sin(t) + y*cos(t))

fig = plt.figure()
ax = fig.add_subplot(projection = "3d")

def animate(n):
    ax.cla()
    Z = g(X, Y, t_[n])
    colorfunction = (X**2+Y**2+Z**2)
    norm = mcolors.Normalize(colorfunction.min(), colorfunction.max())
    ax.plot_surface(
      X, Y, Z, rstride = 1, cstride = 1, facecolors=cm.jet(norm(colorfunction))
    )
    ax.set_zlim(0, 100)
    return fig


anim = FuncAnimation(
  fig = fig, func = animate, frames = len(t_), interval = 1, repeat = False
)
anim.save("Anim.mp4", fps = 10)

(I converted the mp4 to gif thanks to the website ezgif.com but the gif was too big for SO, so I compressed it with gifsicle and there's a loss of quality.)

---

EDIT

Better gif: