torus.py

import matplotlib.pyplot as plt
import numpy as np
tau = 2*np.pi

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

# map from unit square to surface of a torus
def torus(u, v):
    X = (2 + np.cos(tau*v))*np.cos(tau*u)
    Y = (2 + np.cos(tau*v))*np.sin(tau*u)
    Z = np.sin(tau*v)
    return X,Y,Z


# u, v = np.mgrid[0:1:100j, 0:1:100j]
# X, Y, Z = torus(u, v)
# ax.plot_surface(X,Y,Z)

plt.ion()

ln, = ax.plot([],[],[])
ax.set_xlim(-3, 3)
ax.set_ylim(-3, 3)
ax.set_zlim(-2, 2)
for num in range(0, 100, 1):
    m = 1/3
    xs = np.linspace(0, num, num=num*10)
    ys = m * xs

    ln.set_data_3d(*torus(xs, ys))

    plt.pause(0.01)
    plt.savefig(f'frames/frame-{num}.png')
    plt.show()

# to animate run
# ffmpeg -hide_banner -y -framerate 60 -i frames/frame-%1d.png out.gif