nla_11_linear_system_of_equations.ipynb

{"nbformat":4,"nbformat_minor":0,"metadata":{"colab":{"provenance":[],"authorship_tag":"ABX9TyPM5dlYnMm/wvaNgWfgDC0F"},"kernelspec":{"name":"python3","display_name":"Python 3"},"language_info":{"name":"python"}},"cells":[{"cell_type":"markdown","source":["#### Linear system of equations"],"metadata":{"id":"s9JRN0lOHziu"}},{"cell_type":"markdown","source":["Linear system of equations solves $b=Ax$ for $x$, $A\\in \\mathbf{R}^{m \\times m}$, `square and full rank`, we can use Cholesky, LDLT, or LU with partial pivoting"],"metadata":{"id":"LPIeN-tCB2ps"}},{"cell_type":"markdown","source":["#### Cholesky"],"metadata":{"id":"4YjSGPFGlzyu"}},{"cell_type":"markdown","source":["For `positive definite` matrices, we take Cholesky factorization $$A=LL^T$$\n","\n","Then\n","\n","* solve $Ly=b$ using `forward` substitution\n","* solve $L^Tx=y$ using `back` substitution"],"metadata":{"id":"-OSbjoRymnvB"}},{"cell_type":"code","source":["import matplotlib.pyplot as plt\n","import numpy as np\n","np.set_printoptions(formatter={'float': '{: 0.4f}'.format})\n","\n","plt.style.use('dark_background')\n","# color: https://matplotlib.org/stable/gallery/color/named_colors.htm"],"metadata":{"id":"64VKzPhM8fWC"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["def cholesky_factorization(A):\n","    m = A.shape[0]\n","    l_mat = A.copy().astype(float)\n","\n","    for k in range(m):\n","        if l_mat[k, k] <= 0:\n","            raise ValueError('Matrix is not positive definite')\n","\n","        # Follow the first step, iteratively apply to a smaller and smaller K\n","        l_mat[k+1:, k+1:] -= np.outer(l_mat[k+1:, k], l_mat[k+1:, k]) / l_mat[k, k]\n","        l_mat[k:, k] /= np.sqrt(l_mat[k, k])\n","\n","    return np.tril(l_mat)\n","\n","def forward_substitution(L, b):\n","    m, n = L.shape\n","    x = np.zeros(n)\n","    for i in range(n):\n","        x[i] = (b[i] - np.dot(L[i, :i], x[:i])) / L[i, i]\n","    return x\n","\n","def back_substitution(R, b):\n","    m, n = R.shape\n","    x = np.zeros(n)\n","    for i in range(n - 1, -1, -1):\n","        x[i] = (b[i] - np.dot(R[i, i + 1:], x[i + 1:])) / R[i, i]\n","    return x"],"metadata":{"id":"s4WzTQpWAaWU"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["np.random.seed(42)\n","m = 50\n","\n","A = np.random.randn(m, m)\n","A_sym = A.T @ A\n","x = np.random.randn(m)\n","b = A_sym @ x\n","b_2 = A @ x\n","print(np.linalg.cond(A))"],"metadata":{"id":"DeCaeG48pzqu","colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"status":"ok","timestamp":1732901095554,"user_tz":300,"elapsed":110,"user":{"displayName":"Jianjing Zhang","userId":"07583447684872889447"}},"outputId":"8826e3b4-1649-4ab3-def1-a7de8b53f7b9"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["41.591235376905374\n"]}]},{"cell_type":"code","source":["try:\n","    L = cholesky_factorization(A_sym)\n","    y_ch = forward_substitution(L, b)\n","    x_ch = back_substitution(L.T, y_ch)\n","    print(np.linalg.norm(x_ch - x))\n","except Exception as e:\n","    print(e)"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"ozB56xYKPJKA","executionInfo":{"status":"ok","timestamp":1732901095554,"user_tz":300,"elapsed":5,"user":{"displayName":"Jianjing Zhang","userId":"07583447684872889447"}},"outputId":"fb15f7d2-5edb-44db-824c-a5e6ec9c7ebf"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["2.443400533449532e-13\n"]}]},{"cell_type":"markdown","source":["#### LDLT"],"metadata":{"id":"UCBVOr9bNgKB"}},{"cell_type":"markdown","source":["For nonsingular `symmetric` matrices, we can use LDLT factorization\n","\n","* $A=LDL^T$\n","* Forward substitution for $Ly=b$\n","* Scaling for $Dz=y$\n","* Back substitution for $L^Tx=z$"],"metadata":{"id":"ryCx9q32NhyQ"}},{"cell_type":"code","source":["def ldlt_factorization(A):\n","    # Assume A is symmetric and invertible\n","    m = A.shape[0]\n","    l_mat = A.copy().astype(float)\n","    d_mat = np.zeros(m)\n","\n","    for k in range(m):\n","        d_mat[k] = l_mat[k, k]\n","        if l_mat[k, k] == 0:\n","            raise ValueError('Matrix is singular')\n","\n","        l_mat[k+1:, k+1:] -= np.outer(l_mat[k+1:, k], l_mat[k+1:, k]) / l_mat[k, k]\n","        l_mat[k:, k] /= l_mat[k, k]\n","\n","    return np.tril(l_mat), np.diag(d_mat)"],"metadata":{"id":"_QSNMyEbN6I9"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["try:\n","    l_mat, d_mat = ldlt_factorization(A_sym)\n","    y_ldlt = forward_substitution(l_mat, b)\n","    z_ldlt = y_ldlt / np.diag(d_mat)\n","    x_ldlt = back_substitution(l_mat.T, z_ldlt)\n","\n","    print(np.linalg.norm(x_ldlt - x))\n","except Exception as e:\n","    print(e)"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"Rqo02LpJPfa6","executionInfo":{"status":"ok","timestamp":1732901095554,"user_tz":300,"elapsed":4,"user":{"displayName":"Jianjing Zhang","userId":"07583447684872889447"}},"outputId":"c5d7396f-e927-45eb-9653-9b98a1d910b9"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["2.7300175676353613e-13\n"]}]},{"cell_type":"markdown","source":["#### Generic matrices possible using Cholesky and LDLT"],"metadata":{"id":"K9KaoeiGQdMr"}},{"cell_type":"markdown","source":["For well-conditioned generic square matrices $A$, using Cholesky and LDLT is possible by solving\n","\n","$$A^TAx=A^Tb$$"],"metadata":{"id":"9kdYzIPHQgg3"}},{"cell_type":"markdown","source":["##### Cholesky"],"metadata":{"id":"TgPs6B0YRhA3"}},{"cell_type":"code","source":["try:\n","    L = cholesky_factorization(A.T @ A)\n","    y_ch = forward_substitution(L, A.T @ b_2)\n","    x_ch = back_substitution(L.T, y_ch)\n","    print(np.linalg.norm(x_ch - x))\n","except Exception as e:\n","    print(e)"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"dovmwBgFQyZ9","executionInfo":{"status":"ok","timestamp":1732901095554,"user_tz":300,"elapsed":3,"user":{"displayName":"Jianjing Zhang","userId":"07583447684872889447"}},"outputId":"565b241f-d300-4f0e-d6f9-79f90ce520cb"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["1.4616675373045904e-13\n"]}]},{"cell_type":"markdown","source":["##### LDLT"],"metadata":{"id":"uyc57ADZRilk"}},{"cell_type":"code","source":["try:\n","    l_mat, d_mat = ldlt_factorization(A.T @ A)\n","    y_ldlt = forward_substitution(l_mat, A.T @ b_2)\n","    z_ldlt = y_ldlt / np.diag(d_mat)\n","    x_ldlt = back_substitution(l_mat.T, z_ldlt)\n","\n","    print(np.linalg.norm(x_ldlt - x))\n","except Exception as e:\n","    print(e)"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"2OyS34mURFPN","executionInfo":{"status":"ok","timestamp":1732901095554,"user_tz":300,"elapsed":3,"user":{"displayName":"Jianjing Zhang","userId":"07583447684872889447"}},"outputId":"ad13c274-da29-47f1-baae-19422f99c79d"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["1.8680752430985427e-13\n"]}]},{"cell_type":"markdown","source":["#### LU with partial pivoting"],"metadata":{"id":"hOK0FNRNZvnv"}},{"cell_type":"markdown","source":["For generic square matrices, we first find $P$, $L$, $U$ such that $PA=LU$\n","\n","Then, we rewrite $Ax=b$ as $LUx=Pb$ and\n","\n","* solve $Ly=Pb$ using `forward` substitution\n","* solve $Ux=y$ using `back` substitution"],"metadata":{"id":"fH-X1-UaZ0fv"}},{"cell_type":"code","source":["def lu_factorization(A):\n","    m, n = A.shape\n","    u_mat = A.copy().astype(float)\n","    l_mat = np.identity(m)\n","    p_mat = np.identity(m)\n","\n","    for k in range(m-1):\n","        # Find pivot\n","        pivot = np.argmax(np.abs(u_mat[k:, k])) + k\n","\n","        if pivot != k:\n","            # Swap rows in u, p, and l\n","            u_mat[[k, pivot], :] = u_mat[[pivot, k], :]\n","            p_mat[[k, pivot], :] = p_mat[[pivot, k], :]\n","            l_mat[[k, pivot], :k] = l_mat[[pivot, k], :k]\n","\n","        for j in range(k + 1, m):\n","            l_mat[j, k] = u_mat[j, k] / u_mat[k, k]\n","            # Subtract multiply of kth row from jth row\n","            u_mat[j, k:] -= l_mat[j, k] * u_mat[k, k:]\n","\n","    return p_mat, l_mat, u_mat"],"metadata":{"id":"gd5ecpksZnBu"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["p, l, u = lu_factorization(A)\n","y_lu = forward_substitution(l, p @ b_2)\n","x_lu = back_substitution(u, y_lu)\n","\n","print(np.linalg.norm(x_lu - x))"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"IyfQVQGyaWNR","executionInfo":{"status":"ok","timestamp":1732901095654,"user_tz":300,"elapsed":102,"user":{"displayName":"Jianjing Zhang","userId":"07583447684872889447"}},"outputId":"821e521b-016d-46b8-8e22-82a9150c6d57"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["2.5865046977186647e-14\n"]}]},{"cell_type":"markdown","source":["We see that when $A$ is well-conditioned, all is good!"],"metadata":{"id":"oI5SkrSncTwg"}},{"cell_type":"markdown","source":["#### When the matrix is ill-conditioned"],"metadata":{"id":"05tYGO45RQ8r"}},{"cell_type":"markdown","source":["Computing $A^TA$ can amplify the ill-conditionedness with Cholesky and LDLT"],"metadata":{"id":"-sHJqCRkSOiT"}},{"cell_type":"code","source":["np.random.seed(42)\n","\n","A_ill = np.random.randn(m, m)\n","A_ill[:, 0] = A_ill[:, -1] + 1e-8 * np.random.randn(m)\n","x = np.random.randn(m)\n","b_3 = A_ill @ x\n","print(np.linalg.cond(A_ill))"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"rsJtWaGCRTQK","executionInfo":{"status":"ok","timestamp":1732901095654,"user_tz":300,"elapsed":102,"user":{"displayName":"Jianjing Zhang","userId":"07583447684872889447"}},"outputId":"9774e21c-986c-4c96-c717-8c75a6fd7e3c"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["856008450.4491345\n"]}]},{"cell_type":"markdown","source":["##### Cholesky"],"metadata":{"id":"UZ6Yw3ancazd"}},{"cell_type":"code","source":["try:\n","    L = cholesky_factorization(A_ill.T @ A_ill)\n","    y_ch = forward_substitution(L, A_ill.T @ b_3)\n","    x_ch = back_substitution(L.T, y_ch)\n","    print(np.linalg.norm(x_ch - x))\n","except Exception as e:\n","    print(e)"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"YwTX3VsyRfwu","executionInfo":{"status":"ok","timestamp":1732901095654,"user_tz":300,"elapsed":4,"user":{"displayName":"Jianjing Zhang","userId":"07583447684872889447"}},"outputId":"9be2dfb2-a2a3-4939-a495-b30a0e5711ae"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["12.967923732845527\n"]}]},{"cell_type":"markdown","source":["##### LDLT"],"metadata":{"id":"pb0TdPwncb_T"}},{"cell_type":"code","source":["try:\n","    l_mat, d_mat = ldlt_factorization(A_ill.T @ A_ill)\n","    y_ldlt = forward_substitution(l_mat, A_ill.T @ b_3)\n","    z_ldlt = y_ldlt / np.diag(d_mat)\n","    x_ldlt = back_substitution(l_mat.T, z_ldlt)\n","    print(np.linalg.norm(x_ldlt - x))\n","except Exception as e:\n","    print(e)"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"QxcNaeduR30k","executionInfo":{"status":"ok","timestamp":1732901095654,"user_tz":300,"elapsed":3,"user":{"displayName":"Jianjing Zhang","userId":"07583447684872889447"}},"outputId":"b7c2147d-fb90-4364-ff91-ebc009b956fd"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["26.631502576582093\n"]}]},{"cell_type":"markdown","source":["##### LU with partial pivoting"],"metadata":{"id":"6LRFOk7TcdXJ"}},{"cell_type":"code","source":["p, l, u = lu_factorization(A_ill)\n","y_lu = forward_substitution(l, p @ b_3)\n","x_lu = back_substitution(u, y_lu)\n","\n","print(np.linalg.norm(x_lu - x))"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"gSfxCjB7R7Xi","executionInfo":{"status":"ok","timestamp":1732901095654,"user_tz":300,"elapsed":2,"user":{"displayName":"Jianjing Zhang","userId":"07583447684872889447"}},"outputId":"03c9d156-0b2c-4336-9a9e-364471628d3c"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["7.86444716582733e-08\n"]}]},{"cell_type":"markdown","source":["#### On stability of LU with partial pivoting"],"metadata":{"id":"ixmzxoD14Ng0"}},{"cell_type":"markdown","source":["For LU with partial pivoting, a key metric for stability is whether amplification of the entries takes place during decomposition, or the growth factor\n","\n","$$\\rho = \\frac{\\max |u_{ij}|}{\\max |a_{ij}|}$$\n","\n","and we have the following theorem\n","\n","Let $\\tilde{P}, \\tilde{L}, \\tilde{U}$ be computed using LU with partial pivoting, then, they satisfy\n","\n","$$\\tilde{L}\\tilde{U}=\\tilde{P}A+\\delta A, \\frac{\\|\\delta A\\|}{\\|A\\|}=O(\\rho \\epsilon_m)$$\n","\n","We generally want the growth factor to be in the order of 1"],"metadata":{"id":"ITiNONPh7DfZ"}},{"cell_type":"markdown","source":["Below is an example from NLA book where the growth factor is exponential $\\rho=2^{m-1}$, where $m$ is the size of $A$"],"metadata":{"id":"t-l-3uST8dFA"}},{"cell_type":"code","source":["# Example from NLA book\n","A = np.array([[1., 0, 0, 0, 1], [-1, 1, 0, 0, 1], [-1, -1, 1, 0, 1], [-1, -1, -1, 1, 1], [-1, -1, -1, -1, 1]])\n","\n","p_mat, l_mat, u_mat = lu_factorization(A)\n","\n","print('A:\\n', A)\n","print('P:\\n', p_mat)\n","print('L:\\n', l_mat)\n","print('U:\\n', u_mat)\n","\n","print('PA:\\n', p_mat @ A)\n","print('LU:\\n', l_mat @ u_mat)"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"PTiw4kJi5OWD","executionInfo":{"status":"ok","timestamp":1732901095654,"user_tz":300,"elapsed":2,"user":{"displayName":"Jianjing Zhang","userId":"07583447684872889447"}},"outputId":"895282d6-7846-4d16-affc-4bc89ac3b2af"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["A:\n"," [[ 1.0000  0.0000  0.0000  0.0000  1.0000]\n"," [-1.0000  1.0000  0.0000  0.0000  1.0000]\n"," [-1.0000 -1.0000  1.0000  0.0000  1.0000]\n"," [-1.0000 -1.0000 -1.0000  1.0000  1.0000]\n"," [-1.0000 -1.0000 -1.0000 -1.0000  1.0000]]\n","P:\n"," [[ 1.0000  0.0000  0.0000  0.0000  0.0000]\n"," [ 0.0000  1.0000  0.0000  0.0000  0.0000]\n"," [ 0.0000  0.0000  1.0000  0.0000  0.0000]\n"," [ 0.0000  0.0000  0.0000  1.0000  0.0000]\n"," [ 0.0000  0.0000  0.0000  0.0000  1.0000]]\n","L:\n"," [[ 1.0000  0.0000  0.0000  0.0000  0.0000]\n"," [-1.0000  1.0000  0.0000  0.0000  0.0000]\n"," [-1.0000 -1.0000  1.0000  0.0000  0.0000]\n"," [-1.0000 -1.0000 -1.0000  1.0000  0.0000]\n"," [-1.0000 -1.0000 -1.0000 -1.0000  1.0000]]\n","U:\n"," [[ 1.0000  0.0000  0.0000  0.0000  1.0000]\n"," [ 0.0000  1.0000  0.0000  0.0000  2.0000]\n"," [ 0.0000  0.0000  1.0000  0.0000  4.0000]\n"," [ 0.0000  0.0000  0.0000  1.0000  8.0000]\n"," [ 0.0000  0.0000  0.0000  0.0000  16.0000]]\n","PA:\n"," [[ 1.0000  0.0000  0.0000  0.0000  1.0000]\n"," [-1.0000  1.0000  0.0000  0.0000  1.0000]\n"," [-1.0000 -1.0000  1.0000  0.0000  1.0000]\n"," [-1.0000 -1.0000 -1.0000  1.0000  1.0000]\n"," [-1.0000 -1.0000 -1.0000 -1.0000  1.0000]]\n","LU:\n"," [[ 1.0000  0.0000  0.0000  0.0000  1.0000]\n"," [-1.0000  1.0000  0.0000  0.0000  1.0000]\n"," [-1.0000 -1.0000  1.0000  0.0000  1.0000]\n"," [-1.0000 -1.0000 -1.0000  1.0000  1.0000]\n"," [-1.0000 -1.0000 -1.0000 -1.0000  1.0000]]\n"]}]},{"cell_type":"markdown","source":["However, from NLA book\n","\n","\n","\n","> Gaussian elimination with partial pivoting is highly unstable for certain matrices. For instability to occur, however, the column spaces of $A$ must be `skewed in a very special fashion`, one that is exponentially rare in at least one class of random matrices. Decades of computational experience suggest that matrices whose column spaces are skewed in this fashion arise `very rarely` in applications.\n","\n"],"metadata":{"id":"A61VWlU04XzE"}},{"cell_type":"markdown","source":["#### Equations with structured sub-blocks"],"metadata":{"id":"sNal4NPqLte4"}},{"cell_type":"markdown","source":["$$\\begin{bmatrix}A_{11} & A_{12} \\\\ A_{21} & A_{22}\\end{bmatrix}\\begin{bmatrix}x_1  \\\\ x_2\\end{bmatrix}=\\begin{bmatrix}b_1  \\\\ b_2\\end{bmatrix}$$\n","\n","If $A_{11}$ has some structure (e.g., it is easy to compute its inverse), we can solve $x_2$ first\n","\n","* Compute $A_{11}^{-1}$\n","* Form $A_{11}^{-1}A_{12}$ and $A_{11}^{-1}b$, plug $x_1=A_{11}^{-1}(b_1-A_{12}x_2)$ into second block equation\n","* Solve $(A_{22}-A_{21}A_{11}^{-1}A_{12})x_2=b_2-A_{21}A_{11}^{-1}b_1$ for $x_2$\n","* Solve $x_1$ from $x_1=A_{11}^{-1}(b_1-A_{12}x_2)$ for $x_1$\n","\n","This is known as `block elimination`"],"metadata":{"id":"doCdjkK0LwTl"}},{"cell_type":"code","source":[],"metadata":{"id":"Q0T7LD-HR-pD"},"execution_count":null,"outputs":[]}]}