Search code examples
pythonscipyodeodeint

Stop Integrating when Output Reaches 0 in scipy.integrate.odeint


I've written a code which looks at projectile motion of an object with drag. I'm using odeint from scipy to do the forward Euler method. The integration runs until a time limit is reached. I would like to stop integrating either when this limit is reached or when the output for ry = 0 (i.e. the projectile has landed).

def f(y, t):
    # takes vector y(t) and time t and returns function y_dot = F(t,y) as a vector
    # y(t) = [vx(t), vy(t), rx(t), ry(t)]

    # magnitude of velocity calculated
    v = np.sqrt(y[0] ** 2 + y[1] **2)

    # define new drag constant k
    k = cd * rho * v * A / (2 * m)

    return [-k * y[0], -k * y[1] - g, y[0], y[1]] 


def solve_f(v0, ang, h0, tstep, tmax=1000):
    # uses the forward Euler time integration method to solve the ODE using f function
    # create vector for time steps based on args
    t = np.linspace(0, tmax, tmax / tstep)

    # convert angle from degrees to radians
    ang = ang * np.pi / 180

    return odeint(f, [v0 * np.cos(ang), v0 * np.sin(ang), 0, h0], t)

solution = solve_f(v0, ang, h0, tstep)

I've tried several loops and similar to try to stop integrating when ry = 0. And found this question below but have not been able to implement something similar with odeint. solution[:,3] is the output column for ry. Is there a simple way to do this with odeint?


Solution

  • Checkout scipy.integrate.ode here. It is more flexible than odeint and helps with what you want to do.

    A simple example using a vertical shot, integrated until it touches ground:

    from scipy.integrate import ode, odeint
    import scipy.constants as SPC
    
    def f(t, y):
        return [y[1], -SPC.g]
    
    v0 = 10
    y0 = 0
    
    r = ode(f)
    r.set_initial_value([y0, v0], 0)
    
    dt = 0.1
    while r.successful() and r.y[0] >= 0:
        print('time: {:.3f}, y: {:.3f}, vy: {:.3f}'.format(r.t + dt, *r.integrate(r.t + dt)))
    

    Each time you call r.integrate, r will store current time and y value. You can pass them to a list if you want to store them.