diff --git a/src/Dofs/DofHandler.jl b/src/Dofs/DofHandler.jl index 4d50f2611b..295111eaf2 100644 --- a/src/Dofs/DofHandler.jl +++ b/src/Dofs/DofHandler.jl @@ -62,14 +62,13 @@ function SubDofHandler(dh::DH, cellset::AbstractVecOrSet{Int}) where {DH <: Abst error("all cells in a SubDofHandler must be of the same type") end # Make sure this set is disjoint with all other existing - for sdh in dh.subdofhandlers - if !isdisjoint(cellset, sdh.cellset) + if !isdisjoint(cellset, dh.cellset) error("cellset not disjoint with sets in existing SubDofHandlers") - end end # Construct and insert into the parent dh sdh = SubDofHandler{typeof(dh)}(dh, convert_to_orderedset(cellset), Symbol[], Interpolation[], Int[], ScalarWrapper(-1)) push!(dh.subdofhandlers, sdh) + union!(dh.cellset, cellset) return sdh end @@ -95,6 +94,7 @@ end struct DofHandler{dim,G<:AbstractGrid{dim}} <: AbstractDofHandler subdofhandlers::Vector{SubDofHandler{DofHandler{dim, G}}} + cellset::OrderedSet{Int} field_names::Vector{Symbol} # Dofs for cell i are stored in cell_dofs at the range: # cell_dofs_offset[i]:(cell_dofs_offset[i]+ndofs_per_cell(dh, i)-1) @@ -131,7 +131,7 @@ close!(dh) function DofHandler(grid::G) where {dim, G <: AbstractGrid{dim}} ncells = getncells(grid) sdhs = SubDofHandler{DofHandler{dim, G}}[] - DofHandler{dim, G}(sdhs, Symbol[], Int[], zeros(Int, ncells), zeros(Int, ncells), ScalarWrapper(false), grid, ScalarWrapper(-1)) + DofHandler{dim, G}(sdhs, OrderedSet{Int}(), Symbol[], Int[], zeros(Int, ncells), zeros(Int, ncells), ScalarWrapper(false), grid, ScalarWrapper(-1)) end function Base.show(io::IO, mime::MIME"text/plain", dh::DofHandler) diff --git a/src/iterators.jl b/src/iterators.jl index 810fcf5176..9bcac7c78f 100644 --- a/src/iterators.jl +++ b/src/iterators.jl @@ -254,20 +254,30 @@ struct CellIterator{CC<:CellCache, IC<:IntegerCollection} set::IC end -function CellIterator(gridordh::Union{Grid,DofHandler}, +function CellIterator(grid::Grid, set::Union{IntegerCollection,Nothing}=nothing, flags::UpdateFlags=UpdateFlags()) if set === nothing - grid = gridordh isa DofHandler ? get_grid(gridordh) : gridordh set = 1:getncells(grid) end - if gridordh isa DofHandler - # TODO: Since the CellCache is resizeable this is not really necessary to check - # here, but might be useful to catch slow code paths? - _check_same_celltype(get_grid(gridordh), set) + + return CellIterator(CellCache(grid, flags), set) +end + +function CellIterator(dh::DofHandler, + set::Union{IntegerCollection,Nothing}=nothing, + flags::UpdateFlags=UpdateFlags()) + if set === nothing + set = dh.cellset end - return CellIterator(CellCache(gridordh, flags), set) + + # TODO: Since the CellCache is resizeable this is not really necessary to check + # here, but might be useful to catch slow code paths? + _check_same_celltype(get_grid(dh), set) + + return CellIterator(CellCache(dh, flags), set) end + function CellIterator(gridordh::Union{Grid,DofHandler}, flags::UpdateFlags) return CellIterator(gridordh, nothing, flags) end diff --git a/test/test_mixeddofhandler.jl b/test/test_mixeddofhandler.jl index e0a7213217..293d1a497f 100644 --- a/test/test_mixeddofhandler.jl +++ b/test/test_mixeddofhandler.jl @@ -661,11 +661,7 @@ function test_celliterator_on_true_subdomain_smoketest() end for cell in CellIterator(dh) - if cellid(cell) <= 3 - @test length(celldofs(cell)) == getnbasefunctions(ip) - else - @test length(celldofs(cell)) == 0 - end + @test cellid(cell) in [1,2,3] end end