RK integrator

While moving files from my old laptop drive to my new one, I found a nice Runge-Kutta integrator class that I had written ages ago. So long ago, in fact, that I was a little embarrassed at the newbiness of some of my code. So I decided to update my code to get a nice RK class out of it, using list comprehensions instead of “for i in range” loops, and including an integrate method that acts as a generator so that the calling code can cycle through each integration step. As is typical in R-K, the system state is maintained in a vector X, and the calling method must provide a callback function that will return dX/dt.

Here is the class:

class RKIntegrator :
    "Class used to perform Runge-Kutta integration of set of ODE's"

    def __init__( self, dt, derivFunc, degree=0, initConds=None ):
        self.dt = float(dt)
        self.dt_2 = dt / 2.0
        self.t = float(0)
        if not (degree or initConds):
            raise ValueError("must specify degree or initial conditions")
        if initConds is not None:
            self.x = initConds[:]
        else:
            self.x = [0.0] * degree
        self.derivFunc = derivFunc


    def doIntegrationStep( self ):
        dt = self.dt
        dxFunc = self.derivFunc
        t2 = self.t + self.dt_2

        dx = dxFunc( self.t, self.x )
        delx0 = [ dx_i*dt for dx_i in dx ]
        xv = [x_i + delx0_i/2.0 for x_i, delx0_i in zip(self.x, delx0)]
        dx = dxFunc( t2, xv )
        delx1 = [ dx_i*dt for dx_i in dx ]
        xv = [x_i + delx1_i/2.0 for x_i,delx1_i in zip(self.x,delx1)]
        dx = dxFunc( t2, xv )
        delx2 = [ dx_i*dt for dx_i in dx ]
        xv = [x_i + delx1_2 for x_i,delx1_2 in zip(self.x,delx2)]
          
        self.t += dt
        dx = dxFunc(self.t, xv)
        
        self.x = [
            x_i + ( delx0_i + dx_i*dt + 2.0*(delx1_i + delx2_i) ) / 6.0
                for x_i, dx_i, delx0_i, delx1_i, delx2_i in 
                    zip(self.x, dx, delx0, delx1, delx2)
            ]
    
    def integrate(self):
        while True:
            self.doIntegrationStep()
            yield self.t, self.x

Here is an example of finding X with constant acceleration of 4:

def getDX( t, x ):
  return [ 
    x[1], 
    4.0 
    ]

isWhole = lambda x : abs(x-round(x)) < 1e6 

rk = RKIntegrator( dt=0.1, derivFunc=getDX, initConds = [0.0, 0.0] )
for t,x in rk.integrate():
    if t > 10: 
        break
    if isWhole(t):
        print(t, ', '.join('%.2f' % x_i for x_i in x))

Googling for ‘Python runge kutta’, I came across this blog posting:
http://doswa.com/blog/2009/01/02/fourth-order-runge-kutta-numerical-integration/
This does a good job, but hardcodes the vector size to just x, velocity, and acceleration. Here is how my R-K integrator would implement Doswa’s code:

def accel(t,x):
    stiffness = 1
    damping = -0.005
    x,v = x
    return -stiffness*x - damping*v

def getDX(t,x):
    return [
        x[1],
        accel(t,x)
        ]

rk = RKIntegrator( dt=1.0/40.0, derivFunc=getDX, initConds = [50.0, 5.0] )
for t,x in rk.integrate():
    if t > 100.1: 
        break
    if isWhole(t):
        print t,', '.join('%.2f' % x_i for x_i in x)

My results match the posted results to 2 places.

Advertisements
  1. #1 by David on December 28, 2010 - 1:30 am

    Neat. I like the user interface.

  2. #2 by Claudio Hartzstein on June 29, 2015 - 10:14 am

    I changed getDX by
    def getDX( t, x):
    return [
    x[1],
    np.sin(2*t)
    ]
    and got strange results (x is a sinus + a sloped line)
    Can you check this.

    Gracias
    Claudio

    • #3 by ptmcg on March 10, 2017 - 12:45 am

      Claudio –

      Well I have really neglected this blog, and am surprised at how much traffic it has been getting! I ran your code, the results don’t look all that strange to me. What were you expecting?

      — Paul

Leave a Reply

Fill in your details below or click an icon to log in:

WordPress.com Logo

You are commenting using your WordPress.com account. Log Out / Change )

Twitter picture

You are commenting using your Twitter account. Log Out / Change )

Facebook photo

You are commenting using your Facebook account. Log Out / Change )

Google+ photo

You are commenting using your Google+ account. Log Out / Change )

Connecting to %s

%d bloggers like this: