Add lazy BlockVec (#93)

* Add lazy BlockVec

* BlockVec MemoryLayout

* increase cov

* Just use ApplyArray

* simplify muladd!

* Update LazyBandedMatrices.jl

Project.toml

@@ -1,7 +1,7 @@
 name = "LazyBandedMatrices"
 uuid = "d7e5e226-e90b-4449-9968-0f923699bf6f"
 authors = ["Sheehan Olver <solver@mac.com>"]
-version = "0.8.6"
+version = "0.8.7"
 
 [deps]
 ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"

src/LazyBandedMatrices.jl

@@ -8,7 +8,7 @@ import LinearAlgebra
 import MatrixFactorizations: ql, ql!, QLPackedQ, QRPackedQ, reflector!, reflectorApply!,
             QLPackedQLayout, QRPackedQLayout, AdjQLPackedQLayout, AdjQRPackedQLayout
 
-import Base: BroadcastStyle, similar, OneTo, oneto, copy, *, axes, size, getindex, tail, convert, resize!, tuple_type_tail
+import Base: BroadcastStyle, similar, OneTo, oneto, copy, *, axes, size, getindex, tail, convert, resize!, tuple_type_tail, view
 import Base.Broadcast: Broadcasted, broadcasted, instantiate
 import LinearAlgebra: kron, hcat, vcat, AdjOrTrans, AbstractTriangular, BlasFloat, BlasComplex, BlasReal,
                         lmul!, rmul!, checksquare, StructuredMatrixStyle, adjoint, transpose,
@@ -42,7 +42,7 @@ import BlockBandedMatrices: BlockSlice, Block1, AbstractBlockBandedLayout,
                         AbstractBandedBlockBandedLayout, BandedBlockBandedLayout, BandedBlockBandedStyle, BlockBandedStyle,
                         blockcolsupport, BlockRange1, blockrowsupport, BlockIndexRange1,
                         BlockBandedColumnMajor
-import BlockArrays: BlockSlice1, BlockLayout, AbstractBlockStyle, block, blockindex, BlockKron, viewblock, blocks, BlockSlices, AbstractBlockLayout
+import BlockArrays: BlockSlice1, BlockLayout, AbstractBlockStyle, block, blockindex, BlockKron, viewblock, blocks, BlockSlices, AbstractBlockLayout, blockvec
 
 # for bidiag/tridiag
 import Base: -, +, *, /, \, ==, AbstractMatrix, Matrix, Array, size, conj, real, imag, copy,
@@ -343,15 +343,16 @@ function materialize!(M::MatMulVecAdd{<:AllBlockBandedLayout,<:PaddedLayout,<:Pa
     α,A,x,β,y = M.α,M.A,M.B,M.β,M.C
     length(y) == size(A,1) || throw(DimensionMismatch())
     length(x) == size(A,2) || throw(DimensionMismatch())
+    ỹ = paddeddata(y)
 
     if !blockisequal(axes(A,2), axes(x,1))
-        return muladd!(α, A, PseudoBlockVector(x, (axes(A,2),)), β, y)
+        x2 = PseudoBlockVector(x, (axes(A,2),))
+        x̃2 = paddeddata(x)
+        muladd!(α, view(A, axes(ỹ,1), axes(x̃2,1)), x̃2, β, ỹ)
+    else
+        x̃ = paddeddata(x)
+        muladd!(α, view(A, axes(ỹ,1), axes(x̃,1)), x̃, β, ỹ)
     end
-
-    ỹ = paddeddata(y)
-    x̃ = paddeddata(x)
-
-    muladd!(α, view(A, axes(ỹ,1), axes(x̃,1)) , x̃, β, ỹ)
     y
 end
 

src/blockconcat.jl

@@ -522,3 +522,25 @@ end
 MemoryLayout(::Type{BlockBroadcastArray{T,N,FF,Args}}) where {T,N,FF,Args} = blockbroadcastlayout(FF, tuple_type_memorylayouts(Args)...)
 
 resize!(c::BlockBroadcastVector{T,typeof(vcat)}, N::Block{1}) where T = BlockBroadcastVector{T}(vcat, resize!.(c.args, N)...)
+
+####
+# 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])...)),)
+
+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)}, ::PaddedLayout) = PaddedLayout{ApplyLayout{typeof(blockvec)}}()
+paddeddata(b::BlockVec) = BlockVec(paddeddata(b.args[1]))

test/test_blockconcat.jl

@@ -1,7 +1,8 @@
 using LazyBandedMatrices, BlockBandedMatrices, BlockArrays, StaticArrays, FillArrays, LazyArrays, ArrayLayouts, BandedMatrices, Test
-import LazyBandedMatrices: BlockBroadcastArray, ApplyLayout, blockcolsupport, blockrowsupport, arguments, paddeddata, resizedata!
+import LazyBandedMatrices: BlockBroadcastArray, ApplyLayout, blockcolsupport, blockrowsupport, arguments, paddeddata, resizedata!, BlockVec, BlockVecLayout
+import BlockArrays: blockvec
 import LinearAlgebra: Adjoint, Transpose
-import LazyArrays: PaddedArray
+import LazyArrays: PaddedArray, PaddedLayout
 
 @testset "unitblocks" begin
     a = unitblocks(Base.OneTo(5))
@@ -142,7 +143,6 @@ end
     end
 end
 
-
 @testset "BlockHcat" begin
     @testset "vec hcat" begin
         a = BlockHcat(1:5, 10:14)
@@ -406,3 +406,32 @@ end
         @test subblockbandwidths(C) == (0,0)
     end
 end
+
+@testset "BlockVec" begin
+    X = randn(5,4)
+    b = BlockVec(X)
+    @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 PaddedLayout
+    @test paddeddata(c) isa BlockVec
+    @test paddeddata(c) == [2]
+end