Iterative Solution : Jacobi method

jacobi_method.py

import numpy as np

# Iterative solution
# Jacobi method
def jacobi(A, b, x0, k):
    # A = D + L + U
    # D is matrix with diagonal elements
    # L is lower triangular matrix
    # U is upper triangular matrix
    #
    #
    # Explaination to jacobi method :
    #
    # Ax = b
    # (D + L + U)x = b
    # Dx = -(L + U)x + b
    # x = -D^(-1)(L + U)x  +  D^(-1)b
    #
    # x(k+1) = Bjx(k) + dj
    #
    # Bj =  -D^(-1)(L + U)
    # Bj =  +D^(-1)b

    D = np.diagflat(np.diag(A))         # gives ( n x n ) matrix with diagonal elements rest 0
    diagonal_values = np.diag(A)        # gives ( 1 x n ) matrix with diagonal elements
    
    lower_plus_diag = np.tril(A)        # lower_plus_diag is matrix with lower triangular and diagonal elements
    L = lower_plus_diag - D
    U = A - lower_plus_diag

    xk = x0
    for i in range(k):
        # xk1 = (1/diagonal_values)*( b - np.dot(L+U, xk) )
        # or
        xk1 = np.dot(np.linalg.inv(D), b - np.dot(L+U, xk))

        print("\n")
        print("Iteration number: "+str(i+1))
        print("xk+1")
        print(xk1)
        print("xk")
        print(xk)
        print("\n")

        # Checking if xk == xk+1 (upto 4 decimal places)
        if np.array_equal( np.around(xk, decimals=3), np.around(xk1, decimals=3)):
            print("Stopping at iteration number: "+str(i+1))
            break
        xk = xk1
        if i == k-1:
            print("Did not get converge !! Wrong solution\nPrinting last iteration values :\n")
    
    return xk

# ---------------------------------------------------------------------------------- #

print("\n\n --- Jacobi Method --- \n\n")
print("\n Select:")
print("1 : Random input ( Rarely get correct result )")
print("2 : Preset values ( May change in the file yourself )\n")
type = int(input())

# --------------------------------------------------- #
# Random Input

if type == 1:
    print("\nInput size of matrix\n")
    size = int(input())
    print("\nInput number of iteartions to break if it does not converge\n")
    iterations = int(input())

    A = np.random.rand(size, size)
    b = np.random.rand(size)
    x = np.zeros(size)
    k = 10                            # Input iterations yourself

    print("\n --- Creating random array --- \n")

    print("\nA :\n")
    print(A)

    print("\nb :\n")
    print(b)

    print("\nx :\n")
    print(x)

    print("\nIteartions (k) :\n")
    print(k)

# --------------------------------------------------- #
# Self Input

else:
    # Give input array A and b
    A = np.array([[4.0, -2.0, 1.0], [1.0, -3.0, 2.0], [-1.0, 2.0, 6.0]])
    b = [1.0, 2.0, 3.0]

    # Give input x0
    x = [1, 1, 1]

    # k iterations
    k = 100
    # Stop If ( xk == xk+1 ) upto 4 decimal places

    print("\nA :\n")
    print(A)

    print("\nb :\n")
    print(b)

    print("\nx :\n")
    print(x)

    print("\nIteartions (k) :\n")
    print(k)

# --------------------------------------------------- #

print(jacobi(A,b,x,k))