Allows blocks to be <:AbstractVector or <:Tuple

- This is a breaking change, as the type param `V` is now different,
  being the type of the `blocks` collection, rather than the `eltype`
  of that collection.
- This allows allows the collection to be something other than a
  `Vector` e.g.
  - a `CuArray` of matrices (allowing use on GPU hopefully)
  - a `Tuple` with different concrete subtypes of `AbstractMatrix{T}`
    (avoiding heterogeneous block to being abstractly typed
    i.e. a way to avoid `Vector{AbstractMatrix{T}}`)

Author: nickrobinson251

Project.toml

@@ -1,7 +1,7 @@
 name = "BlockDiagonals"
 uuid = "0a1fb500-61f7-11e9-3c65-f5ef3456f9f0"
 authors = ["Invenia Technical Computing Corporation"]
-version = "0.1.14"
+version = "0.2.0"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

src/blockdiagonal.jl

@@ -1,19 +1,22 @@
 # Core functionality for the `BlockDiagonal` type
 
 """
-    BlockDiagonal{T, V<:AbstractMatrix{T}} <: AbstractMatrix{T}
+    BlockDiagonal{T, V} <: AbstractMatrix{T}
+    BlockDiagonal(blocks::V) -> BlockDiagonal{T,V}
 
 A matrix with matrices on the diagonal, and zeros off the diagonal.
-"""
-struct BlockDiagonal{T, V<:AbstractMatrix{T}} <: AbstractMatrix{T}
-    blocks::Vector{V}
 
