Archive for the ‘snippet’ Tag

Pure Python implementation of Runge Kutta

I was asked (or forced, if you like) to implement the Hodgkin-Huxley model (a model that describes how action potentials in neurons are initiated and propagated) using Runge Kutta for a system of differential equations. The following snippet is the Runge Kutta 4th order, for an arbitrary number of differential equation.


class RungeKutta(object):
	"""Runge Kutta 4th order for a system of differential equations."""
	"""https://testparticle.wordpress.com"""
	__author__ = 'woci'
	__version__ = 0.1b
	__url__ = """https://testparticle.wordpress.com"""
	def __init__(self, functions, initConditions, dh, save=True):
		self.Trajectory, self.save   = [initConditions], save
		self.functions    = [lambda *args: 1.0]  + functions
		[self.N, self.dh] = [len(initConditions), dh]
		self.coeff        = [1.0/6.0, 2.0/6.0, 2.0/6.0, 1.0/6.0]
		self.InArgCoeff   = [0.0, 0.5, 0.5, 1.0]
	def iterate(self):
		step   = self.Trajectory[-1][:]
		istep, iac  = step[:], self.InArgCoeff
		k, ktmp = self.N * [0.0], self.N * [0.0]
		for ic, c in enumerate(self.coeff):
			for if_, f in enumerate(self.functions):
				arguments  = [ (x + k[i]*iac[ic]) \
						for i, x in enumerate(istep)] 
				ktmp[if_]  = self.dh * f(*arguments)
			k = ktmp[:]
			step = [s + c*k[ik] for ik,s in enumerate(step)]
		if self.save: self.Trajectory += [step]
		else: self.Trajectory = [step]
	def solve(self, finishtime):
		while (self.Trajectory[-1][0] < finishtime ):
			self.iterate()
		return None
	def series(self):
			return zip(*self.Trajectory)
&#91;/sourcecode&#93;
You define the functions, the initial conditions, the timestep and you have yourself a solver.
The functions can be defined using the standard <strong>def</strong> or by just using lambda.

def dxdt(t):
	return cos(t)
dxdt(t) = lambda t: cos(x)

The following source solves the «Lorenz Oscillator» equations and plots the timeseries if you have matplotlib at hand.

if __name__ == '__main__':
	print "Solving Lorenz Equation"
	dxdt = lambda t, x, y, z: 10.0*(y-x)
	dydt = lambda t, x, y, z: 28.0*x - x*z - y
	dzdt = lambda t, x, y, z: x*y - (8.0/3.0)*z
	funs = [dxdt, dydt, dzdt]
	init_cond = [0., 1, 2, 3]
	dt = 0.01
	RK4 = RungeKutta(funs, init_cond, dt)
	RK4.solve(10.0)
	t, x, y, z = RK4.series()
	pylab.plot(t, x)
	pylab.plot(t, y)
	pylab.plot(t, z)
	pylab.show()

Lorenz Solution
Please refer if you use this script.