Better support for BlockRange indexing (#69)

* BlockRange indexing produces (Banded)BlockBandedMatrix

* v0.8.2

* generalise sublayout

Author: dlfivefifty

Project.toml

@@ -1,6 +1,6 @@
 name = "BlockBandedMatrices"
 uuid = "ffab5731-97b5-5995-9138-79e8c1846df0"
-version = "0.8.1"
+version = "0.8.2"
 
 [deps]
 ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"

src/AbstractBlockBandedMatrix.jl

@@ -1,14 +1,24 @@
-####
-# Matrix memory layout traits
-#
-# if MemoryLayout(A) returns BandedColumnMajor, you must override
-# pointer and leadingdimension
-# in addition to the banded matrix interface
-####
 
+"""
+   AbstractBlockBandedLayout
+
+isa a `MemoryLayout` that indicates that the array implements the block-banded
+interface.
+"""
 abstract type AbstractBlockBandedLayout <: AbstractBlockLayout end
+
+"""
+   AbstractBandedBlockBandedLayout
+
+isa a `MemoryLayout` that indicates that the array implements the banded-block-banded
+interface.
+"""
 abstract type AbstractBandedBlockBandedLayout <: AbstractBlockBandedLayout end
 
+
+struct BandedBlockBandedLayout <: AbstractBandedBlockBandedLayout end
+struct BlockBandedLayout <: AbstractBlockBandedLayout end
+
 struct BandedBlockBandedColumns{LAY} <: AbstractBandedBlockBandedLayout end
 struct BandedBlockBandedRows{LAY} <: AbstractBandedBlockBandedLayout end
 struct BlockBandedColumns{LAY} <: AbstractBlockBandedLayout end

src/BandedBlockBandedMatrix.jl

@@ -316,13 +316,18 @@ end
 end
 
 ## structured matrix methods ##
-function Base.replace_in_print_matrix(A::BandedBlockBandedMatrix, i::Integer, j::Integer, s::AbstractString)
+function _bandedblockbanded_replace_in_print_matrix(A, i, j, s)
     bi = findblockindex.(axes(A), (i,j))
     I,J = block.(bi)
     i,j = blockindex.(bi)
-    -A.l ≤ Int(J-I) ≤ A.u && -A.λ ≤ j-i ≤ A.μ ? s : Base.replace_with_centered_mark(s)
+    l,u = blockbandwidths(A)
+    λ,μ = subblockbandwidths(A)
+    -l ≤ Int(J-I) ≤ u && -λ ≤ j-i ≤ μ ? s : Base.replace_with_centered_mark(s)
 end
 
+Base.replace_in_print_matrix(A::BandedBlockBandedMatrix, i::Integer, j::Integer, s::AbstractString) = 
+    _bandedblockbanded_replace_in_print_matrix(A, i, j, s)
+
 
 ############
 # Indexing #
@@ -411,32 +416,43 @@ sublayout(::BandedBlockBandedColumnMajor, ::Type{<:Tuple{BlockSlice{Block1},Bloc
 sublayout(::BandedBlockBandedColumnMajor, ::Type{<:Tuple{BlockSlice{BlockRange1},BlockSlice{Block1}}}) = BandedBlockBandedColumnMajor()
 sublayout(::BandedBlockBandedColumnMajor, ::Type{<:Tuple{BlockSlice{BlockRange1},BlockSlice{BlockIndexRange1}}}) = BandedBlockBandedColumnMajor()
 
+sublayout(::AbstractBandedBlockBandedLayout, ::Type{<:Tuple{BlockSlice{Block1},BlockSlice{Block1}}}) = BandedLayout()
+sublayout(::AbstractBandedBlockBandedLayout, ::Type{<:Tuple{BlockSlice{BlockIndexRange1},BlockSlice{BlockIndexRange1}}}) = BandedLayout()
+sublayout(::AbstractBandedBlockBandedLayout, ::Type{<:Tuple{BlockSlice{BlockIndexRange1},BlockSlice{Block1}}}) = BandedLayout()
+sublayout(::AbstractBandedBlockBandedLayout, ::Type{<:Tuple{BlockSlice{Block1},BlockSlice{BlockIndexRange1}}}) = BandedLayout()
+sublayout(::AbstractBandedBlockBandedLayout, ::Type{<:Tuple{BlockSlice{BlockRange1},BlockSlice{BlockRange1}}}) = BandedBlockBandedLayout()
+sublayout(::AbstractBandedBlockBandedLayout, ::Type{<:Tuple{BlockSlice{Block1},BlockSlice{BlockRange1}}}) = BandedBlockBandedLayout()
+sublayout(::AbstractBandedBlockBandedLayout, ::Type{<:Tuple{BlockSlice{BlockRange1},BlockSlice{Block1}}}) = BandedBlockBandedLayout()
+sublayout(::AbstractBandedBlockBandedLayout, ::Type{<:Tuple{BlockSlice{BlockRange1},BlockSlice{BlockIndexRange1}}}) = BandedBlockBandedLayout()
+
+
+sub_materialize(::AbstractBandedBlockBandedLayout, V, _) = BandedBlockBandedMatrix(V)
+
 isbanded(A::SubArray{<:Any,2,<:BandedBlockBandedMatrix}) = MemoryLayout(typeof(A)) == BandedColumnMajor()
 isbandedblockbanded(A::SubArray{<:Any,2,<:BandedBlockBandedMatrix}) = MemoryLayout(typeof(A)) == BandedBlockBandedColumnMajor()
 
 
-subblockbandwidths(V::SubBandedBlockBandedMatrix) = subblockbandwidths(parent(V))
+subblockbandwidths(V::SubArray{<:Any,2,<:Any,<:Tuple{<:BlockSlice{BlockRange1},<:BlockSlice{BlockRange1}}}) = 
+    subblockbandwidths(parent(V))
 
-function blockbandwidths(V::SubBandedBlockBandedMatrix{<:Any,BlockRange1,Block1})
+function blockbandwidths(V::SubArray{<:Any,2,<:Any,<:Tuple{<:BlockSlice{BlockRange1},<:BlockSlice{Block1}}})
     A = parent(V)
-
     KR = parentindices(V)[1].block.indices[1]
     J = parentindices(V)[2].block
     shift = Int(KR[1])-Int(J)
     blockbandwidth(A,1) - shift, blockbandwidth(A,2) + shift
 end
 
-function blockbandwidths(V::SubBandedBlockBandedMatrix{<:Any,Block1,BlockRange1})
+function blockbandwidths(V::SubArray{<:Any,2,<:Any,<:Tuple{<:BlockSlice{<:Block1},<:BlockSlice{<:BlockRange1}}})
     A = parent(V)
-
     K = parentindices(V)[1].block
     JR = parentindices(V)[2].block.indices[1]
     shift = Int(K)-Int(JR[1])
-
-    blockbandwidth(A,1) - shift, blockbandwidth(A,2) + shift
+    l,u = blockbandwidths(A)
+    l - shift, u + shift
 end
 
-function blockbandwidths(V::SubBandedBlockBandedMatrix{<:Any,BlockRange1,BlockRange1})
+function blockbandwidths(V::SubArray{<:Any,2,<:Any,<:Tuple{<:BlockSlice{<:BlockRange1},<:BlockSlice{<:BlockRange1}}})
     A = parent(V)
 
     KR = parentindices(V)[1].block.indices[1]

src/BlockBandedMatrices.jl

@@ -23,7 +23,7 @@ import SparseArrays: sparse
 
 import ArrayLayouts: BlasMatLmulVec,
                     triangularlayout, UpperTriangularLayout, TriangularLayout, MatLdivVec,
-                    triangulardata, sublayout, 
+                    triangulardata, sublayout, sub_materialize,
                     AbstractColumnMajor, DenseColumnMajor, ColumnMajor,
                     DiagonalLayout, MulAdd, mul, colsupport, rowsupport,
                     _qr, _factorize, _copyto!
@@ -38,7 +38,7 @@ import BandedMatrices: isbanded, bandwidths, bandwidth, banded_getindex, colrang
                         inbands_setindex!, inbands_getindex, banded_setindex!,
                         banded_generic_axpy!,
                         BlasFloat, banded_dense_axpy!, MemoryLayout,
-                        BandedColumnMajor,
+                        BandedLayout, BandedColumnMajor,
                         BandedSubBandedMatrix, bandeddata, tribandeddata,
                         _BandedMatrix, colstart, colstop, rowstart, rowstop,
                         BandedStyle, _fill_lmul!, bandshift,
@@ -50,7 +50,7 @@ export BandedBlockBandedMatrix, BlockBandedMatrix, BlockSkylineMatrix, blockband
 
 const Block1 = Block{1,Int}
 const BlockRange1 = BlockRange{1,Tuple{UnitRange{Int}}}
-const BlockIndexRange1 = BlockIndexRange{1,Tuple{UnitRange{Int}}}    
+const BlockIndexRange1 = BlockIndexRange{1,Tuple{UnitRange{Int}}}
 
 blockcolrange(A...) = blockcolsupport(A...)
 blockrowrange(A...) = blockrowsupport(A...)

src/linalg.jl

@@ -77,8 +77,9 @@ function blockbandwidths(V::SubBlockSkylineMatrix{<:Any,LL,UU,BlockRange1,BlockR
     KR = parentindices(V)[1].block.indices[1]
     JR = parentindices(V)[2].block.indices[1]
 
-    @assert KR[1] == JR[1] == 1
-    blockbandwidths(A)
+    shift =  KR[1] - JR[1]
+    l,u = blockbandwidths(A)
+    l-shift, u+shift
 end
 
 
@@ -92,6 +93,10 @@ sublayout(::BlockBandedColumnMajor, ::Type{<:Tuple{<:BlockSlice{BlockRange1}, <:
 sublayout(::BlockBandedColumnMajor, ::Type{<:Tuple{<:BlockSlice{BlockRange1}, <:BlockSlice{BlockIndexRange1}}}) = ColumnMajor()
 sublayout(::BlockBandedColumnMajor, ::Type{<:Tuple{<:BlockSlice{BlockIndexRange1}, <:BlockSlice{BlockIndexRange1}}}) = ColumnMajor()
 
+isblockbanded(V::SubArray{<:Any,2,<:Any,<:Tuple{<:BlockSlice{BlockRange1}, <:BlockSlice{BlockRange1}}}) = 
+    isblockbanded(parent(V))
+
+sub_materialize(::AbstractBlockBandedLayout, V, _) = BlockBandedMatrix(V)
 
 strides(V::SubBlockSkylineMatrix{<:Any,LL,UU,<:Union{BlockRange1,Block1},Block1}) where {LL,UU} =
     (1,parent(V).block_sizes.block_strides[Int(parentindices(V)[2].block)])

test/test_bandedblockbanded.jl

@@ -298,7 +298,9 @@ import BlockBandedMatrices: _BandedBlockBandedMatrix, blockcolsupport, blockrows
         @test B isa BandedBlockBandedMatrix
         @test blockbandwidths(V) == blockbandwidths(B) == (2,0)
         @test subblockbandwidths(V) == subblockbandwidths(B) == subblockbandwidths(A) == (λ,μ)
-        @test B == V
+        @test B == V == A[Block.(2:3), Block.(3:4)] 
+
+        @test A[Block.(2:3), Block.(3:4)] isa BandedBlockBandedMatrix
 
         x = randn(size(B,2))
         y = similar(x, size(B,1))

test/test_blockbanded.jl

@@ -11,7 +11,7 @@ import Base.Broadcast: materialize!
         cols = rows = 1:N
 
         @test Matrix(BlockBandedMatrix(Zeros(sum(rows),sum(cols)), rows, cols, (l,u))) ==
-            Array(BlockBandedMatrix(Zeros(sum(rows),sum(cols)), rows, cols, (l,u))) == 
+            Array(BlockBandedMatrix(Zeros(sum(rows),sum(cols)), rows, cols, (l,u))) ==
             zeros(Float64, 10, 10)
 
         @test Matrix(BlockBandedMatrix{Int}(Zeros(sum(rows),sum(cols)), rows,cols, (l,u))) ==
@@ -35,13 +35,11 @@ import Base.Broadcast: materialize!
         N = M = 4
         cols = rows = 1:N
         A = BlockBandedMatrix{Int}(undef, rows,cols, (l,u))
-            A.data .= 1:length(A.data)
+        A.data .= 1:length(A.data)
 
         @test A[1,1] == 1
         @test A[1,3] == 10
 
-
-
         @test blockbandwidth(A,1)  == 1
         @test blockbandwidths(A) == (l,u)
 
@@ -92,12 +90,34 @@ import Base.Broadcast: materialize!
         @test ret[1,2] == 0
     end
 
+    @testset "block-banded matrix interface for blockranges" begin
+        l , u = 1,1
+        N = M = 4
+        cols = rows = 1:N
+        A = BlockBandedMatrix{Int}(undef, rows,cols, (l,u))
+        A.data .= 1:length(A.data)
+
+        V = view(A, Block.(2:3), Block.(3:4))
+        @test isblockbanded(V)
+
+        B = BlockBandedMatrix(V)
+        @test B isa BlockBandedMatrix
+        @test blockbandwidths(V) == blockbandwidths(B) == (2,0)
+        @test B == V == A[Block.(2:3), Block.(3:4)]
+
+        @test A[Block.(2:3), Block.(3:4)] isa BlockBandedMatrix
+
+        x = randn(size(B,2))
+        y = similar(x, size(B,1))
+        @test all((similar(y) .= MulAdd(B, x)) .=== (similar(y) .= MulAdd(V,x)))
+    end
+
     @testset "BlockBandedMatrix indexing" begin
         l , u = 1,1
         N = M = 10
         cols = rows = 1:N
         A = BlockBandedMatrix{Float64}(undef, rows,cols, (l,u))
-            A.data .= 1:length(A.data)
+        A.data .= 1:length(A.data)
 
         A[1,1] = 5
         @test A[1,1] == 5