Как изменить программу для численного анализа траектории снаряда под действием силы тяжести и силы перетаскивания для использования функций в определенном формате

Question

Привет, это мой код, который я использовал и изменил с веб-сайта проблем моего дополнительного класса, как описано в заголовке, я получил его, чтобы дать желаемые результаты, но я должен изменить его, чтобы следовать определенному конкретному использованиюФункции вместо того, чтобы быть в одном большом блоке.

import math
import matplotlib.pyplot as plt

M = 1.0                                                         # Mass of projectile in kg
g = 9.81                                                        # Acceleration due to gravity (m/s^2)
V = float(input("Enter value for initial velocity"))            # Initial velocity in m/s
ang = float(input("Enter value for angle of initial velocity")) # Angle of initial velocity in degrees
Cd = float(input("Drag co-effiecent here"))                     # Drag coefficient
dt = float(input("Time step value here"))                       #Time step in seconds

t = [0]                                                         # list to keep track of time
vx = [V*math.cos(ang/180*math.pi)]                              # list for velocity x and y components
vy = [V*math.sin(ang/180*math.pi)]
x = [0]                                                         # list for x and y position
y = [0]


drag = Cd*V**2                                                  # Calculate Drag force based of Cd value 

                                                            # Acceleration components
ax = [-(drag*math.cos(ang/180*math.pi))/M ]          
ay = [-g-(drag*math.sin(ang/180*math.pi)/M) ]

counter = 0
while (y[counter] >= 0):                                        # Check that the last value of y is     >= 0 so that the object reaches ground again
    t.append(t[counter]+dt)                                     # increment by dt and add to the list of time 
                                                            # Update velocity
    vx.append(vx[counter]+dt*ax[counter])  
    vy.append(vy[counter]+dt*ay[counter])
                                                            # Update position
    x.append(x[counter]+dt*vx[counter])    
    y.append(y[counter]+dt*vy[counter])
    vel = math.sqrt(vx[counter+1]**2 + vy[counter+1]**2)        # magnitude of velocity
    drag = Cd*vel**2                                            # drag force based on requested     input Cd
    ax.append(-(drag*math.cos(ang/180*math.pi))/M)     
    ay.append(-g-(drag*math.sin(ang/180*math.pi)/M))


    counter = (counter +1)


plt.plot(x,y,'ro')                                              #Plot graph
plt.ylabel("y displacement(m)")
plt.xlabel("x displacement(m)")
plt.show()

Я должен изменить его на формат, который выглядит следующим образом:

def set_initial(v_initial, theta):
"""   Set inital conditions
"""
return x,y,vx,vy

def acceleration(vx, vy, beta):
""" Calculate the acceleration 
"""
return ax, ay

def step_forward(x, y, vx, vy, beta, delta_t):
"""   Do a forward step
""""
return x, y, vx, vy

def main():
"""    Main program to read in from terminal, do itteration, and plot out graph
"""