This notebook illustrates how to use the scipy.sparse module to define a tridiagonal matrix and solve a linear system with such a matrix.
See http://docs.scipy.org/doc/scipy/reference/sparse.html for more documentation on sparse matrix routines.
Note that it would be inefficient to set up a full matrix that only has three nonzero diagonals and then call the general numpy.linalg.solve routine to solve the system. This would perform Gauss Elimination on the full matrix without taking advantage of the fact that most of the lower triangular part is already full of 0's. So it would take $O(n^3)$ flops to solve an $n\times n$ system. By using the sparse routines, the same system can be solved in $O(n)$ flops.
In [1]:
%pylab inline
In [2]:
from scipy import sparse # to define sparse matrices
First a simple example to illustrate how data is specified in a banded matrix. For a tridiagonal matrix only the diagonals numbered:
are nonzero, and the data is specified by providing three 1-dimensional arrays of data giving values along these three diagonals.
In [3]:
d_main = linspace(1,5,5) # values that will go on main diagonal
d_sub = 10*d_main # values that will go on subdiagonal
d_super = 100*d_main # values that will go on superdiagonal
data = [d_sub, d_main, d_super] # list of all the data
diags = [-1,0,1] # which diagonal each vector goes into
A = sparse.spdiags(data,diags,5,5,format='csc') # create the matrix
print A.shape
Note that although $A$ is a $5 \times 5$ matrix, the full matrix is not stored. Only the non-zero elements are stored along with information about which elements these are. The format 'csc' means it is stored in compressed sparse column format,and so if you print out $A$ you will see this information by columns:
In [4]:
print A
If you want to see the full matrix, you can use the function todense associated with any sparse matrix, which fills in all the 0's and returns a numpy ndarray:
In [5]:
print A.todense()
Note that only the first 4 elements of the "d_sub" vector were used since only these columns needed to be filled. Also only the second through 5th elements of the "d_super" vector were used since only these columns needed to be filled. More generally, the elements of a vector that are used in filling the matrix correspond to the column being filled.
Here's another example with the same data, but filling different diagonals:
In [6]:
data = [d_sub, d_main, d_super] # list of all the data
diags = [-2,0,3] # which diagonal each vector goes into
B = sparse.spdiags(data,diags,5,5,format='csc') # create the matrix
print B.todense()
Sparse matrices are a special matrix class for which the * operation is overloaded with matrix or matrx-vector multiplication, so you can compute A*x or A*B:
In [7]:
x = ones((5,1))
print "A*x = \n",A*x
print "B*x = \n",B*x
C = A*B
print "A*B = \n",C.todense()
In [8]:
from scipy.sparse.linalg import spsolve # to solve sparse systems
b = A*ones(5) # create right hand side
print "b = \n",b
x = spsolve(A,b) # solve for x, should be all 1's
print "x = \n",x
Recall that for the steady-state diffusion equation in one space dimension, $u''(x) = -f(x)$ for $a < x < b$, we need to solve a system involving a tridiagonal matrix with $-2$ on the diagonal and 1 on the subdiagonal and superdiagonal. We can easily set up such a system:
In [9]:
n = 10
d1 = ones(n)
d0 = -2 * ones(n)
A = sparse.spdiags([d1,d0,d1], [-1,0,1],n,n,format='csc')
print A.todense()
The final project will involve working with this matrix!
In [10]:
A = sparse.csc_matrix(ones((3,3)))
print "A =\n", A.todense()
b = A*ones(3)
print "b =\n", b
x = spsolve(A,b) # should be all 1's
print "x = \n",x
It should instead give a warning like the dense solver does:
In [11]:
from numpy.linalg import solve
A = A.todense()
x = solve(A,b)
Also remember that a matrix can be very close to singular (ill-conditioned) and you may get answers with few digits of accuracy numerically even if it the matrix is not singular.