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

Author: jagot

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,132 @@ 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, 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

@@ -110,17 +110,3 @@ for op in (:-, :+)
         end
     end
 end
-
-function replace_in_print_matrix(A::BlockDiagonal, i::Integer, j::Integer, s::AbstractString)
-    bi = findblockindex.(axes(A), (i,j))
-    I,J = block.(bi)
-    i,j = blockindex.(bi)
-    Int(J-I) == 0 ? s : Base.replace_with_centered_mark(s)
-end
-
-function replace_in_print_matrix(A::BlockTridiagonal, i::Integer, j::Integer, s::AbstractString)
-    bi = findblockindex.(axes(A), (i,j))
-    I,J = block.(bi)
-    i,j = blockindex.(bi)
-    -1 ≤ Int(J-I) ≤ 1 ? s : Base.replace_with_centered_mark(s)
-end

test/test_blockskyline.jl

@@ -104,4 +104,48 @@ 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(Matrix(A), Matrix(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) = [(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