BlockSkylineMatrix–Diagonal arithmetic, fix BlockDiagonal (#58)

* Fixed BlockDiagonal

* Added BlockSkylineMatrix–{,Bi,Tri,SymTri}Diagonal arithmetic (fixes #57)

Author: jagot

Project.toml

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

src/BlockSkylineMatrix.jl

@@ -52,7 +52,7 @@ struct BlockSkylineSizes{BS<:NTuple{2,AbstractUnitRange{Int}}, LL<:AbstractVecto
     u::UU
 end
 
-const BlockBandedSizes = BlockSkylineSizes{NTuple{2,BlockedUnitRange{Vector{Int}}}, Fill{Int,1,Tuple{OneTo{Int}}}, Fill{Int,1,Tuple{OneTo{Int}}}, 
+const BlockBandedSizes = BlockSkylineSizes{NTuple{2,BlockedUnitRange{Vector{Int}}}, Fill{Int,1,Tuple{OneTo{Int}}}, Fill{Int,1,Tuple{OneTo{Int}}},
                                             BandedMatrix{Int,Matrix{Int},OneTo{Int}}, Vector{Int}}
 
 
@@ -247,7 +247,7 @@ BlockBandedMatrix(A::AbstractMatrix) = convert(BlockBandedMatrix, A)
 similar(A::BlockSkylineMatrix, T::Type=eltype(A), bs::BlockSkylineSizes=A.block_sizes) =
     BlockSkylineMatrix{T}(undef, bs)
 
-axes(A::BlockSkylineMatrix) = A.block_sizes.axes    
+axes(A::BlockSkylineMatrix) = A.block_sizes.axes
 
 ################################
 # BlockSkylineMatrix Interface #
@@ -390,7 +390,7 @@ const BlockBandedBlock{T} = SubArray{T,2,<:BlockSkylineMatrix,<:Tuple{<:BlockSli
 
 
 # gives the columns of parent(V).data that encode the block
-_parent_blocks(V::BlockBandedBlock)::Tuple{Int,Int} = 
+_parent_blocks(V::BlockBandedBlock)::Tuple{Int,Int} =
     first(first(parentindices(V)).block.n),first(last(parentindices(V)).block.n)
 
 ######################################
@@ -437,3 +437,147 @@ end
         throw(BandError(A, J-K))
     end
 end
+
+"""
+    copy_accommodating_diagonals(A::BlockSkylineMatrix, diagonals::UnitRange{<:Integer})
+
+Return copy of `A`, ensuring blocks are present covering the
+`diagonals` as well.
+"""
+function copy_accommodating_diagonals(A::BlockSkylineMatrix, diagonals::UnitRange{<:Integer}, ::Type{T}=eltype(A)) where T
+    checksquareblocks(A)
+    bs = A.block_sizes
+    l,u = bs.l,bs.u
+    ax = first(bs.axes)
+
+    if 0 ∈ diagonals
+        l = max.(l, 0)
+        u = max.(u, 0)
+    end
+
+    for d = extrema(diagonals)
+        d == 0 && continue # Already taken care of above
+        v = d > 0 ? u : l
+        # The j:th element of rows is the row index which the diagonal
+        # d covers in the j:th column.
+        rows = clamp.(ax .- d, 1, last(ax))
+        for (j,i) in enumerate(rows)
+            # First we find which block the j:th element of the main
+            # diagonal would occupy.
+            md_block = findfirst(≥(j), ax.lasts)
+
+            # Next we find which block covers row i
+            d_block = findfirst(≥(i), ax.lasts)
+
+            # Finally, we increase the block-bandwidth as necessary
+            v[md_block] = max(v[md_block], abs(d_block-md_block))
+        end
+    end
+
+    BlockSkylineMatrix{T}(A, BlockSkylineSizes((ax,ax), l, u))
+end
+
+for op in (:-, :+)
+    @eval begin
+        function $op(A::BlockSkylineMatrix, I::UniformScaling)
+            B = copy_accommodating_diagonals(A, 0:0, Base._return_type(+, Tuple{eltype(A), eltype(I)}))
+            @inbounds for i in axes(A, 1)
+                B[i,i] = $op(B[i,i], I.λ)
+            end
+            B
+        end
+        function $op(I::UniformScaling, A::BlockSkylineMatrix)
+            B = copy_accommodating_diagonals($op(A), 0:0, Base._return_type(+, Tuple{eltype(A), eltype(I)}))
+            @inbounds for i in axes(A, 1)
+                B[i,i] += I.λ
+            end
+            B
+        end
+
+        function $op(A::BlockSkylineMatrix, D::Diagonal)
+            B = copy_accommodating_diagonals(A, 0:0, Base._return_type(+, Tuple{eltype(A), eltype(D)}))
+            @inbounds for i in axes(A, 1)
+                B[i,i] = $op(B[i,i], D.diag[i])
+            end
+            B
+        end
+        function $op(D::Diagonal, A::BlockSkylineMatrix)
+            B = copy_accommodating_diagonals($op(A), 0:0, Base._return_type(+, Tuple{eltype(A), eltype(D)}))
+            @inbounds for i in axes(A, 1)
+                B[i,i] += D.diag[i]
+            end
+            B
+        end
+
+        function $op(A::BlockSkylineMatrix, Bd::Bidiagonal)
+            B = copy_accommodating_diagonals(A, Bd.uplo == 'U' ? (0:1) : (-1:0),
+                                             Base._return_type(+, Tuple{eltype(A), eltype(Bd)}))
+            @inbounds for i in axes(A, 1)
+                B[i,i] = $op(B[i,i], Bd.dv[i])
+            end
+            @inbounds for i in 1:size(A, 1)-1
+                Bd.uplo == 'U' && (B[i,i+1] = $op(B[i,i+1], Bd.ev[i]))
+                Bd.uplo == 'L' && (B[i+1,i] = $op(B[i+1,i], Bd.ev[i]))
+            end
+            B
+        end
+        function $op(Bd::Bidiagonal, A::BlockSkylineMatrix)
+            B = copy_accommodating_diagonals($op(A), Bd.uplo == 'U' ? (0:1) : (-1:0),
+                                             Base._return_type(+, Tuple{eltype(A), eltype(Bd)}))
+            @inbounds for i in axes(A, 1)
+                B[i,i] += Bd.dv[i]
+            end
+            @inbounds for i in 1:size(A, 1)-1
+                Bd.uplo == 'U' && (B[i,i+1] += Bd.ev[i])
+                Bd.uplo == 'L' && (B[i+1,i] += Bd.ev[i])
+            end
+            B
+        end
+
+        function $op(A::BlockSkylineMatrix, T::Tridiagonal)
+            B = copy_accommodating_diagonals(A, -1:1, Base._return_type(+, Tuple{eltype(A), eltype(T)}))
+            @inbounds for i in axes(A, 1)
+                B[i,i] = $op(B[i,i], T.d[i])
+            end
+            @inbounds for i in 1:size(A, 1)-1
+                B[i,i+1] = $op(B[i,i+1], T.du[i])
+                B[i+1,i] = $op(B[i+1,i], T.dl[i])
+            end
+            B
+        end
+        function $op(T::Tridiagonal, A::BlockSkylineMatrix)
+            B = copy_accommodating_diagonals($op(A), -1:1, Base._return_type(+, Tuple{eltype(A), eltype(T)}))
+            @inbounds for i in axes(A, 1)
+                B[i,i] += T.d[i]
+            end
+            @inbounds for i in 1:size(A, 1)-1
+                B[i,i+1] += T.du[i]
+                B[i+1,i] += T.dl[i]
+            end
+            B
+        end
+
+        function $op(A::BlockSkylineMatrix, T::SymTridiagonal)
+            B = copy_accommodating_diagonals(A, -1:1, Base._return_type(+, Tuple{eltype(A), eltype(T)}))
+            @inbounds for i in axes(A, 1)
+                B[i,i] = $op(B[i,i], T.dv[i])
+            end
+            @inbounds for i in 1:size(A, 1)-1
+                B[i,i+1] = $op(B[i,i+1], T.ev[i])
+                B[i+1,i] = $op(B[i+1,i], T.ev[i])
+            end
+            B
+        end
+        function $op(T::SymTridiagonal, A::BlockSkylineMatrix)
+            B = copy_accommodating_diagonals($op(A), -1:1, Base._return_type(+, Tuple{eltype(A), eltype(T)}))
+            @inbounds for i in axes(A, 1)
+                B[i,i] += T.dv[i]
+            end
+            @inbounds for i in 1:size(A, 1)-1
+                B[i,i+1] += T.ev[i]
+                B[i+1,i] += T.ev[i]
+            end
+            B
+        end
+    end
+end

src/interfaceimpl.jl

@@ -25,15 +25,15 @@ const BlockDiagonal{T,VT<:Matrix{T}} = BlockMatrix{T,<:Diagonal{VT}}
 
 BlockDiagonal(A) = mortar(Diagonal(A))
 
-function sizes_from_blocks(A::Diagonal, _) 
+function sizes_from_blocks(A::Diagonal, _)
     # for k = 1:length(A.du)
     #     size(A.du[k],1) == sz[1][k] || throw(ArgumentError("block sizes of upper diagonal inconsisent with diagonal"))
     #     size(A.du[k],2) == sz[2][k+1] || throw(ArgumentError("block sizes of upper diagonal inconsisent with diagonal"))
     #     size(A.dl[k],1) == sz[1][k+1] || throw(ArgumentError("block sizes of lower diagonal inconsisent with diagonal"))
     #     size(A.dl[k],2) == sz[2][k] || throw(ArgumentError("block sizes of lower diagonal inconsisent with diagonal"))
     # end
     (size.(A.diag, 1), size.(A.diag,2))
-end    
+end
 
 
 # Block Bi/Tridiagonal
@@ -43,7 +43,7 @@ const BlockBidiagonal{T,VT<:Matrix{T}} = BlockMatrix{T,<:Bidiagonal{VT}}
 BlockTridiagonal(A,B,C) = mortar(Tridiagonal(A,B,C))
 BlockBidiagonal(A, B, uplo) = mortar(Bidiagonal(A,B,uplo))
 
-function sizes_from_blocks(A::Tridiagonal, _) 
+function sizes_from_blocks(A::Tridiagonal, _)
     # for k = 1:length(A.du)
     #     size(A.du[k],1) == sz[1][k] || throw(ArgumentError("block sizes of upper diagonal inconsisent with diagonal"))
     #     size(A.du[k],2) == sz[2][k+1] || throw(ArgumentError("block sizes of upper diagonal inconsisent with diagonal"))
@@ -53,7 +53,7 @@ function sizes_from_blocks(A::Tridiagonal, _)
     (size.(A.d, 1), size.(A.d,2))
 end
 
-function sizes_from_blocks(A::Bidiagonal, _) 
+function sizes_from_blocks(A::Bidiagonal, _)
     # for k = 1:length(A.du)
     #     size(A.du[k],1) == sz[1][k] || throw(ArgumentError("block sizes of upper diagonal inconsisent with diagonal"))
     #     size(A.du[k],2) == sz[2][k+1] || throw(ArgumentError("block sizes of upper diagonal inconsisent with diagonal"))
@@ -83,21 +83,30 @@ checksquareblocks(A) = blockisequal(axes(A)...) || throw(DimensionMismatch("bloc
 
 for op in (:-, :+)
     @eval begin
-        function $op(A::BlockTridiagonal, λ::UniformScaling) 
+        function $op(A::BlockDiagonal, λ::UniformScaling)
+            checksquareblocks(A)
+            mortar(Diagonal(broadcast($op, A.blocks.diag, Ref(λ))))
+        end
+        function $op(λ::UniformScaling, A::BlockDiagonal)
+            checksquareblocks(A)
+            mortar(Diagonal(broadcast($op, Ref(λ), A.blocks.diag)))
+        end
+
+        function $op(A::BlockTridiagonal, λ::UniformScaling)
             checksquareblocks(A)
             mortar(Tridiagonal(A.blocks.dl, broadcast($op, A.blocks.d, Ref(λ)), A.blocks.du))
         end
-        function $op(λ::UniformScaling, A::BlockTridiagonal) 
+        function $op(λ::UniformScaling, A::BlockTridiagonal)
             checksquareblocks(A)
             mortar(Tridiagonal(broadcast($op,A.blocks.dl), broadcast($op, Ref(λ), A.blocks.d), broadcast($op,A.blocks.du)))
         end
-        function $op(A::BlockBidiagonal, λ::UniformScaling) 
+        function $op(A::BlockBidiagonal, λ::UniformScaling)
             checksquareblocks(A)
             mortar(Bidiagonal(broadcast($op, A.blocks.dv, Ref(λ)), A.blocks.ev, A.blocks.uplo))
         end
-        function $op(λ::UniformScaling, A::BlockBidiagonal) 
+        function $op(λ::UniformScaling, A::BlockBidiagonal)
             checksquareblocks(A)
             mortar(Bidiagonal(broadcast($op, Ref(λ), A.blocks.dv), broadcast($op,A.blocks.ev), A.blocks.uplo))
         end
     end
-end
+end

test/test_blockskyline.jl

@@ -104,4 +104,50 @@ Random.seed!(0)
         @test Matrix(AC) == Matrix(A) + Matrix(C)
         @test blockisequal(axes(AC),axes(C))
     end
+
+    @testset "BlockSkylineMatrix arithmetic" begin
+        t = Int
+        rows = cols = 1:4
+        N = sum(rows)
+        D = Diagonal(ones(t, N))
+        Bu = Bidiagonal(ones(t, N), 2ones(t, N-1), 'U')
+        Bl = Bidiagonal(ones(t, N), 2ones(t, N-1), 'L')
+        T = Tridiagonal(ones(t, N-1), 2ones(t, N), 3ones(t, N-1))
+        ST = SymTridiagonal(ones(t, N), 2ones(t, N-1))
+
+        function check_blockskylinematrix_arithmetic(A, op, B, (l,u))
+            C = op(A, B)
+            Cdense = op((A isa BlockSkylineMatrix ? Matrix(A) : A),
+                        (B isa BlockSkylineMatrix ? Matrix(B) : B))
+            @test C == Cdense
+
+            bs = (A isa BlockSkylineMatrix ? A : B).block_sizes
+            expl = max.(l, bs.l)
+            expu = max.(u, bs.u)
+
+            Cbs = C.block_sizes
+            @test all(Cbs.l .== expl)
+            @test all(Cbs.u .== expu)
+        end
+
+        for (l,u) = [([0,0,0,0],[0,0,0,0]),
+                     ([-1,-1,-1,-1],[1,2,2,2]),
+                     ([-1,-2,-2,-2],[1,1,2,2]),
+                     ([2,2,1,1],[-1,-1,-1,-1]),
+                     ([2,2,1,1],[-2,-2,-2,-1])]
+            M = BlockSkylineMatrix{t}(undef, rows, cols, (l,u))
+            M.data .= 1
+            for (A,expected_bws) = [(3I, ([0,0,0,0],[0,0,0,0])),
+                                    (D, ([0,0,0,0],[0,0,0,0])),
+                                    (Bu, ([0,0,0,0],[0,1,1,1])),
+                                    (Bl, ([1,1,1,0],[0,0,0,0])),
+                                    (T, ([1,1,1,0],[0,1,1,1])),
+                                    (ST, ([1,1,1,0],[0,1,1,1]))]
+                for op = (+, -)
+                    check_blockskylinematrix_arithmetic(M, op, A, expected_bws)
+                    check_blockskylinematrix_arithmetic(A, op, M, expected_bws)
+                end
+            end
+        end
+    end
 end

test/test_misc.jl

@@ -21,6 +21,17 @@ import BandedMatrices: bandwidths, AbstractBandedMatrix, BandedStyle, bandeddata
     @test BandedMatrix(V) == V
 end
 
+@testset "Block Diagonal" begin
+    A = BlockBandedMatrices.BlockDiagonal(fill([1 2],3))
+    @test blockbandwidths(A) == (0,0)
+    @test isblockbanded(A)
+    @test A[Block(1,1)] == [1 2]
+    @test @inferred(getblock(A,1,2)) == @inferred(A[Block(1,2)]) == [0 0]
+    @test_throws DimensionMismatch A+I
+    A = BlockBandedMatrices.BlockDiagonal(fill([1 2; 1 2],3))
+    @test A+I == I+A == mortar(Diagonal(fill([2 2; 1 3],3))) == Matrix(A) + I
+end
+
 @testset "Block Bidiagonal" begin
     Bu = BlockBidiagonal(fill([1 2],4), fill([3 4],3), :U)
     Bl = BlockBidiagonal(fill([1 2],4), fill([3 4],3), :L)
@@ -52,4 +63,4 @@ end
     @test A+I == I+A == mortar(Tridiagonal(fill([1 2; 1 2],3), fill([4 4; 3 5],4), fill([4 5; 4 5],3))) == Matrix(A) + I
     @test A-I == mortar(Tridiagonal(fill([1 2; 1 2],3), fill([2 4; 3 3],4), fill([4 5; 4 5],3))) == Matrix(A) - I
     @test I-A == mortar(Tridiagonal(fill(-[1 2; 1 2],3), fill([-2 -4; -3 -3],4), fill(-[4 5; 4 5],3))) == I - Matrix(A)
-end
+end