How to concatenate matrices in diagonal form in Julia

user_11446037 0
Viewed 593

I have a matrix A and I want to build a block diagonal matrix M as

A 0 ⋯ 0
0 A ⋯ 0
⋮ ⋮ ⋱ ⋮
0 0 ⋯ A

where the 0s are matrices of the same size as A filled with zeros. Is there a convenient way to do that?

In other words, is there an equivalent to Python's linalg.block_diag (from scipy)?

diagonal julia concatenation matrix
3 Answers

You can use blockdiag from the SparseArrays.jl standard library. For example:

julia> using SparseArrays

julia> A = sparse(rand(3,2)) # A must be sparse for blockdiag
3×2 SparseMatrixCSC{Float64, Int64} with 6 stored entries:
 0.465226  0.473656
 0.230678  0.391924
 0.928409  0.772551

julia> M = blockdiag(A, A, A)
9×6 SparseMatrixCSC{Float64, Int64} with 18 stored entries:
 0.465226  0.473656   ⋅         ⋅         ⋅         ⋅
 0.230678  0.391924   ⋅         ⋅         ⋅         ⋅
 0.928409  0.772551   ⋅         ⋅         ⋅         ⋅
  ⋅         ⋅        0.465226  0.473656   ⋅         ⋅
  ⋅         ⋅        0.230678  0.391924   ⋅         ⋅
  ⋅         ⋅        0.928409  0.772551   ⋅         ⋅
  ⋅         ⋅         ⋅         ⋅        0.465226  0.473656
  ⋅         ⋅         ⋅         ⋅        0.230678  0.391924
  ⋅         ⋅         ⋅         ⋅        0.928409  0.772551

And if you really need M to be dense, you can just convert it back with

julia> Matrix(M)
9×6 Matrix{Float64}:
 0.465226  0.473656  0.0       0.0       0.0       0.0
 0.230678  0.391924  0.0       0.0       0.0       0.0
 0.928409  0.772551  0.0       0.0       0.0       0.0
 0.0       0.0       0.465226  0.473656  0.0       0.0
 0.0       0.0       0.230678  0.391924  0.0       0.0
 0.0       0.0       0.928409  0.772551  0.0       0.0
 0.0       0.0       0.0       0.0       0.465226  0.473656
 0.0       0.0       0.0       0.0       0.230678  0.391924
 0.0       0.0       0.0       0.0       0.928409  0.772551
user_56021490

You can do this with cat, giving it two dimensions:

julia> A = [1 2; 3 4]
2×2 Matrix{Int64}:
 1  2
 3  4

julia> cat(A, A', reverse(A); dims=(1,2))
6×6 Matrix{Int64}:
 1  2  0  0  0  0
 3  4  0  0  0  0
 0  0  1  3  0  0
 0  0  2  4  0  0
 0  0  0  0  4  3
 0  0  0  0  2  1
user_137067350

If you don't want to allocate the extra space for the copies of the matrix and the 0s, you can use the following struct that I wrote, which emulates a block-diagonal matrix.

struct BlockDiagonalMatrix{T} <: AbstractMatrix{T}
    matrix::AbstractMatrix{T}
    repeats::Integer

    BlockDiagonalMatrix(mat::AbstractMatrix{T}, repeats) where {T} = new{T}(mat, repeats)
end

Base.IndexStyle(::Type{BlockDiagonalMatrix}) = IndexLinear()
Base.size(bdm::BlockDiagonalMatrix) = size(bdm.matrix) .* bdm.repeats
function Base.getindex(bdm::BlockDiagonalMatrix{T}, i) where {T}
    bdm_height = size(bdm, 1)

    i -= 1  # Math is slightly simpler when indexing is 0-based
    i_col, i_row = divrem(i, bdm_height)

    block_height, block_width = size(bdm.matrix)
    block_col = fld(i_col, block_width)
    block_row = fld(i_row, block_height)

    if block_row == block_col
        inner_mat_col = 1 + (i_col % block_width)  # convert back to 1-based
        inner_mat_row = 1 + (i_row % block_height)
        return bdm.matrix[inner_mat_row, inner_mat_col]
    else
        return zero(T)
    end
end

Now, this emulates the full block-diagonal matrix while using just the storage of the one array. In fact, even the pretty-printing comes for free.

julia> bdm = BlockDiagonalMatrix([1 2; 3 4; 5 6], 4)
12×8 BlockDiagonalMatrix{Int64}:
 1  2  0  0  0  0  0  0
 3  4  0  0  0  0  0  0
 5  6  0  0  0  0  0  0
 0  0  1  2  0  0  0  0
 0  0  3  4  0  0  0  0
 0  0  5  6  0  0  0  0
 0  0  0  0  1  2  0  0
 0  0  0  0  3  4  0  0
 0  0  0  0  5  6  0  0
 0  0  0  0  0  0  1  2
 0  0  0  0  0  0  3  4
 0  0  0  0  0  0  5  6

julia> bdm = BlockDiagonalMatrix([1.0 2.0], 4)
4×8 BlockDiagonalMatrix{Float64}:
 1.0  2.0  0.0  0.0  0.0  0.0  0.0  0.0
 0.0  0.0  1.0  2.0  0.0  0.0  0.0  0.0
 0.0  0.0  0.0  0.0  1.0  2.0  0.0  0.0
 0.0  0.0  0.0  0.0  0.0  0.0  1.0  2.0

You can learn more about how to implement a custom array type here – it takes surprisingly little work to achieve the full breadth of functionality offered by the built-in Array.

user_54964330
Related