Support BlockRange in Mul , add ApplyBandedBlockBandedLayout (#16)

* Add ApplyBandedBlockBandedLayout

* Supoort BlockRange in Mul

* Update Project.toml

* Update LazyBandedMatrices.jl

* Update runtests.jl

Project.toml

@@ -1,7 +1,7 @@
 name = "LazyBandedMatrices"
 uuid = "d7e5e226-e90b-4449-9968-0f923699bf6f"
 authors = ["Sheehan Olver <solver@mac.com>"]
-version = "0.2.11"
+version = "0.2.12"
 
 [deps]
 ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
@@ -15,13 +15,13 @@ MatrixFactorizations = "a3b82374-2e81-5b9e-98ce-41277c0e4c87"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 
 [compat]
-ArrayLayouts = "0.3.3"
-BandedMatrices = "0.15.10"
-BlockArrays = "0.12.8"
-BlockBandedMatrices = "0.8.5"
-FillArrays = "0.8.10"
-LazyArrays = "0.16.12"
-MatrixFactorizations = "0.4.1"
+ArrayLayouts = "0.3.6"
+BandedMatrices = "0.15.15"
+BlockArrays = "0.12.10"
+BlockBandedMatrices = "0.8.10"
+FillArrays = "0.8.13"
+LazyArrays = "0.16.15"
+MatrixFactorizations = "0.4.1, 0.5"
 julia = "1.2"
 
 [extras]

src/LazyBandedMatrices.jl

@@ -11,15 +11,16 @@ import LinearAlgebra: kron, hcat, vcat, AdjOrTrans, AbstractTriangular, BlasFloa
                         lmul!, rmul!, checksquare, StructuredMatrixStyle
 
 import ArrayLayouts: materialize!, colsupport, rowsupport, MatMulVecAdd, require_one_based_indexing,
-                    sublayout, transposelayout, _copyto!, MemoryLayout, AbstractQLayout
+                    sublayout, transposelayout, _copyto!, MemoryLayout, AbstractQLayout, _mul
 import LazyArrays: LazyArrayStyle, combine_mul_styles, mulapplystyle, PaddedLayout,
                         broadcastlayout, applylayout, arguments, _arguments, call,
                         LazyArrayApplyStyle, ApplyArrayBroadcastStyle, ApplyStyle,
                         LazyLayout, ApplyLayout, BroadcastLayout, FlattenMulStyle, CachedVector,
-                        _mul_args_rows, _mul_args_cols, paddeddata, sub_materialize,
+                        _mat_mul_arguments, paddeddata, sub_materialize,
                         MulMatrix, Mul, CachedMatrix, CachedArray, cachedlayout, _cache,
                         resizedata!, applybroadcaststyle, 
-                        LazyMatrix, LazyVector, LazyArray, MulAddStyle
+                        LazyMatrix, LazyVector, LazyArray, MulAddStyle,
+                        _mul_args_colsupport, _mul_args_rowsupport
 import BandedMatrices: bandedcolumns, bandwidths, isbanded, AbstractBandedLayout,
                         prodbandwidths, BandedStyle, BandedColumns, BandedRows,
                         AbstractBandedMatrix, BandedSubBandedMatrix, BandedStyle, _bnds,
@@ -31,7 +32,7 @@ import BlockBandedMatrices: BlockSlice, Block1, AbstractBlockBandedLayout,
                         BandedBlockBandedColumns, BlockBandedColumns,
                         subblockbandwidths, BandedBlockBandedMatrix, BlockBandedMatrix,
                         AbstractBandedBlockBandedLayout, BandedBlockBandedStyle,
-                        blockcolsupport, BlockRange1
+                        blockcolsupport, BlockRange1, blockrowsupport
 import BlockArrays: blockbroadcaststyle, BlockSlice1, BlockLayout
 
 export DiagTrav, KronTrav, blockkron
@@ -51,6 +52,7 @@ for LazyLay in (:(BandedColumns{LazyLayout}), :(BandedRows{LazyLayout}),
     @eval begin
         combine_mul_styles(::$LazyLay) = LazyArrayApplyStyle()
         mulapplystyle(::AbstractQLayout, ::$LazyLay) = LazyArrayApplyStyle()
+        _mul(::$LazyLay, _, A, B) = apply(*, A, B)
     end
 end
 
@@ -174,16 +176,29 @@ isbanded(M::MulMatrix) = isbanded(Applied(M))
 
 struct ApplyBandedLayout{F} <: AbstractBandedLayout end
 struct ApplyBlockBandedLayout{F} <: AbstractBlockBandedLayout end
+struct ApplyBandedBlockBandedLayout{F} <: AbstractBlockBandedLayout end
 
 arguments(::ApplyBandedLayout{F}, A) where F = arguments(ApplyLayout{F}(), A)
 sublayout(::ApplyBandedLayout{F}, A) where F = sublayout(ApplyLayout{F}(), A)
 
+
+# The following catches the arguments machinery to work for BlockRange
+# see LazyArrays.jl/src/mul.jl
+
+_mul_args_colsupport(a, kr::BlockRange) = blockcolsupport(a, kr)
+_mul_args_rowsupport(a, kr::BlockRange) = blockrowsupport(a, kr)
+_mul_args_colsupport(a, kr::Block) = blockcolsupport(a, kr)
+_mul_args_rowsupport(a, kr::Block) = blockrowsupport(a, kr)
+_mat_mul_arguments(args, (kr,jr)::Tuple{BlockSlice,BlockSlice}) = _mat_mul_arguments(args, (kr.block, jr.block))
+
 arguments(::ApplyBlockBandedLayout{F}, A) where F = arguments(ApplyLayout{F}(), A)
 sublayout(::ApplyBlockBandedLayout{F}, A) where F = sublayout(ApplyLayout{F}(), A)
+arguments(::ApplyBandedBlockBandedLayout{F}, A) where F = arguments(ApplyLayout{F}(), A)
+sublayout(::ApplyBandedBlockBandedLayout{F}, A) where F = sublayout(ApplyLayout{F}(), A)
 
 applylayout(::Type{typeof(*)}, ::AbstractBandedLayout...) = ApplyBandedLayout{typeof(*)}()
 applylayout(::Type{typeof(*)}, ::AllBlockBandedLayout...) = ApplyBlockBandedLayout{typeof(*)}()
-
+applylayout(::Type{typeof(*)}, ::AbstractBandedBlockBandedLayout...) = ApplyBandedBlockBandedLayout{typeof(*)}()
 
 applybroadcaststyle(::Type{<:AbstractMatrix}, ::ApplyBandedLayout{typeof(*)}) = BandedStyle()
 
@@ -305,20 +320,13 @@ end
 # sub materialize
 ###
 
-function arguments(::ApplyBandedLayout{typeof(*)}, V::SubArray)
-    P = parent(V)
-    kr, jr = parentindices(V)
-    as = arguments(P)
-    kjr = intersect.(_mul_args_rows(kr, as...), _mul_args_cols(jr, reverse(as)...))
-    view.(as, (kr, kjr...), (kjr..., jr))
-end
-
 @inline sub_materialize(::ApplyBandedLayout{typeof(*)}, V, _) = BandedMatrix(V)
 @inline sub_materialize(::BroadcastBandedLayout, V, _) = BandedMatrix(V)
 @inline sub_materialize(::BandedColumns{LazyLayout}, V, _) = V
 @inline sub_materialize(::BandedColumns{LazyLayout}, V, ::Tuple{<:OneTo,<:OneTo}) = BandedMatrix(V)
 
 @inline sub_materialize(::ApplyBlockBandedLayout{typeof(*)}, V, _) = BlockBandedMatrix(V)
+@inline sub_materialize(::ApplyBandedBlockBandedLayout{typeof(*)}, V, _) = BandedBlockBandedMatrix(V)
 
 ###
 # copyto!
@@ -342,16 +350,24 @@ for op in (:+, :-, :*)
     end
 end
 
+_mulbanded_copyto!(dest, a) = copyto!(dest, a)
+_mulbanded_copyto!(dest::AbstractArray{T}, a, b) where T = muladd!(one(T), a, b, zero(T), dest)
+_mulbanded_copyto!(dest::AbstractArray{T}, a, b, c, d...) where T = _mulbanded_copyto!(dest, apply(*,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))
 
 _copyto!(::AbstractBandedLayout, ::ApplyBandedLayout{typeof(*)}, dest::AbstractMatrix, src::AbstractMatrix) =
     _mulbanded_copyto!(dest, map(_mulbanded_BandedMatrix,arguments(src))...)
 
-_mulbanded_copyto!(dest, a) = copyto!(dest, a)
-_mulbanded_copyto!(dest::AbstractArray{T}, a, b) where T = muladd!(one(T), a, b, zero(T), dest)
-_mulbanded_copyto!(dest::AbstractArray{T}, a, b, c, d...) where T = _mulbanded_copyto!(dest, apply(*,a,b), c, d...)
+_mulbanded_BandedBlockBandedMatrix(A, _) = A
+_mulbanded_BandedBlockBandedMatrix(A, ::NTuple{2,OneTo{Int}}) = BandedBlockBandedMatrix(A)
+_mulbanded_BandedBlockBandedMatrix(A) = _mulbanded_BandedBlockBandedMatrix(A, axes(A))
+
+_copyto!(::AbstractBandedBlockBandedLayout, ::ApplyBandedBlockBandedLayout{typeof(*)}, dest::AbstractMatrix, src::AbstractMatrix) =
+    _mulbanded_copyto!(dest, map(_mulbanded_BandedBlockBandedMatrix,arguments(src))...)
+
 
 _broadcast_sub_arguments(::BroadcastBandedLayout{F}, A, V) where F = arguments(BroadcastLayout{F}(), V)
 _broadcast_sub_arguments(::BroadcastBandedBlockBandedLayout{F}, A, V) where F = arguments(BroadcastLayout{F}(), V)
@@ -370,6 +386,7 @@ transposelayout(b::BroadcastBandedLayout) = b
 arguments(b::BroadcastBandedLayout, A::AdjOrTrans) where F = arguments(BroadcastLayout(b), A)
 
 sublayout(M::ApplyBlockBandedLayout{typeof(*)}, ::Type{<:Tuple{BlockSlice{BlockRange1},BlockSlice{BlockRange1}}}) = M
+sublayout(M::ApplyBandedBlockBandedLayout{typeof(*)}, ::Type{<:Tuple{BlockSlice{BlockRange1},BlockSlice{BlockRange1}}}) = M
 
 
 ######

test/runtests.jl

@@ -2,7 +2,7 @@ using LazyBandedMatrices, BlockBandedMatrices, BandedMatrices, LazyArrays, Block
             ArrayLayouts, MatrixFactorizations, LinearAlgebra, Random, Test
 import LazyArrays: Applied, resizedata!, FillLayout, MulAddStyle, arguments, colsupport, rowsupport, LazyLayout, ApplyStyle, PaddedLayout, paddeddata, call, ApplyLayout
 import LazyBandedMatrices: VcatBandedMatrix, BroadcastBlockBandedLayout, BroadcastBandedLayout, 
-                    ApplyBandedLayout, ApplyBlockBandedLayout, BlockKron, LazyBandedLayout, BroadcastBandedBlockBandedLayout
+                    ApplyBandedLayout, ApplyBlockBandedLayout, ApplyBandedBlockBandedLayout, BlockKron, LazyBandedLayout, BroadcastBandedBlockBandedLayout
 import BandedMatrices: BandedStyle, _BandedMatrix, AbstractBandedMatrix, BandedRows, BandedColumns
 import ArrayLayouts: StridedLayout
 
@@ -161,6 +161,24 @@ end
         @test arguments(V) == (A[Block.(1:2),Block.(1:2)], B[Block.(1:2),Block.(1:2)])
         @test M[Block.(1:2), Block.(1:2)] isa BlockBandedMatrix
     end
+    @testset "MulBandedBlockBanded" begin
+        A = BandedBlockBandedMatrix{Float64}(undef, 1:4, 1:4, (1,0), (1,0)); A.data .= randn.();
+        B = BandedBlockBandedMatrix{Float64}(undef, 1:4, 1:4, (1,1), (1,1)); B.data .= randn.();
+        M = ApplyMatrix(*, A, B)
+        @test blockbandwidths(M) == (2,1)
+        @test subblockbandwidths(M) == (2,1)
+        @test MemoryLayout(M) isa ApplyBandedBlockBandedLayout{typeof(*)}
+        @test BandedBlockBandedMatrix(M) ≈ BlockBandedMatrix(M) ≈ A*B
+        @test arguments(M) == (A,B)
+        V = view(M, Block.(1:2), Block.(1:2))
+        @test MemoryLayout(V) isa ApplyBandedBlockBandedLayout{typeof(*)}
+        @test arguments(V) == (A[Block.(1:2),Block.(1:2)], B[Block.(1:2),Block.(1:2)])
+        @test M[Block.(1:2), Block.(1:2)] isa BandedBlockBandedMatrix
+        V = view(M, 1:3, 1:3)
+        @test MemoryLayout(V) isa ApplyLayout{typeof(*)}
+        @test arguments(V) == (A[1:3,1:3], B[1:3,1:3])
+        @test M[1:3, 1:3] ≈ (A*B)[1:3,1:3]
+    end
     @testset "Psuedo Mul" begin
         A = PseudoBandedMatrix(rand(5, 4), 1, 2)
         B = PseudoBandedMatrix(rand(4, 4), 2, 3)