adapted parametrization from MixedDofHandler to DofHandler, tests pas…

…s, still need to test it on our subtyped grid
Ferrite-FEM · Dec 17, 2020 · 5db9b04 · 5db9b04
2 parents adcf645 + 96ae0b6
commit 5db9b04
Show file tree

Hide file tree

Showing 5 changed files with 42 additions and 42 deletions.
diff --git a/src/Dofs/ConstraintHandler.jl b/src/Dofs/ConstraintHandler.jl
@@ -259,8 +259,8 @@ end
 
 # for nodes
 function _update!(values::Vector{Float64}, f::Function, nodes::Set{Int}, field::Symbol, nodeidxs::Vector{Int}, globaldofs::Vector{Int},
-                  components::Vector{Int}, dh::DofHandler{Grid{dim,C,M},M}, facevalues::BCValues,
-                  dofmapping::Dict{Int,Int}, time::Float64) where {dim,C,M}
+                  components::Vector{Int}, dh::DofHandler{dim, M, G}, facevalues::BCValues,
+                  dofmapping::Dict{Int,Int}, time::Float64) where {dim,M,G<:AbstractGrid}
     counter = 1
     for (idx, nodenumber) in enumerate(nodeidxs)
         x = dh.grid.nodes[nodenumber].x

diff --git a/src/Dofs/DofHandler.jl b/src/Dofs/DofHandler.jl
@@ -11,7 +11,7 @@ abstract type AbstractDofHandler end
 
 Construct a `DofHandler` based on the grid `grid`.
 """
-struct DofHandler{G <: AbstractGrid,T} <: AbstractDofHandler
+struct DofHandler{dim,T,G <: AbstractGrid{dim}} <: AbstractDofHandler
     field_names::Vector{Symbol}
     field_dims::Vector{Int}
     # TODO: field_interpolations can probably be better typed: We should at least require
@@ -26,14 +26,18 @@ struct DofHandler{G <: AbstractGrid,T} <: AbstractDofHandler
 end
 
 function DofHandler(grid::Grid{dim,C,T}) where {dim,C,T}
-    DofHandler{Grid{dim,C,T}, T}(Symbol[], Int[], Interpolation[], BCValues{Float64}[], Int[], Int[], ScalarWrapper(false), grid, JuAFEM.ScalarWrapper(-1))
+    DofHandler{dim,T,Grid{dim,C,T}}(Symbol[], Int[], Interpolation[], BCValues{Float64}[], Int[], Int[], ScalarWrapper(false), grid, JuAFEM.ScalarWrapper(-1))
+end
+
+function DofHandler(grid::AbstractGrid{dim}) where dim
+    DofHandler{dim,Float64,typeof(grid)}(Symbol[], Int[], Interpolation[], BCValues{Float64}[], Int[], Int[], ScalarWrapper(false), grid, JuAFEM.ScalarWrapper(-1))
 end
 
 function Base.show(io::IO, ::MIME"text/plain", dh::DofHandler)
     println(io, "DofHandler")
     println(io, "  Fields:")
     for i in 1:nfields(dh)
-        println(io, "    ", repr(dh.field_names[i]), ", interpolation: ", dh.field_interpolations[i],", dim: ", dh.field_dims[i])
+        println(io, "    ", repr(dh.field_names[i]), ", interpolation: ", dh.field_interpolations[i], ", dim: ", dh.field_dims[i])
     end
     if !isclosed(dh)
         print(io, "  Not closed!")
@@ -44,21 +48,21 @@ function Base.show(io::IO, ::MIME"text/plain", dh::DofHandler)
 end
 
 ndofs(dh::AbstractDofHandler) = dh.ndofs[]
-ndofs_per_cell(dh::AbstractDofHandler, cell::Int=1) = dh.cell_dofs_offset[cell+1] - dh.cell_dofs_offset[cell]
+ndofs_per_cell(dh::AbstractDofHandler, cell::Int=1) = dh.cell_dofs_offset[cell + 1] - dh.cell_dofs_offset[cell]
 isclosed(dh::AbstractDofHandler) = dh.closed[]
 nfields(dh::AbstractDofHandler) = length(dh.field_names)
 getfieldnames(dh::AbstractDofHandler) = dh.field_names
 ndim(dh::AbstractDofHandler, field_name::Symbol) = dh.field_dims[find_field(dh, field_name)]
 function find_field(dh::DofHandler, field_name::Symbol)
-    j = findfirst(i->i == field_name, dh.field_names)
+    j = findfirst(i -> i == field_name, dh.field_names)
     j == 0 && error("did not find field $field_name")
     return j
 end
 
 # Calculate the offset to the first local dof of a field
 function field_offset(dh::DofHandler, field_name::Symbol)
     offset = 0
-    for i in 1:find_field(dh, field_name)-1
+    for i in 1:find_field(dh, field_name) - 1
         offset += getnbasefunctions(dh.field_interpolations[i])::Int * dh.field_dims[i]
     end
     return offset
@@ -89,7 +93,7 @@ function dof_range(dh::DofHandler, field_name::Symbol)
     f = find_field(dh, field_name)
     offset = field_offset(dh, field_name)
     n_field_dofs = getnbasefunctions(dh.field_interpolations[f])::Int * dh.field_dims[f]
-    return (offset+1):(offset+n_field_dofs)
+    return (offset + 1):(offset + n_field_dofs)
 end
 
 function Base.push!(dh::DofHandler, name::Symbol, dim::Int, ip::Interpolation=default_interpolation(getcelltype(dh.grid)))
@@ -123,8 +127,7 @@ function close!(dh::DofHandler)
 end
 
 # close the DofHandler and distribute all the dofs
-function __close!(dh::DofHandler{G, T}) where {G<:AbstractGrid,T}
-    dim = getdim(dh.grid)
+function __close!(dh::DofHandler{dim, T, G}) where {dim,T, G<:AbstractGrid{dim}}
     C = getcelltype(dh.grid)
     @assert !isclosed(dh)
 
@@ -158,7 +161,7 @@ function __close!(dh::DofHandler{G, T}) where {G<:AbstractGrid,T}
     end
 
     # not implemented yet: more than one facedof per face in 3D
-    dim == 3 && @assert(!any(x->x.nfacedofs > 1, interpolation_infos))
+    dim == 3 && @assert(!any(x -> x.nfacedofs > 1, interpolation_infos))
 
     nextdof = 1 # next free dof to distribute
     push!(dh.cell_dofs_offset, 1) # dofs for the first cell start at 1
@@ -176,8 +179,8 @@ function __close!(dh::DofHandler{G, T}) where {G<:AbstractGrid,T}
                     if token > 0 # haskey(vertexdicts[fi], vertex) # reuse dofs
                         reuse_dof = vertexdicts[fi].vals[token] # vertexdicts[fi][vertex]
                         for d in 1:dh.field_dims[fi]
-                            @debug println("      reusing dof #$(reuse_dof + (d-1))")
-                            push!(dh.cell_dofs, reuse_dof + (d-1))
+                            @debug println("      reusing dof #$(reuse_dof + (d - 1))")
+                            push!(dh.cell_dofs, reuse_dof + (d - 1))
                         end
                     else # token <= 0, distribute new dofs
                         for vertexdof in 1:interpolation_info.nvertexdofs
@@ -201,7 +204,7 @@ function __close!(dh::DofHandler{G, T}) where {G<:AbstractGrid,T}
                             startdof, olddir = edgedicts[fi].vals[token] # edgedicts[fi][sedge] # first dof for this edge (if dir == true)
                             for edgedof in (dir == olddir ? (1:interpolation_info.nedgedofs) : (interpolation_info.nedgedofs:-1:1))
                                 for d in 1:dh.field_dims[fi]
-                                    reuse_dof = startdof + (d-1) + (edgedof-1)*dh.field_dims[fi]
+                                    reuse_dof = startdof + (d - 1) + (edgedof - 1) * dh.field_dims[fi]
                                     @debug println("      reusing dof#$(reuse_dof)")
                                     push!(dh.cell_dofs, reuse_dof)
                                 end
@@ -228,7 +231,7 @@ function __close!(dh::DofHandler{G, T}) where {G<:AbstractGrid,T}
                         startdof = facedicts[fi].vals[token] # facedicts[fi][sface]
                         for facedof in interpolation_info.nfacedofs:-1:1 # always reverse (YOLO)
                             for d in 1:dh.field_dims[fi]
-                                reuse_dof = startdof + (d-1) + (facedof-1)*dh.field_dims[fi]
+                                reuse_dof = startdof + (d - 1) + (facedof - 1) * dh.field_dims[fi]
                                 @debug println("      reusing dof#$(reuse_dof)")
                                 push!(dh.cell_dofs, reuse_dof)
                             end
@@ -257,7 +260,7 @@ function __close!(dh::DofHandler{G, T}) where {G<:AbstractGrid,T}
             end
         end # field loop
         # push! the first index of the next cell to the offset vector
-        push!(dh.cell_dofs_offset, length(dh.cell_dofs)+1)
+        push!(dh.cell_dofs_offset, length(dh.cell_dofs) + 1)
     end # cell loop
     dh.ndofs[] = maximum(dh.cell_dofs)
     dh.closed[] = true
@@ -281,7 +284,7 @@ function cellnodes!(global_nodes::Vector{Int}, grid::Grid{dim,C}, i::Int) where
     return global_nodes
 end
 
-function cellnodes!(global_nodes::Vector{Int}, grid::G, i::Int) where {G<:AbstractGrid}
+function cellnodes!(global_nodes::Vector{Int}, grid::G, i::Int) where {G <: AbstractGrid}
     C = getcelltype(grid)
     dim = getdim(grid)
     @assert length(global_nodes) == nnodes(C)
@@ -291,14 +294,14 @@ function cellnodes!(global_nodes::Vector{Int}, grid::G, i::Int) where {G<:Abstra
     return global_nodes
 end
 
-function cellcoords!(global_coords::Vector{Vec{dim,T}}, grid::AbstractGrid, i::Int) where {dim, T}
+function cellcoords!(global_coords::Vector{Vec{dim,T}}, grid::AbstractGrid, i::Int) where {dim,T}
     @assert dim == getdim(grid)
     C = getcelltype(grid)
     @assert length(global_coords) == nnodes(C)
-    #for j in 1:nnodes(C)
+    # for j in 1:nnodes(C)
     #    nodeid = getcells(grid, i).nodes[j]
     #    global_coords[j] = getnodes(grid, nodeid).x
-    #end
+    # end
     global_coords .= getcoordinates(grid, i)
     return global_coords
 end
@@ -348,7 +351,7 @@ See the [Sparsity Pattern](@ref) section of the manual.
 function _create_sparsity_pattern(dh::DofHandler, sym::Bool)
     ncells = getncells(dh.grid)
     n = ndofs_per_cell(dh)
-    N = sym ? div(n*(n+1), 2) * ncells : n^2 * ncells
+    N = sym ? div(n * (n + 1), 2) * ncells : n^2 * ncells
     N += ndofs(dh) # always add the diagonal elements
     I = Int[]; resize!(I, N)
     J = Int[]; resize!(J, N)

diff --git a/src/Dofs/MixedDofHandler.jl b/src/Dofs/MixedDofHandler.jl
@@ -26,18 +26,18 @@ function Base.getindex(elvec::CellVector, el::Int)
     return elvec.values[offset:offset + elvec.length[el]-1]
  end
 
-struct MixedDofHandler{dim,C,T} <: JuAFEM.AbstractDofHandler
+struct MixedDofHandler{dim,T,G<:AbstractGrid{dim}} <: JuAFEM.AbstractDofHandler 
     fieldhandlers::Vector{FieldHandler}
     cell_dofs::CellVector{Int}
     cell_nodes::CellVector{Int}
     cell_coords::CellVector{Vec{dim,T}}
     closed::ScalarWrapper{Bool}
-    grid::Grid{dim,C,T}
+    grid::G
     ndofs::ScalarWrapper{Int}
 end
 
 function MixedDofHandler(grid::Grid{dim,C,T}) where {dim,C,T}
-    MixedDofHandler{dim,C,T}(FieldHandler[], CellVector(Int[],Int[],Int[]), CellVector(Int[],Int[],Int[]), CellVector(Vec{dim,T}[],Int[],Int[]), JuAFEM.ScalarWrapper(false), grid, JuAFEM.ScalarWrapper(-1))
+    MixedDofHandler{dim,T,typeof(grid)}(FieldHandler[], CellVector(Int[],Int[],Int[]), CellVector(Int[],Int[],Int[]), CellVector(Vec{dim,T}[],Int[],Int[]), JuAFEM.ScalarWrapper(false), grid, JuAFEM.ScalarWrapper(-1))
 end
 
 getfieldnames(fh::FieldHandler) = [field.name for field in fh.fields]

diff --git a/src/Grid/grid.jl b/src/Grid/grid.jl
@@ -55,12 +55,12 @@ struct FaceIndex
     idx::Tuple{Int,Int} # cell and side
 end
 
-abstract type AbstractGrid end
+abstract type AbstractGrid{dim} end
 
 """
 A `Grid` is a collection of `Cells` and `Node`s which covers the computational domain, together with Sets of cells, nodes and faces.
 """
-mutable struct Grid{dim,C<:AbstractCell,T<:Real} <: AbstractGrid
+mutable struct Grid{dim,C<:AbstractCell,T<:Real} <: AbstractGrid{dim}
     cells::Vector{C}
     nodes::Vector{Node{dim,T}}
     # Sets
@@ -83,7 +83,7 @@ end
 ##########################
 # Grid utility functions #
 ##########################
-@inline getdim(grid::Grid{dim}) where {dim} = dim
+@inline getdim(::AbstractGrid{dim}) where {dim} = dim
 @inline getcells(grid::AbstractGrid) = grid.cells
 @inline getcells(grid::AbstractGrid, v::Union{Int, Vector{Int}}) = grid.cells[v]
 @inline getcells(grid::AbstractGrid, set::String) = grid.cells[collect(grid.cellsets[set])]

diff --git a/src/iterators.jl b/src/iterators.jl
@@ -23,27 +23,26 @@ for cell in CellIterator(grid)
 end
 ```
 """
