Support for tags on entities for all meshes

src/core_types/core_types.jl

@@ -466,6 +466,7 @@ An entity for something that isn't associated with an entity
 """
 struct NoEntity <: JutulEntity end
 
+
 # Sim model
 
 """
@@ -1039,3 +1040,133 @@ function renumber!(x, im::IndexRenumerator)
         x[i] = im[v]
     end
 end
+
+struct EntityTags{T}
+    tags::Dict{Symbol, Dict{Symbol, Vector{T}}}
+    count::Int
+end
+
+function EntityTags(n; tag_type = Int)
+    data = Dict{Symbol, Dict{Symbol, Vector{tag_type}}}()
+    return EntityTags{tag_type}(data, n)
+end
+
+Base.haskey(et::EntityTags, k) = Base.haskey(et.tags, k)
+Base.keys(et::EntityTags) = Base.keys(et.tags)
+Base.getindex(et::EntityTags, arg...) = Base.getindex(et.tags, arg...)
+Base.setindex!(et::EntityTags, arg...) = Base.setindex!(et.tags, arg...)
+
+export set_mesh_entity_tag!, get_mesh_entity_tag
+struct MeshEntityTags{T}
+    tags::Dict{JutulEntity, EntityTags{T}}
+end
+
+Base.haskey(et::MeshEntityTags, k) = Base.haskey(et.tags, k)
+Base.keys(et::MeshEntityTags) = Base.keys(et.tags)
+Base.getindex(et::MeshEntityTags, arg...) = Base.getindex(et.tags, arg...)
+Base.setindex!(et::MeshEntityTags, arg...) = Base.setindex!(et.tags, arg...)
+
+function Base.show(io::IO, t::MIME"text/plain", options::MeshEntityTags{T}) where T
+    println(io, "MeshEntityTags stored as $T:")
+    for (k, v) in pairs(options.tags)
+        kv = keys(v)
+        if length(kv) == 0
+            kv = "<no tags>"
+        else
+            s = map(x -> "x $(keys(v[x]))", collect(kv))
+            kv = join(s, ",")
+        end
+        println(io, "    $k: $(kv)")
+    end
+end
+
+function MeshEntityTags(g::JutulMesh; kwarg...)
+    return MeshEntityTags(declare_entities(g); kwarg...)
+end
+
+function MeshEntityTags(entities = missing; tag_type = Int)
+    tags = Dict{JutulEntity, EntityTags{tag_type}}()
+    if !ismissing(entities)
+        for epair in entities
+            tags[epair.entity] = EntityTags(epair.count, tag_type = tag_type)
+        end
+    end
+    return MeshEntityTags{tag_type}(tags)
+end
+
+function initialize_entity_tags!(g::JutulMesh)
+    tags = mesh_entity_tags(g)
+    for epair in declare_entities(g)
+        tags[epair.entity] = EntityTags(epair.count)
+    end
+    return g
+end
+
+function mesh_entity_tags(x)
+    return x.tags
+end
+
+function set_mesh_entity_tag!(m::JutulMesh, arg...; kwarg...)
+    set_mesh_entity_tag!(mesh_entity_tags(m), arg...; kwarg...)
+    return m
+end
+
+function set_mesh_entity_tag!(met::MeshEntityTags{T}, entity::JutulEntity, tag_group::Symbol, tag_value::Symbol, ix::Vector{T}; allow_merge = true, allow_new = true) where T
+    tags = met.tags[entity]
+    tags::EntityTags{T}
+    if !haskey(tags, tag_group)
+        @assert allow_new "allow_new = false and tag group $tag_group for entity $entity already exists."
+        tags[tag_group] = Dict{Symbol, Vector{T}}()
+    end
+    @assert maximum(ix) <= tags.count "Tag value must not exceed $(tags.count) for $entity"
+    @assert minimum(ix) > 0 "Tags must have positive indices."
+
+    tg = tags[tag_group]
+    if haskey(tg, tag_value)
+        @assert allow_merge "allow_merge = false and tag $tag_value in group $tag_group for entity $entity already exists."
+        vals = tg[tag_value]
+        for i in ix
+            push!(vals, i)
+        end
+    else
+        tg[tag_value] = copy(ix)
+    end
+    vals = tg[tag_value]
+    sort!(vals)
+    unique!(vals)
+    return met
+end
+
+
+"""
+    get_mesh_entity_tag(met::JutulMesh, entity::JutulEntity, tag_group::Symbol, tag_value = missing; throw = true)
+
+Get the indices tagged for `entity` in group `tag_group`, optionally for the
+specific `tag_value`. If `ismissing(tag_value)`, the Dict containing the tag
+group will be returned.
+"""
+function get_mesh_entity_tag(m::JutulMesh, arg...; kwarg...)
+    return get_mesh_entity_tag(mesh_entity_tags(m), arg...; kwarg...)
+end
+
+function get_mesh_entity_tag(met::MeshEntityTags, entity::JutulEntity, tag_group::Symbol, tag_value = missing; throw = true)
+    out = missing
+    tags = met.tags[entity]
+    if haskey(tags, tag_group)
+        tg = tags[tag_group]
+        if ismissing(tag_value)
+            out = tg
+        elseif haskey(tg, tag_value)
+            tag_value::Symbol
+            out = tg[tag_value]
+        end
+    end
+    if ismissing(out) && throw
+        if ismissing(tag_value)
+            error("Tag $tag_group not found in $entity.")
+        else
+            error("Tag $tag_group.$tag_value not found in $entity.")
+        end
+    end
+    return out
+end

