From 4dfccd9de1d823b156c12e30b56a506d42dcff10 Mon Sep 17 00:00:00 2001 From: Fredrik Ekre Date: Fri, 16 Dec 2022 22:53:33 +0100 Subject: [PATCH] Add package extension for BlockArrays.jl. This patch adds support for global block matrices. A global block matrix can be obtained by `create_sparsity_pattern(BlockMatrix, dh, ch)` or by `create_sparsity_pattern(A::BlockMatrix, dh, ch)`. The former infers the block sizes by looking at the fields of the DofHandler, and the latters uses the blocksizes of the (partially) initilized BlockMatrix (i.e. by `undef_blocks`). For global assembly, `start_assemble(::BlockMatrix, ...)` returns a `BlockAssembler` which pre allocates a vector for pre-computing the (block number, local index) pairs for each DoF. BlockArrays only support local condensation of constraints currently and you have to use `apply_assemble!` instead of `assemble!` + `apply!`. --- Project.toml | 2 + ext/BlockArraysExt.jl | 149 ++++++++++++++++++++++++++++++++++++++++++ test/blockarrays.jl | 87 ++++++++++++++++++++++++ test/runtests.jl | 1 + 4 files changed, 239 insertions(+) create mode 100644 ext/BlockArraysExt.jl create mode 100644 test/blockarrays.jl diff --git a/Project.toml b/Project.toml index 4bd79e1ead..e875010279 100644 --- a/Project.toml +++ b/Project.toml @@ -13,9 +13,11 @@ Tensors = "48a634ad-e948-5137-8d70-aa71f2a747f4" WriteVTK = "64499a7a-5c06-52f2-abe2-ccb03c286192" [weakdeps] +BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e" Metis = "2679e427-3c69-5b7f-982b-ece356f1e94b" [extensions] +BlockArraysExt = "BlockArrays" MetisExt = "Metis" [compat] diff --git a/ext/BlockArraysExt.jl b/ext/BlockArraysExt.jl new file mode 100644 index 0000000000..8489004284 --- /dev/null +++ b/ext/BlockArraysExt.jl @@ -0,0 +1,149 @@ +module BlockArraysExt + +using BlockArrays: BlockArray, BlockIndex, BlockMatrix, BlockVector, block, blockaxes, + blockindex, blocks, findblockindex +using Ferrite +using Ferrite: addindex!, fillzero! + +# TODO: Move into Ferrite +function global_dof_range(dh::DofHandler, f::Symbol) + set = Set{Int}() + frange = dof_range(dh, f) + for cc in CellIterator(dh) + union!(set, @view celldofs(cc)[frange]) + end + dofmin, dofmax = extrema(set) + r = dofmin:dofmax + if length(set) != length(r) + error("renumber by blocks you donkey") + end + return r +end + +################################### +## Creating the sparsity pattern ## +################################### + +# Note: +# Creating the full unblocked matrix and then splitting into blocks inside the BlockArray +# constructor (i.e. by `getindex(::SparseMatrixCSC, ::UnitRange, ::UnitRange)`) is +# consistently faster than creating individual blocks directly. However, the latter approach +# uses less than half of the memory (measured for a 2x2 block system and various problem +# sizes), so might be useful in the future to provide an option on what algorithm to use. + +# TODO: Could potentially extract the element type and matrix type for the individual blocks +# by allowing e.g. create_sparsity_pattern(BlockMatrix{Float32}, ...) but that is not +# even supported by regular pattern right now. +function Ferrite.create_sparsity_pattern(::Type{<:BlockMatrix}, dh, ch; kwargs...) + K = create_sparsity_pattern(dh, ch; kwargs...) + # Infer block sizes from the fields in the DofHandler + global_ranges = [global_dof_range(dh, f) for f in dh.field_names] + block_sizes = length.(global_ranges) + return BlockArray(K, block_sizes, block_sizes) +end + +function Ferrite.create_sparsity_pattern(B::BlockMatrix, dh, ch; kwargs...) + if !(size(B, 1) == size(B, 2) == ndofs(dh)) + error("size of input matrix ($(size(B))) does not match number of dofs ($(ndofs(dh)))") + end + K = create_sparsity_pattern(dh, ch; kwargs...) + ax = axes(B) + for block_j in blockaxes(B, 2), block_i in blockaxes(B, 1) + range_j = ax[2][block_j] + range_i = ax[1][block_i] + B[block_i, block_j] = K[range_i, range_j] + end + return B +end + +########################################### +## BlockAssembler and associated methods ## +########################################### + +struct BlockAssembler{BM, Bv} <: Ferrite.AbstractSparseAssembler + K::BM + f::Bv + blockindices::Vector{BlockIndex{1}} +end + +Ferrite.matrix_handle(ba::BlockAssembler) = ba.K +Ferrite.vector_handle(ba::BlockAssembler) = ba.f + +function Ferrite.start_assemble(K::BlockMatrix, f; fillzero::Bool=true) + fillzero && (fillzero!(K); fillzero!(f)) + return BlockAssembler(K, f, BlockIndex{1}[]) +end + +# Split into the block and the local index +splindex(idx::BlockIndex{1}) = (block(idx), blockindex(idx)) + +function Ferrite.assemble!(assembler::BlockAssembler, dofs::AbstractVector{<:Integer}, ke::AbstractMatrix, fe::AbstractVector) + K = assembler.K + f = assembler.f + blockindices = assembler.blockindices + + @assert blockaxes(K, 1) == blockaxes(K, 2) + @assert axes(K, 1) == axes(K, 2) == axes(f, 1) + @boundscheck checkbounds(K, dofs, dofs) + @boundscheck checkbounds(f, dofs) + + # Update the cached the block indices + resize!(blockindices, length(dofs)) + @inbounds for (i, I) in pairs(dofs) + blockindices[i] = findblockindex(axes(K, 1), I) + end + + # Assemble matrix entries + @inbounds for (j, blockindex_j) in pairs(blockindices) + Bj, lj = splindex(blockindex_j) + for (i, blockindex_i) in pairs(blockindices) + Bi, li = splindex(blockindex_i) + KB = @view K[Bi, Bj] + addindex!(KB, ke[i, j], li, lj) + end + end + + # Assemble vector entries + if blockaxes(f) == blockaxes(K, 1) + # If f::BlockVector with the same axes the same blockindex cache can be used... + @inbounds for (i, blockindex_i) in pairs(blockindices) + Bi, li = splindex(blockindex_i) + fB = @view f[Bi] + fB[li] += fe[i] + end + else + # ... otherwise, use regular indexing in fallback assemble! + @inbounds assemble!(f, dofs, fe) + end + return +end + +function Ferrite.apply!(::BlockMatrix, ::AbstractVector, ::ConstraintHandler) + error( + "Condensation of constraints with `apply!` after assembling not supported yet " * + "for BlockMatrix, use local condensation with `apply_assemble!` instead." + ) +end + + +####################################################### +## Overloaded assembly pieces from src/arrayutils.jl ## +####################################################### + +function Ferrite.addindex!(B::BlockMatrix{Tv}, v::Tv, i::Int, j::Int) where Tv + @boundscheck checkbounds(B, i, j) + Bi, li = splindex(findblockindex(axes(B, 1), i)) + Bj, lj = splindex(findblockindex(axes(B, 2), j)) + BB = @view B[Bi, Bj] + @inbounds addindex!(BB, v, li, lj) + return B +end + +function Ferrite.fillzero!(B::Union{BlockVector,BlockMatrix}) + for blk in blocks(B) + fillzero!(blk) + end + return B +end + +end # module BlockArraysExt diff --git a/test/blockarrays.jl b/test/blockarrays.jl new file mode 100644 index 0000000000..4415069446 --- /dev/null +++ b/test/blockarrays.jl @@ -0,0 +1,87 @@ +using Ferrite, BlockArrays, Test + +@testset "BlockArrays.jl extension" begin + grid = generate_grid(Triangle, (10, 10)) + + dh = DofHandler(grid) + push!(dh, :u, 2) + push!(dh, :p, 1) + close!(dh) + renumber!(dh, DofOrder.FieldWise()) + nd = ndofs(dh) ÷ 3 + + ch = ConstraintHandler(dh) + periodic_faces = collect_periodic_faces(grid, "top", "bottom") + add!(ch, PeriodicDirichlet(:u, periodic_faces)) + add!(ch, Dirichlet(:u, union(getfaceset(grid, "left"), getfaceset(grid, "top")), (x, t) -> [0, 0])) + add!(ch, Dirichlet(:p, getfaceset(grid, "left"), (x, t) -> 0)) + close!(ch) + update!(ch, 0) + + K = create_sparsity_pattern(dh, ch) + f = zeros(axes(K, 1)) + KB = create_sparsity_pattern(BlockMatrix, dh, ch) + @test KB isa BlockMatrix + @test blocksize(KB) == (2, 2) + @test size(KB[Block(1), Block(1)]) == (2nd, 2nd) + @test size(KB[Block(2), Block(1)]) == (1nd, 2nd) + @test size(KB[Block(1), Block(2)]) == (2nd, 1nd) + @test size(KB[Block(2), Block(2)]) == (1nd, 1nd) + fB = similar(KB, axes(KB, 1)) + + # Test the pattern + fill!(K.nzval, 1) + foreach(x -> fill!(x.nzval, 1), blocks(KB)) + @test K == KB + + # Zeroing out in start_assemble + assembler = start_assemble(K, f) + @test iszero(K) + @test iszero(f) + block_assembler = start_assemble(KB, fB) + @test iszero(KB) + @test iszero(fB) + + # Assembly procedure + npc = ndofs_per_cell(dh) + for cc in CellIterator(dh) + ke = rand(npc, npc) + fe = rand(npc) + dofs = celldofs(cc) + # Standard assemble + assemble!(assembler, dofs, ke, fe) + assemble!(block_assembler, dofs, ke, fe) + # Assemble with local condensation of constraints + let ke = copy(ke), fe = copy(fe) + apply_assemble!(assembler, ch, dofs, ke, fe) + end + let ke = copy(ke), fe = copy(fe) + apply_assemble!(block_assembler, ch, dofs, ke, fe) + end + end + @test K ≈ KB + @test f ≈ fB + + # Global application of BC not supported yet + @test_throws ErrorException apply!(KB, fB, ch) + + # Custom blocking by passing a partially initialized matrix + perm = invperm([ch.free_dofs; ch.prescribed_dofs]) + renumber!(dh, ch, perm) + nfree = length(ch.free_dofs) + npres = length(ch.prescribed_dofs) + K = create_sparsity_pattern(dh, ch) + block_sizes = [nfree, npres] + KBtmp = BlockArray(undef_blocks, SparseMatrixCSC{Float64, Int}, block_sizes, block_sizes) + KB = create_sparsity_pattern(KBtmp, dh, ch) + @test KBtmp === KB + @test blocksize(KB) == (2, 2) + @test size(KB[Block(1), Block(1)]) == (nfree, nfree) + @test size(KB[Block(2), Block(1)]) == (npres, nfree) + @test size(KB[Block(1), Block(2)]) == (nfree, npres) + @test size(KB[Block(2), Block(2)]) == (npres, npres) + # Test the pattern + fill!(K.nzval, 1) + foreach(x -> fill!(x.nzval, 1), blocks(KB)) + @test K == KB +end diff --git a/test/runtests.jl b/test/runtests.jl index edfb54ccf2..3c41f179f8 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -28,5 +28,6 @@ include("test_pointevaluation.jl") # include("test_notebooks.jl") include("test_apply_rhs.jl") include("test_apply_analytical.jl") +HAS_EXTENSIONS && include("blockarrays.jl") include("test_examples.jl") @test all(x -> isdefined(Ferrite, x), names(Ferrite)) # Test that all exported symbols are defined