Tech Blog
Python ODE
Adaptive time-step integration of ODEs in Python
Overview
The ability to solve a system of ordinary differential equations accurately is fundamental to any work involving dynamic systems. In this post we look at how to solve such systems using the routines available in scipy.integrate
.
Adaptive time-step integration allows us to capture the different time scales in our system. If the system is changing rapidly the time step will reduce whereas when the system is evolving slowly a larger time step will be used, keeping the overall simulation time to a minimum. In scipy
the 4/5th order Runge-Kutta method of Dormand and Prince has been implemented under the moniker dopri5
. This method is that used in Matlab’s ode45
routine.
Spring-mass-damper system
As a motivating example, lets solve the classic spring-mass-damper system. This system can be described by the second-order differential equation
My″ + Cy′ + ky = 0,
where a prime denotes the derivative with respect to time. Our task is to solve the equation for the position of the body y.
Solving systems of ODEs
In general, ODE solvers solve systems of the form y′(t) = f(t, y), meaning we need to create a function which has arguments t, y and returns y′. The solver then integrates y′ to give us the desired y. We therefore need to write our second-order ODE as a series of first order ODES: to do this, start by introducing a new variable u = y′, so that from our system above
u′ = ( − Cy′ − ky)/M.
In order to solve the system we are going to use the scipy.integrate.ode()
method which takes as its only (required) argument the function that we want to integrate. We can define this function as follows
import scipy.integrate
import matplotlib.pyplot as plt
import numpy as np
def smd_ode(t, Y, M, C, k):
'''Spring-mass-damper system as a set of 1st order ODEs
Inputs:
t - time
Y - state vector of position and velocity [y, y']
M - mass of body
C - damping coefficient
k - spring constant
'''
# Unpack the state vector
y = Y[0]
yp = Y[1]
# Our ODE system
u = yp
up = (-C*yp - k*y) / M
# Return our system
return [u, up]
Example solution
As an example, take M = 1000 kg, C = 150 Ns/m and k = 250 N/m for a system with an initial displacement of y0 = 5m and zero initial velocity. The solution has the following steps:
- Define the problem parameters
- Create the integration object using the function we want to integrate
- Set the type of integrator we want to use
- Set the initial values of the problem
- Set the additional parameters we want to pass to the function we want to integrate
In code, this sequence takes the form
# Simulation parameters
M = 1000 # kg
C = 150 # Ns/m
k = 250 # N/m
y0 = 5 # m
yp0 = 0 # m/s
tsim = 50 # s
# Set up the ode-solver object
intobj = scipy.integrate.ode(smd_ode)
intobj.set_integrator('dopri5')
intobj.set_initial_value([y0, yp0], 0.0)
intobj.set_f_params(M, C, k)
# Function to call at each timestep
intobj.set_solout(solout)
In order to retrieve our solution we make use of the set_solout
method which takes a function to be called at the end of each time step. We can therefore save the time and solution at each time using:
sol = []
def solout(t, y):
sol.append([t, *y])
The “*” in front of the y
is critical to unpack y
from a sequence. Further information is here.
Having set up the solution, we can finally call the integrate
method with the simulation time and plot the results.
# Perform the integration
intobj.integrate(tsim)
# Convert our solution into a numpy array for convenience
asol = np.array(sol)
# Plot everything
plt.figure()
plt.plot(asol[:,0], asol[:,1], 'b.-', markerfacecolor='b')
plt.xlabel('t (s)')
plt.ylabel('y (m)')
plt.grid()
plt.figure()
plt.plot(asol[:,0], asol[:,2], 'b.-', markerfacecolor='b')
plt.xlabel('t (s)')
plt.ylabel('y\' (m/s)')
plt.grid()
plt.show()
Results
The figures below show the results for this particular problem. We can see that the position starts at 5m as we specified then decays over time due to the damping is the system.
Comments
Please let us know if there are other topics related to scientific computing in python that you would like us to cover – just send us an email or leave a comment.