# Math 551: uMass Amherst Fall 2022
%load_ext autoreload
%autoreload 2

import os
import numpy as np
import scipy as sp
import scipy.sparse as sps
import scipy.linalg as LA
import scipy.sparse.linalg as LAS
import matplotlib.pyplot as plt

# For output
homedir = os.path.expanduser('~')

The autoreload extension is already loaded. To reload it, use:
  %reload_ext autoreload

LUQ decomposition¶

from scipy.sparse import csr_array

def luq(A, tol=1e-8):
    """
        Computes LUQ decomposition that is
        
        A = L U Q 
        
        where 
            - L is square lower triangular
            - U 
            - Q 
    """
    n, m = A.shape
    
    # Special casese
    if n == 0:
        L = sps.eye(n)
        U = A
        Q = sps.eye(m)
        return L, U, Q
    
    if m == 0:
        L = sps.eye(n)
        U = A
        Q = sps.eye(m)
        return L, U, Q
    
    # Call LU decomposition
    P, L, U = LA.lu(A)
    Q = sps.eye(m).tocsr()
    
    # Assemble some stuff
    p  = n - L.shape[1]
    if p > 0:
        L  = sps.hstack([L, P[n-p+1:n, :].T])
        U  = sps.vstack([U, csr_array((p, m))])
    
    # Find rows with zero and non-zero elements on the diagonal
    I     = np.where(np.abs(np.diag(U))>tol)[0]
    Jl    = np.delete(np.arange(n), I)
    Jq    = np.delete(np.arange(m), I)
    
    Ubar1 = U[np.ix_(I, I)]
    Ubar2 = U[np.ix_(Jl, Jq)]
    Qbar1 = Q[I, :]
    Lbar1 = L[:, I]
    
    # Eliminates non-zero elements below and on the right
    # of the invertible block of the matrix U
    if I.size:
        Utmp  = U[np.ix_(I, Jq)]
        X     = LA.solve_triangular(Ubar1.T, U[np.ix_(Jl, I)].T)
        Ubar2 = Ubar2 - X.T @ Utmp
        Lbar1 = Lbar1 + L[:, Jl] @ X.T
        X     = LA.solve_triangular(Ubar1, Utmp)
        Qbar1 = Qbar1 + X @ Q[Jq, :]
        Utmp  = np.zeros((0, 0))
        X     = np.zeros((0, 0))
    
    # If Ubar2 is empty we are done!
    if Ubar2.shape[0] == 0 and Ubar2.shape[1] == 0:
        return L, U, Q
    
    # Find rows and cols with only zero elements
    I2 = np.where(np.max(np.abs(Ubar2), axis=1)>tol)[0]
    I5 = np.where(np.max(np.abs(Ubar2), axis=0)>tol)[0]
    
    # These are the cols / rows with zero diagonal but some non-zero elements
    I3 = Jl[I2]
    I4 = Jq[I5]
    
    # These are the row/cols with zero diagonal only zero elements
    Jq = np.delete(Jq, I5)
    Jl = np.delete(Jl, I2)
    U  = np.zeros((0, 0)) 
    
    # Find a part of the matrix U which is not in the required form
    A = Ubar2[np.ix_(I2, I5)]
    
    # Performs LUQ decomposition of the matrix A
    L1, U1, Q1 = luq(A, tol=tol)
    
    # Update matrices L, U, Q
    Lbar2 = L[:, I3] @ L1
    Qbar2 = Q1 @ Q[I4, :]
    L     = np.hstack([Lbar1, Lbar2, L[:, Jl]])
    Q     = sps.vstack([Qbar1, Qbar2, Q[Jq, :]])
    
    # Assemble U
    n1    = I.size
    n2    = I3.size
    m2    = I4.size
    U     = sps.vstack([sps.hstack([Ubar1, csr_array((n1, m - n1))]),
                        sps.hstack([csr_array((n2, n1)), U1, csr_array((n2, m-n1-m2))]),
                        sps.hstack([csr_array((n-n1-n2, m))])])
    
    return L, U.todense(), Q.todense(), P

A = np.array([[2, 1, 3], 
              [4, 2, 1],
              [6, 3, 9]])

L, U, Q, P = luq(A)

print(L)
print(U)
print(Q)

print(P @ L @ U @ Q)

[[1.         0.         0.        ]
 [0.66666667 1.         0.        ]
 [0.33333333 0.         1.        ]]
[[ 6.  0.  0.]
 [ 0. -5.  0.]
 [ 0.  0.  0.]]
[[1.  0.5 1.5]
 [0.  0.  1. ]
 [0.  1.  0. ]]
[[2. 1. 3.]
 [4. 2. 1.]
 [6. 3. 9.]]

A = np.ones((2, 2))
A[1, :] *= 2
print(A)
L, U, Q, P = luq(A)
print(L)
print(U)
print(Q)

[[1. 1.]
 [2. 2.]]
[[1.  0. ]
 [0.5 1. ]]
[[2. 0.]
 [0. 0.]]
[[1. 1.]
 [0. 1.]]

P @ L @ U @ Q

matrix([[1., 1.],
        [2., 2.]])

A = np.array([[2, 1, 3], 
              [4, 2, 1],
              [6, 3, 9]])

np.argmax(A, axis=1)

array([2, 0, 2])

np.argmax(A, axis=0)

array([2, 2, 2])