-struct CellIterator{dim,C,T,G<:AbstractGrid}
+struct CellIterator{dim,C,T,DH<:Union{AbstractDofHandler,Nothing}}
     flags::UpdateFlags
-    grid::G
+    grid::AbstractGrid
     current_cellid::ScalarWrapper{Int}
     nodes::Vector{Int}
     coords::Vector{Vec{dim,T}}
     cellset::Union{Vector{Int},Nothing}
-    dh::Union{DofHandler{G,T}, MixedDofHandler{dim,C,T}} #Future: remove DofHandler and rename MixedDofHandler->DofHandler
+    dh::Union{DH,Nothing}
     celldofs::Vector{Int}
 
-    function CellIterator(dh::Union{DofHandler{G,T}, MixedDofHandler}, cellset::Union{AbstractVector{Int},Nothing}, flags::UpdateFlags) where {G<:AbstractGrid, T}
+    function CellIterator(dh::Union{DofHandler{dim,T,G},MixedDofHandler{dim,T,G}}, cellset::Union{AbstractVector{Int},Nothing}, flags::UpdateFlags) where {dim,T,G<:AbstractGrid}
         C = getcelltype(dh.grid)
-        dim = getdim(dh.grid)
         isconcretetype(C) || _check_same_celltype(dh.grid, cellset)
         N = nnodes_per_cell(dh.grid, cellset === nothing ? 1 : first(cellset))
         cell = ScalarWrapper(0)
         nodes = zeros(Int, N)
         coords = zeros(Vec{dim,T}, N)
         n = ndofs_per_cell(dh, cellset === nothing ? 1 : first(cellset))
         celldofs = zeros(Int, n)