src/meshes/cart.jl

@@ -25,10 +25,15 @@ julia> CartesianMesh((2, 3), ([1.0, 2.0], [0.1, 3.0, 2.5]))
 CartesianMesh (3D) with 3x5x2=30 cells
 ```
 """
-struct CartesianMesh{D, Δ, O} <: FiniteVolumeMesh
-    dims::D   # Tuple of dimensions (nx, ny, [nz])
-    deltas::Δ # Either a tuple of scalars (uniform grid) or a tuple of vectors (non-uniform grid)
-    origin::O # Coordinate of lower left corner
+struct CartesianMesh{D, Δ, O, T} <: FiniteVolumeMesh
+    "Tuple of dimensions (nx, ny, [nz])"
+    dims::D
+    "Either a tuple of scalars (uniform grid) or a tuple of vectors (non-uniform grid)"
+    deltas::Δ
+    "Coordinate of lower left corner"
+    origin::O
+    "Tags on cells/faces/nodes"
+    tags::MeshEntityTags{T}
     function CartesianMesh(dims::Tuple, deltas_or_size::Union{Nothing, Tuple} = nothing; origin = nothing)
         dim = length(dims)
         if isnothing(deltas_or_size)
@@ -55,7 +60,10 @@ struct CartesianMesh{D, Δ, O} <: FiniteVolumeMesh
         end
         @assert length(deltas_or_size) == dim
         deltas = generate_deltas(deltas_or_size)
-        return new{typeof(dims), typeof(deltas), typeof(origin)}(dims, deltas, origin)
+        tags = MeshEntityTags()
+        g = new{typeof(dims), typeof(deltas), typeof(origin), Int}(dims, deltas, origin, tags)
+        initialize_entity_tags!(g)
+        return g
     end
 end
 Base.show(io::IO, g::CartesianMesh) = print(io, "CartesianMesh ($(dim(g))D) with $(join(grid_dims_ijk(g), "x"))=$(number_of_cells(g)) cells")

src/meshes/coarse.jl

@@ -8,6 +8,7 @@ struct CoarseMesh{G, T} <: FiniteVolumeMesh
     coarse_boundary_to_fine::T
     face_neighbors::Vector{Tuple{Int, Int}}
     boundary_cells::Vector{Int}
+    tags::MeshEntityTags{Int}
 end
 
 """
@@ -79,8 +80,10 @@ function CoarseMesh(G, p)
         return (IndirectionMap(fmap, pos), neigh)
     end
     cfaces, coarse_neighbors = to_indir_facemap(coarse_faces)
-    # cboundary, coarse_boundary_cells = to_indir_facemap(coarse_boundary)
-    return CoarseMesh(G, p, cells_p, cfaces, cboundary, coarse_neighbors, coarse_boundary_cells)
+    tags = MeshEntityTags()
+    cg = CoarseMesh(G, p, cells_p, cfaces, cboundary, coarse_neighbors, coarse_boundary_cells, tags)
+    initialize_entity_tags!(cg)
+    return cg
 end
 
 function Base.show(io::IO, t::MIME"text/plain", g::CoarseMesh)

src/meshes/mrst.jl

@@ -3,6 +3,7 @@ struct MRSTWrapMesh <: FiniteVolumeMesh
     N::Matrix
     nc::Int
     nf::Int
+    tags::MeshEntityTags{Int}
 end
 
 function MRSTWrapMesh(G, N = nothing)
@@ -19,7 +20,10 @@ function MRSTWrapMesh(G, N = nothing)
     end
     nf = size(N, 2)
     nc = G.cells.num
-    return MRSTWrapMesh(G, N, nc, nf)
+    tags = MeshEntityTags()
+    g = MRSTWrapMesh(G, N, nc, nf, tags)
+    initialize_entity_tags!(g)
+    return g
 end
 
 function grid_dims_ijk(g::MRSTWrapMesh)

src/meshes/unstructured/types.jl

@@ -22,7 +22,7 @@ struct FaceMap{M, N}
     end
 end
 
-struct UnstructuredMesh{D, S, IM, IF, M, F, BM, NM} <: FiniteVolumeMesh
+struct UnstructuredMesh{D, S, IM, IF, M, F, BM, NM, T} <: FiniteVolumeMesh
     structure::S
     faces::FaceMap{M, Tuple{Int, Int}}
     boundary_faces::FaceMap{M, Int}
@@ -31,6 +31,8 @@ struct UnstructuredMesh{D, S, IM, IF, M, F, BM, NM} <: FiniteVolumeMesh
     face_map::IF
     boundary_map::BM
     node_map::NM
+    "Tags on cells/faces/nodes"
+    tags::MeshEntityTags{T}
 end
 
 function convert_coord_points(points::AbstractMatrix{F}) where F
@@ -249,7 +251,20 @@ function UnstructuredMesh(
     bnd = FaceMap(cells_to_bnd, bnd_to_nodes, boundary_cells)
     @assert length(face_neighbors) == length(faces_to_nodes)
     @assert length(boundary_cells) == length(bnd_to_nodes)
-    return UnstructuredMesh{dim, S, IM, IF, T, F, BM, NM}(structure, faces, bnd, node_points, cell_map, face_map, boundary_map, node_map)
+    tags = MeshEntityTags()
+    g = UnstructuredMesh{dim, S, IM, IF, T, F, BM, NM, Int}(
+        structure,
+        faces,
+        bnd,
+        node_points,
+        cell_map,
+        face_map,
+        boundary_map,
+        node_map,
+        tags
+    )
+    initialize_entity_tags!(g)
+    return g
 end
 
 function UnstructuredMesh(G::UnstructuredMesh)

test/utils.jl

@@ -237,3 +237,22 @@ end
     @test isdir(jutul_output_path())
     @test last(splitdir(jutul_output_path("testname"))) == "testname"
 end
+@testset "mesh tags" begin
+    for i in 1:2
+        g = CartesianMesh((10,1,1))
+        if i == 2
+            g = UnstructuredMesh(g)
+        end
+        set_mesh_entity_tag!(g, Cells(), :group, :tag1, [1, 2, 3])
+        set_mesh_entity_tag!(g, Cells(), :group, :tag2, [5, 4, 6])
+        set_mesh_entity_tag!(g, Cells(), :group2, :tag1, [1, 2, 3])
+        set_mesh_entity_tag!(g, Cells(), :group2, :tag1, [3, 7, 1])
+
+        @test_throws "Tag value must not exceed 10 for Cells()" set_mesh_entity_tag!(g, Cells(), :group2, :tag1, [21])
+        @test get_mesh_entity_tag(g, Cells(), :group, :tag1) == [1, 2, 3]
+        @test get_mesh_entity_tag(g, Cells(), :group, :tag2) == [4, 5, 6] # Sorted.
+        @test get_mesh_entity_tag(g, Cells(), :group2, :tag1) == [1, 2, 3, 7] # Sorted.
+        @test_throws "Tag group2.tag3 not found in Cells()." get_mesh_entity_tag(g, Cells(), :group2, :tag3)
+        @test ismissing(get_mesh_entity_tag(g, Cells(), :group2, :tag3, throw = false))
+    end
+end