How can I animate object A thrown upwards from an given height then free-falling with matplotlib?

Using Python 3.12, Matplotlib 3.9.2,

From height h=20(m), throw object A upwards with a given velocity v and drop object B to let it free-falling. Air resistance is neglectable. Calculate v so that object A falls to the ground dt =2 seconds(s) after object B and animate. g = 9.81 (m/s2)

I've been trying to animate said problem with matplotlib using 'set_data', but the results are unsatisfying for previous positions of A is kept which made the animation overlap itself after a while. I want the previous data to be removed, but I couldn't find any functions to help it.

Here is the code that I've used:

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

#Input Variables
h, g, dt= 20, 9.81, 2

#Calculated Variables
tmp = np.sqrt(2*h/g) + dt
v = 0.5*g*tmp - h/tmp 

#Arrays 
t = np.linspace(0, tmp , 100)
hA = h + v*t - 0.5*g*t**2
hB = h - 0.5*g*t**2

#Plot
fig, ax = plt.subplots()
plt.axis([0, 10, 0, 100])
ax.set_title("Physic 1")
plt.grid()

#Define object A and object B
A, = ax.plot([],[],"o",markersize=4,color="red")
B, = ax.plot([],[],"o",markersize=4,color="blue")

#Animation
def animate(frame):
    A.set_data([2],hA[:frame])
    B.set_data([6],hB[:frame])
    return A,B

ani=FuncAnimation(fig, animate, frames=len(t)+2, interval=25, repeat =False)
plt.show()

Solution

When you use hA[:frame] you're passing all the points up to index frame. Since you only want a single index, you should do hA[frame]. You will need to wrap that result in a list so that y is a sequence (otherwise you get the error mentioned in the comments). In the end, you'll have this:

def animate(frame):
    A.set_data([2], [hA[frame]])
    B.set_data([6], [hB[frame]])
    return A, B

This can also be done using scatter plots by changing the corresponding lines to this:

# the zorder argument puts the scatter points on top of the grid markings
A = ax.scatter([], [], s=16, color="red", zorder=2)
B = ax.scatter([], [], s=16, color="blue", zorder=2)

def animate(frame):
    A.set_offsets([2, hA[frame]])
    B.set_offsets([6, hB[frame]])
    return A, B

Lastly, your hA and hB arrays are the same length as t, so the number of frames should be len(t), not len(t)+2.