Make copying a view of a KronTrav faster

Project.toml

@@ -1,7 +1,7 @@
 name = "LazyBandedMatrices"
 uuid = "d7e5e226-e90b-4449-9968-0f923699bf6f"
 authors = ["Sheehan Olver <[email protected]>"]
-version = "0.8.9"
+version = "0.8.10"
 
 [deps]
 ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"

src/LazyBandedMatrices.jl

@@ -41,7 +41,7 @@ import BlockBandedMatrices: BlockSlice, Block1, AbstractBlockBandedLayout,
                         subblockbandwidths, BandedBlockBandedMatrix, BlockBandedMatrix, BlockBandedLayout,
                         AbstractBandedBlockBandedLayout, BandedBlockBandedLayout, BandedBlockBandedStyle, BlockBandedStyle,
                         blockcolsupport, BlockRange1, blockrowsupport, BlockIndexRange1,
-                        BlockBandedColumnMajor
+                        BlockBandedColumnMajor, blockbanded_copyto!
 import BlockArrays: BlockSlice1, BlockLayout, AbstractBlockStyle, block, blockindex, BlockKron, viewblock, blocks, BlockSlices, AbstractBlockLayout, blockvec
 
 # for bidiag/tridiag
@@ -533,10 +533,10 @@ _broadcastarray2broadcasted(::StructuredBroadcastLayouts{F}, A) where F = _broad
 _broadcastarray2broadcasted(::StructuredBroadcastLayouts{F}, A::BroadcastArray) where F = _broadcastarray2broadcasted(BroadcastLayout{F}(), A)
 
 _copyto!(::AbstractBandedLayout, ::BroadcastBandedLayout, dest::AbstractMatrix, bc::AbstractMatrix) =
-    copyto!(dest, _broadcastarray2broadcasted(bc))
+    copyto!(dest, _broadcastarray2broadcasted(MemoryLayout(bc), bc))
 
 _copyto!(_, ::BroadcastBandedLayout, dest::AbstractMatrix, bc::AbstractMatrix) =
-    copyto!(dest, _broadcastarray2broadcasted(bc))
+    copyto!(dest, _broadcastarray2broadcasted(MemoryLayout(bc), bc))
 
 _banded_broadcast!(dest::AbstractMatrix, f, (A,B)::Tuple{AbstractMatrix{T},AbstractMatrix{V}}, _, ::Tuple{<:Any,ApplyBandedLayout{typeof(*)}}) where {T,V} =
     broadcast!(f, dest, BandedMatrix(A), BandedMatrix(B))