-        return new{dim,C,T,G}(flags, dh.grid, cell, nodes, coords, cellset, dh, celldofs)
+        return new{dim,C,T,typeof(dh)}(flags, dh.grid, cell, nodes, coords, cellset, dh, celldofs)
     end
 
     function CellIterator(grid::Grid{dim,C,T}, cellset::Union{AbstractVector{Int},Nothing}, flags::UpdateFlags) where {dim,C,T}
@@ -52,7 +51,7 @@ struct CellIterator{dim,C,T,G<:AbstractGrid}
         cell = ScalarWrapper(0)
         nodes = zeros(Int, N)
         coords = zeros(Vec{dim,T}, N)
-        return new{dim,C,T, Grid}(flags, grid, cell, nodes, coords, cellset)
+        return new{dim,C,T,Nothing}(flags, grid, cell, nodes, coords, cellset)
     end
 
     function CellIterator(grid::G, cellset::Union{AbstractVector{Int},Nothing}, flags::UpdateFlags) where {G<:AbstractGrid}
@@ -64,15 +63,13 @@ struct CellIterator{dim,C,T,G<:AbstractGrid}
         cell = ScalarWrapper(0)
         nodes = zeros(Int, N)
         coords = zeros(Vec{dim,T}, N)
-        return new{dim,C,T,G}(flags, grid, cell, nodes, coords, cellset)
+        return new{dim,C,T,Nothing}(flags, grid, cell, nodes, coords, cellset)
     end
 end
 
 CellIterator(grid::G; cellset::Union{AbstractVector{Int},Nothing}=nothing, flags::UpdateFlags=UpdateFlags()) where {G<:AbstractGrid} =
     CellIterator(grid, cellset, flags)
