Using Python's ODE Suite

In [2]:
import numpy as np
from scipy.integrate import solve_ivp
import matplotlib.pyplot as plt
%matplotlib inline
np.seterr(divide='ignore');

Simple Usage

In [3]:
# Dynamics function
def simple_golf(t, z):
    '''
        z[0] = x(t)
        z[1] = y(t)
        z[3] = y'(t)
    '''
    # Vx is defined globally
    return [Vx, z[2] , -9.81]
In [4]:
# IVP parameters
tspan = [0, 10]
y0 = [0, 0, 15]
Vx = 25.
In [5]:
# Call the numerical solver
sol = solve_ivp(simple_golf, tspan, y0, max_step=1)
In [6]:
sol.y[0]
Out[6]:
array([0.00000000e+00, 8.57435763e-04, 9.43179339e-03, 9.51753697e-02,
       9.52611133e-01, 9.52696876e+00, 3.45269688e+01, 5.95269688e+01,
       8.45269688e+01, 1.09526969e+02, 1.34526969e+02, 1.59526969e+02,
       1.84526969e+02, 2.09526969e+02, 2.34526969e+02, 2.50000000e+02])
In [7]:
# Plot the solution
plt.plot(sol.y[0], sol.y[1], 'ro-');
plt.title('Golf Ball Trajectory')
plt.xlabel('Distance (m)')
plt.ylabel('Height (m)');
In [8]:
sol.y[:,-1]
Out[8]:
array([ 250. , -340.5,  -83.1])

Passing additional parameters

In [9]:
# Suppose we want to pass Vx as a parameter to our dynamics function, like this.
def simple_golf2(t, z, Vx):
    '''
        z[0] = x(t)
        z[1] = y(t)
        z[3] = y'(t)
    '''
    return [Vx, z[2] , -9.81]
In [10]:
# We can create a simple wrapper for our dynamics function, like this:
fun = lambda t,x: simple_golf2(t,x,20)  # Vx will be passed through as a constant
sol = solve_ivp(fun, tspan, y0, max_step=1.)
In [11]:
# Plot
plt.plot(sol.y[0], sol.y[1], 'ro-');
plt.title('Golf Ball Trajectory')
plt.xlabel('Distance (m)')
plt.ylabel('Height (m)');

ODE Solver Controls

In [12]:
# Maximum step size (max_step)
sol = solve_ivp(simple_golf, tspan, y0, max_step=0.5)
plt.plot(sol.y[0], sol.y[1], 'ro-')
plt.xlabel('x'); plt.ylabel('y');
In [13]:
# Choosing solution timestamps (t_eval)
tt = np.linspace(tspan[0], tspan[1], 41)
sol = solve_ivp(simple_golf, tspan, y0, t_eval=tt)
plt.plot(sol.y[0], sol.y[1], 'ro-')
plt.xlabel('x'); plt.ylabel('y');
In [14]:
# Choosing error tolerances (rtol, atol)
fun = (lambda t,z: 2*z[0]*(1-z[0])/(1+t))
sol1 = solve_ivp(fun, [0,10], [0.3], rtol=1.e-2, atol=1.e-2)
sol2 = solve_ivp(fun, [0,10], [0.3], rtol=1.e-7, atol=1.e-7)
plt.plot(sol1.t, sol1.y[0], 'ro-')
plt.plot(sol2.t, sol2.y[0], 'bo-')
plt.xlabel('x'); plt.ylabel('y');

Event Functions

In [15]:
# Define the event function
def my_event(t,z):
    return z[1]
# We want integration to terminate
my_event.terminal = True
# And we want only positive-to-negative crossings.
my_event.direction = -1
In [16]:
fun = lambda t,x: simple_golf2(t,x,10)
sol = solve_ivp(fun, tspan, y0, max_step=0.6, events=my_event)
In [17]:
plt.plot(sol.y[0], sol.y[1], 'ro-');
plt.title('Golf Ball Trajectory')
plt.xlabel('Distance (m)')
plt.ylabel('Height (m)');
In [ ]: