Static Condensation of structured array

StructuredStaticCondensation.jl

using LinearAlgebra, SparseArrays

struct StaticCondensationMatrix{T, M<:AbstractMatrix{T}} <: AbstractMatrix{T}
  A::M
  indices::Vector{UnitRange{Int}}
  nlocalblocks::Int
  ncouplingblocks::Int
  reducedlocalindices::Vector{UnitRange{Int}}
  reducedcoupledindices::Vector{UnitRange{Int}}
  reducedlhs::M
end

function StaticCondensationMatrix(A::AbstractMatrix{T}, blockindices) where T
  n = size(A, 1)
  @assert isodd(length(blockindices))
  ncouplingblocks = (length(blockindices) - 1) ÷ 2
  nlocalblocks = ncouplingblocks + 1

  reducedlocalindices = Vector{UnitRange{Int}}()
  push!(reducedlocalindices, blockindices[1])
  for (c, i) in enumerate(3:2:length(blockindices))
    inds = (blockindices[i] .- blockindices[i][1] .+ 1) .+ reducedlocalindices[c][end]
    push!(reducedlocalindices, inds) 
  end
  reducedcoupledindices = Vector{UnitRange{Int}}()
  push!(reducedcoupledindices, blockindices[2] .- blockindices[2][1] .+ 1)
  for (c, i) in enumerate(4:2:length(blockindices))
    inds = (blockindices[i] .- blockindices[i][1] .+ 1) .+ reducedcoupledindices[c][end]
    push!(reducedcoupledindices, inds) 
  end

  totalcouplingblocksize = sum(length(i) for i in reducedcoupledindices)
  # allocate reduced lhs
  reducedlhs = similar(A, totalcouplingblocksize, totalcouplingblocksize)
  fill!(reducedlhs, 0)

  return StaticCondensationMatrix(A, blockindices, nlocalblocks, ncouplingblocks, reducedlocalindices, reducedcoupledindices, reducedlhs)
end

function StaticCondensationMatrix(A::AbstractMatrix{T}, localblocksize, couplingblocksize) where T
  n = size(A, 1)
  ncouplingblocks = (n - localblocksize) ÷ (localblocksize + couplingblocksize)
  nlocalblocks = ncouplingblocks + 1

  indices = Vector{UnitRange{Int}}()
  a = 1
  for i = 1:nlocalblocks-1
    inds = a:a + localblocksize - 1
    push!(indices, inds)
    a = a + localblocksize
    inds = a:a + couplingblocksize - 1
    push!(indices, inds)
    a = a + couplingblocksize
  end
  push!(indices, a:a + localblocksize - 1)
  @assert indices[end][end] == size(A, 1) == size(A, 2)
  return StaticCondensationMatrix(A, indices)
end
Base.size(A::StaticCondensationMatrix) = (size(A.A, 1), size(A.A, 2))
Base.size(A::StaticCondensationMatrix, i) = size(A.A, i)
islocalblock(i) = isodd(i)
iscouplingblock(i) = !islocalblock(i)
function enumeratelocalindices(A::StaticCondensationMatrix)
  return zip(1:A.nlocalblocks, 1:2:length(A.indices), A.indices[1:2:end])
end
function enumeratecouplingindices(A::StaticCondensationMatrix)
  return zip(1:A.ncouplingblocks, 2:2:length(A.indices), A.indices[2:2:end-1])
end

function factoriselocals(A::StaticCondensationMatrix{T}) where T
  lis = A.indices
  localfact = lu!(view(A.A, lis[1], lis[1])) # can't use a view
  d = Dict{Int, typeof(localfact)}(1=>localfact)
  for (i_1, li) in enumerate(lis[2:end]) # parallelisable
    i = i_1 + 1
    iscouplingblock(i) && continue
    d[i] = lu!(view(A.A, li, li)) # can't use a view
  end
  return d
end

function calculatecouplings(A::StaticCondensationMatrix{T,M}, localfactors) where {T,M}
  d = Dict{Tuple{Int, Int}, M}()
  for (c, i, li) in enumeratecouplingindices(A) # parallelisable
    if i - 1 >= 1
      lim = A.indices[i-1]
      d[(i-1, i)] = localfactors[i-1] \ A.A[lim, li]
    end
    if i + 1 <= length(A.indices)
      lip = A.indices[i+1]
      d[(i+1, i)] = localfactors[i+1] \ A.A[lip, li]
    end
  end
  return d
end

function solvelocalparts(A::StaticCondensationMatrix{T,M}, b, localfactors) where {T,M}
  d = Dict{Int, M}()
  for (c, i, li) in enumeratelocalindices(A) # parallelisable
    d[i] = localfactors[i] \ b[li, :]
  end
  return d
end

function couplingblockindices(A::StaticCondensationMatrix, i)
  @assert i > 1
  return A.indices[i] .- A.indices[i-1][1] .+ 1
end

function localblockindices(A::StaticCondensationMatrix, i)
  @assert i > 1
  return A.indices[i] .- A.indices[i-1][1] .+ 1
end