-CellIterator(dh::Union{DofHandler{Grid{dim,C,T},T}, MixedDofHandler{dim,C,T}}; cellset::Union{AbstractVector{Int},Nothing}=nothing, flags::UpdateFlags=UpdateFlags()) where {dim,C,T} =
-    CellIterator(dh, cellset, flags)
-CellIterator(dh::DofHandler{G,T}; cellset::Union{AbstractVector{Int},Nothing}=nothing, flags::UpdateFlags=UpdateFlags()) where {G<:AbstractGrid, T} =
+CellIterator(dh::Union{DofHandler{dim, T, G}, MixedDofHandler{dim,T,G}}; cellset::Union{AbstractVector{Int},Nothing}=nothing, flags::UpdateFlags=UpdateFlags()) where {dim,T,G<:AbstractGrid} =
     CellIterator(dh, cellset, flags)
 
 # iterator interface
@@ -102,21 +99,21 @@ function reinit!(ci::CellIterator{dim,C}, i::Int) where {dim,C}
     ci.current_cellid[] = ci.cellset === nothing ? i : i.cellset[i]
 
     if ci.flags.nodes
-        if isdefined(ci, :dh) && ci.dh isa MixedDofHandler
+        if ci.dh !== nothing && ci.dh isa MixedDofHandler
             cellnodes!(ci.nodes, ci.dh, ci.current_cellid[])
         else
             cellnodes!(ci.nodes, ci.grid, ci.current_cellid[])
         end
     end
     if ci.flags.coords
-        if isdefined(ci, :dh) && ci.dh isa MixedDofHandler
+        if ci.dh !== nothing && ci.dh isa MixedDofHandler
             cellcoords!(ci.coords, ci.dh, ci.current_cellid[])
         else
             cellcoords!(ci.coords, ci.grid, ci.current_cellid[])
         end
     end
 
-    if isdefined(ci, :dh) && ci.flags.celldofs # update celldofs
+    if ci.dh !== nothing && ci.flags.celldofs # update celldofs
         celldofs!(ci.celldofs, ci.dh, ci.current_cellid[])
     end
     return ci