Support lu factorization for strided arrays (#78)

* support lu factorization for strided

* add general ldiv! for LU

* fix tests, default to Eye in UniformScaling addition

* test LU

* zeroeltype to support special cases for vector mul

* fix cholesky tests

* Tests pass on Julia v1.7

* increase coverage

Author: dlfivefifty

Project.toml

@@ -1,7 +1,7 @@
 name = "ArrayLayouts"
 uuid = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
 authors = ["Sheehan Olver <[email protected]>"]
-version = "0.7.8"
+version = "0.7.7"
 
 [deps]
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"

src/ArrayLayouts.jl

@@ -36,7 +36,7 @@ import Base.Broadcast: BroadcastStyle, AbstractArrayStyle, Broadcasted, broadcas
 
 import LinearAlgebra: AbstractTriangular, AbstractQ, checksquare, pinv, fill!, tilebufsize, Abuf, Bbuf, Cbuf, factorize, qr, lu, cholesky,
                         norm2, norm1, normInf, normMinusInf, qr, lu, qr!, lu!, AdjOrTrans, HermOrSym, copy_oftype,
-                        AdjointAbsVec, TransposeAbsVec, cholcopy
+                        AdjointAbsVec, TransposeAbsVec, cholcopy, checknonsingular, _apply_ipiv_rows!, ipiv2perm, RealHermSymComplexHerm, chkfullrank
 
 import LinearAlgebra.BLAS: BlasFloat, BlasReal, BlasComplex
 
@@ -53,6 +53,13 @@ export materialize, materialize!, MulAdd, muladd!, Ldiv, Rdiv, Lmul, Rmul, Dot,
         colsupport, rowsupport, layout_getindex, QLayout, LayoutArray, LayoutMatrix, LayoutVector,
         RangeCumsum
 
+if VERSION < v"1.7-"
+    const ColumnNorm = Val{true}
+    const RowMaximum = Val{true}
+    const NoPivot = Val{false}
+end
+        
+
 struct ApplyBroadcastStyle <: BroadcastStyle end
 @inline function copyto!(dest::AbstractArray, bc::Broadcasted{ApplyBroadcastStyle})
     @assert length(bc.args) == 1

src/factorizations.jl

@@ -15,10 +15,21 @@ factors are stored with layout SLAY and τ stored with layout TLAY
 """
 struct QRPackedLayout{SLAY,TLAY} <: AbstractQRLayout end
 
+
+"""
+    LULayout{SLAY}()
+
+represents a Packed QR factorization whose 
+factors are stored with layout SLAY and τ stored with layout TLAY
+"""
+struct LULayout{SLAY} <: AbstractQRLayout end
+
 MemoryLayout(::Type{<:LinearAlgebra.QRCompactWY{<:Any,MAT}}) where MAT = 
     QRCompactWYLayout{typeof(MemoryLayout(MAT)),DenseColumnMajor}()
 MemoryLayout(::Type{<:LinearAlgebra.QR{<:Any,MAT}}) where MAT = 
     QRPackedLayout{typeof(MemoryLayout(MAT)),DenseColumnMajor}()
+MemoryLayout(::Type{<:LinearAlgebra.LU{<:Any,MAT}}) where MAT = 
+    LULayout{typeof(MemoryLayout(MAT))}()
 
 function materialize!(L::Ldiv{<:QRCompactWYLayout,<:Any,<:Any,<:AbstractVector})
     A,b = L.A, L.B
@@ -32,6 +43,15 @@ function materialize!(L::Ldiv{<:QRCompactWYLayout,<:Any,<:Any,<:AbstractMatrix})
     B
 end
 
+materialize!(L::Ldiv{<:LULayout{<:AbstractColumnMajor},<:AbstractColumnMajor,<:LU{T},<:AbstractVecOrMat{T}}) where {T<:BlasFloat} =
+    LAPACK.getrs!('N', L.A.factors, L.A.ipiv, L.B)
+
+function materialize!(L::Ldiv{<:LULayout})
+    A,B = L.A,L.B
+    _apply_ipiv_rows!(A, B)
+    ldiv!(UpperTriangular(A.factors), ldiv!(UnitLowerTriangular(A.factors), B))
+end
+
 # Julia implementation similar to xgelsy
 function materialize!(Ldv::Ldiv{<:QRPackedLayout,<:Any,<:Any,<:AbstractMatrix{T}}) where T
     A,B = Ldv.A,Ldv.B
@@ -287,16 +307,69 @@ _lu(layout, axes, A, pivot::P; kwds...) where P = Base.invoke(lu, Tuple{Abstract
 _lu!(layout, axes, A, args...; kwds...) = error("Overload _lu!(::$(typeof(layout)), axes, A)")
 _cholesky(layout, axes, A, ::Val{false}=Val(false); check::Bool = true) = cholesky!(cholcopy(A); check = check)
 _cholesky(layout, axes, A, ::Val{true}; tol = 0.0, check::Bool = true) = cholesky!(cholcopy(A), Val(true); tol = tol, check = check)
-_cholesky!(layout, axes, A, v::Val{tf}; kwds...) where tf = Base.invoke(cholesky!, Tuple{LinearAlgebra.RealHermSymComplexHerm,Val{tf}}, A, v; kwds...)
 _factorize(layout, axes, A) = qr(A) # Default to QR
 
+
+_factorize(::AbstractStridedLayout, axes, A) = lu(A)
+function _lu!(::AbstractColumnMajor, axes, A::AbstractMatrix{T}, pivot::Union{NoPivot, RowMaximum} = RowMaximum();
+            check::Bool = true) where T<:BlasFloat
+    if pivot === NoPivot()
+        return generic_lufact!(A, pivot; check = check)
+    end
+    lpt = LAPACK.getrf!(A)
+    check && checknonsingular(lpt[3])
+    return LU{T,typeof(A)}(lpt[1], lpt[2], lpt[3])
+end
+
+# for some reason only defined for StridedMatrix in LinearAlgebra
+function getproperty(F::LU{T,<:LayoutMatrix}, d::Symbol) where T
+    m, n = size(F)
+    if d === :L
+        L = tril!(getfield(F, :factors)[1:m, 1:min(m,n)])
+        for i = 1:min(m,n); L[i,i] = one(T); end
+        return L
+    elseif d === :U
+        return triu!(getfield(F, :factors)[1:min(m,n), 1:n])
+    elseif d === :p
+        return ipiv2perm(getfield(F, :ipiv), m)
+    elseif d === :P
+        return Matrix{T}(I, m, m)[:,invperm(F.p)]
+    else
+        getfield(F, d)
+    end
+end
+
+
 # Cholesky factorization without pivoting (copied from stdlib/LinearAlgebra).
-function _cholesky!(layout, axes, A::LinearAlgebra.RealHermSymComplexHerm, ::Val{false}; check::Bool = true)
-    C, info = LinearAlgebra._chol!(A.data, A.uplo == 'U' ? UpperTriangular : LowerTriangular)
+
+# _chol!. Internal methods for calling unpivoted Cholesky
+## BLAS/LAPACK element types
+function _chol!(::SymmetricLayout{<:AbstractColumnMajor}, A::AbstractMatrix{<:BlasFloat}, ::Type{UpperTriangular})
+    C, info = LAPACK.potrf!('U', A)
+    return UpperTriangular(C), info
+end
+function _chol!(::SymmetricLayout{<:AbstractColumnMajor}, A::AbstractMatrix{<:BlasFloat}, ::Type{LowerTriangular})
+    C, info = LAPACK.potrf!('L', A)
+    return LowerTriangular(C), info
+end
+
+_chol!(_, A, UL) = LinearAlgebra._chol!(A, UL)
+
+function _cholesky!(layout, axes, A::RealHermSymComplexHerm, ::Val{false}; check::Bool = true)
+    C, info = _chol!(layout, A.data, A.uplo == 'U' ? UpperTriangular : LowerTriangular)
     check && LinearAlgebra.checkpositivedefinite(info)
     return Cholesky(C.data, A.uplo, info)
 end
 
+function _cholesky!(::SymmetricLayout{<:AbstractColumnMajor}, axes, A::AbstractMatrix{<:BlasReal},
+    ::Val{true}; tol = 0.0, check::Bool = true)
+    AA, piv, rank, info = LAPACK.pstrf!(A.uplo, A.data, tol)
+    C = CholeskyPivoted{eltype(AA),typeof(AA)}(AA, A.uplo, piv, rank, tol, info)
+    check && chkfullrank(C)
+    return C
+end
+
+
 _inv_eye(_, ::Type{T}, axs::NTuple{2,Base.OneTo{Int}}) where T = Matrix{T}(I, map(length,axs)...)
 function _inv_eye(A, ::Type{T}, (rows,cols)) where T
     dest = zero!(similar(A, T, (cols,rows)))
@@ -318,14 +391,16 @@ end
 macro _layoutfactorizations(Typ)
     esc(quote
         LinearAlgebra.cholesky(A::$Typ, args...; kwds...) = ArrayLayouts._cholesky(ArrayLayouts.MemoryLayout(A), axes(A), A, args...; kwds...)
-        LinearAlgebra.cholesky!(A::$Typ, v::Val{false}=Val(false); check::Bool = true) = ArrayLayouts._cholesky!(ArrayLayouts.MemoryLayout(A), axes(A), A, v; check=check)
+        LinearAlgebra.cholesky!(A::RealHermSymComplexHerm{<:Real,<:$Typ}, v::Val{false}=Val(false); check::Bool = true) = ArrayLayouts._cholesky!(ArrayLayouts.MemoryLayout(A), axes(A), A, v; check=check)
+        LinearAlgebra.cholesky!(A::RealHermSymComplexHerm{<:Real,<:$Typ}, v::Val{true}; check::Bool = true, tol = 0.0) = ArrayLayouts._cholesky!(ArrayLayouts.MemoryLayout(A), axes(A), A, v; check=check, tol=tol)
         LinearAlgebra.qr(A::$Typ, args...; kwds...) = ArrayLayouts._qr(ArrayLayouts.MemoryLayout(A), axes(A), A, args...; kwds...)
         LinearAlgebra.qr!(A::$Typ, args...; kwds...) = ArrayLayouts._qr!(ArrayLayouts.MemoryLayout(A), axes(A), A, args...; kwds...)
-        LinearAlgebra.lu(A::$Typ, pivot::Union{Val{false}, Val{true}}; kwds...) = ArrayLayouts._lu(ArrayLayouts.MemoryLayout(A), axes(A), A, pivot; kwds...)
+        LinearAlgebra.lu(A::$Typ, pivot::Union{NoPivot,RowMaximum}; kwds...) = ArrayLayouts._lu(ArrayLayouts.MemoryLayout(A), axes(A), A, pivot; kwds...)
         LinearAlgebra.lu(A::$Typ{T}; kwds...) where T = ArrayLayouts._lu(ArrayLayouts.MemoryLayout(A), axes(A), A; kwds...)
         LinearAlgebra.lu!(A::$Typ, args...; kwds...) = ArrayLayouts._lu!(ArrayLayouts.MemoryLayout(A), axes(A), A, args...; kwds...)
         LinearAlgebra.factorize(A::$Typ) = ArrayLayouts._factorize(ArrayLayouts.MemoryLayout(A), axes(A), A)
         LinearAlgebra.inv(A::$Typ) = ArrayLayouts._inv(ArrayLayouts.MemoryLayout(A), axes(A), A)
+        LinearAlgebra.ldiv!(L::LU{<:Any,<:$Typ}, B) = ArrayLayouts.ldiv!(L, B)
     end)
 end
 

src/mul.jl

@@ -30,9 +30,13 @@ axes(M::Mul) = _mul_axes(axes(M.A), axes(M.B))
 
 # The following design is to support QuasiArrays.jl where indices
 # may not be `Int`
+
+zeroeltype(M) = zero(eltype(M)) # allow special casing where we know more about zero
+zeroeltype(M::Mul{<:Any,<:Any,<:SubArray}) = zeroeltype(Mul(parent(M.A), M.B))
+
 function _getindex(::Type{Tuple{AA}}, M::Mul, (k,)::Tuple{AA}) where AA
     A,B = M.A, M.B
-    ret = zero(eltype(M))
+    ret = zeroeltype(M)
     for j = rowsupport(A, k) ∩ colsupport(B,1)
         ret += A[k,j] * B[j]
     end
@@ -41,7 +45,7 @@ end
 
 function _getindex(::Type{Tuple{AA,BB}}, M::Mul, (k, j)::Tuple{AA,BB}) where {AA,BB}
     A,B = M.A,M.B
-    ret = zero(eltype(M))
+    ret = zeroeltype(M)
     @inbounds for ℓ in (rowsupport(A,k) ∩ colsupport(B,j))
         ret += A[k,ℓ] * B[ℓ,j]
     end
@@ -292,8 +296,9 @@ LinearAlgebra.dot(x::AbstractVector, A::Symmetric{<:Real,<:LayoutMatrix}, y::Abs
 
 # allow overloading for infinite or lazy case
 @inline _power_by_squaring(_, _, A, p) = Base.invoke(Base.power_by_squaring, Tuple{AbstractMatrix,Integer}, A, p)
-@inline _apply(_, _, op, A::AbstractMatrix, Λ::UniformScaling) = Base.invoke(op, Tuple{AbstractMatrix,UniformScaling}, A, Λ)
-@inline _apply(_, _, op, Λ::UniformScaling, A::AbstractMatrix) = Base.invoke(op, Tuple{UniformScaling,AbstractMatrix}, Λ, A)
+# TODO: Remove unnecessary _apply
+_apply(_, _, op, Λ::UniformScaling, A::AbstractMatrix) = op(Diagonal(Fill(Λ.λ,size(A,1))), A)
+_apply(_, _, op, A::AbstractMatrix, Λ::UniformScaling) = op(A, Diagonal(Fill(Λ.λ,size(A,1))))
 
 for Typ in (:LayoutMatrix, :(Symmetric{<:Any,<:LayoutMatrix}), :(Hermitian{<:Any,<:LayoutMatrix}),
             :(Adjoint{<:Any,<:LayoutMatrix}), :(Transpose{<:Any,<:LayoutMatrix}))

test/test_layoutarray.jl

@@ -1,10 +1,6 @@
 using ArrayLayouts, LinearAlgebra, FillArrays, Base64, Test
-import ArrayLayouts: sub_materialize
+import ArrayLayouts: sub_materialize, MemoryLayout, ColumnNorm, RowMaximum
 
-if VERSION < v"1.7-"
-    ColumnNorm() = Val(true)
-    RowMaximum() = Val(true)
-end
 
 struct MyMatrix <: LayoutMatrix{Float64}
     A::Matrix{Float64}
@@ -16,6 +12,7 @@ Base.size(A::MyMatrix) = size(A.A)
 Base.strides(A::MyMatrix) = strides(A.A)
 Base.unsafe_convert(::Type{Ptr{T}}, A::MyMatrix) where T = Base.unsafe_convert(Ptr{T}, A.A)
 MemoryLayout(::Type{MyMatrix}) = DenseColumnMajor()
+Base.copy(A::MyMatrix) = MyMatrix(copy(A.A))
 
 struct MyVector <: LayoutVector{Float64}
     A::Vector{Float64}
@@ -93,14 +90,24 @@ MemoryLayout(::Type{MyVector}) = DenseColumnMajor()
             @test lu(A).factors ≈ lu(A.A).factors
             @test lu(A,RowMaximum()).factors ≈ lu(A.A,RowMaximum()).factors
             @test_throws ErrorException qr!(A)
-            @test_throws ErrorException lu!(A)
+            @test lu!(copy(A)).factors ≈ lu(A.A).factors
+            b = randn(5)
+            @test all(A \ b .≡ A.A \ b)
+            @test all(lu(A).L .≡ lu(A.A).L)
+            @test all(lu(A).U .≡ lu(A.A).U)
+            @test lu(A).p == lu(A.A).p
+            @test lu(A).P == lu(A.A).P
 
             @test qr(A) isa LinearAlgebra.QRCompactWY
             @test inv(A) ≈ inv(A.A)
 
             S = Symmetric(MyMatrix(reshape(inv.(1:25),5,5) + 10I))
             @test cholesky(S).U ≈ @inferred(cholesky!(deepcopy(S))).U
             @test cholesky(S,Val(true)).U ≈ cholesky(Matrix(S),Val(true)).U
+
+            S = Symmetric(MyMatrix(reshape(inv.(1:25),5,5) + 10I),:L)
+            @test cholesky(S).U ≈ @inferred(cholesky!(deepcopy(S))).U
+            @test cholesky(S,Val(true)).U ≈ cholesky(Matrix(S),Val(true)).U
         end
         Bin = randn(5,5)
         B = MyMatrix(copy(Bin))
@@ -208,7 +215,7 @@ MemoryLayout(::Type{MyVector}) = DenseColumnMajor()
         C = randn(ComplexF64,5,5)
         @test ArrayLayouts.lmul!(2, Hermitian(copy(C))) == ArrayLayouts.rmul!(Hermitian(copy(C)), 2) == 2Hermitian(C)
 
-        
+
         @test ldiv!(2, deepcopy(b)) == rdiv!(deepcopy(b), 2) == 2\b
         @test ldiv!(2, deepcopy(A)) == rdiv!(deepcopy(A), 2) == 2\A
         @test ldiv!(2, deepcopy(A)') == rdiv!(deepcopy(A)', 2) == 2\A'