function assemblecoupledrhs(A::StaticCondensationMatrix, B, localsolutions, couplings)
  b = similar(A.reducedlhs, size(A.reducedlhs, 1), size(B, 2))
  @views for (c, i, li) in enumeratecouplingindices(A) # parallelisable
    rows = A.reducedcoupledindices[c]
    b[rows, :] .= B[li, :]
    b[rows, :] .-= A.A[li, A.indices[i-1]] * localsolutions[i-1]
    b[rows, :] .-= A.A[li, A.indices[i+1]] * localsolutions[i+1]
  end
  return b
end

function assemblecoupledlhs!(A::StaticCondensationMatrix, couplings; assignblocks=false, applycouplings=false)
  M = A.reducedlhs
  @views for (c, i, li) in enumeratecouplingindices(A) # parallelisable
    rows = A.reducedcoupledindices[c]
    assignblocks && (M[rows, rows] .= A.A[li, li])
    aim = A.A[li, A.indices[i-1]]
    aip = A.A[li, A.indices[i+1]]
    applycouplings && (M[rows, rows] .-= aim * couplings[(i - 1, i)])
    applycouplings && (M[rows, rows] .-= aip * couplings[(i + 1, i)])
    if c + 1 <= A.ncouplingblocks
      right = A.reducedcoupledindices[c + 1]
      assignblocks && (M[rows, right] .= A.A[li, A.indices[i + 2]])
      applycouplings && (M[rows, right] .-= aip * couplings[(i + 1, i + 2)])
    end
    if c - 1 >= 1
      left = A.reducedcoupledindices[c - 1]
      assignblocks && (M[rows, left] .= A.A[li, A.indices[i - 2]])
      applycouplings && (M[rows, left] .-= aim * couplings[(i - 1, i - 2)])
    end
  end
  return M
end

function coupledx!(x, A::StaticCondensationMatrix{T}, b, localsolutions, couplings) where {T}
  Ac = assemblecoupledlhs!(A, couplings; applycouplings=true)
  bc = assemblecoupledrhs(A, b, localsolutions, couplings)
  xc = Ac \ bc
  for (c, i, ind) in enumeratecouplingindices(A)
    x[ind, :] .= xc[A.reducedcoupledindices[c], :]
  end
  return x
end

function localx!(x, A::StaticCondensationMatrix{T}, b, localsolutions, couplings) where T
  for (c, i, li) in enumeratelocalindices(A) # parallelisable
    rows = A.reducedlocalindices[c]
    x[li, :] .= localsolutions[i]
    for j in (i + 1, i - 1)
      0 < j <= length(A.indices) || continue
      x[li, :] .-= couplings[(i, j)] * x[A.indices[j], :]
    end
  end
  return x
end

function solve(A::StaticCondensationMatrix{T}, b) where T
  assemblecoupledlhs!(A, nothing; assignblocks=true)
  localfactors = factoriselocals(A)
  localsolutions = solvelocalparts(A, b, localfactors)
  couplings = calculatecouplings(A, localfactors)

  x = zeros(T, size(b))
  x = coupledx!(x, A, b, localsolutions, couplings)
  x = localx!(x, A, b, localsolutions, couplings)
  return x
end

using Random, Test
Random.seed!(0)
@testset "StaticCondensationCondensation" begin
for (L, C) in ((2, 1), (3, 2), (4, 2), (16, 4), (5, 7), (128, 64), (256, 128), (1024, 512))
  A11, A33, A55, A77 = (rand(L, L) for _ in 1:4)
  t12, t32, t34, t54, t56, t76 = (rand(L, C) for _ in 1:6)
  s21, s23, s43, s45, s65, s67 = (rand(C, L) for _ in 1:6)
  a22, a44, a66, a24, a42, a46, a64 = (rand(C, C) for _ in 1:7)
  zLL = zeros(L, L)
  zCL = zeros(C, L)
  zLC = zeros(L, C)
  zCC = zeros(C, C)
  
  A = [A11 t12 zLL;
       s21 a22 s23;
       zLL t32 A33]
  b = rand(size(A, 1))
  x = A \ b
  SCM = StaticCondensationMatrix(A, L, C)
  @test solve(SCM, b) ≈ x
  
  A = [A11 t12 zLL zLC zLL;
       s21 a22 s23 a24 zCL;
       zLL t32 A33 t34 zLL;
       zCL a42 s43 a44 s45;
       zLL zLC zLL t54 A55]
  b = rand(size(A, 1))
  x = A \ b
  SCM = StaticCondensationMatrix(A, L, C)
  @test  solve(SCM, b) ≈ x
  
  A = [A11 t12 zLL zLC zLL zLC zLL;
       s21 a22 s23 a24 zCL zCC zCL;
       zLL t32 A33 t34 zLL zLC zLL;
       zCL a42 s43 a44 s45 a46 zCL;
       zLL zLC zLL t54 A55 t56 zLL;
       zCL zCC zCL a64 s65 a66 s67;
       zLL zLC zLL zLC zLL t76 A77]
  
  b = rand(size(A, 1))
  x = A \ b
  SCM = StaticCondensationMatrix(A, L, C)
  @test solve(SCM, b) ≈ x
end
end