import numpy as np
This function performs LU decomposition of given matrix A with pivoting to obtain $$ PA = LU $$ where $P$ is a permutation matrix. Here we dont store $P$ as a matrix, but only as a vector. We initialize $$ P = [0,1,2,\ldots,n-1] $$ and then swap entries of $P$ whenever two rows are swapped during pivoting.
def LU(A):
n = A.shape[0]
L = np.identity(n)
P = np.arange(n,dtype=int) # Permutation matrix
U = np.array(A)
for k in range(n-1):
i = np.argmax(np.abs(U[k:n,k])) + k
U[[k,i],k:n] = U[[i,k],k:n] # swap row i and k
L[[k,i],0:k] = L[[i,k],0:k] # swap row i and k
P[[k,i]] = P[[i,k]] # swap row i and k
for j in range(k+1,n):
L[j,k] = U[j,k]/U[k,k]
U[j,k:n] = U[j,k:n] - L[j,k]*U[k,k:n]
return L,U,P
This performs solution of $$ LUx = Pb $$ in two steps. In first step, solve $$ Ly = Pb $$ In second step, solve $$ Ux = y $$
def LUSolve(L,U,P,b):
n = L.shape[0]
# solve Ly = Pb
pb = b[P]
y = np.empty_like(b)
for i in range(n):
y[i] = (pb[i] - L[i,0:i].dot(y[0:i]))/L[i,i]
#solve Ux = y
x = np.empty_like(b)
for i in range(n-1,-1,-1):
x[i] = (y[i] - U[i,i+1:n].dot(x[i+1:n]))/U[i,i]
return x
Now we test the above function for LU decomposition.
n = 3
A = np.random.rand(n,n)
L,U,P = LU(A)
print("A =\n",A)
print("L =\n",L)
print("U =\n",U)
print("P =\n",P)
# Create a permutation matrix from the P vector
Pm = np.zeros((n,n))
for i in range(n):
Pm[i,P[i]] = 1.0
print("Pm=\n",Pm)
print("PA-LU =\n",Pm.dot(A) - L.dot(U))
A = [[0.32201298 0.7874649 0.11501008] [0.35560336 0.28107935 0.29992026] [0.62139609 0.78687268 0.56771417]] L = [[ 1. 0. 0. ] [ 0.5182089 1. 0. ] [ 0.57226521 -0.44566839 1. ]] U = [[ 0.62139609 0.78687268 0.56771417] [ 0. 0.37970048 -0.17918445] [ 0. 0. -0.10481965]] P = [2 0 1] Pm= [[0. 0. 1.] [1. 0. 0.] [0. 1. 0.]] PA-LU = [[0. 0. 0.] [0. 0. 0.] [0. 0. 0.]]
Solve the linear system.
b = np.random.rand(n)
x = LUSolve(L,U,P,b)
print(A.dot(x) - b)
[ 0.00000000e+00 0.00000000e+00 -1.11022302e-16]