jacobi method in python

1import numpy as np
2from numpy.linalg import *
3
4def jacobi(A, b, x0, tol, maxiter=200):
5    """
6    Performs Jacobi iterations to solve the line system of
7    equations, Ax=b, starting from an initial guess, ``x0``.
8
9    Terminates when the change in x is less than ``tol``, or
10    if ``maxiter`` [default=200] iterations have been exceeded.
11
12    Returns 3 variables:
13        1.  x, the estimated solution
14        2.  rel_diff, the relative difference between last 2
15            iterations for x
16        3.  k, the number of iterations used.  If k=maxiter,
17            then the required tolerance was not met.
18    """
19    n = A.shape[0]
20    x = x0.copy()
21    x_prev = x0.copy()
22    k = 0
23    rel_diff = tol * 2
24
25    while (rel_diff > tol) and (k < maxiter):
26
27        for i in range(0, n):
28            subs = 0.0
29            for j in range(0, n):
30                if i != j:
31                    subs += A[i,j] * x_prev[j]
32
33            x[i] = (b[i] - subs ) / A[i,i]
34        k += 1
35
36        rel_diff = norm(x - x_prev) / norm(x)
37        print(x, rel_diff)
38        x_prev = x.copy()
39
40    return x, rel_diff, k
41
42# Main code starts here
43# ---------------------
44GL = 1.6
45d = 0.8
46A = np.array([
47    [1.0,      0,      0,      0,   0],
48    [ GL, -(d+1),    1.0,      0,   0],
49    [  0,      d, -(d+1),    1.0,   0],
50    [  0,      0,      d, -(d+1), 1.0],
51    [  0,      0,      0,      0, 1.0]])
52b = [0.5,      0,      0,      0, 0.1]
53x0 = np.zeros(5);
54
55tol = 1E-9
56maxiter = 200
57x, rel_diff, k = jacobi(A, b, x0, tol, maxiter)
58if k == maxiter:
59    print(('WARNING: the Jacobi iterations did not '
60           'converge within the required tolerance.'))
61print(('The solution is %s; within a tolerance of %g, '
62        'using %d iterations.' % (x, rel_diff, k)))
63print('Solution error = norm(Ax-b) = %g' % \
64            norm(np.dot(A,x)-b)) 
65print('Condition number of A = %0.5f' % cond(A))
66print('Solution from built-in functions = %s' % solve(A, b))
67