Increase support for padded QR (#381)

* Increase support for padded QR

* Move over BlockVec

* Move over blockvec

* Work around padded data being a Blocked Array

* Update padded.jl

* increase coverage

* Update paddedtests.jl

* Update paddedtests.jl

* increase coverage

* Update padded.jl

Author: dlfivefifty

Project.toml

@@ -1,6 +1,6 @@
 name = "LazyArrays"
 uuid = "5078a376-72f3-5289-bfd5-ec5146d43c02"
-version = "2.7"
+version = "2.8"
 
 [deps]
 ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
@@ -23,7 +23,7 @@ LazyArraysStaticArraysExt = "StaticArrays"
 
 [compat]
 Aqua = "0.8"
-ArrayLayouts = "1.10.1"
+ArrayLayouts = "1.12"
 BandedMatrices = "1.6.2"
 Base64 = "1"
 BlockArrays = "1.0"

ext/LazyArraysBlockArraysExt.jl

@@ -4,13 +4,14 @@ using BlockArrays
 using LazyArrays
 using LazyArrays.ArrayLayouts
 using LazyArrays.FillArrays
+using LazyArrays.LinearAlgebra
 import LazyArrays: resizedata!, paddeddata, paddeddata_axes, arguments, call,
                     LazyArrayStyle, CachedVector, AbstractPaddedLayout, PaddedLayout, PaddedRows, PaddedColumns, BroadcastLayout,
                     AbstractCachedMatrix, AbstractCachedArray, setindex, applybroadcaststyle,
-                    ApplyLayout, cache_layout
+                    ApplyLayout, cache_layout, applied_eltype, applylayout, applied_ndims, broadcast_deblock
 import ArrayLayouts: sub_materialize
-import Base: getindex, BroadcastStyle, broadcasted, OneTo
-import BlockArrays: AbstractBlockStyle, AbstractBlockedUnitRange, blockcolsupport, blockrowsupport, BlockSlice, BlockIndexRange, AbstractBlockLayout
+import Base: getindex, setindex!, BroadcastStyle, broadcasted, OneTo, axes, size, view, resize!
+import BlockArrays: AbstractBlockStyle, AbstractBlockedUnitRange, blockcolsupport, blockrowsupport, BlockSlice, BlockIndexRange, AbstractBlockLayout, blockvec
 
 BlockArrays._broadcaststyle(S::LazyArrays.LazyArrayStyle{1}) = S
 
@@ -233,5 +234,33 @@ function arguments(lay::ApplyLayout{typeof(hcat)}, V::SubArray{<:Any,2,<:Any,<:T
 end
 
 
+####
+# BlockVec
+####
+const BlockVec{T, M<:AbstractMatrix{T}} = ApplyVector{T, typeof(blockvec), <:Tuple{M}}
+
+BlockVec{T}(M::AbstractMatrix{T}) where T = ApplyVector{T}(blockvec, M)
+BlockVec(M::AbstractMatrix{T}) where T = BlockVec{T}(M)
+axes(b::BlockVec) = (blockedrange(Fill(size(b.args[1])...)),)
+size(b::BlockVec) = (length(b.args[1]),)
+applied_eltype(::typeof(blockvec), A) = eltype(A)
+applied_ndims(::typeof(blockvec), A) = 1
+
+view(b::BlockVec, K::Block{1}) = view(b.args[1], :, Int(K))
+Base.@propagate_inbounds getindex(b::BlockVec, k::Int) = b.args[1][k]
+Base.@propagate_inbounds setindex!(b::BlockVec, v, k::Int) = setindex!(b.args[1], v, k)
+
+_resize!(A::AbstractMatrix, m, n) = A[1:m, 1:n]
+_resize!(At::Transpose, m, n) = transpose(transpose(At)[1:n, 1:m])
+_resize!(Ac::Adjoint, m, n) = (Ac')[1:n, 1:m]'
+resize!(b::BlockVec, K::Block{1}) = BlockVec(_resize!(b.args[1], size(b.args[1],1), Int(K)))
+
+applylayout(::Type{typeof(blockvec)}, ::AbstractPaddedLayout) = PaddedColumns{ApplyLayout{typeof(blockvec)}}()
+paddeddata(b::BlockVec) = BlockVec(paddeddata(b.args[1]))
+
+# work around case where padded data is blocked
+broadcast_deblock(op, A, B::BlockedArray) = broadcast(op, A, B.blocks)
+broadcast_deblock(op, A::BlockedArray, B) = broadcast(op, A.blocks, B)
+broadcast_deblock(op, A::BlockedArray, B::BlockedArray) = broadcast(op, A.blocks, B.blocks)
 
 end

src/LazyArrays.jl

@@ -40,7 +40,7 @@ import ArrayLayouts: AbstractQLayout, Dot, Dotu, Ldiv, Lmul, MatMulMatAdd, MatMu
                      sublayout, symmetriclayout, symtridiagonallayout, transposelayout, triangulardata,
                      triangularlayout, tridiagonallayout, zero!, transtype, OnesLayout,
                      diagonaldata, subdiagonaldata, supdiagonaldata, MemoryLayout, MatLmulVec, MatLmulMat,
-                     AdjQRCompactWYQLayout
+                     AdjQRCompactWYQLayout, QRCompactWYQLayout, MatLmulMat, MatRmulMat
 
 import FillArrays: AbstractFill, getindex_value
 

src/linalg/mul.jl

@@ -162,6 +162,7 @@ _mul(A,B,C...) = lazymaterialize(*,A,B,C...)
 _mul_colsupport(j) = j
 _mul_colsupport(j, Z::Number, Y...) = _mul_colsupport(j, Y...) # scalar mul doesn't do anything
 _mul_colsupport(j, Z::AbstractArray, Y...) = _mul_colsupport(colsupport(Z,j), Y...)
+_mul_colsupport(j, Z::AbstractQ, Y...) = _mul_colsupport(colsupport(Z,j), Y...)
 
 colsupport(B::Applied{<:Any,typeof(*)}, j) = _mul_colsupport(j, reverse(B.args)...)
 colsupport(B::MulArray, j) = _mul_colsupport(j, reverse(B.args)...)
@@ -247,6 +248,7 @@ _vec_mul_arguments(args, (kr,jr)::Tuple{AbstractVector,Number}) =
 
 # this is a row-vector view
 _transposeifnumber(a::AbstractArray{<:Number}) = transpose(a)
+_transposeifnumber(a::AbstractQ{<:Number}) = transpose(a)
 _transposeifnumber(a) = permutedims(a)
 
 _vec_mul_arguments(args, (kr,jr)::Tuple{Number,AbstractVector}) =

src/padded.jl

@@ -81,13 +81,15 @@ _hvcat_paddeddata(N, A, B::Zeros...) = A
 paddeddata(A::ApplyMatrix{<:Any,typeof(hvcat)}) = _hvcat_paddeddata(A.args...)
 
 
+paddeddata(A::Array) = A # support padded interface for strided arrays
 
 
 #####
 # conversion
 #####
 
 cache_layout(::AbstractPaddedLayout, O::AbstractArray) = CachedArray(copy(paddeddata(O)), Zeros{eltype(O)}(axes(O)))
+cache_layout(::AbstractPaddedLayout, O::CachedArray) = CachedArray(copy(O.data), O.array, O.datasize)
 
 
 #####
@@ -252,48 +254,51 @@ end
 
 makearray(a::Number, ::Tuple{Any}) = [a]
 makearray(a::Number, ::Tuple{Any,Any}) = [a ;;]
+makearray(a::AbstractVector, ::Tuple{Any,Any}) = reshape(a, :, 1)
 makearray(a, _) = convert(Array, a)
 
+broadcast_deblock(op, A, B) = broadcast(op, A, B)
+
 function layout_broadcasted(_, ::AbstractPaddedLayout, op, A::Union{Ref,Number}, B::AbstractArray)
     b = paddeddata(B)
     m = length(b)
     zB = Zeros{eltype(B)}(size(B)...)
-    CachedArray(makearray(broadcast(op, A, b), size(B)), broadcast(op, A, zB))
+    CachedArray(makearray(broadcast_deblock(op, A, b), size(B)), broadcast(op, A, zB))
 end
 
 function layout_broadcasted(::AbstractPaddedLayout, _, op, A::AbstractArray, B::Union{Ref,Number})
     a = paddeddata(A)
     n = length(a)
     zA = Zeros{eltype(A)}(size(A)...)
-    CachedArray(makearray(broadcast(op, a, B), size(A)), broadcast(op, zA, B))
+    CachedArray(makearray(broadcast_deblock(op, a, B), size(A)), broadcast(op, zA, B))
 end
 
 function layout_broadcasted(_, ::AbstractPaddedLayout, op, A::AbstractVector, B::AbstractVector)
     b = paddeddata(B)
     m = length(b)
     zB = Zeros{eltype(B)}(size(B)...)
-    CachedArray(makearray(broadcast(op, view(A,1:m), b), size(B)), broadcast(op, A, zB))
+    CachedArray(makearray(broadcast_deblock(op, view(A,1:m), b), size(B)), broadcast(op, A, zB))
 end
 
 function layout_broadcasted(::AbstractPaddedLayout, _, op, A::AbstractVector, B::AbstractVector)
     a = paddeddata(A)
     n = length(a)
     zA = Zeros{eltype(A)}(size(A)...)
-    CachedArray(makearray(broadcast(op, a, view(B,1:n)), size(A)), broadcast(op, zA, B))
+    CachedArray(makearray(broadcast_deblock(op, a, view(B,1:n)), size(A)), broadcast(op, zA, B))
 end
 
 function layout_broadcasted(_, ::AbstractPaddedLayout, op, A::AbstractVector, B::AbstractMatrix)
     b = paddeddata(B)
     m = size(b,1)
     zB = Zeros{eltype(B)}(size(B)...)
-    CachedArray(makearray(broadcast(op, view(A,1:m), b), size(B)), broadcast(op, A, zB))
+    CachedArray(makearray(broadcast_deblock(op, view(A,1:(size(B,1) == 1 ? size(A,1) : m)), b), size(B)), broadcast(op, A, zB))
 end
 
 function layout_broadcasted(::AbstractPaddedLayout, _, op, A::AbstractMatrix, B::AbstractVector)
     a = paddeddata(A)
     n = size(a,1)
     zA = Zeros{eltype(A)}(size(A)...)
-    CachedArray(makearray(broadcast(op, a, view(B,1:n)), size(A)), broadcast(op, zA, B))
+    CachedArray(makearray(broadcast_deblock(op, a, view(B,1:(size(A,1) == 1 ? size(B,1) : n))), size(A)), broadcast(op, zA, B))
 end
 
 
@@ -724,10 +729,15 @@ function ArrayLayouts.qr!_layout(::AbstractPaddedLayout, ax, A)
     padqr(F, ax)
 end
 
+colsupport(::QRCompactWYQLayout{<:AbstractPaddedLayout}, Q, k) = colsupport(paddeddata(Q.factors), k)
+rowsupport(::QRCompactWYQLayout{<:AbstractPaddedLayout}, Q, k) = colsupport(paddeddata(Q.factors), k)
+colsupport(::AdjQRCompactWYQLayout{<:AbstractPaddedLayout}, Q, k) = colsupport(paddeddata(Q'.factors), k)
+rowsupport(::AdjQRCompactWYQLayout{<:AbstractPaddedLayout}, Q, k) = colsupport(paddeddata(Q'.factors), k)
+
 similar(M::Lmul{<:AdjQRCompactWYQLayout{<:PaddedColumns}, <:PaddedColumns}, ::Type{T}, ax) where T = CachedArray(Zeros{T}(ax))
 
 
-function materialize!(L::MatLmulVec{<:AdjQRCompactWYQLayout{<:PaddedColumns}, <:PaddedColumns})
+function materialize!(L::MatLmulVec{<:AdjQRCompactWYQLayout{<:PaddedColumns}, <:Union{AbstractStridedLayout,PaddedColumns}})
     Q,b = L.A',L.B
     F = paddeddata(Q.factors)
     m = size(F,1)
@@ -737,3 +747,49 @@ function materialize!(L::MatLmulVec{<:AdjQRCompactWYQLayout{<:PaddedColumns}, <:
     lmul!(Q̃', view(p_b,oneto(m)))
     b
 end
+
+function materialize!(L::MatLmulMat{<:AdjQRCompactWYQLayout{<:PaddedColumns}, <:Union{AbstractStridedLayout,AbstractPaddedLayout}})
+    Q,b = L.A',L.B
+    F = paddeddata(Q.factors)
+    m = size(F,1)
+    resizedata!(b, m, size(b,2))
+    Q̃ = LinearAlgebra.QRCompactWYQ(F, Q.T)
+    p_b = paddeddata(b)
+    lmul!(Q̃', view(p_b,oneto(m),:))
+    b
+end
+
+
+
+function materialize!(L::MatLmulVec{<:QRCompactWYQLayout{<:PaddedColumns}, <:Union{AbstractStridedLayout,PaddedColumns}})
+    Q,b = L.A,L.B
+    F = paddeddata(Q.factors)
+    m = size(F,1)
+    resizedata!(b, m)
+    Q̃ = LinearAlgebra.QRCompactWYQ(F, Q.T)
+    p_b = paddeddata(b)
+    lmul!(Q̃, view(p_b,oneto(m)))
+    b
+end
+
+function materialize!(L::MatLmulMat{<:QRCompactWYQLayout{<:PaddedColumns}, <:Union{AbstractStridedLayout,AbstractPaddedLayout}})
+    Q,b = L.A,L.B
+    F = paddeddata(Q.factors)
+    m = size(F,1)
+    resizedata!(b, m, size(b,2))
+    Q̃ = LinearAlgebra.QRCompactWYQ(F, Q.T)
+    p_b = paddeddata(b)
+    lmul!(Q̃, view(p_b,oneto(m),:))
+    b
+end
+
+function materialize!(L::MatRmulMat{<:Union{AbstractStridedLayout,AbstractPaddedLayout}, <:QRCompactWYQLayout{<:PaddedColumns}})
+    A,Q = L.A,L.B
+    F = paddeddata(Q.factors)
+    m = size(F,1)
+    resizedata!(A, size(A,1), m)
+    Q̃ = LinearAlgebra.QRCompactWYQ(F, Q.T)
+    p_A = paddeddata(A)
+    rmul!(view(p_A,:, oneto(m)), Q̃)
+    A
+end

test/blocktests.jl

@@ -1,7 +1,8 @@
 module LazyArraysBlockArraysTest
 using LazyArrays, ArrayLayouts, BlockArrays, FillArrays, Test
-using LazyArrays: LazyArrayStyle, PaddedLayout, PaddedColumns, PaddedRows, paddeddata
-using BlockArrays: blockcolsupport, blockrowsupport
+using LazyArrays: LazyArrayStyle, PaddedLayout, PaddedColumns, PaddedRows, paddeddata, ApplyLayout
+using BlockArrays: blockcolsupport, blockrowsupport, blockvec
+const BlockVec{T, M<:AbstractMatrix{T}} = ApplyVector{T, typeof(blockvec), <:Tuple{M}}
 
 @testset "Lazy BlockArrays" begin
     @testset "LazyBlock" begin
@@ -143,5 +144,44 @@ using BlockArrays: blockcolsupport, blockrowsupport
         B = BlockedArray(Vcat([1 2 3; 4 5 6], Ones(3,3), fill(4, 2, 3))', [3], [2,3,2])
         @test ApplyArray(view(B, Block.(1:1), Block(1))) == [1 4; 2 5; 3 6]
     end
+
+    @testset "BlockVec" begin
+        X = randn(5,4)
+        b = BlockVec(X)
+        @test size(b) == (20,)
+        @test length(b) == 20
+        @test MemoryLayout(b) isa ApplyLayout{typeof(blockvec)}
+        @test b == vec(X)
+        @test view(b, Block(3)) ≡ view(X, :, 3)
+        @test b[Block(3)] isa Vector
+        b[5] = 6
+        @test X[5] == 6
+        @test resize!(b, Block(2)) == b[Block.(1:2)]
+
+        c = BlockVec(X')
+        @test c == vec(X')
+        @test view(c, Block(3)) ≡ view(X', :, 3)
+        @test resize!(c, Block(2)) == c[Block.(1:2)]
+
+        c = BlockVec(transpose(X))
+        @test c == vec(transpose(X))
+        @test view(c, Block(3)) ≡ view(transpose(X), :, 3)
+        @test resize!(c, Block(2)) == c[Block.(1:2)]
+
+        X = cache(Zeros(5,6));
+        X[1,1] = 2
+        c = BlockVec(X);
+        @test MemoryLayout(c) isa PaddedColumns
+        @test paddeddata(c) isa BlockVec
+        @test paddeddata(c) == [2]
+
+        a = ApplyArray(blockvec, (1:5)')
+        b = BlockedArray(Vcat([1,2], Zeros(3)), (axes(a,1),))
+        
+        @test a + b == b + a
+        @test a - b == -(b-a)
+        @test a + a == 2a
+        @test b + b == 2b
+    end
 end
 end

test/paddedtests.jl

@@ -457,22 +457,52 @@ paddeddata(a::PaddedPadded) = a
         @test B*p == Matrix(B)*p
         @test simplifiable(*,B,p) == Val(true)
     end
+    
+    @testset "cached broadcasted" begin
+        z = Zeros(5)
+        a = cache(z); a[1] = 3;
+        
+        @test z ./ a ≡ a .\ z ≡ z
+    end
+
+    @testset "Number * padded matrix - of - vector" begin
+        A = Vcat([1,2], Zeros(3,1))
+        @test 2 * A == A * 2 == 2 * Matrix(A)
+    end
 
     @testset "QR" begin
         A = Vcat(randn(5,5), Zeros(4,5))
         F = qr!(copy(A))
         F̃ = qr!(Matrix(A))
         @test F.R ≈ F̃.R
-        b = Vcat([1,2], Zeros(7))
-        @test F.Q'b ≈ F̃.Q'b
+
+        Q = F.Q
+        Q̃ = F̃.Q
+
+        b = cache(Vcat([1,2], Zeros(7)));
+        B = cache(Vcat([1 2;3 4], Zeros(7,2)));
+
+        @test Q'b ≈ Q̃'b
         @test ldiv!(F,cache(b)) ≈ ldiv!(F̃,Vector(b))
-    end
+        @test lmul!(Q, deepcopy(b)) ≈ lmul!(Q, Vector(b)) ≈ Q*b
+        @test lmul!(Q', deepcopy(b)) ≈ lmul!(Q', Vector(b)) ≈ Q'b
 
-    @testset "cached broadcasted" begin
-        z = Zeros(5)
-        a = cache(z); a[1] = 3;
-        
-        @test z ./ a ≡ a .\ z ≡ z
+
+        @test lmul!(Q, deepcopy(B)) ≈ lmul!(Q, Matrix(B)) ≈ Q*B
+        @test lmul!(Q', deepcopy(B)) ≈ lmul!(Q', Matrix(B)) ≈ Q'B
+
+        @test Matrix(Q) ≈ [Q[k,j] for k in axes(A,1), j in axes(A,2)] ≈ Matrix(Q̃)
+        @test colsupport(Q,3) == rowsupport(Q,3) == colsupport(Q',3) == rowsupport(Q',3) == Base.OneTo(5)
+
+        @test Q'A ≈ [F.R; zeros(4,5)]
+        @test A'*Q ≈ [F.R; zeros(4,5)]'
+
+        M = ApplyArray(*, Q, b)
+        @test colsupport(M) == Base.OneTo(5)
+
+        M = ApplyArray(*, Q, B)
+        @test colsupport(M) == Base.OneTo(5)
+        @test M[1,:] ≈ (Q*B)[1,:]
     end
 end
 end # module