#!/usr/bin/env python
# coding: utf-8

# # Gaussian Elimination

# In[7]:


import numpy as np


# This function performs LU decomposition of given matrix A.

# In[8]:


def LU(A):
    n = A.shape[0]
    L = np.identity(n)
    U = np.array(A)
    for i in range(n-1):
        L[i+1:n,i] = U[i+1:n,i]/U[i,i]
        for j in range(i+1,n):
            U[j,i+1:n] = U[j,i+1:n] - L[j,i] * U[i,i+1:n]
        U[i+1:n,i] = 0.0
    return L,U


# This performs performs solution of
# $$
# LUx=b
# $$
# in two steps. In first step, solve
# $$
# Ly = b
# $$
# using
# $$
# y_i = \frac{1}{L_{ii}} \left[b_i - \sum_{j=0}^{i-1} L_{ij} y_j\right], \qquad i=0,1,\ldots,n-1
# $$
# Note that $L_{ii}=1$, so we can drop the division step. In second step, solve
# $$
# Ux = y
# $$
# using
# $$
# x_i = \frac{1}{U_{ii}}\left[y_i - \sum_{j=i+1}^{n-1} U_{ij} x_j \right], \qquad i=n-1,n-2,\ldots,0
# $$

# In[9]:


def LUSolve(L,U,b):
    n = L.shape[0]
    # solve Ly = b
    y = np.empty_like(b)
    for i in range(n):
        y[i] = b[i] - L[i,0:i].dot(y[0: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. We initialize $A$ to a random matrix and compute its LU decomposition.

# In[10]:


n = 3
A = np.random.rand(n,n)
L,U = LU(A)
print("A =\n",A)
print("L =\n",L)
print("U =\n",U)
print(A - L.dot(U))


# Solve the linear system.

# In[11]:


b = np.random.rand(n)
x = LUSolve(L,U,b)
print(A.dot(x) - b)