Beginner-Make a frame based video in python from figures

Question

I plot figures in a for loop which is a loop for my time, basically at each time step I plot a surf out of my data as below:

for time_step in range(0,nt):
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    x = np.arange(xmin, xmax+dx, dx)
    z = np.arange(zmin, zmax+dz, dz)
    X, Z = np.meshgrid(x, z)
    ax.plot_surface(X, Z, w1[time_step])
    plt.show()

Suppose that w1[time_step] changes in the loop and is sth different at each time step, all other assumptions you can have. I plot but don't know only how to make them into a video.

I have done it matlab, but I want to do sth similar in Python

Answer 1

Matplotlib as some animation features you might want to use. Check the following recipe (that I collected from here):

    from mpl_toolkits.mplot3d import axes3d
    import matplotlib.pyplot as plt
    import numpy as np
    import time


    def generate(X, Y, phi):
        R = 1 - np.sqrt(X**2 + Y**2)
        return np.cos(2 * np.pi * X + phi) * R

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

    xs = np.linspace(-1, 1, 50)
    ys = np.linspace(-1, 1, 50)
    X, Y = np.meshgrid(xs, ys)
    Z = generate(X, Y, 0.0)

    wframe = None
    tstart = time.time()
    for phi in np.linspace(0, 360 / 2 / np.pi, 100):

        oldcol = wframe

        Z = generate(X, Y, phi)
        wframe = ax.plot_wireframe(X, Y, Z, rstride=2, cstride=2)

        # Remove old line collection before drawing
        if oldcol is not None:
            ax.collections.remove(oldcol)

        plt.pause(.001)

    print('FPS: %f' % (100 / (time.time() - tstart)))

Just replace the wireframe plot for whatever you want (and also use your data obviously) and you should have what you are looking for.