Adaptive time-step integration of ODEs in Python
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
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
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 yp = Y # Our ODE system u = yp up = (-C*yp - k*y) / M # Return our system return [u, up]
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()
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.
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.