Tridiagonal Matrix Algorithm solver in Python with numpy array

	try:
	import numpypy as np # for compatibility with numpy in pypy
	except:
	import numpy as np # if using numpy in cpython
	## Tri Diagonal Matrix Algorithm(a.k.a Thomas algorithm) solver
	def TDMAsolver(a, b, c, d):
	'''
	TDMA solver, a b c d can be NumPy array type or Python list type.
	refer to http://en.wikipedia.org/wiki/Tridiagonal_matrix_algorithm
	'''
	nf = len(a) # number of equations
	ac, bc, cc, dc = map(np.array, (a, b, c, d)) # copy the array
	for it in xrange(1, nf):
	mc = ac[it]/bc[it-1]
	bc[it] = bc[it] - mc*cc[it-1]
	dc[it] = dc[it] - mc*dc[it-1]
	xc = ac
	xc[-1] = dc[-1]/bc[-1]
	for il in xrange(nf-2, -1, -1):
	xc[il] = (dc[il]-cc[il]*xc[il+1])/bc[il]
	del bc, cc, dc # delete variables from memory
	return xc

view raw TDMAsolver.py hosted with ❤ by GitHub