Overload lmul!, rmul! for AbstractQ (#264)

Author: dlfivefifty

src/ArrayLayouts.jl

@@ -17,7 +17,7 @@ import Base.Broadcast: BroadcastStyle, broadcastable, instantiate, materialize,
 
 using LinearAlgebra: AbstractQ, AbstractTriangular, AdjOrTrans, AdjointAbsVec, HermOrSym, HessenbergQ, QRCompactWYQ,
                      QRPackedQ, RealHermSymComplexHerm, TransposeAbsVec, _apply_ipiv_rows!, checknonsingular,
-                     checksquare, chkfullrank, cholcopy, ipiv2perm
+                     checksquare, chkfullrank, cholcopy, ipiv2perm, QRCompactWYQ, AdjointQ, QRPackedQ
 
 using LinearAlgebra.BLAS: BlasComplex, BlasFloat, BlasReal
 
@@ -81,6 +81,10 @@ end
 abstract type LayoutArray{T,N} <: AbstractArray{T,N} end
 const LayoutMatrix{T} = LayoutArray{T,2}
 const LayoutVector{T} = LayoutArray{T,1}
+const LayoutVecOrMat{T} = Union{LayoutVector{T},LayoutMatrix{T}}
+const LayoutMatrices{T} = Union{LayoutMatrix{T}, SubArray{T,2,<:LayoutMatrix}}
+const LayoutVecOrMats{T} = Union{LayoutVecOrMat{T}, SubArray{T,1,<:LayoutVecOrMat}, SubArray{T,2,<:LayoutMatrix}}
+
 
 ## TODO: Following are type piracy which may be removed in Julia v1.5
 ## No, it can't, because strides(::AdjointAbsMat) is defined only for real eltype!
@@ -374,9 +378,6 @@ end
 # printing
 ###
 
-const LayoutVecOrMat{T} = Union{LayoutVector{T},LayoutMatrix{T}}
-const LayoutVecOrMats{T} = Union{LayoutVecOrMat{T},SubArray{T,1,<:LayoutVecOrMat},SubArray{T,2,<:LayoutVecOrMat}}
-
 layout_replace_in_print_matrix(_, A, i, j, s) =
     i in colsupport(A,j) ? s : Base.replace_with_centered_mark(s)
 

src/factorizations.jl

@@ -134,6 +134,8 @@ mulreduce(M::Mul{<:AbstractQLayout}) = Lmul(M)
 mulreduce(M::Mul{<:Any,<:AbstractQLayout}) = Rmul(M)
 mulreduce(M::Mul{<:TriangularLayout,<:AbstractQLayout}) = Rmul(M)
 mulreduce(M::Mul{<:AbstractQLayout,<:TriangularLayout}) = Lmul(M)
+mulreduce(M::Mul{<:DiagonalLayout, <:AbstractQLayout}) = Lmul(M)
+mulreduce(M::Mul{<:AbstractQLayout, <:DiagonalLayout}) = Rmul(M)
 
 function copyto!(dest::AbstractArray{T}, M::Lmul{<:AbstractQLayout}) where T
     A,B = M.A,M.B
@@ -152,8 +154,8 @@ function copyto!(dest::AbstractArray, M::Ldiv{<:AbstractQLayout})
     ldiv!(A,dest)
 end
 
-materialize!(M::Lmul{LAY}) where LAY<:AbstractQLayout = error(LazyString("Overload materialize!(::Lmul{", LAY, "})"))
-materialize!(M::Rmul{LAY}) where LAY<:AbstractQLayout = error(LazyString("Overload materialize!(::Rmul{", LAY, "})"))
+materialize!(M::Lmul{LAY, LAY2}) where {LAY<:AbstractQLayout,LAY2} = error(LazyString("Overload materialize!(::Lmul{", LAY, ", ", LAY2, "})"))
+materialize!(M::Rmul{LAY1, LAY}) where {LAY<:AbstractQLayout,LAY1} = error(LazyString("Overload materialize!(::Rmul{", LAY1, ", ", LAY, "})"))
 
 materialize!(M::Ldiv{<:AbstractQLayout}) = materialize!(Lmul(M.A',M.B))
 

src/lmul.jl

@@ -28,13 +28,13 @@ end
 
 
 const MatLmulVec{StyleA,StyleB} = Lmul{StyleA,StyleB,<:Union{AbstractMatrix,AbstractQ},<:AbstractVector}
-const MatLmulMat{StyleA,StyleB} = Lmul{StyleA,StyleB,<:AbstractMatrix,<:AbstractMatrix}
+const MatLmulMat{StyleA,StyleB} = Lmul{StyleA,StyleB,<:Union{AbstractMatrix,AbstractQ},<:Union{AbstractMatrix,AbstractQ}}
 
 const BlasMatLmulVec{StyleA,StyleB,T<:BlasFloat} = Lmul{StyleA,StyleB,<:Union{AbstractMatrix{T},AbstractQ{T}},<:AbstractVector{T}}
-const BlasMatLmulMat{StyleA,StyleB,T<:BlasFloat} = Lmul{StyleA,StyleB,<:AbstractMatrix{T},<:AbstractMatrix{T}}
+const BlasMatLmulMat{StyleA,StyleB,T<:BlasFloat} = Lmul{StyleA,StyleB,<:Union{AbstractMatrix{T},AbstractQ{T}},<:Union{AbstractMatrix{T},AbstractQ{T}}}
 
-const MatRmulMat{StyleA,StyleB} = Rmul{StyleA,StyleB,<:AbstractMatrix,<:AbstractMatrix}
-const BlasMatRmulMat{StyleA,StyleB,T<:BlasFloat} = Rmul{StyleA,StyleB,<:AbstractMatrix{T},<:AbstractMatrix{T}}
+const MatRmulMat{StyleA,StyleB} = Rmul{StyleA,StyleB,<:Union{AbstractMatrix,AbstractQ},<:Union{AbstractMatrix,AbstractQ}}
+const BlasMatRmulMat{StyleA,StyleB,T<:BlasFloat} = Rmul{StyleA,StyleB,<:Union{AbstractMatrix{T},AbstractQ{T}},<:Union{AbstractMatrix{T},AbstractQ{T}}}
 
 
 ####
@@ -153,3 +153,22 @@ macro layoutrmul(Typ)
         ArrayLayouts.@_layoutrmul UnitLowerTriangular{T, <:Transpose{T,<:$Typ{T}}} where T        
     end)
 end
+
+
+# TODO: uncomment out. commented for now due to ambiguities
+# LinearAlgebra.lmul!(Q::AbstractQ, v::LayoutVecOrMats) = lmul!(Q, v)
+# LinearAlgebra.rmul!(A::LayoutMatrices, Q::AbstractQ) = rmul!(A, Q)
+
+for Typ in (:QRCompactWYQ, :QRPackedQ)
+    @eval begin
+        LinearAlgebra.lmul!(Q::$Typ{<:Any,<:LayoutMatrix}, v::LayoutVecOrMats) = lmul!(Q, v)
+        LinearAlgebra.lmul!(Q::$Typ{<:Any,<:LayoutMatrix}, v::AbstractVecOrMat) = lmul!(Q, v)
+        LinearAlgebra.lmul!(Q::AdjointQ{<:Any,<:$Typ{<:Any,<:LayoutMatrix}}, v::LayoutVecOrMats) = lmul!(Q, v)
+        LinearAlgebra.lmul!(Q::AdjointQ{<:Any,<:$Typ{<:Any,<:LayoutMatrix}}, v::AbstractVecOrMat) = lmul!(Q, v)
+        LinearAlgebra.rmul!(V::LayoutMatrices, Q::$Typ{<:Any,<:LayoutMatrix}) = rmul!(V, Q)
+        LinearAlgebra.rmul!(V::AbstractMatrix, Q::$Typ{<:Any,<:LayoutMatrix}) = rmul!(V, Q)
+        LinearAlgebra.rmul!(V::LayoutMatrices, Q::AdjointQ{<:Any,<:$Typ{<:Any,<:LayoutMatrix}}) = rmul!(V, Q)
+        LinearAlgebra.rmul!(V::AbstractMatrix, Q::AdjointQ{<:Any,<:$Typ{<:Any,<:LayoutMatrix}}) = rmul!(V, Q)
+        LinearAlgebra.rmul!(V::StridedVecOrMat{T}, Q::AdjointQ{<:Any,<:$Typ{T,<:LayoutMatrix}}) where T<:BlasReal = rmul!(V, Q)
+    end
+end

test/test_layoutarray.jl

@@ -602,6 +602,37 @@ Base.copy(A::MyVector) = MyVector(copy(A.A))
             @test mul!(copy(B), D, A, 2.0, 3.0) ≈ mul!(copy(B), D, V, 2.0, 3.0) ≈ 2D * A + 3B
             @test mul!(copy(B), A, D, 2.0, 3.0) ≈ mul!(copy(B), V, D, 2.0, 3.0) ≈ 2A * D + 3B
     end
+
+    @testset "Diagonal * Q" begin
+        A = randn(10,5)
+        Q̃ = qr(A).Q
+        Q = LinearAlgebra.QRCompactWYQ(MyMatrix(Q̃.factors), Q̃.T)
+        D = Diagonal(1:10)
+        @test ArrayLayouts.mul(D, Q) ≈ ArrayLayouts.mul(D, Q̃) ≈ D*Q
+        @test ArrayLayouts.mul(Q, D) ≈ ArrayLayouts.mul(Q̃, D) ≈ Q*D
+    end
+
+    @testset "QR" begin
+        A = randn(10,5)
+        Q̃ = qr(A).Q
+        Q = LinearAlgebra.QRCompactWYQ(MyMatrix(Q̃.factors), Q̃.T)
+        b = randn(10)
+        B = randn(10,3)
+        @test Q*b ≈ Q̃*b ≈ lmul!(Q, copy(b)) ≈ lmul!(Q, MyVector(copy(b)))
+        @test Q'b ≈ Q̃'b ≈ lmul!(Q', copy(b)) ≈ lmul!(Q', MyVector(copy(b)))
+        @test Q*B ≈ Q̃*B ≈ lmul!(Q, copy(B)) ≈ lmul!(Q, MyMatrix(copy(B)))
+        @test Q'B ≈ Q̃'B ≈ lmul!(Q', copy(B)) ≈ lmul!(Q', MyMatrix(copy(B)))
+
+        @test B'Q ≈ B'Q̃ ≈ rmul!(copy(B'),Q) ≈ rmul!(MyMatrix(copy(B')), Q)
+        @test B'Q' ≈ B'Q̃' ≈ rmul!(copy(B'),Q') ≈ rmul!(MyMatrix(copy(B')), Q')
+
+        @test_throws ErrorException lmul!(Q, big.(B))
+        @test_throws ErrorException lmul!(Q', big.(B))
+        @test_throws ErrorException rmul!(big.(B'),Q')
+        @test_throws ErrorException rmul!(big.(B'),Q)
+
+        @test_broken Q*b ≈ lmul!(Q̃, MyVector(copy(b))) # broken due to commented out code for ambiguity with MatrixFactorizations.jl
+    end
 end
 
 struct MyUpperTriangular{T} <: AbstractMatrix{T}