Iterate only over cells in subdomain of dof handler

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)

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

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