Jacobi Method in Python and NumPy

Jacobi Method in Python and NumPy

This article will discuss the Jacobi Method in Python. We've already looked at some other numerical linear algebra implementations in Python, including three separate matrix decomposition methods: LU Decomposition, Cholesky Decomposition and QR Decomposition. The Jacobi method is a matrix iterative method used to solve the equation $Ax=b$ for a known square matrix $A$ of size $n\times n$ and known vector $b$ or length $n$.

Jacobi's method is used extensively in finite difference method (FDM) calculations, which are a key part of the quantitative finance landscape. The Black-Scholes PDE can be formulated in such a way that it can be solved by a finite difference technique. The Jacobi method is one way of solving the resulting matrix equation that arises from the FDM.

The algorithm for the Jacobi method is relatively straightforward. We begin with the following matrix equation:

\begin{eqnarray*} Ax = b \end{eqnarray*}

$A$ is split into the sum of two separate matrices, $D$ and $R$, such that $A=D+R$. $D_{ii} = A_{ii}$, but $D_{ij}=0$, for $i\neq j$. $R$ is essentially the opposite. $R_{ii} = 0$, but $R_{ij} = A_{ij}$ for $i \neq j$. The solution to the equation, i.e. the value of $x$, is given by the following iterative equation:

\begin{eqnarray*} x^{(k+1)} = D^{-1}(b-Rx^{(k)}). \end{eqnarray*}

We will make use of the NumPy library to speed up the calculation of the Jacobi method. NumPy is significantly more efficient than writing an implementation in pure Python. The iterative nature of the Jacobi method means that any increases in speed within each iteration can have a large impact on the overall calculation.

Note that this implementation uses a predetermined number of steps when converging upon the correct solution. Depending upon your needs in production, you may wish to use a residual tolerance method. For that you will need to take a look at the spectral radius.

Here is the implementation via NumPy:

from pprint import pprint
from numpy import array, zeros, diag, diagflat, dot

def jacobi(A,b,N=25,x=None):
    """Solves the equation Ax=b via the Jacobi iterative method."""
    # Create an initial guess if needed                                                                                                                                                            
    if x is None:
        x = zeros(len(A[0]))

    # Create a vector of the diagonal elements of A                                                                                                                                                
    # and subtract them from A                                                                                                                                                                     
    D = diag(A)
    R = A - diagflat(D)

    # Iterate for N times                                                                                                                                                                          
    for i in range(N):
        x = (b - dot(R,x)) / D
    return x

A = array([[2.0,1.0],[5.0,7.0]])
b = array([11.0,13.0])
guess = array([1.0,1.0])

sol = jacobi(A,b,N=25,x=guess)

print "A:"
pprint(A)

print "b:"
pprint(b)

print "x:"
pprint(sol)

The output from the NumPy implementation is given below:

A:
array([[ 2.,  1.],
       [ 5.,  7.]])
b:
array([ 11.,  13.])
x:
array([ 7.11110202, -3.22220342])

We can see that after 25 iterations, the output is as given by the Wikipedia article on the Jacobi Method, up to the difference in significant figures.