##^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^Lance Simms, Stanford University 2009 #PURPOSE: # To decompose a tri-diagonal matrix M into upper and lower triangular # matrices, U and L, respectively, so that M = LU. The matrix should # have the upper band as a, the middle band as b, and the lower band as c. # # | a_0 c_0 0 0 0 ... | # M = | b_0 a_1 c_1 0 0 ... | # | 0 b_1 a_2 c_2 0 ... | # | 0 0 b_2 a_3 c_3 ... | # | .. .. .. .. .. ... | # #INPUTS: # a: int array # The middle band of the matrix # b: int array # The lower band of the matrix # c: int array # The upper band of the matrix # #KEYWORDS: # FullReturn: int # 0 - Don't return the full L and U matrices, just the diagonals # 1 - Return all the values # #OUTPUTS: # L: 2d matrix # The lower triangular matrix # U: 2d matrix # The upper triangular matrix # alpha: 1d array # The main diagonal of L # gamma: 1d array # The upper diagonal of U # ##EXAMPLE: ## a=ones(6) ## b=ones(5) ## c=ones(5) ## L, U, alpha, gamma = LU_decomposition(a,b,c,JL) ## ######################################################################### from numpy import * from pylab import * import pdb import matplotlib.pylab as mplot def LU_decomposition(a, b, c, FullReturn=0): #Form the diagonals for the output matrices JL = len(a) alpha=zeros(JL,dtype=double) gamma=zeros(JL-1,dtype=double) alpha[0] = a[0] if alpha[0] != 0: gamma[0] = c[0]/alpha[0] else: gamma[0]=0 for i in arange(JL-2)+1: alpha[i] = a[i]-b[i-1]*gamma[i-1] gamma[i] = c[i]/alpha[i] alpha[JL-1] = a[JL-1]-b[JL-2]*gamma[JL-2] if FullReturn == 0: L = 0 ; U = 0 else: L = diag(alpha)+diag(b,-1) U = diag(ones(JL))+diag(gamma,1) return L, U, alpha, gamma ######################################################################### ##LU_solve ## ##PURPOSE: ## To solve for a vector after a LU decomposition has been made. The form ## of the equation involving the vector should be: ## ## L * U * PHI = f ## ## The procedure follows with two loops, one to solve the L * PHI# = f, the ## next to solve for U PHI = PHI* ## ##INPUTS: ## alpha: int array ## The main diagonal of the lower triangular matrix L ## b: int array ## The lower diagonal of the lower triangular matrix L ## gamma: int array ## The upper diagonal of the upper triangular matrix U ## f: int array ## The vector on the rhs of the equation ## ##OUTPUTS: ## Phi: int array ## The solution matrix ## ########################################################### def LU_solve(alpha, b, gamma, f): #Set up the vectors that will hold the solution and temp. solution JL = len(alpha) Phi = zeros(JL, dtype = double) PhiStar = zeros(JL, dtype=double) ###FIRST LOOP #### L*PhiStar = f ----------FORWARD SWEEP PhiStar[0] = f[0]/alpha[0] for i in arange(JL-1)+1: PhiStar[i] = (f[i]-b[i-1]*PhiStar[i-1])/alpha[i] ###SECOND LOOP ### U*Phi = PhiStar Phi[JL-1] = PhiStar[JL-1] iarr = arange(JL-1) for i in iarr[::-1]: Phi[i] = PhiStar[i]-gamma[i]*Phi[i+1] return Phi ########################################################################### ##3x3 block diagonal solution methods #PURPOSE: # To decompose an array of 3x3 matrices into an array of upper and lower # triangular matrix elements. # # A will be a set of 3x3 matrices with length N. IndexRange will be a range of # indices (say, arange[1:N-1])+1. The returned arrays will be the arrays of # a specific element of the upper and lower triangular matrices. # #INPUTS: # A: int array (N, 3,3) # The Nx3x3 matrix # IndexRange: int array (N-2) # The range of indices #OUTPUTS: # L11, L12, L22, L31, L32, L33 : int array (of length IndexRange) # The elements of the lower triangular matrix # U12, U13, U23: N int array (of length IndexRange) # The elements of the upper triangular matrix # V1, V2, V3: int array (of length IndexRange) # 1/ The lower elements of the upper triangular matrix to. # i.e. V1 = 1/L11 ; V2 = 1/L22 ; V3 = 1/L33 # ierr: # 0 - No error # 1 - first diagonal element was zero # 2 - failure due to singular matrix # ##EXAMPLE: ## a=ones([3,3]) ## b=ones(5) ## c=ones(5) ## L, U, alpha, gamma = LU_decomposition(a,b,c,JL) ## ######################################################################### from numpy import * import pdb def LU_decomposition_3x3(A,IndexRange): #Simplify the notation IR = IndexRange #Do the top row ierr = 0 L11 = A[IR,0,0] if len(where(L11==0)[0]) != 0: ierr = -1 return 0,0,0,0,0,0,0,0,0,0,0,0, ierr #Do the second row V1 = 1./L11 U12 = V1*A[IR,0,1] U13 = V1*A[IR,0,2] L21 = A[IR,1,0] L22 = A[IR,1,1]-L21*U12 if len(where(L22==0)[0]) !=0 : ierr = -2 return 0,0,0,0,0,0,0,0,0,0,0,0, ierr #Do the third row V2 = 1./L22 U23 = V2*(A[IR,1,2]-L21*U13) L31 = A[IR,2,0] L32 = A[IR,2,1]-L31*U12 L33 = A[IR,2,2]-L31*U13-L32*U23 if len(where(L33==0)[0])!=0: ierr = -2 return 0,0,0,0,0,0,0,0,0,0,0,0, ierr V3 = 1./L33 return L11, L21, L22, L31, L32, L33, V1, V2, V3, U12, U13, U23, ierr ######################################################################### ##LU_solve_3x3_1d ## ##PURPOSE: ## To solve for a vector after a LU decomposition has been made. The form ## of the equation involving the vector should be: ## ## B*del_y_i-1+C*del_y+D*del_y+1 = f ## ## The methodology employed is from Chapter 11, page 3 of Bob MackCormack's ## notes. We first decompose our block matrix into an upper block ## triangular matrix, U, and a lower block triangular matrix, L. ## Note that we have made the interchanges ## ## A -> C C-> D ## ## The matrix L has alpha as its main diagonal and B as its lower diagonal. ## The matrix U has identitify as its main diagonal and gamma as its upper. ## Gamma and Alpha should only have DL-1 elements, where DL is the range total ## range of indices. ## ##INPUTS: ## B: (N,3,3) int array ## The lower diagonal of the block tri-diagonal matrix ## C: (N,3,3) int array ## The middle diagonal of the block tri-diagonal matrix ## D: (N,3,3) int array ## The upper diagonal of the block tri-diagonal matrix ## R: int array ## The vector on the rhs of the equation ## ##KEYWORDS: ## PBC: int array ## 0 - Don't use periodic boundary conditions ## B and D should have one less element than C since ## they are off diagonal ## 1 - Use periodic boundary conditions ## B and D should have the same number of elements as ## C. They are off diagonal, but the first and last ## row of the matrix have an additional element in the ## corners. E.g ## ## | a_0 c_0 0 0 0 ... b_0 | ## M = | b_0 a_1 c_1 0 0 ... 0 | ## | 0 b_1 a_2 c_2 0 ... 0 | ## | 0 0 b_2 a_3 c_3 ... 0 | ## | .. .. .. .. .. ... 0 | ## | c_I .. .. .. .. b_I a_I | ## ## DEBUG: int ## 0 - Don't debug; 1 - debug ## ##OUTPUTS: ## del_F: int array ## The solution matrix ## ##NOTES: ## The original version of the code used LU decomposition to do the inversion ## of all the 3x3 matrices. However, pylab's inv function does the job just ## fine (and a lot faster!), so the code was switched to use this. ## ## To reload this module: ## import LU_decomposition; reload(LU_decomposition); from LU_decomposition import * ## ########################################################### def LU_solve_3x3_1d(IL, IU, B, C, D, R, PBC=0, Debug=0): #Set up the vectors that will hold the solution and temp. solution I = IL #Lowest index (starting index for decomposition) FR = [0] #Range used for LU decomposition; meant for 2-d DL = IU-IL #Full range of indices for decomposition if Debug == 1: pdb.set_trace() #Make a copy of the Right hand side? Maybe I did this because F was being affected outside? F = R.copy() del R ##Create arrays for the LU matrices (only necessary to hold one alpha value) Alpha = zeros([1,3,3],dtype=double) Gamma = zeros([DL, 3, 3], dtype=double) ##Create the matrices for the periodic case if the keyword is set if PBC == 1: e = zeros([DL, 3, 3], dtype=double) #Corresponds to little e in MacCormack's notes d = zeros([DL, 3, 3], dtype=double) #Corresponds to little d in MacCormack's notes #########################FORWARD SWEEP###################################### ##Get alpha_0 of L Alpha[0,:,:] = C[I,:,:] ##Solve for the first element of temp vector : R -> del_F* F[I,:] = dot(inv(Alpha[0,:,:]),F[I,:]) ##Now solve for gamma_0 of U: Gamma[I,:,:] = dot(inv(Alpha[0,:,:]), D[I,:,:]) ##Solve for d_0 if we are using periodic boundary conditions if PBC == 1: e[0,:,:] = D[DL,:,:] d[0,:,:] = dot(inv(Alpha[0,:,:]), B[I,:,:]) ##Do the forward sweep for I = 1 to I = IL for I in arange(DL)+1: ##Periodic Boundary Conditions and I=DL require special attention if not(PBC == 1 and I == DL): #Calculate the new Alpha_I and its inverse Alpha_I_Inv: Alpha[0,:,:] = C[I,:,:] - dot(B[I,:,:],Gamma[I-1,:,:]) Alpha_Inv = inv(Alpha[0,:,:]) #Solve for del_F** = B_i*del_F (temporary val) F[I,:] = F[I,:] - dot(B[I,:,:],F[I-1,:]) #Now solve for del_F*_i (Use pylab's matrix ops instead of LU decomposition) F[I,:] = dot(Alpha_Inv,F[I,:]) #Now solve for Gamma_i = Alpha_I^-1 * D_I. Only do this for DL-1 entries if I < DL : Gamma[I,:,:] = dot(Alpha_Inv, D[I,:,:]) #Solve for d_I and e_I if periodic boundary conditions are used if (PBC == 1 and I <= (DL-1)): d[I,:,:] = - dot(Alpha_Inv,dot(B[I,:,:],d[I-1,:,:])) e[I,:,:] = - dot(e[I-1,:,:],Gamma[I-1,:,:]) #The last two elements require special attention for PBC if (PBC == 1 and I == (DL-1)): Gamma[I,:,:] = Gamma[I,:,:] + d[I,:,:] #Gamma_3 #Need to handle the last element Alpha_2, B_2 differently for PBC if (PBC == 1 and I == DL): B[I,:,:] = B[I,:,:] + e[I-1,:,:] Alpha[0,:,:] = C[I,:,:] - dot(B[I,:,:],Gamma[I-1,:,:]) F[I,:] = F[I,:] - dot(B[I,:,:],F[I-1,:]) IRarr = arange(DL-1) for IR in IRarr[::-1]: Alpha[0,:,:] = Alpha[0,:,:] - dot(e[IR,:,:],d[IR,:,:]) F[I,:] = F[I,:]-dot(e[IR,:,:],F[IR,:]) Alpha_Inv = inv(Alpha[0,:,:]) F[I,:] = dot(Alpha_Inv, F[I,:]) if Debug: pdb.set_trace() #######################BACKWARD SWEEP############################## # dF*[1]=dF**[1] Iarr = arange(DL) for I in Iarr[::-1]: if (I == (DL-1) or PBC == 0): F[I,:] = F[I,:]-dot(Gamma[I,:,:],F[I+1,:]) else: F[I,:] = F[I,:]-dot(Gamma[I,:,:],F[I+1,:])-dot(d[I,:,:],F[DL,:]) if Debug: pdb.set_trace() return F ######################################################################### ##LU_solve_5x5_2d ## ##PURPOSE: ## To solve for a penta-diagonal matrix where the vectors lie in separate ## dimensions (x and y) after a LU decomposition has been made. The form ## of the equation involving the vector should be: ## ## A*del_U(j,i-1)+B*del_U(j-1,i)+C*del_U(j,i)+D*del_U(j+1,i)+E*del_U(j,i+1) = f(j,i) ## ## The methodology employed is from Chapter 11, page 9 of Bob MackCormack's ## notes from the course AA 214B at Stanford University. ## We first decompose our matrix into the product of 3 matrices ## according to the diagonally dominant method: ## ## Ma del_U(i,j) = Del_U(i,j) ## ## Ma = Ma_x*D^-1*Ma_y ## ## The X dimension has M elements and the Y dimension has N elements. ## ##INPUTS: ## A: (N,M,3,3) float array ## The distant, lower diagonal of the block penta-diagonal matrix ## B: (N,M,3,3) float array ## The close, lower diagonal of the block penta-diagonal matrix ## C: (N,M,3,3) float array ## The middle diagonal of the block penta-diagonal matrix ## D: (N,M,3,3) float array ## The close, upper diagonal of the block penta-diagonal matrix ## E: (N,M,3,3) float ## The distant, upper diagonal of the block penta-diagonal matrix ## ## R: (N,M,3) float array ## The vector on the rhs of the equation ## ##KEYWORDS: ## PBC_X: int array ## 0 - Don't use periodic boundary conditions in the X direction ## 1 - Use periodic boundary conditions in the X direction ## PBC_Y: int array ## 0 - Don't use periodic boundary conditions in the Y direction ## 1 - Use periodic boundary conditions in the Y direction ## ##OUTPUTS: ## del_F: int array ## The solution matrix ## ##To reload this module: ## import LU_decomposition; reload(LU_decomposition); from LU_decomposition import * ##Debug plot command: ## mplot.plot(dU_k[:,10,2]) ########################################################### def LU_solve_5x5_2d(IL, IU, JL, JU, A_Arr, B_Arr, C_Arr, D_Arr, E_Arr, F_Arr, \ PBC_X=0, PBC_Y=0, Order=0 , Debug=0): ##Get the dimensionality of the matrices YDim = A_Arr.shape[0] #The length of the Y axis = J XDim = A_Arr.shape[1] #The length of the X axis = I dU = zeros([YDim, XDim, 3], dtype=float64) #The full solution matrix dU_Err = zeros([YDim, XDim, 3], dtype=float64) #The error in the full solution matrix dX = zeros([YDim, XDim, 3], dtype=float64) #The 1st temp solution matrix dY = zeros([YDim, XDim, 3], dtype=float64) #The 2nd temp solution matrix ##Determine which order to use if Order == 0: ##FIRST: Solve the equation for all vectors Y_j along the y axis ## D^-1*My*del_U(i,j) = Mx^-1*Del_U(j,i) = X(j,i) for x in arange(XDim): dX[:,x,:] = LU_solve_3x3_1d(0, IU, B_Arr[:,x,:,:], C_Arr[:,x,:,:], \ D_Arr[:,x,:,:], F_Arr[:,x,:], PBC=PBC_Y, \ Debug=Debug) ##SECOND: Solve the equation ## My*del_U(i,j) = D*X(j,i) = Y(j,i) for x in arange(JU+1): for y in arange(IU+1): dY[y,x,:] = dot(C_Arr[y,x,:,:],dX[y,x,:]) ##THIRD: Solve the equation ## del_U(i,j) = My^-1 *Y(i,j) for y in arange(YDim): dU[y,:,:] = LU_solve_3x3_1d(0, JU, A_Arr[y,:,:,:], C_Arr[y,:,:,:], \ E_Arr[y,:,:,:], dY[y,:,:], PBC=PBC_X, \ Debug=Debug) else: ##FIRST: Solve the equation for all vectors X_i along the x axis ## D^-1*My*del_U(i,j) = Mx^-1*Del_U(j,i) = X(j,i) for y in arange(YDim): dX[y,:,:] = LU_solve_3x3_1d(0, JU, A_Arr[y,:,:,:], C_Arr[y,:,:,:], \ E_Arr[y,:,:,:], F_Arr[y,:,:], PBC=PBC_X) ##SECOND: Solve the equation ## My*del_U(i,j) = D*X(j,i) = Y(j,i) for x in arange(JU+1): for y in arange(IU+1): dY[y,x,:] = dot(C_Arr[y,x,:,:],dX[y,x,:]) ##THIRD: Solve the equation ## del_U(i,j) = My^-1 *Y(i,j) for x in arange(XDim): dU[:,x,:] = LU_solve_3x3_1d(0, IU, B_Arr[:,x,:,:], C_Arr[:,x,:,:], \ D_Arr[:,x,:,:], dY[:,x,:], PBC=PBC_Y) ##CALCULATE DECOMPOSITION ERROR for x in arange(XDim): ##Apply periodic boundary conditions on X if PBC_X == 1: if x == 0: xL = XDim-1 ; xR = 1 elif x == (XDim-1): xL = XDim-2 ; xR = 0 else: xL = x-1 ; xR = x+1 else: xL = x-1 ; xR = x+1 ##Avoid change in y-boundaries for y in arange(YDim-2)+1: dU_Err[y,x,:] = dot(A_Arr[y,x,:,:],dU[y,xL,:])+dot(B_Arr[y,x,:,:],dU[y-1,x,:])+\ dot(C_Arr[y,x,:,:],dU[y,x,:])+\ dot(D_Arr[y,x,:,:],dU[y+1,x,:])+dot(E_Arr[y,x,:,:],dU[y,xR,:]) return dU, dU_Err