Python: How can I substitute this for loop with a while loop? (No negative values)

Question

When defining this trajectory of a cannonball function, I would like for the loop to stop at negative values of y. IE the cannonball should not continue moving once hitting the ground.

I tried using while >= 0, but y points is a list and I don't know how to do that. Any help?

def trajectory(v_0=1, mass=1, theta=np.pi/2, t_0=0, t_f=2, n=100):
"""
:param v_0: initial velocity, in meters per second
:param mass: Mass of the object in kg, set to 1kg as default
:param theta: Angle of the trajectory's shot, set to pi/2 as default
:param t_0: initial time, in seconds, default as 0
:param t_f: final time, in seconds, default as 2
:param n: Number of points
:return: array with x and y coordinates of the trajectory
"""
h = (t_f - t_0)/n
t_points = np.arange(t_0, t_f, h)
x_points = []
y_points = []
r = np.array([0, v_0 * np.cos(theta), 0, v_0 * np.sin(theta)], float)

for t in t_points:
    x_points.append(r[0])
    y_points.append(r[2])
    k1 = h * F(r, t, mass)
    k2 = h * F(r + 0.5 * k1, t + 0.5 * h, mass)
    k3 = h * F(r + 0.5 * k2, t + 0.5 * h, mass)
    k4 = h * F(r + k3, t + h, mass)
    r += (k1 + 2 * k2 + 2 * k3 + k4) / 6
return np.array(x_points, float), np.array(y_points, float)

when graphing the trajectory, I get a graph that contains negative values of y, which I would like to prevent from being calculated in the first place in order to not affect the code's performance.

Answer 1

You don't need to convert to a while loop. The simplest way I can think of is to put an exit condition in the for loop:

for t in t_points:
    if r[2] < 0:
        break
    ...

This will exit the loop when the y position is less than zero. You might consider making the condition r[2] < 0 and r[3] < 0 so that if you have one starting below the ground, it makes its way up before coming down and colliding with the ground.

If you really have your heart set on a while loop, you can make an iterator variable, and then use it to iterate through t_points :

iterator_variable = 0
while r[2] < 0 and iterator_variable < len(t_points):
    t = t_points[iterator_variable]
    ...
return np.array(x_points, float), np.array(y_points, float)

While I don't know what your function F() does, I think it might be easier to go without defining n , h , or t_points . You would start at the starting point, and calculate each next point until you hit the ground. This strategy lends itself nicely to a while loop.

while r[2] > 0:
    //calculate next position and velocity
    //add that position to the list of x_points and y_points

Instead of n as an input to your function which dictates the number of points you calculate, you could have it be a point density measure, or a maximum number of points.