Implement coupling support for MixedDofHandler (#650)

This patch enables creating sparsity patterns with the field/component coupling specified just like for DofHandler. Also enables the Metis reordering to use coupling information when renumbering dofs in a MixedDofHandler. Closes #630.
Ferrite-FEM · Mar 28, 2023 · 9ca56f3 · 9ca56f3
1 parent a12d87d
commit 9ca56f3
Show file tree

Hide file tree

Showing 4 changed files with 153 additions and 74 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -10,6 +10,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
  - Support reordering dofs of a `MixedDofHandler` by the built-in orderings `FieldWise` and
    `ComponentWise`. This includes support for reordering dofs of fields on subdomains.
    ([#645][github-645])
+ - Support specifying the coupling between fields in a `MixedDofHandler` when creating the
+   sparsity pattern. ([#650][github-650])
+ - Support Metis dof reordering with coupling information for `MixedDofHandler`.
+   ([#650][github-650])
 
 ## [0.3.13] - 2023-03-23
 ### Added
@@ -376,6 +380,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 [github-621]: https://github.com/Ferrite-FEM/Ferrite.jl/pull/621
 [github-633]: https://github.com/Ferrite-FEM/Ferrite.jl/pull/633
 [github-645]: https://github.com/Ferrite-FEM/Ferrite.jl/pull/645
+[github-650]: https://github.com/Ferrite-FEM/Ferrite.jl/pull/650
 
 [Unreleased]: https://github.com/Ferrite-FEM/Ferrite.jl/compare/v0.3.13...HEAD
 [0.3.13]: https://github.com/Ferrite-FEM/Ferrite.jl/compare/v0.3.12...v0.3.13

diff --git a/ext/FerriteMetis.jl b/ext/FerriteMetis.jl
@@ -5,7 +5,6 @@ module FerriteMetis
 if VERSION >= v"1.10.0-DEV.90"
 
 using Ferrite
-using Ferrite: AbstractDofHandler
 using Metis.LibMetis: idx_t
 using Metis: Metis
 using SparseArrays: sparse
@@ -30,42 +29,51 @@ function DofOrder.Ext{Metis}(;
 end
 
 function Ferrite.compute_renumber_permutation(
-    dh::AbstractDofHandler,
+    dh::Union{DofHandler,MixedDofHandler},
     ch::Union{ConstraintHandler,Nothing},
     order::DofOrder.Ext{Metis}
 )
 
     # Expand the coupling matrix to size ndofs_per_cell × ndofs_per_cell
     coupling = order.coupling
-    if coupling === nothing
-        n = ndofs_per_cell(dh)
-        entries_per_cell = n * (n - 1)
-    else # coupling !== nothing
+    if coupling !== nothing
         # Set sym = true since Metis.permutation requires a symmetric graph.
         # TODO: Perhaps just symmetrize it: coupling = coupling' .| coupling
-        coupling = Ferrite._coupling_to_local_dof_coupling(dh, coupling, #= sym =# true)
-        # Compute entries per cell, subtract diagonal elements
-        entries_per_cell =
-            count(coupling[i, j] for i in axes(coupling, 1), j in axes(coupling, 2) if i != j)
+        couplings = Ferrite._coupling_to_local_dof_coupling(dh, coupling, #= sym =# true)
     end
 
     # Create the CSR (CSC, but pattern is symmetric so equivalent) using
     # Metis.idx_t as the integer type
-    L = entries_per_cell * getncells(dh.grid)
-    I = Vector{idx_t}(undef, L)
-    J = Vector{idx_t}(undef, L)
+    buffer_length = 0
+    for (fhi, fh) in pairs(dh isa DofHandler ? (dh, ) : dh.fieldhandlers)
+        set = fh isa DofHandler ? (1:getncells(dh.grid)) : fh.cellset
+        n = ndofs_per_cell(dh, first(set)) # TODO: ndofs_per_cell(fh)
+        entries_per_cell = if coupling === nothing
+            n * (n - 1)
+        else
+            count(couplings[fhi][i, j] for i in 1:n, j in 1:n if i != j)
+        end
+        buffer_length += entries_per_cell * length(set)
+    end
+    I = Vector{idx_t}(undef, buffer_length)
+    J = Vector{idx_t}(undef, buffer_length)
     idx = 0
-    @inbounds for cc in CellIterator(dh)
-        dofs = celldofs(cc)
-        for (i, dofi) in pairs(dofs), (j, dofj) in pairs(dofs)
-            dofi == dofj && continue # Metis doesn't want the diagonal
-            coupling === nothing || coupling[i, j] || continue
-            idx += 1
-            I[idx] = dofi
-            J[idx] = dofj
+
+    for (fhi, fh) in pairs(dh isa DofHandler ? (dh, ) : dh.fieldhandlers)
+        coupling === nothing || (coupling_fh = couplings[fhi])
+        set = fh isa DofHandler ? (1:getncells(dh.grid)) : fh.cellset
+        for cc in CellIterator(dh, set)
+            dofs = celldofs(cc)
+            for (j, dofj) in pairs(dofs), (i, dofi) in pairs(dofs)
+                dofi == dofj && continue # Metis doesn't want the diagonal
+                coupling === nothing || coupling_fh[i, j] || continue
+                idx += 1
+                I[idx] = dofi
+                J[idx] = dofj
+            end
         end
     end
-    @assert idx == L
+    @assert length(I) == length(J) == idx
     N = ndofs(dh)
     # TODO: Use spzeros! in Julia 1.10.
     S = sparse(I, J, zeros(Float32, length(I)), N, N)

diff --git a/src/Dofs/sparsity_pattern.jl b/src/Dofs/sparsity_pattern.jl
@@ -56,83 +56,106 @@ end
 # ii) (ncomponents × ncomponents) specifying coupling between components, or iii)
 # (ndofs_per_cell × ndofs_per_cell) specifying coupling between all local dofs, i.e. a
 # "template" local matrix.
-function _coupling_to_local_dof_coupling(dh::DofHandler, coupling::AbstractMatrix{Bool}, sym::Bool)
-    out = zeros(Bool, ndofs_per_cell(dh), ndofs_per_cell(dh))
+function _coupling_to_local_dof_coupling(dh::Union{DofHandler,MixedDofHandler}, coupling::AbstractMatrix{Bool}, sym::Bool)
     sz = size(coupling, 1)
     sz == size(coupling, 2) || error("coupling not square")
     sym && (issymmetric(coupling) || error("coupling not symmetric"))
-    dof_ranges = [dof_range(dh, f) for f in dh.field_names]
-    if sz == length(dh.field_names) # Coupling given by fields
-        for (j, jrange) in pairs(dof_ranges), (i, irange) in pairs(dof_ranges)
-            out[irange, jrange] .= coupling[i, j]
-        end
-    elseif sz == sum(dh.field_dims) # Coupling given by components
-        component_offsets = pushfirst!(cumsum(dh.field_dims), 0)
-        for (jf, jrange) in pairs(dof_ranges), (j, J) in pairs(jrange)
-            jc = mod1(j, dh.field_dims[jf]) + component_offsets[jf]
-            for (i_f, irange) in pairs(dof_ranges), (i, I) in pairs(irange)
-                ic = mod1(i, dh.field_dims[i_f]) + component_offsets[i_f]
-                out[I, J] = coupling[ic, jc]
+
+    # Return one matrix per (potential) sub-domain
+    outs = Matrix{Bool}[]
+    field_dims = map(fieldname -> getfielddim(dh, fieldname), dh.field_names)
+
+    for fh in (dh isa DofHandler ? (dh,) : dh.fieldhandlers)
+        ci = dh isa DofHandler ? 1 : first(fh.cellset)
+        out = zeros(Bool, ndofs_per_cell(dh, ci), ndofs_per_cell(dh, ci))
+        push!(outs, out)
+
+        field_names = fh isa DofHandler ? fh.field_names : (f.name for f in fh.fields)
+        dof_ranges = [dof_range(fh, f) for f in field_names]
+        global_idxs = [findfirst(x -> x === f, dh.field_names) for f in field_names]
+
+        if sz == length(dh.field_names) # Coupling given by fields
+            for (j, jrange) in pairs(dof_ranges), (i, irange) in pairs(dof_ranges)
+                out[irange, jrange] .= coupling[global_idxs[i], global_idxs[j]]
             end
+        elseif sz == sum(field_dims) # Coupling given by components
+            component_offsets = pushfirst!(cumsum(field_dims), 0)
+            for (jf, jrange) in pairs(dof_ranges), (j, J) in pairs(jrange)
+                jc = mod1(j, field_dims[global_idxs[jf]]) + component_offsets[global_idxs[jf]]
+                for (i_f, irange) in pairs(dof_ranges), (i, I) in pairs(irange)
+                    ic = mod1(i, field_dims[global_idxs[i_f]]) + component_offsets[global_idxs[i_f]]
+                    out[I, J] = coupling[ic, jc]
+                end
+            end
+        elseif sz == ndofs_per_cell(dh, ci) # Coupling given by template local matrix
+            # TODO: coupling[fieldhandler_idx] if different template per subddomain
+            out .= coupling
+        else
+            error("could not create coupling")
         end
-    elseif sz == ndofs_per_cell(dh) # Coupling given by template local matrix
-        out .= coupling
-    else
-        error("could not create coupling")
     end
-    return out
+    return outs
 end
 
 function _create_sparsity_pattern(dh::AbstractDofHandler, ch#=::Union{ConstraintHandler, Nothing}=#, sym::Bool, keep_constrained::Bool, coupling::Union{AbstractMatrix{Bool},Nothing})
     @assert isclosed(dh)
     if !keep_constrained
         @assert ch !== nothing && isclosed(ch)
     end
-    ncells = getncells(dh.grid)
     if coupling !== nothing
         # Extend coupling to be of size (ndofs_per_cell × ndofs_per_cell)
-        coupling = _coupling_to_local_dof_coupling(dh, coupling, sym)
+        couplings = _coupling_to_local_dof_coupling(dh, coupling, sym)
     end
-    # Compute approximate size for the buffers using the dofs in the first element
-    n = ndofs_per_cell(dh)
-    N = (coupling === nothing ? (sym ? div(n*(n+1), 2) : n^2) : count(coupling)) * ncells
-    N += ndofs(dh) # always add the diagonal elements
-    I = Int[]; resize!(I, N)
-    J = Int[]; resize!(J, N)
-    global_dofs = zeros(Int, n)
+
+    # Allocate buffers. Compute an upper bound for the buffer length and allocate it all up
+    # front since they will become large and expensive to re-allocate. The bound is exact
+    # when keeping constrained dofs (default) and if not it only over-estimates with number
+    # of entries eliminated by constraints.
+    max_buffer_length = ndofs(dh) # diagonal elements
+    for (fhi, fh) in pairs(dh isa DofHandler ? (dh, ) : dh.fieldhandlers)
+        set = fh isa DofHandler ? (1:getncells(dh.grid)) : fh.cellset
+        n = ndofs_per_cell(dh, first(set)) # TODO: ndofs_per_cell(fh)
+        entries_per_cell = if coupling === nothing
+            sym ? div(n * (n + 1), 2) : n^2
+        else
+            count(couplings[fhi][i, j] for i in 1:n for j in (sym ? i : 1):n)
+        end
+        max_buffer_length += entries_per_cell * length(set)
+    end
+    I = Vector{Int}(undef, max_buffer_length)
+    J = Vector{Int}(undef, max_buffer_length)
+    global_dofs = Int[]
     cnt = 0
-    for element_id in 1:ncells
-        # MixedDofHandler might have varying number of dofs per element
-        resize!(global_dofs, ndofs_per_cell(dh, element_id))
-        celldofs!(global_dofs, dh, element_id)
-        @inbounds for j in eachindex(global_dofs), i in eachindex(global_dofs)
-            coupling === nothing || coupling[i, j] || continue
-            dofi = global_dofs[i]
-            dofj = global_dofs[j]
-            sym && (dofi > dofj && continue)
-            !keep_constrained && (haskey(ch.dofmapping, dofi) || haskey(ch.dofmapping, dofj)) && continue
-            cnt += 1
-            if cnt > length(J)
-                resize!(I, trunc(Int, length(I) * 1.5))
-                resize!(J, trunc(Int, length(J) * 1.5))
-            end
-            I[cnt] = dofi
-            J[cnt] = dofj
 
+    for (fhi, fh) in pairs(dh isa DofHandler ? (dh, ) : dh.fieldhandlers)
+        coupling === nothing || (coupling_fh = couplings[fhi])
+        set = fh isa DofHandler ? (1:getncells(dh.grid)) : fh.cellset
+        n = ndofs_per_cell(dh, first(set)) # TODO: ndofs_per_cell(fh)
+        resize!(global_dofs, n)
+        @inbounds for element_id in set
+            celldofs!(global_dofs, dh, element_id)
+            for j in eachindex(global_dofs), i in eachindex(global_dofs)
+                coupling === nothing || coupling_fh[i, j] || continue
+                dofi = global_dofs[i]
+                dofj = global_dofs[j]
+                sym && (dofi > dofj && continue)
+                !keep_constrained && (haskey(ch.dofmapping, dofi) || haskey(ch.dofmapping, dofj)) && continue
+                cnt += 1
+                I[cnt] = dofi
+                J[cnt] = dofj
+            end
         end
     end
+
+    # Always add diagonal entries
+    resize!(I, cnt + ndofs(dh))
+    resize!(J, cnt + ndofs(dh))
     @inbounds for d in 1:ndofs(dh)
         cnt += 1
-        if cnt > length(J)
-            resize!(I, trunc(Int, length(I) + ndofs(dh)))
-            resize!(J, trunc(Int, length(J) + ndofs(dh)))
-        end
         I[cnt] = d
         J[cnt] = d
     end
-
-    resize!(I, cnt)
-    resize!(J, cnt)
+    @assert length(I) == length(J) == cnt
 
     K = spzeros!!(Float64, I, J, ndofs(dh), ndofs(dh))
 

diff --git a/test/test_dofs.jl b/test/test_dofs.jl
@@ -397,6 +397,11 @@ end
         @test_throws ErrorException renumber!(dh, ch, DofOrder.Ext{Metis}())
         renumber!(dh, DofOrder.Ext{Metis}(coupling=[true true; true false]))
         @test_throws ErrorException renumber!(dh, ch, DofOrder.Ext{Metis}(coupling=[true true; true false]))
+        dh, ch = testdhch(MixedDofHandler)
+        renumber!(dh, DofOrder.Ext{Metis}())
+        @test_throws ErrorException renumber!(dh, ch, DofOrder.Ext{Metis}())
+        renumber!(dh, DofOrder.Ext{Metis}(coupling=[true true; true false]))
+        @test_throws ErrorException renumber!(dh, ch, DofOrder.Ext{Metis}(coupling=[true true; true false]))
     end
 end
 
@@ -492,4 +497,42 @@ end
     @test_throws ErrorException("coupling not square") create_sparsity_pattern(dh; coupling=[true true])
     @test_throws ErrorException("coupling not symmetric") create_symmetric_sparsity_pattern(dh; coupling=[true true; false true])
     @test_throws ErrorException("could not create coupling") create_symmetric_sparsity_pattern(dh; coupling=falses(100, 100))
+
+    # Test coupling with subdomains
+    grid = generate_grid(Quadrilateral, (1, 2))
+    dh = MixedDofHandler(grid)
+    fh1 = FieldHandler(
+        [Field(:u, Lagrange{2,RefCube,1}(), 2), Field(:p, Lagrange{2,RefCube,1}(), 2)],
+        Set(1)
+    )
+    add!(dh, fh1)
+    fh2 = FieldHandler(
+        [Field(:u, Lagrange{2,RefCube,1}(), 2)],
+        Set(2)
+    )
+    add!(dh, fh2)
+    close!(dh)
+    K = create_sparsity_pattern(dh; coupling = [true true; true false])
+    KS = create_symmetric_sparsity_pattern(dh; coupling = [true true; true false])
+    # Subdomain 1: u and p
+    udofs = celldofs(dh, 1)[dof_range(fh1, :u)]
+    pdofs = celldofs(dh, 1)[dof_range(fh1, :p)]
+    for j in udofs, i in Iterators.flatten((udofs, pdofs))
+        @test is_stored(K, i, j)
+        @test is_stored(KS, i, j) == (i <= j)
+    end
+    for j in pdofs, i in udofs
+        @test is_stored(K, i, j)
+        @test is_stored(KS, i, j)
+    end
+    for j in pdofs, i in pdofs
+        @test is_stored(K, i, j) == (i == j)
+        @test is_stored(KS, i, j) == (i == j)
+    end
+    # Subdomain 2: u
+    udofs = celldofs(dh, 2)[dof_range(fh2, :u)]
+    for j in udofs, i in udofs
+        @test is_stored(K, i, j)
+        @test is_stored(KS, i, j) == (i <= j)
+    end
 end