Tech Blog

Python ODE

MAB_headshot_circle

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:

  1. Define the problem parameters
  2. Create the integration object using the function we want to integrate
  3. Set the type of integrator we want to use
  4. Set the initial values of the problem
  5. 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.


Position of the body in time Position of the body in time

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.