From a6eb9cefee30f2b7c6f40dd4d914caf2f4121ae8 Mon Sep 17 00:00:00 2001 From: Fredrik Ekre Date: Thu, 19 Sep 2024 16:47:57 +0200 Subject: [PATCH] Extend COOAssembler to also handle vector assembly --- docs/src/literate-tutorials/stokes-flow.jl | 4 +- src/assembler.jl | 138 +++++++++++++++------ test/test_assemble.jl | 50 ++++++-- 3 files changed, 137 insertions(+), 55 deletions(-) diff --git a/docs/src/literate-tutorials/stokes-flow.jl b/docs/src/literate-tutorials/stokes-flow.jl index 70c5da2e16..e5e97c6fd5 100644 --- a/docs/src/literate-tutorials/stokes-flow.jl +++ b/docs/src/literate-tutorials/stokes-flow.jl @@ -284,7 +284,7 @@ end # example](https://www.dealii.org/current/doxygen/deal.II/step_11.html). function setup_mean_constraint(dh, fvp) - assembler = start_assemble() + assembler = Ferrite.COOAssembler() ## All external boundaries set = union( getfacetset(dh.grid, "Γ1"), @@ -313,7 +313,7 @@ function setup_mean_constraint(dh, fvp) ## Assemble to row 1 assemble!(assembler, [1], element_dofs_p, Ce) end - C = finish_assemble(assembler) + C, _ = finish_assemble(assembler) ## Create an AffineConstraint from the C-matrix _, J, V = findnz(C) _, constrained_dof_idx = findmax(abs2, V) diff --git a/src/assembler.jl b/src/assembler.jl index e79521dc4b..6c0ffc5203 100644 --- a/src/assembler.jl +++ b/src/assembler.jl @@ -2,58 +2,98 @@ abstract type AbstractAssembler end abstract type AbstractCSCAssembler <: AbstractAssembler end """ -This assembler creates a COO (**coo**rdinate format) representation of a sparse matrix during -assembly and converts it into a `SparseMatrixCSC{T}` on finalization. + struct COOAssembler{Tv, Ti} + +This assembler creates a COO (**coo**rdinate format) representation of a sparse matrix +during assembly and converts it into a `SparseMatrixCSC{Tv, Ti}` on finalization. """ -struct COOAssembler{T} # <: AbstractAssembler - I::Vector{Int} - J::Vector{Int} - V::Vector{T} +struct COOAssembler{Tv, Ti} # <: AbstractAssembler + nrows::Int + ncols::Int + f::Vector{Tv} + I::Vector{Ti} + J::Vector{Ti} + V::Vector{Tv} end -COOAssembler(N) = COOAssembler(Float64, N) - -function COOAssembler(::Type{T}, N) where T +function COOAssembler{Tv, Ti}(nrows::Int, ncols::Int; sizehint::Int = 0) where {Tv, Ti} I = Int[] J = Int[] - V = T[] - sizehint!(I, N) - sizehint!(J, N) - sizehint!(V, N) - - COOAssembler(I, J, V) + V = Tv[] + sizehint!(I, sizehint) + sizehint!(J, sizehint) + sizehint!(V, sizehint) + f = Tv[] + return COOAssembler{Tv, Ti}(nrows, ncols, f, I, J, V) end """ - start_assemble([N=0]) -> COOAssembler - -Create an `Assembler` object which can be used to assemble element contributions to the -global sparse matrix. Use [`assemble!`](@ref) for each element, and [`finish_assemble`](@ref), -to finalize the assembly and return the sparse matrix. + COOAssembler(nrows::Int, ncols::Int; sizehint::Int = 0) -Note that giving a sparse matrix as input can be more efficient. See below and -as described in the [manual](@ref man-assembly). - -!!! note - When the same matrix pattern is used multiple times (for e.g. multiple time steps or - Newton iterations) it is more efficient to create the sparse matrix **once** and reuse - the same pattern. See the [manual section](@ref man-assembly) on assembly. +Create a new assembler. """ -function start_assemble(N::Int=0) - return start_assemble(Float64, N) -end - -function start_assemble(t::Type{T}, N::Int=0) where T - return COOAssembler(t, N) +function COOAssembler(nrows::Int, ncols::Int; sizehint::Int = 0) + return COOAssembler{Float64, Int}(nrows, ncols; sizehint = sizehint) end +COOAssembler(; sizehint::Int = 0) = COOAssembler(-1, -1; sizehint = sizehint) + +# """ +# start_assemble(; sizehint::Int = 0) -> COOAssembler +# start_assemble(nrows::Int, ncols::Int; sizehint::Int = 0) -> COOAssembler + +# Create an `COOAssembler` which can be used to assemble element contributions to a global +# sparse matrix and vector. `nrows` is the number of rows in the final matrix and `ncols` the +# number of columns. If `nrows` and `ncols` are not passed they are inferred from the +# maximum indices added to the assembler during assembly. `sizehint` is a hint for how many +# entries in total will be added to the assembler and can be passed to optimize for +# allocations. + +# Use [`assemble!`](@ref) to insert element contributions, and [`finish_assemble`](@ref), to +# finalize the assembly and return the sparse matrix (and optionally vector). + +# Note that allocating a sparse matrix and assemble into it is generally preferred. See below +# and the [manual section on assembly](@ref man-assembly). + +# !!! note +# When the same matrix pattern is used multiple times (for e.g. multiple time steps or +# Newton iterations) it is more efficient to create the sparse matrix **once** and reuse +# the same pattern. See the [manual section](@ref man-assembly) on assembly. +# """ +# function start_assemble(nrows::Int, ncols::Int; sizehint::Int = 0) +# return COOAssembler{Float64, Int}(nrows, ncols; sizehint = sizehint) +# end + +# function start_assemble(; sizehint::Int = 0) +# return COOAssembler{Float64, Int}(-1, -1; sizehint = sizehint) +# end """ assemble!(a::COOAssembler, dofs, Ke) + assemble!(a::COOAssembler, dofs, Ke, fe) -Assembles the element matrix `Ke` into `a`. +Assembles the element matrix `Ke` and element vector `fe` into `a`. """ -function assemble!(a::COOAssembler{T}, dofs::AbstractVector{Int}, Ke::AbstractMatrix{T}) where {T} +function assemble!(a::COOAssembler{T}, dofs::AbstractVector{Int}, Ke::AbstractMatrix{T}, fe::Union{AbstractVector{T}, Nothing} = nothing) where {T} assemble!(a, dofs, dofs, Ke) + if fe !== nothing + # If the final number of rows is unknown we grow the vector lazily, + # otherwise we resize it directly to nrows. + if a.nrows == -1 + m = maximum(dofs; init = 0) + lf = length(a.f) + if lf < m + resize!(a.f, m) + for i in (lf + 1):m + a.f[i] = 0 + end + end + elseif isempty(a.f) + resize!(a.f, a.nrows) + fill!(a.f, 0) + end + assemble!(a.f, dofs, fe) + end + return end """ @@ -78,13 +118,31 @@ function assemble!(a::COOAssembler{T}, rowdofs::AbstractVector{Int}, coldofs::Ab end """ - finish_assemble(a::Assembler) -> K + finish_assemble(a::COOAssembler) -> K, f -Finalizes an assembly. Returns a sparse matrix with the -assembled values. Note that this step is not necessary for `AbstractAssembler`s. +Finalize the assembly and return the sparse matrix `K::SparseMatrixCSC` and vector +`f::Vector`. If the assembler have not been used for vector assembly, `f` is an empty +vector. """ function finish_assemble(a::COOAssembler) - return sparse(a.I, a.J, a.V) + # Create the matrix + nrows = a.nrows == -1 ? maximum(a.I) : a.nrows + ncols = a.ncols == -1 ? maximum(a.J) : a.ncols + K = sparse(a.I, a.J, a.V, nrows, ncols) + # Finalize the vector + f = a.f + if !isempty(f) + # There have been things assembled, make sure it is resized correctly + lf = length(f) + @assert lf <= nrows + if lf < nrows + resize!(f, nrows) + for i in (lf + 1):nrows + f[i] = 0 + end + end + end + return K, f end """ @@ -173,7 +231,7 @@ function start_assemble(K::Symmetric{T,<:SparseMatrixCSC}, f::Vector=T[]; fillze end function finish_assemble(a::Union{CSCAssembler, SymmetricCSCAssembler}) - return a.K, isempty(a.f) ? nothing : a.f + return a.K, a.f end """ diff --git a/test/test_assemble.jl b/test/test_assemble.jl index ae8c1b3a68..d1a586d3eb 100644 --- a/test/test_assemble.jl +++ b/test/test_assemble.jl @@ -1,31 +1,55 @@ @testset "assemble" begin dofs = [1, 3, 5, 7] + maxd = maximum(dofs) - # residual + # Vector assembly ge = rand(4) g = zeros(8) assemble!(g, dofs, ge) - @test g[1] == ge[1] - @test g[3] == ge[2] - @test g[5] == ge[3] - @test g[7] == ge[4] + @test g[dofs] == ge - # stiffness - a = start_assemble() + # COOAssembler: matrix only, inferred size + a = Ferrite.COOAssembler() Ke = rand(4, 4) assemble!(a, dofs, Ke) - K = finish_assemble(a) - @test K[1,1] == Ke[1,1] - @test K[1,5] == Ke[1,3] - @test K[5,1] == Ke[3,1] + K, f = finish_assemble(a) + @test K[dofs, dofs] == Ke + @test size(K) == (maxd, maxd) + @test isempty(f) + + # COOAssembler: matrix only, given size + a = Ferrite.COOAssembler(10, 10) + assemble!(a, dofs, Ke) + K, f = finish_assemble(a) + @test K[dofs, dofs] == Ke + @test size(K) == (10, 10) + @test isempty(f) + + # COOAssembler: matrix and vector, inferred size + a = Ferrite.COOAssembler() + assemble!(a, dofs, Ke, ge) + K, f = finish_assemble(a) + @test K[dofs, dofs] == Ke + @test f[dofs] == ge + @test size(K) == (maxd, maxd) + @test length(f) == maxd + + # COOAssembler: matrix and vector, given size + a = Ferrite.COOAssembler(10, 10) + assemble!(a, dofs, Ke, ge) + K, f = finish_assemble(a) + @test K[dofs, dofs] == Ke + @test f[dofs] == ge + @test size(K) == (10, 10) + @test length(f) == 10 # assemble with different row and col dofs rdofs = [1,4,6] cdofs = [1,7] - a = start_assemble() + a = Ferrite.COOAssembler() Ke = rand(length(rdofs), length(cdofs)) assemble!(a, rdofs, cdofs, Ke) - K = finish_assemble(a) + K, _ = finish_assemble(a) @test (K[rdofs,cdofs] .== Ke) |> all # SparseMatrix assembler