''' Solve a tri-diagonal system of equations ''' import numpy as np def trisolve(alpha,beta,gamma,b): n=len(alpha) # we will do LU factorization in place and # store the matrix L in the array beta for i in range(1,n): beta[i] /= alpha[i-1] alpha[i] -= beta[i]*gamma[i-1] # now do forward-backward substitutions # this can also be done in place, using only one array x x=np.zeros_like(b) x[0]=b[0] for i in range(1,n): x[i]=b[i]-beta[i]*x[i-1] x[n-1]=x[n-1]/alpha[n-1] for i in range(n-2,-1,-1): x[i] = (x[i]-gamma[i]*x[i+1])/alpha[i] return x if __name__=='__main__': # test the algorithm on random matrices n = 10000 beta = np.random.randn(n) gamma = np.random.randn(n) beta[0] = 0 gamma[n-1] = 0 # make alpha such that the matrix is strictly diagonally dominate alpha = 0.1 + np.abs(np.random.randn(n))+np.abs(beta)+np.abs(gamma) x_exact = np.ones(n) b = beta+alpha+gamma x=trisolve(alpha,beta,gamma,b) print('Relative forward error: %e' % (np.linalg.norm(x-x_exact,np.Inf)/np.linalg.norm(x_exact,np.Inf))) # generate a new problem and compare with the scipy sparse solver beta = np.random.randn(n) gamma = np.random.randn(n) beta[0] = 0 gamma[n-1] = 0 # make alpha such that the matrix is strictly diagonally dominate alpha = 0.1 + np.abs(np.random.randn(n))+np.abs(beta)+np.abs(gamma) import scipy.sparse as sp import scipy.sparse.linalg as sla data = np.zeros((3,n),dtype='f8') data[0,:] = alpha[:] data[1,:-1] = beta[1:] data[2,1:] = gamma[:-1] A = sp.dia_matrix((data,[0,-1,1]),(n,n)) # now convert this matrix to compressed column format, such that the sparse # solver can be used A=A.tocsc() x_scipy = sla.spsolve(A,b) x = trisolve(alpha,beta,gamma,b) print('Relative forward error: %e' % (np.linalg.norm(x-x_scipy,np.Inf)/np.linalg.norm(x_scipy,np.Inf)))