{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Gaussian Elimination with row pivoting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This function performs LU decomposition of given matrix A with pivoting to obtain\n",
    "$$\n",
    "PA = LU\n",
    "$$\n",
    "where $P$ is a permutation matrix. Here we dont store $P$ as a matrix, but only as a vector. We initialize\n",
    "$$\n",
    "P = [0,1,2,\\ldots,n-1]\n",
    "$$\n",
    "and then swap entries of $P$ whenever two rows are swapped during pivoting."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def LU(A):\n",
    "    n = A.shape[0]\n",
    "    L = np.identity(n)\n",
    "    P = np.arange(n,dtype=int) # Permutation matrix\n",
    "    U = np.array(A)\n",
    "    for k in range(n-1):\n",
    "        i = np.argmax(np.abs(U[k:n,k])) + k\n",
    "        U[[k,i],k:n] = U[[i,k],k:n] # swap row i and k\n",
    "        L[[k,i],0:k] = L[[i,k],0:k] # swap row i and k\n",
    "        P[[k,i]] = P[[i,k]]         # swap row i and k\n",
    "        for j in range(k+1,n):\n",
    "            L[j,k] = U[j,k]/U[k,k]\n",
    "            U[j,k:n] = U[j,k:n] - L[j,k]*U[k,k:n]\n",
    "    return L,U,P"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This performs solution of\n",
    "$$\n",
    "LUx = Pb\n",
    "$$\n",
    "in two steps. In first step, solve\n",
    "$$\n",
    "Ly = Pb\n",
    "$$\n",
    "In second step, solve\n",
    "$$\n",
    "Ux = y\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def LUSolve(L,U,P,b):\n",
    "    n = L.shape[0]\n",
    "    # solve Ly = Pb\n",
    "    pb = b[P]\n",
    "    y = np.empty_like(b)\n",
    "    for i in range(n):\n",
    "        y[i] = (pb[i] - L[i,0:i].dot(y[0:i]))/L[i,i]\n",
    "    #solve Ux = y\n",
    "    x = np.empty_like(b)\n",
    "    for i in range(n-1,-1,-1):\n",
    "        x[i] = (y[i] - U[i,i+1:n].dot(x[i+1:n]))/U[i,i]\n",
    "    return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we test the above function for LU decomposition."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "A =\n",
      " [[0.32201298 0.7874649  0.11501008]\n",
      " [0.35560336 0.28107935 0.29992026]\n",
      " [0.62139609 0.78687268 0.56771417]]\n",
      "L =\n",
      " [[ 1.          0.          0.        ]\n",
      " [ 0.5182089   1.          0.        ]\n",
      " [ 0.57226521 -0.44566839  1.        ]]\n",
      "U =\n",
      " [[ 0.62139609  0.78687268  0.56771417]\n",
      " [ 0.          0.37970048 -0.17918445]\n",
      " [ 0.          0.         -0.10481965]]\n",
      "P =\n",
      " [2 0 1]\n",
      "Pm=\n",
      " [[0. 0. 1.]\n",
      " [1. 0. 0.]\n",
      " [0. 1. 0.]]\n",
      "PA-LU =\n",
      " [[0. 0. 0.]\n",
      " [0. 0. 0.]\n",
      " [0. 0. 0.]]\n"
     ]
    }
   ],
   "source": [
    "n = 3\n",
    "A = np.random.rand(n,n)\n",
    "L,U,P = LU(A)\n",
    "print(\"A =\\n\",A)\n",
    "print(\"L =\\n\",L)\n",
    "print(\"U =\\n\",U)\n",
    "print(\"P =\\n\",P)\n",
    "\n",
    "# Create a permutation matrix from the P vector\n",
    "Pm = np.zeros((n,n))\n",
    "for i in range(n):\n",
    "    Pm[i,P[i]] = 1.0\n",
    "print(\"Pm=\\n\",Pm)\n",
    "print(\"PA-LU =\\n\",Pm.dot(A) - L.dot(U))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Solve the linear system."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ 0.00000000e+00  0.00000000e+00 -1.11022302e-16]\n"
     ]
    }
   ],
   "source": [
    "b = np.random.rand(n)\n",
    "x = LUSolve(L,U,P,b)\n",
    "print(A.dot(x) - b)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}