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 for i in range(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, 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:
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, 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.