import numpy as np import numpy.linalg as la import matplotlib.pyplot as plt from scipy import sparse from scipy.sparse.linalg import spsolve from mpl_toolkits.mplot3d import Axes3D # For 3-d plot from matplotlib import cm import time newparams = {'figure.figsize': (8.0, 4.0), 'axes.grid': True, 'lines.markersize': 8, 'lines.linewidth': 2, 'font.size': 14} plt.rcParams.update(newparams) def A_laplace(M): # Construct the discrete laplacian, that is # A = blocktridiag(I,T,I) is a M^2xM^2 matrix where # T = tridiag(1,-4,1) is a MxM matrix M2 = M**2 A = -4*np.eye(M2) # The diagonal matrix for i in range(M2-1): A[i,i+1] = 1 A[i+1,i] = 1 for i in range(M-1,M2-1,M): A[i,i+1] = 0 A[i+1,i] = 0 for i in range(M2-M): # The block sub- and sup diagonal A[i,i+M] = 1 A[i+M,i] = 1 return A def plot_solution(x, y, U, txt='Solution', nfig = 1): # Plot the solution of the heat equation on the unit square fig = plt.figure(nfig) plt.clf() ax = fig.gca(projection='3d') X,Y = np.meshgrid(x, y) ax.plot_surface(X, Y, U, cmap=cm.coolwarm) ax.view_init(azim=30) # Rotate the figure plt.xlabel('x') plt.ylabel('y') plt.title(txt); plt.show() def poisson_dirichlet(M): # Solve the Poisson equation by the 5-point formula on the unit square, with # Dirichlet boundary conditions. # Input: M, the number of interval # Set up the problem # Boundary conditions def gs(x): # y=0 return x**3 def gn(x): # y=1 return 2+x**3 def gw(y): # x=0 return 2*y**2 def ge(y): # x=1 return 1+2*y**2 # Source function def f(x,y): return 6*x+4 # The grid x = np.linspace(0, 1, M+1) y = np.linspace(0, 1, M+1) h = 1/M # Inner grid xi = x[1:-1] yi = y[1:-1] Xi, Yi = np.meshgrid(xi, yi) Mi = M-1 # Number of inner gridpoints Mi2 = Mi**2 # Total number of inner gridpoints (and unknowns) A = A_laplace(M-1) # The right hand side b = np.zeros(Mi2) # Boundary conditions b[0:Mi] = b[0:Mi]-gs(xi) # y=0 b[Mi2-Mi:Mi2] = b[Mi2-Mi:Mi2] - gn(xi) # y=1 b[0:Mi2:Mi] = b[0:Mi2:Mi] - gw(yi) # x=0 b[Mi-1:Mi2:Mi] = b[Mi-1:Mi2:Mi] - ge(yi) # x=1 # Source function b = b+h**2*f(Xi,Yi).flatten() # Solve the linear system Ui = la.solve(A, b) # Make an (M+1)x(M+1) array to store the solution, # including boundary U = np.zeros((M+1, M+1)) # Reshape the solution vector, and insert into the solution matrix U[1:-1,1:-1] = np.reshape(Ui, (Mi,Mi)) # include the boundaries U[0,:] = gs(x) U[M,:] = gn(x) U[:,0] = gw(y) U[:,M] = ge(y) return x, y, U def sparse_demo(M): # Demonstate the benefits of using a sparse solver, by constructing the # discrete laplacian A, solve the system Ax=b for some vector b, and # compare the time used for solving the system by a full and a sparse # solver. # NB! This solves a system of M^2 equations. A = A_laplace(M) b = np.ones(M**2) # Or some other arbitrary b-vector start = time.time() U = la.solve(A, b) finish= time.time() print('Time used for solving a full system: {:8.4f} sec'.format(finish-start)) # Sparse solver As = sparse.csr_matrix(A) # Convert the matrix to a sparse format (csr) start = time.time() U = spsolve(As, b) # Use a sparse linear solver finish = time.time() print('Time used for solving a sparse system: {:8.4f} sec'.format(finish-start)) if __name__ == "__main__": x, y, U = poisson_dirichlet(10) X, Y = np.meshgrid(x,y) U_exact = X**3+2*Y**2 plot_solution(x, y, U) plot_solution(x, y, U-U_exact, txt='Error', nfig=2)