I am just playing around the matplotlib's animation module. I wanted to create an animated plot of the projectile motion. I managed to create it with the following code:

import matplotlib.animation as ani
import matplotlib.pyplot as plt
import numpy as np

plt.rcParams['figure.dpi'] = 150

## Inputs
angle = 30  # deg
angleRad = angle*np.pi/180
vel = 2     # m/s
deltat = 0.005    # s (timestep)
acc = -9.81     # m/s^2 (gravitational acc)


## Calculations
# Initialisations
ux0 = vel*np.cos(angleRad)
uy0 = vel*np.sin(angleRad)
U = [ux0]
V = [uy0]

# Acceleration
ax = 0
ay = acc

# Velocity
def calcVel(time):
    Uvel = ux0   ax*time
    Vvel = uy0   ay*time
    return (Uvel, Vvel)

# Location
x0 = 0
y0 = 0
x = [x0]
y = [y0]
   
def location(time):
    locx = U[0]*time   0.5*ax*time**2
    locy = V[0]*time   0.5*ay*time**2
    return (locx, locy)
   
yn = y0
t0 = 0
n = 1
count = [0]
while yn>=0:
    t = t0   n*deltat
    xVel, yVel = calcVel(t)
    posx, posy = location(t)
    yn = posy
    if yn>=0:
        U.append(xVel)
        V.append(yVel)
        x.append(posx)
        y.append(posy)
        count.append(t)
        n  =1


def buildframes(i=int):
    line.set_data(x[:i], y[:i])
    points.set_offsets([x[i-1], y[i-1]])
    index.set_text('time = %.3f s' % count[i])
    plt.pause(0.001)


## Plotting
fig = plt.figure()
axes = fig.add_subplot(1,1,1)
line, = axes.plot([], []) # , 'bo', lw=2, scalex = True, scaley = True
points = plt.scatter([], [], c='red', s=15)
axes.set_ylim(min(y), np.true_divide(np.ceil(max(y) * 10**2), 10**2))
axes.set_xlim(min(x), np.true_divide(np.ceil(max(x) * 10**2), 10**2))
plt.ylabel('Height')
plt.xlabel('Distance')
index = axes.annotate('', xy=(0.75, 0.95), xycoords='axes fraction')

animator = ani.FuncAnimation(fig, buildframes, interval = 500)
#animator.save("filename")
plt.show()
animator.save('animation.gif')

The title question could easily have been a generic "how do I improve this code" but those aren't encouraged here. While the code produces the animation, there are several finer issues with this code. I have listed them in decreasing order of importance:

1. The "list index out of range" error: this is due to the the way the line and the scattered points are plotted. For lines there should be two points already. How do I avoid this and have a cleaner animator function? (It still plots - just that this error remains and it is annoying me.)
2. The line tip is always over the point (notice the tiny blue line segment over red ball). I have tried to change the sequence but that didn't help to have the red point always be above the line. How do I do it?
3. When run in Spyder it seems to animate it twice in the plotting window (with differing speeds). I'm not sure why. The gif created is always with the later speed (slower). It might be the case that when saving it overwrites the frames with the slower speed and hence the observation. Which would mean the animator function is being called twice? How do I fix this? This doesn't happen if I run it on Linux python environment.
4. Generic: Is there any smarter way to do the same animation? e.g. I feel like the code is too long, is there any way to make it compact and maybe more efficient?

CodePudding user response：

1. Set the number of frames in your animation to be equal to the length of your variables that you are plotting. For example: ani.FuncAnimation(fig, buildframes, frames=len(x), interval=500)
2. You can use zorder to force things to be displayed with an appropriate order. Here I used ``zorder=1for the line andzorder=2` for the scatter, so the scatter is going to be visualized on top of the line.
3. I can't comment on that because I don't use Spyder.
4. When using animations, don't place plt.pause inside your animation function. Set an appropriate interval when calling ani.FuncAnimation. The animation function is good as it is. In the code below I have done the same computation by using Numpy's "vectorization", so no need to use loops (this saves some lines of code).

import matplotlib.animation as ani
import matplotlib.pyplot as plt
import numpy as np

plt.rcParams['figure.dpi'] = 150

## Inputs
angle = 30  # deg
angleRad = angle*np.pi/180
vel = 2     # m/s
deltat = 0.005    # s (timestep)
acc = -9.81     # m/s^2 (gravitational acc)

# Initialisations
x0 = 0
y0 = 0
ux0 = vel*np.cos(angleRad)
uy0 = vel*np.sin(angleRad)
ax = 0
ay = acc

# https://stackoverflow.com/a/64083631/2329968
def linspace(start, stop, step=1.):
    """
    Like np.linspace but uses step instead of num
    This is inclusive to stop, so if start=1, stop=3, step=0.5
    Output is: array([1., 1.5, 2., 2.5, 3.])
    """
    return np.linspace(start, stop, int((stop - start) / step   1))

t0 = 0
tf = 1 # arbitrary value
# discretize the time with a given deltat
time = linspace(t0, tf, deltat)
# compute x-y locations
x = ux0*time   0.5*ax*time**2
y = uy0*time   0.5*ay*time**2
# now there is a change that y < 0 for some x.
# We need to only keep the locations and times where y >= 0.
idx = y >= 0
x = x[idx]
y = y[idx]
time = time[idx]
# compute velocities
U = ux0   ax*time
V = uy0   ay*time

def buildframes(i):
    line.set_data(x[:i 1], y[:i 1])
    points.set_offsets([x[i], y[i]])
    index.set_text('time = %.3f s' % time[i])


## Plotting
fig = plt.figure()
axes = fig.add_subplot(1,1,1)
line, = axes.plot([], [], zorder=1) # , 'bo', lw=2, scalex = True, scaley = True
points = plt.scatter([], [], c='red', s=15, zorder=2)
axes.set_ylim(min(y), np.true_divide(np.ceil(max(y) * 10**2), 10**2))
axes.set_xlim(min(x), np.true_divide(np.ceil(max(x) * 10**2), 10**2))
plt.ylabel('Height')
plt.xlabel('Distance')
index = axes.annotate('', xy=(0.75, 0.95), xycoords='axes fraction')

animator = ani.FuncAnimation(fig, buildframes, frames=len(x), interval=500)
# plt.show()
animator.save('animation.gif')