Implementation of Fitzhug-Nagumo model with Runge-Kutta method

#-*- coding:utf-8 -*-
import numpy as np
import matplotlib.pylab as plt

#node
def phase_func(x, t, a=-1., b=0., c=0., d=-2.):
	return np.array([a*x[0]+b*x[1], c*x[0] + d*x[1]])

#focus
def phase_focus(x, t, a=1., b=2., c=-2., d=1.):
	return np.array([a*x[0]+b*x[1], c*x[0] + d*x[1]])

#saddle
def phase_saddle(x, t, a=-1., b=-1., c=4., d=1.):
	return np.array([a*x[0]+b*x[1], c*x[0] + d*x[1]])

#Fitz_Nagumo
def fitz_nagumo(x, t, a=0.7, b=0.8, c=3, z=-0.00):
	return np.array([c*(x[1] + x[0] - x[0]**3/3. + z), -1./c * (x[0] -a + b*x[1])])

def runge_kutta(t0 = 0.0, x0 = np.array([1.0]), t1 = 100., h = 0.01, ng = None):
	T = np.arange(t0, t1, h)
	N = np.size(T)
	X = np.zeros((N, np.size(x0)))
	X[0] = x0
	
	for i in xrange(1, N):
		x = X[i-1]; t = T[i-1]
		w1 = h * ng(x, t)
		w2 = h * ng(x + w1 * 0.5, t + 0.5 * h)
		w3 = h * ng(x + w2 * 0.5, t + 0.5 * h)
		w4 = h * ng(x + w3, t + h)
		X[i] = X[i-1] + 1./6. * (w1 + 2*w2 + 2*w3 + w4)
	return X,T

def main():
	plt.figure()
	
	#Node
	#for i in xrange(-5, 5):
	#	for j in xrange(-5, 5):
	#		X = runge_kutta(x0 = np.array([i, j]), ng = phase_func)
	#		plt.plot(X[:, 0], X[:, 1])
	#plt.xlabel('x1')
	#plt.ylabel('x2')
	#plt.title('Phase Plane(node) p^2 >4q, q >0')
	#plt.show()
	
	#focus
	#for i in xrange(-5, 5)	:
	#	for j in xrange(-5, 5):
	#		X = runge_kutta(x0 = np.array([i, j]), ng = phase_focus)
	#		plt.plot(X[:, 0], X[:, 1])
	#plt.xlabel('x1')
	#plt.ylabel('x2')
	#plt.title('Phase Plane(focus) p^2 < 4q, q >0')
	#plt.show()
	
	#saddle
	#for i in xrange(-5, 5)	:
	#	for j in xrange(-5, 5):
	#		X = runge_kutta(x0 = np.array([i, j]), ng = phase_saddle)
	#		plt.plot(X[:, 0], X[:, 1])
	#plt.xlabel('x1')
	#plt.ylabel('x2')
	#plt.title('Phase Plane(saddle) q < 0')
	#plt.show()
	
	#Fit_Nagumo_model
	#for i in xrange(-5, 5):
	#	for j in xrange(-5, 5):
	X, T = runge_kutta(x0 = np.array([-1., -1.]), ng = fitz_nagumo)
	#		plt.plot(X[:, 0], X[:, 1]);
	plt.plot(T, X[:, 0])
	plt.xlabel('x1')
	plt.ylabel('x2')
	ax = plt.gca()
	ax.invert_yaxis()
	plt.title('Fitzhuge-Nagumo model(a=0.7, b=0.8, c=3, z=0.0)')
	plt.show()

if __name__ == "__main__":
	main()