-    function BlockDiagonal{T, V}(blocks::Vector{V}) where {T, V<:AbstractMatrix{T}}
-        return new{T, V}(blocks)
-    end
+!!! info "`V` type"
+    `blocks::V` should be a `Tuple` or `AbstractVector` where each component (each block) is
+    `<:AbstractMatrix{T}` for some common element type `T`.
+"""
+struct BlockDiagonal{T, V} <: AbstractMatrix{T}
+    blocks::V
 end
 
-function BlockDiagonal(blocks::Vector{V}) where {T, V<:AbstractMatrix{T}}
+function BlockDiagonal(blocks::V) where {
+    T, V<:Union{Tuple{Vararg{<:AbstractMatrix{T}}}, AbstractVector{<:AbstractMatrix{T}}}
+}
     return BlockDiagonal{T, V}(blocks)
 end
 
@@ -151,7 +154,9 @@ function _block_indices(B::BlockDiagonal, i::Integer, j::Integer)
         p += 1
         j -= ncols[p]
     end
-    i -= sum(nrows[1:(p-1)])
+    if !isempty(nrows[1:(p-1)])
+        i -= sum(nrows[1:(p-1)])
+    end
     # if row `i` outside of block `p`, set `p` to place-holder value `-1`
     if i <= 0 || i > nrows[p]
         p = -1

src/chainrules.jl

@@ -1,5 +1,5 @@
 # constructor
-function ChainRulesCore.rrule(::Type{<:BlockDiagonal}, blocks::Vector{V}) where {V}
+function ChainRulesCore.rrule(::Type{<:BlockDiagonal}, blocks::V) where {V<:AbstractVector}
     BlockDiagonal_pullback(Δ::Composite) = (NO_FIELDS, Δ.blocks)
     return BlockDiagonal(blocks), BlockDiagonal_pullback
 end
@@ -27,7 +27,7 @@ function ChainRulesCore.rrule(
         ::typeof(*),
         bm::BlockDiagonal{T, V},
         v::StridedVector{T}
-    ) where {T<:Union{Real, Complex}, V<:Matrix{T}}
+    ) where {T<:Union{Real, Complex}, V<:Vector{Matrix{T}}}
 
     y = bm * v
 

src/linalg.jl

@@ -76,7 +76,7 @@ svdvals_blockwise(B::BlockDiagonal) = mapreduce(svdvals, vcat, blocks(B))
 LinearAlgebra.svdvals(B::BlockDiagonal) = sort!(svdvals_blockwise(B); rev=true)
 
 # `B = U * Diagonal(S) * Vt` with `U` and `Vt` `BlockDiagonal` (`S` only sorted block-wise).
-function svd_blockwise(B::BlockDiagonal{T}; full::Bool=false) where T
+function svd_blockwise(B::BlockDiagonal{T, <:AbstractVector}; full::Bool=false) where T
     U = Matrix{float(T)}[]
     S = Vector{float(T)}()
     Vt = Matrix{float(T)}[]
@@ -88,6 +88,17 @@ function svd_blockwise(B::BlockDiagonal{T}; full::Bool=false) where T
     end
     return BlockDiagonal(U), S, BlockDiagonal(Vt)
 end
+function svd_blockwise(B::BlockDiagonal{T, <:Tuple}; full::Bool=false) where T
+    S = Vector{float(T)}()
+    U_Vt = ntuple(length(blocks(B))) do i
+        F = svd(getblock(B, i), full=full)
+        append!(S, F.S)
+        (F.U, F.Vt)
+    end
+    U = first.(U_Vt)
+    Vt = last.(U_Vt)
+    return BlockDiagonal(U), S, BlockDiagonal(Vt)
+end
 
 function LinearAlgebra.svd(B::BlockDiagonal; full::Bool=false)::SVD
     U, S, Vt = svd_blockwise(B, full=full)

test/base_maths.jl

@@ -7,9 +7,14 @@ using Test
     rng = MersenneTwister(123456)
     N1, N2, N3 = 3, 4, 5
     N = N1 + N2 + N3
-    b1 = BlockDiagonal([rand(rng, N1, N1), rand(rng, N2, N2), rand(rng, N3, N3)])
-    b2 = BlockDiagonal([rand(rng, N1, N1), rand(rng, N3, N3), rand(rng, N2, N2)])
-    b3 = BlockDiagonal([rand(rng, N1, N1), rand(rng, N2, N2), rand(rng, N2, N2)])
+    blocks1 = [rand(rng, N1, N1), rand(rng, N2, N2), rand(rng, N3, N3)]
+    blocks2 = [rand(rng, N1, N1), rand(rng, N3, N3), rand(rng, N2, N2)]
+    blocks3 = [rand(rng, N1, N1), rand(rng, N2, N2), rand(rng, N2, N2)]
+
+    @testset "$T" for (T, (b1, b2, b3)) in (
+        Tuple => (BlockDiagonal(Tuple(blocks1)), BlockDiagonal(Tuple(blocks2)), BlockDiagonal(Tuple(blocks3))),
+        Vector => (BlockDiagonal(blocks1), BlockDiagonal(blocks2), BlockDiagonal(blocks3)),
+    )
     A = rand(rng, N, N + N1)
     B = rand(rng, N + N1, N + N2)
     A′, B′ = A', B'
@@ -127,8 +132,8 @@ using Test
         end
 
         @testset "Non-Square BlockDiagonal * Non-Square BlockDiagonal" begin
-    	    b4 = BlockDiagonal([ones(2, 4), 2 * ones(3, 2)])
-            b5 = BlockDiagonal([3 * ones(2, 2), 2 * ones(4, 1)])
+            b4 = BlockDiagonal(T([ones(2, 4), 2 * ones(3, 2)]))
+            b5 = BlockDiagonal(T([3 * ones(2, 2), 2 * ones(4, 1)]))
 
             @test b4 * b5 isa Array
             @test b4 * b5 == [6 * ones(2, 2) 4 * ones(2, 1); zeros(3, 2) 8 * ones(3, 1)]
@@ -138,3 +143,4 @@ using Test
         end
     end  # Multiplication
 end
+end

test/blockdiagonal.jl

@@ -7,9 +7,14 @@ using Test
     rng = MersenneTwister(123456)
     N1, N2, N3 = 3, 4, 5
     N = N1 + N2 + N3
-    b1 = BlockDiagonal([rand(rng, N1, N1), rand(rng, N2, N2), rand(rng, N3, N3)])
-    b2 = BlockDiagonal([rand(rng, N1, N1), rand(rng, N3, N3), rand(rng, N2, N2)])
-    b3 = BlockDiagonal([rand(rng, N1, N1), rand(rng, N2, N2), rand(rng, N2, N2)])
+    blocks1 = [rand(rng, N1, N1), rand(rng, N2, N2), rand(rng, N3, N3)]
+    blocks2 = [rand(rng, N1, N1), rand(rng, N3, N3), rand(rng, N2, N2)]
+    blocks3 = [rand(rng, N1, N1), rand(rng, N2, N2), rand(rng, N2, N2)]
+
+    @testset "$T" for (T, (b1, b2, b3)) in (
+        Tuple => (BlockDiagonal(Tuple(blocks1)), BlockDiagonal(Tuple(blocks2)), BlockDiagonal(Tuple(blocks3))),
+        Vector => (BlockDiagonal(blocks1), BlockDiagonal(blocks2), BlockDiagonal(blocks3)),
+    )
     A = rand(rng, N, N + N1)
     B = rand(rng, N + N1, N + N2)
     A′, B′ = A', B'
@@ -42,8 +47,10 @@ using Test
         end
 
         @testset "parent" begin
-            @test parent(b1) isa Vector{<:AbstractMatrix}
+            @test parent(b1) isa Union{Tuple,AbstractVector}
+            @test eltype(parent(b1)) <: AbstractMatrix
             @test parent(BlockDiagonal([X, Y])) == [X, Y]
+            @test parent(BlockDiagonal((X, Y))) == (X, Y)
         end
 
         @testset "similar" begin
@@ -117,3 +124,4 @@ using Test
         @test_throws DimensionMismatch copy!(b2, b1)
     end
 end
+end

test/linalg.jl

@@ -8,9 +8,14 @@ using Test
     rng = MersenneTwister(123456)
     N1, N2, N3 = 3, 4, 5
     N = N1 + N2 + N3
-    b1 = BlockDiagonal([rand(rng, N1, N1), rand(rng, N2, N2), rand(rng, N3, N3)])
-    b2 = BlockDiagonal([rand(rng, N1, N1), rand(rng, N3, N3), rand(rng, N2, N2)])
-    b3 = BlockDiagonal([rand(rng, N1, N1), rand(rng, N2, N2), rand(rng, N2, N2)])
+    blocks1 = [rand(rng, N1, N1), rand(rng, N2, N2), rand(rng, N3, N3)]
+    blocks2 = [rand(rng, N1, N1), rand(rng, N3, N3), rand(rng, N2, N2)]
+    blocks3 = [rand(rng, N1, N1), rand(rng, N2, N2), rand(rng, N2, N2)]
+
+    @testset "$T" for (T, (b1, b2, b3)) in (
+        Tuple => (BlockDiagonal(Tuple(blocks1)), BlockDiagonal(Tuple(blocks2)), BlockDiagonal(Tuple(blocks3))),
+        Vector => (BlockDiagonal(blocks1), BlockDiagonal(blocks2), BlockDiagonal(blocks3)),
+    )
     A = rand(rng, N, N + N1)
     B = rand(rng, N + N1, N + N2)
     A′, B′ = A', B'
@@ -143,7 +148,7 @@ using Test
 
         @testset "eigvals on LinearAlgebra types" begin
             # `eigvals` has different methods for different types, e.g. Hermitian
-            b_herm = BlockDiagonal([Hermitian(rand(rng, 3, 3) + I) for _ in 1:3])
+            b_herm = BlockDiagonal(T(Hermitian(rand(rng, 3, 3) + I) for _ in 1:3))
             @test eigvals(b_herm) ≈ eigvals(Matrix(b_herm))
             @test eigvals(b_herm, 1.0, 2.0) ≈ eigvals(Hermitian(Matrix(b_herm)), 1.0, 2.0)
         end
@@ -163,21 +168,21 @@ using Test
               0.0 1.0  5.0
               0.0 0.0  3.0]
 
-        B = BlockDiagonal([X, X])
+        B = BlockDiagonal(T([X, X]))
         C = cholesky(B)
         @test C isa Cholesky{Float64, <:BlockDiagonal{Float64}}
         @test C.U ≈ cholesky(Matrix(B)).U
-        @test C.U ≈ BlockDiagonal([U, U])
-        @test C.L ≈ BlockDiagonal([U', U'])
+        @test C.U ≈ BlockDiagonal(T([U, U]))
+        @test C.L ≈ BlockDiagonal(T([U', U']))
         @test C.UL ≈ C.U
         @test C.uplo === 'U'
         @test C.info == 0
 
         M = BlockDiagonal(map(Matrix, blocks(C.L)))
         C = Cholesky(M, 'L', 0)
         @test C.U ≈ cholesky(Matrix(B)).U
-        @test C.U ≈ BlockDiagonal([U, U])
-        @test C.L ≈ BlockDiagonal([U', U'])
+        @test C.U ≈ BlockDiagonal(T([U, U]))
+        @test C.L ≈ BlockDiagonal(T([U', U']))
         @test C.UL ≈ C.L
         @test C.uplo === 'L'
         @test C.info == 0
@@ -187,7 +192,7 @@ using Test
         X = [  4  12 -16
               12  37 -43
              -16 -43  98]
-        B = BlockDiagonal([X, X])
+        B = BlockDiagonal(T([X, X]))
 
         @testset "full=$full" for full in (true, false)
 
@@ -230,3 +235,4 @@ using Test
         end
     end  # SVD
 end
+end