Solving Block Diagonal Systems

A block diagonal matrix is one whose nonzeros sit entirely in square blocks along the diagonal, so the system it defines is really a batch of smaller, independent systems stacked together:

\[\begin{bmatrix} A_1 & & \\ & A_2 & \\ & & A_3 \end{bmatrix} \begin{bmatrix} u_1 \\ u_2 \\ u_3 \end{bmatrix} = \begin{bmatrix} b_1 \\ b_2 \\ b_3 \end{bmatrix}\]

These show up whenever a problem decouples into independent subsystems: per-element systems in a finite element assembly, per-sample systems in a batched inverse problem, or the diagonal approximation of a larger operator.

BlockDiagonals.jl provides a BlockDiagonal type that stores only the blocks, and LinearSolve.jl accepts it directly as the A of a LinearProblem.

Setting Up the Problem

import LinearSolve as LS
import LinearAlgebra as LA
using BlockDiagonals

blocks = [rand(3, 3) + 3LA.I for _ in 1:4]
A = BlockDiagonal(blocks)
b = rand(size(A, 1))

prob = LS.LinearProblem(A, b)
sol = LS.solve(prob)
sol.u

12-element Vector{Float64}:
  0.2223992953294319
  0.050051774323911835
 -0.053779641099343106
  0.01403996367050804
  0.01471783776089492
  0.2615490732387348
  0.07406233920676945
  0.12814578795703338
  0.13763330863002882
  0.0033388005792773516
  0.134663696772687
  0.10176765007947158

The blocks do not have to be the same size:

A_mixed = BlockDiagonal([rand(n, n) + n * LA.I for n in (2, 3, 4)])
b_mixed = rand(size(A_mixed, 1))

LS.solve(LS.LinearProblem(A_mixed, b_mixed)).u

9-element Vector{Float64}:
  0.10779738382434123
  0.09030068997905083
  0.3130743523786977
 -0.04887786445787903
 -0.044732456170131066
  0.14870134800719995
 -0.005607246185176176
  0.10541889156249867
  0.03460674741857196

Which Solvers Work

All of the standard dense algorithms accept a BlockDiagonal:

u_ref = Matrix(A) \ b

for alg in (LS.LUFactorization(), LS.GenericFactorization(),
        LS.QRFactorization(LA.NoPivot()), LS.SimpleGMRES(), LS.KrylovJL_GMRES())
    u = LS.solve(LS.LinearProblem(A, b), alg).u
    println(rpad(nameof(typeof(alg)), 22), " residual = ",
        LA.norm(Matrix(A) * u - b))
end

LUFactorization        residual = 1.6883057536160649e-16
GenericFactorization   residual = 2.0955000055363631e-16
QRFactorization        residual = 2.4552312978479334e-16
SimpleGMRES            residual = 4.687735131751618e-16
KrylovJL               residual = 1.9215096052097404e-9

Sparse-only algorithms such as UMFPACKFactorization and KLUFactorization expect a SparseMatrixCSC, so convert with sparse(A) first if you want those.

Blockwise Factorization

LUFactorization and QRFactorization do not treat the matrix as a dense N x N system. Because the blocks are independent, each one is factorized on its own, which turns O(N^3) work into O(\sum m_i^3) and lets every per-block call land on BLAS. For k blocks of size m that is roughly a factor of k^2 less arithmetic.

The cache stores the per-block factorizations rather than one big dense one:

cache_lu = LS.init(LS.LinearProblem(A, b), LS.LUFactorization())
LS.solve!(cache_lu)

(n_block_factorizations = length(cache_lu.cacheval.facts),
    block_factorization = nameof(typeof(first(cache_lu.cacheval.facts))))

(n_block_factorizations = 4, block_factorization = :LU)

Two cases deliberately keep the generic dense path, because they are not a batch of independent square systems or need machinery the blockwise path does not carry:

blocks that are not square, which can still add up to a square matrix but do not decompose into independent subsystems, and
LUFactorization(residualsafety = true), which uses the a-posteriori residual check of the standard LU solve.

Both still give correct answers, just without the structural speedup.

The `SimpleGMRES` Specialization

The iterative side has its own specialization. When every block is the same size, SimpleGMRES solves the batch of subsystems together instead of running one Krylov iteration over the whole stacked system. LinearSolve.jl detects this automatically when the extension is loaded, so no keyword is needed:

cache_uniform = LS.init(LS.LinearProblem(A, b), LS.SimpleGMRES())
LS.solve!(cache_uniform).u

12-element Vector{Float64}:
  0.22239929532943184
  0.050051774323911814
 -0.053779641099343085
  0.014039963670508034
  0.014717837760894914
  0.26154907323873483
  0.07406233920676948
  0.12814578795703338
  0.13763330863002884
  0.003338800579277334
  0.13466369677268697
  0.10176765007947156

You can see the specialization in the cache that init builds. The first type parameter of the cache records whether the batched path was selected, and the blocksize field records the detected block size:

cache_mixed = LS.init(LS.LinearProblem(A_mixed, b_mixed), LS.SimpleGMRES())

(uniform_blocksize = cache_uniform.cacheval.blocksize,
    mixed_blocksize = cache_mixed.cacheval.blocksize)

(uniform_blocksize = 3, mixed_blocksize = 0)

The uniform case reports the block size it found, while the mixed case reports 0, meaning it fell back to the generic path. Both give the right answer; only the uniform case gets the batched treatment.

If your matrix is stored as a plain dense Matrix but you know it is block diagonal with uniform blocks, you can request the same specialization explicitly with the blocksize keyword:

LS.solve(LS.LinearProblem(Matrix(A), b), LS.SimpleGMRES(; blocksize = 3)).u

12-element Vector{Float64}:
  0.22239929532943184
  0.050051774323911793
 -0.05377964109934309
  0.01403996367050804
  0.014717837760894914
  0.26154907323873483
  0.07406233920676948
  0.12814578795703338
  0.13763330863002884
  0.003338800579277345
  0.13466369677268697
  0.1017676500794716

Reusing the Factorization

The caching interface works with BlockDiagonal exactly as it does for any other matrix, which matters when the same block structure is solved against many right hand sides:

cache = LS.init(LS.LinearProblem(A, b), LS.QRFactorization(LA.NoPivot()))
sol1 = LS.solve!(cache)

cache.b = rand(size(A, 1))
sol2 = LS.solve!(cache)

sol2.u

12-element Vector{Float64}:
  0.012107813524070681
  0.05685676136100502
  0.24998630609394593
  0.07056200468581512
 -0.05768707551528127
  0.29668405092084965
  0.22346766596839007
  0.2201080957272611
 -0.039562494263392936
  0.22338859039357933
 -0.025575181383727528
  0.07145824262745717

The second solve reuses the stored per-block factorizations and only applies them to the new right hand side.