@@ -582,15 +582,15 @@ _mulbanded_copyto!(dest::AbstractArray{T}, a, b) where T = muladd!(one(T), a, b,
 _mulbanded_copyto!(dest::AbstractArray{T}, a, b, c, d...) where T = _mulbanded_copyto!(dest, mul(a,b), c, d...)
 
 _mulbanded_BandedMatrix(A, _) = A
-_mulbanded_BandedMatrix(A, ::NTuple{2,OneTo{Int}}) = BandedMatrix(A)
-_mulbanded_BandedMatrix(A) = _mulbanded_BandedMatrix(A, axes(A))
+_mulbanded_BandedMatrix(A, ::NTuple{2,Int}) = BandedMatrix(A)
+_mulbanded_BandedMatrix(A) = _mulbanded_BandedMatrix(A, size(A))
 
 _copyto!(::AbstractBandedLayout, ::ApplyBandedLayout{typeof(*)}, dest::AbstractMatrix, src::AbstractMatrix) =
     _mulbanded_copyto!(dest, map(_mulbanded_BandedMatrix,arguments(src))...)
 
 _mulbanded_BandedBlockBandedMatrix(A, _) = A
-_mulbanded_BandedBlockBandedMatrix(A, ::NTuple{2,OneTo{Int}}) = BandedBlockBandedMatrix(A)
-_mulbanded_BandedBlockBandedMatrix(A) = _mulbanded_BandedBlockBandedMatrix(A, axes(A))
+_mulbanded_BandedBlockBandedMatrix(A, ::NTuple{2,Int}) = BandedBlockBandedMatrix(A)
+_mulbanded_BandedBlockBandedMatrix(A) = _mulbanded_BandedBlockBandedMatrix(A, size(A))
 
 _copyto!(::AbstractBandedBlockBandedLayout, ::ApplyBandedBlockBandedLayout{typeof(*)}, dest::AbstractMatrix, src::AbstractMatrix) =
     _mulbanded_copyto!(dest, map(_mulbanded_BandedBlockBandedMatrix,arguments(src))...)
@@ -881,7 +881,8 @@ BandedLazyLayouts = Union{AbstractLazyBandedLayout, BandedColumns{LazyLayout}, B
                 TriangularLayout{UPLO,UNIT,BandedRows{LazyLayout}} where {UPLO,UNIT},
                 TriangularLayout{UPLO,UNIT,BandedColumns{LazyLayout}} where {UPLO,UNIT},
                 SymTridiagonalLayout{LazyLayout}, BidiagonalLayout{LazyLayout}, TridiagonalLayout{LazyLayout},
-                SymmetricLayout{BandedColumns{LazyLayout}}, HermitianLayout{BandedColumns{LazyLayout}}}
+                SymmetricLayout{BandedColumns{LazyLayout}}, HermitianLayout{BandedColumns{LazyLayout}},
+                BidiagonalLayout{LazyLayout,LazyLayout}}
 
 StructuredLazyLayouts = Union{BandedLazyLayouts,
                 BlockBandedColumns{LazyLayout}, BandedBlockBandedColumns{LazyLayout},
@@ -958,9 +959,13 @@ copy(M::Mul{BroadcastLayout{typeof(*)}, <:StructuredLazyLayouts}) = lazymaterial
 ## padded copy
 mulreduce(M::Mul{<:StructuredLazyLayouts, <:PaddedLayout}) = MulAdd(M)
 mulreduce(M::Mul{<:StructuredApplyLayouts{F}, D}) where {F,D<:PaddedLayout} = Mul{ApplyLayout{F},D}(M.A, M.B)
+mulreduce(M::Mul{<:PaddedLayout, <:StructuredLazyLayouts}) = MulAdd(M)
+mulreduce(M::Mul{D, <:StructuredApplyLayouts{F}}) where {F,D<:PaddedLayout} = Mul{D,ApplyLayout{F}}(M.A, M.B)
 # need to overload copy due to above
 copy(M::Mul{<:StructuredLazyLayouts, <:PaddedLayout}) = copy(mulreduce(M))
+copy(M::Mul{<:PaddedLayout, <:StructuredLazyLayouts}) = copy(mulreduce(M))
 simplifiable(::Mul{<:StructuredLazyLayouts, <:PaddedLayout}) = Val(true)
+simplifiable(::Mul{<:PaddedLayout, <:StructuredLazyLayouts}) = Val(true)
 
 
 copy(L::Ldiv{ApplyBandedLayout{typeof(*)}, Lay}) where Lay = copy(Ldiv{ApplyLayout{typeof(*)},Lay}(L.A, L.B))

src/blockkron.jl

@@ -64,6 +64,9 @@ function colsupport(A::DiagTrav{<:Any,2}, _)
     axes(A,1)[bs]
 end
 
+copy(A::DiagTrav) = DiagTrav(copy(A.array))
+==(A::DiagTrav, B::PseudoBlockVector) = PseudoBlockVector(A) == B
+==(A::PseudoBlockVector, B::DiagTrav) = A == PseudoBlockVector(B)
 
 function getindex(A::DiagTrav, K::Block{1})
     @boundscheck checkbounds(A, K)
@@ -201,6 +204,7 @@ function _krontrav_getindex(Kin::Block{2}, A, B, C)
 end
 
 getindex(M::KronTrav{<:Any,2}, K::Block{2}) = _krontrav_getindex(K, M.args...)
+# getindex(M::KronTrav{<:Any,2}, K::BlockIndex{2}) = _krontrav_getindex(K, M.args...)
 
 
 getindex(A::KronTrav{<:Any,N}, kj::Vararg{Int,N}) where N =
@@ -223,6 +227,7 @@ isbandedblockbanded(A::KronTrav) = isblockbanded(A) && length(A.args) == 2
 convert(::Type{B}, A::KronTrav{<:Any,2}) where B<:BandedBlockBandedMatrix = convert(B, BandedBlockBandedMatrix(A))
 
 struct KronTravBandedBlockBandedLayout <: AbstractBandedBlockBandedLayout end
+struct SubKronTravBandedBlockBandedLayout <: AbstractBandedBlockBandedLayout end
 struct KronTravLayout{M} <: AbstractBlockLayout end
 
 
@@ -243,6 +248,16 @@ sublayout(::KronTravBandedBlockBandedLayout, ::Type{<:NTuple{2,BlockSlice{BlockR
 sub_materialize(::Union{KronTravLayout,KronTravBandedBlockBandedLayout}, V) = KronTrav(map(sub_materialize, krontravargs(V))...)
 
 
+_copyto!(_, ::KronTravBandedBlockBandedLayout, dest::AbstractMatrix, src::AbstractMatrix) = blockbanded_copyto!(dest, sub_materialize(src))
+function _copyto!(_, ::SubKronTravBandedBlockBandedLayout, dest::AbstractMatrix, src::AbstractMatrix)
+    KR, JR = parentindices(src)
+    # use Base.OneTo range which is a Kron Trav
+    KR2,JR2 = Block.(Base.OneTo(Int(last(KR.block)))),Block.(Base.OneTo(Int(last(JR.block))))
+    # materialize parent with KR2,JR2 ranges then copyto!
+    copyto!(dest, view(BandedBlockBandedMatrix(view(parent(src),KR2, JR2)), KR, JR))
+end
+
+
 krontravargs(K::KronTrav) = K.args
 function krontravargs(V::SubArray)
     KR,JR = parentindices(V)

test/test_banded.jl

@@ -1,4 +1,4 @@
-
+using LazyBandedMatrices, LazyArrays, BandedMatrices, Test
 
 struct PseudoBandedMatrix{T} <: AbstractMatrix{T}
     data::Array{T}
@@ -67,12 +67,19 @@ LinearAlgebra.lmul!(β::Number, A::PseudoBandedMatrix) = (lmul!(β, A.data); A)
         b = Vcat(randn(2), Zeros(3))
         @test A*b ≈ Matrix(A)b
         @test B*b ≈ Matrix(B)b
+
+        P = Vcat(randn(2,5), Zeros(3,5))
+        @test A*P ≈ Matrix(A)P
+        @test P*A ≈ P*Matrix(A)
     end
 
     @testset "Apply * Banded" begin
         B = brand(5,5,2,1)
         A = ApplyArray(*, B, B)
         @test A * Vcat([1,2], Zeros(3)) ≈ B*B*[1,2,0,0,0]
+        P = Vcat(randn(2,5), Zeros(3,5))
+        @test A*P ≈ Matrix(A)P
+        @test P*A ≈ P*Matrix(A)
     end
 
     @testset "Banded Perturbed" begin