%pylab inline
import numpy as np
import matplotlib.pyplot as pl
from scipy.sparse import spdiags,linalg
Here's a function to construct the 5-point Laplacian finite difference formula
$$\frac{1}{h^2}\left(U_{i-1,j}+U_{i+1,j}+U_{i,j-1}+U_{i,j+1}-4U_{ij}\right)$$
as a matrix $A$, so that the product $AU$ gives the formula above when $U$ is constructed using
the natural row-wise ordering:
def five_pt_laplacian(m):
e=np.ones(m**2)
e2=([0]+[1]*(m-1))*m
h=1./(m+1)
A=np.diag(-4*e,0)+np.diag(e2[1:],-1)+np.diag(e2[1:],1)+np.diag(e[m:],m)+np.diag(e[m:],-m)
A/=h**2
return A
Matplotlib has a nice spy() command, just like MATLAB's, which shows us the structure of $A$ (note here the structure of the +1 and -1 diagonals):
A=five_pt_laplacian(4)
plt.spy(A)
Notice that each row has at most 5 non-zero entries, and at least $m^2-5$ entries equal to zero. The function above constructs a full matrix, meaning all those zeros are stored in memory. This will take up a lot of memory! We might as well just store the non-zero entries. This is done with the following code (take a look at the help for spdiags):
def five_pt_laplacian_sparse(m):
e=np.ones(m**2)
e2=([1]*(m-1)+[0])*m
e3=([0]+[1]*(m-1))*m
h=1./(m+1)
A=spdiags([-4*e,e2,e3,e,e],[0,-1,1,-m,m],m**2,m**2)
A/=h**2
return A.tocsr()
Another tool we need is a way to transform a matrix whose $i,j$ entry corresponds to the point $x_i,y_j$ on the grid into a vector using the natural row-wise ordering. Like MATLAB, numpy has a reshape() command that does just this:
U=np.array([[1,2,3],[4,5,6],[7,8,9]])
U.reshape([9])
Now let's solve the 2D BVP
$$\nabla^2 u = -(20y^3+9\pi^2(y-y^5))\sin(3\pi x)$$on the unit square $[0,1]\times[0,1]$ with homogeneous Dirichlet boundary conditions (i.e., $u(x,y)=0$ everywhere on the boundary). You can do this as follows:
Now to plot your solution, you'll need to reshape $U$ to be the shape of the grid; i.e, $m\times m$. Then you can plot it using the function plt.pcolor().
U=U.reshape([m,m])
plt.clf()
plt.pcolor(X,Y,U)
plt.colorbar()
How accurate is this solution? In fact, the problem above was obtained by taking
$$u(x,y) = (y-y^5)\sin(3\pi x)$$
and computing the Laplacian of $u$. Notice that $u=0$ on the boundaries.
Use this to compute the error in your solution. Try changing $m$. What is the rate of convergence?