From eae7317773e9140b02564111c8fb59474110dffd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Olav=20M=C3=B8yner?= <olav.moyner@gmail.com>
Date: Wed, 20 Dec 2023 17:04:22 +0100
Subject: [PATCH 01/16] Add throw kw to cell_index

---
 src/meshes/cart.jl                      |  6 ++++--
 src/meshes/unstructured/unstructured.jl | 10 ++++++++--
 2 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/src/meshes/cart.jl b/src/meshes/cart.jl
index a0fc9e47..217f029a 100644
--- a/src/meshes/cart.jl
+++ b/src/meshes/cart.jl
@@ -93,13 +93,15 @@ coord_offset(pos, δ::AbstractVector) = sum(δ[1:(pos-1)])
 
 Get linear (scalar) index of mesh cell from provided IJK tuple `pos`.
 """
-function cell_index(g, pos::Tuple)
+function cell_index(g, pos::Tuple; throw = true)
     nx, ny, nz = grid_dims_ijk(g)
     x, y, z = cell_ijk(g, pos)
     return (z-1)*nx*ny + (y-1)*nx + x
 end
 
-cell_index(g, pos::Integer) = pos
+function cell_index(g, pos::Integer; throw = true)
+    return pos
+end
 
 function lower_corner_3d(g, index)
     pos = cell_ijk(g, index)
diff --git a/src/meshes/unstructured/unstructured.jl b/src/meshes/unstructured/unstructured.jl
index c2772d6d..ff71486a 100644
--- a/src/meshes/unstructured/unstructured.jl
+++ b/src/meshes/unstructured/unstructured.jl
@@ -65,7 +65,7 @@ function cell_ijk(g::UnstructuredMesh{D, CartesianIndex{D}}, index::Integer) whe
     return (x, y, z)
 end
 
-function cell_index(g::UnstructuredMesh, pos::Tuple)
+function cell_index(g::UnstructuredMesh, pos::Tuple; throw = true)
     nx, ny, nz = grid_dims_ijk(g)
     x, y, z = cell_ijk(g, pos)
     @assert x > 0 && x <= nx
@@ -77,7 +77,13 @@ function cell_index(g::UnstructuredMesh, pos::Tuple)
         t = index
     else
         t = findfirst(isequal(index), g.cell_map)
-        @assert !isnothing(t) "Cell $pos not found in active set"
+        if isnothing(t)
+            if throw
+                error("Cell $pos not found in active set")
+            else
+                return nothing
+            end
+        end
     end
     return t
 end

From ebd3b9bb46bd2d1fff499f390dead8392541b147 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Olav=20M=C3=B8yner?= <olav.moyner@gmail.com>
Date: Thu, 21 Dec 2023 14:27:48 +0100
Subject: [PATCH 02/16] plot_mesh: Avoid points that are not used

---
 ext/JutulMakieExt/mesh_plots.jl | 25 +++++++++++++++++++++----
 1 file changed, 21 insertions(+), 4 deletions(-)

diff --git a/ext/JutulMakieExt/mesh_plots.jl b/ext/JutulMakieExt/mesh_plots.jl
index 819c6cac..f34b04eb 100644
--- a/ext/JutulMakieExt/mesh_plots.jl
+++ b/ext/JutulMakieExt/mesh_plots.jl
@@ -20,11 +20,25 @@ function Jutul.plot_mesh_impl!(ax, m; cells = nothing, is_depth = true, outer =
             keep[i] = cell_ix[tri[i, 1]] in cells
         end
         tri = tri[keep, :]
+        tri, pts = remove_unused_points(tri, pts)
     end
     f = mesh!(ax, pts, tri; color = color, kwarg...)
     return f
 end
 
+function remove_unused_points(tri, pts)
+    unique_pts_ix = unique(vec(tri))
+    renum = Dict{Int, Int}()
+    for (i, ix) in enumerate(unique_pts_ix)
+        renum[ix] = i
+    end
+    pts = pts[unique_pts_ix, :]
+    for i in eachindex(tri)
+        tri[i] = renum[tri[i]]
+    end
+    return (tri, pts)
+end
+
 function Jutul.plot_cell_data_impl(m, data;
         colorbar = :horizontal,
         resolution = default_jutul_resolution(),
@@ -55,11 +69,13 @@ function Jutul.plot_cell_data_impl!(ax, m, data::AbstractVecOrMat; cells = nothi
             @assert length(cells) == nc
             cells = findall(cells)
         end
-        new_data = zeros(nc)
-        @. new_data = NaN
+        new_data = fill(NaN, nc)
         if length(data) == length(cells)
-            new_data[cells] = data
+            for (i, j) in enumerate(cells)
+                new_data[j] = data[i]
+            end
         else
+            @assert length(data) == nc
             for i in cells
                 new_data[i] = data[i]
             end
@@ -67,5 +83,6 @@ function Jutul.plot_cell_data_impl!(ax, m, data::AbstractVecOrMat; cells = nothi
         data = new_data
     end
     @assert length(data) == nc
-    return mesh!(ax, pts, tri; color = mapper.Cells(data), kwarg...)
+    color = mapper.Cells(data)
+    return mesh!(ax, pts, tri; color = color, kwarg...)
 end

From 292cdfefb9a062cc98124e8371de9de1aceb1c64 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Olav=20M=C3=B8yner?= <olav.moyner@gmail.com>
Date: Tue, 26 Dec 2023 22:25:55 +0100
Subject: [PATCH 03/16] Update interpolation.jl

---
 src/interpolation.jl | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/interpolation.jl b/src/interpolation.jl
index 95d07fad..e71bdd9f 100644
--- a/src/interpolation.jl
+++ b/src/interpolation.jl
@@ -80,6 +80,8 @@ function get_1d_interpolator(xs, ys; method = LinearInterpolant, cap_endpoints =
             ys = vcat(ys, ys[end])
         end
     end
+    @assert length(xs) > 1
+    @assert length(ys) == length(xs)
     return method(xs, ys)
 end
 

From d0eb5339c1bc78dc6b1b0c4f295dbe89efea4bca Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Olav=20M=C3=B8yner?= <olav.moyner@gmail.com>
Date: Tue, 26 Dec 2023 22:26:11 +0100
Subject: [PATCH 04/16] Better normal behavior for degenerate grids

---
 src/meshes/unstructured/unstructured.jl | 25 ++++++++++++++++++++-----
 1 file changed, 20 insertions(+), 5 deletions(-)

diff --git a/src/meshes/unstructured/unstructured.jl b/src/meshes/unstructured/unstructured.jl
index ff71486a..3e1b93bc 100644
--- a/src/meshes/unstructured/unstructured.jl
+++ b/src/meshes/unstructured/unstructured.jl
@@ -25,11 +25,26 @@ function face_normal(G::UnstructuredMesh, f, e = Faces())
     get_nodes(::BoundaryFaces) = G.boundary_faces
     nodes = get_nodes(e).faces_to_nodes[f]
     pts = G.node_points
-    a = pts[nodes[1]]
-    b = pts[nodes[2]]
-    c = pts[nodes[3]]
-
-    normal = cross(c - b, a - b)
+    n = length(nodes)
+    # If the geometry is well defined it would be sufficient to take the first
+    # triplet and use that to generate the normals. We assume it isn't and
+    # create a weighted sum where each weight corresponds to the areal between
+    # the triplets.
+    normal = zero(eltype(pts))
+    for i in 1:n
+        if i == 1
+            a = pts[nodes[n]]
+        else
+            a = pts[nodes[i-1]]
+        end
+        b = pts[nodes[i]]
+        if i == n
+            c = pts[nodes[1]]
+        else
+            c = pts[nodes[i+1]]
+        end
+        normal += cross(c - b, a - b)
+    end
     normal /= norm(normal, 2)
     return normal
 end

From de1189dde44c1cb6e8fbbc15ed95c60cab61f609 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Olav=20M=C3=B8yner?= <olav.moyner@gmail.com>
Date: Thu, 28 Dec 2023 20:47:27 +0100
Subject: [PATCH 05/16] Slight refactoring of trans. No functional changes.

---
 src/discretization/finite-volume.jl | 26 +++++++++-----------------
 1 file changed, 9 insertions(+), 17 deletions(-)

diff --git a/src/discretization/finite-volume.jl b/src/discretization/finite-volume.jl
index c8ba7294..957317d6 100644
--- a/src/discretization/finite-volume.jl
+++ b/src/discretization/finite-volume.jl
@@ -13,29 +13,19 @@ the symbol of some data defined on `Cells()`, a vector of numbers for each cell
 or a matrix with number of columns equal to the number of cells.
 """
 function compute_half_face_trans(g::DataDomain, perm)
-    compute_half_face_trans(g[:cell_centroids], g[:face_centroids], g[:normals], g[:areas], perm, g[:neighbors])
+    return compute_half_face_trans(g[:cell_centroids], g[:face_centroids], g[:normals], g[:areas], perm, g[:neighbors])
 end
 
 function compute_half_face_trans(g::TwoPointFiniteVolumeGeometry, perm)
-    compute_half_face_trans(g.cell_centroids, g.face_centroids, g.normals, g.areas, perm, g.neighbors)
+    return compute_half_face_trans(g.cell_centroids, g.face_centroids, g.normals, g.areas, perm, g.neighbors)
 end
 
 function compute_half_face_trans(cell_centroids, face_centroids, face_normals, face_areas, perm, N)
     nc = size(cell_centroids, 2)
     @assert size(N) == (2, length(face_areas))
     faces, facepos = get_facepos(N, nc)
-    facesigns = similar(faces)
-    for c in 1:nc
-        for ix in facepos[c]:(facepos[c+1]-1)
-            f = faces[ix]
-            if N[2, f] == c
-                facesigns[ix] = -1
-            else
-                facesigns[ix] = 1
-            end
-        end
-    end
-    compute_half_face_trans(cell_centroids, face_centroids, face_normals, face_areas, perm, faces, facepos, facesigns)
+    facesigns = get_facesigns(N, faces, facepos, nc)
+    return compute_half_face_trans(cell_centroids, face_centroids, face_normals, face_areas, perm, faces, facepos, facesigns)
 end
 
 
@@ -72,9 +62,11 @@ function compute_half_face_trans(cell_centroids, face_centroids, face_normals, f
     for cell = 1:nc
         for fpos = facepos[cell]:(facepos[cell+1]-1)
             face = faces[fpos]
+            cc = cell_centroids[:, cell]
+            fc = face_centroids[:, face]
             A = face_areas[face]
             K = expand_perm(perm[:, cell], dim)
-            C = face_centroids[:, face] - cell_centroids[:, cell]
+            C = fc - cc
             Nn = facesigns[fpos]*face_normals[:, face]
             T_hf[fpos] = compute_half_face_trans(A, K, C, Nn)
         end
@@ -115,9 +107,9 @@ function compute_face_trans(T_hf, N)
     nf = size(N, 2)
     T = zeros(nf)
     for i in eachindex(faces)
-        T[faces[i]] += 1/T_hf[i]
+        T[faces[i]] += 1.0/T_hf[i]
     end
-    T = 1 ./T
+    @. T = 1.0 /T
     return T
 end
 

From ef61ad592fff3f31a062cbccc544cc47635c154e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Olav=20M=C3=B8yner?= <olav.moyner@gmail.com>
Date: Fri, 29 Dec 2023 14:31:10 +0100
Subject: [PATCH 06/16] Add utility

---
 src/utils.jl | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/src/utils.jl b/src/utils.jl
index f7cd0b5b..22cada39 100644
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -753,6 +753,22 @@ function get_facepos(N, arg...)
     return (faces, facePos)
 end
 
+function get_facesigns(N, faces, facepos, nc)
+    facesigns = similar(faces)
+    for c in 1:nc
+        for ix in facepos[c]:(facepos[c+1]-1)
+            f = faces[ix]
+            if N[2, f] == c
+                facesigns[ix] = -1
+            else
+                @assert N[1, f] == c
+                facesigns[ix] = 1
+            end
+        end
+    end
+    return facesigns
+end
+
 export timing_breakdown
 function timing_breakdown(report)
     D = Dict{Symbol, Any}(:total => report[:total_time])

From 9eff4a15ac23b45f11fd7b485258093b66daf56d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Olav=20M=C3=B8yner?= <olav.moyner@gmail.com>
Date: Fri, 29 Dec 2023 14:31:24 +0100
Subject: [PATCH 07/16] Type assert for unit

---
 src/units/interface.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/units/interface.jl b/src/units/interface.jl
index 8060f2fb..382d762d 100644
--- a/src/units/interface.jl
+++ b/src/units/interface.jl
@@ -84,7 +84,7 @@ julia> si_unit(:day) # Get days represented as seconds
 ```
 """
 function si_unit(uname::Symbol)
-    return si_unit(Val(uname))
+    return si_unit(Val(uname))::Float64
 end
 
 function si_unit(uname::String)

From 489517c3a602ae906275f9e5de14810208255492 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Olav=20M=C3=B8yner?= <olav.moyner@gmail.com>
Date: Fri, 29 Dec 2023 18:25:40 +0100
Subject: [PATCH 08/16] Add tags for meshes

---
 src/core_types/core_types.jl     | 131 +++++++++++++++++++++++++++++++
 src/meshes/cart.jl               |  18 +++--
 src/meshes/mrst.jl               |   5 +-
 src/meshes/unstructured/types.jl |  19 ++++-
 test/utils.jl                    |  19 +++++
 5 files changed, 184 insertions(+), 8 deletions(-)

diff --git a/src/core_types/core_types.jl b/src/core_types/core_types.jl
index 77e143e4..d3e19086 100644
--- a/src/core_types/core_types.jl
+++ b/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
diff --git a/src/meshes/cart.jl b/src/meshes/cart.jl
index 217f029a..1ab847a2 100644
--- a/src/meshes/cart.jl
+++ b/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")
diff --git a/src/meshes/mrst.jl b/src/meshes/mrst.jl
index 06e0d554..891db950 100644
--- a/src/meshes/mrst.jl
+++ b/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,9 @@ function MRSTWrapMesh(G, N = nothing)
     end
     nf = size(N, 2)
     nc = G.cells.num
-    return MRSTWrapMesh(G, N, nc, nf)
+    g = MRSTWrapMesh(G, N, nc, nf, tags)
+    initialize_entity_tags!(g)
+    return g
 end
 
 function grid_dims_ijk(g::MRSTWrapMesh)
diff --git a/src/meshes/unstructured/types.jl b/src/meshes/unstructured/types.jl
index 95b89128..1ee05c21 100644
--- a/src/meshes/unstructured/types.jl
+++ b/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)
diff --git a/test/utils.jl b/test/utils.jl
index dc2d6482..a6f10e02 100644
--- a/test/utils.jl
+++ b/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

From 56b9e54607580119d0ea8cb58263b12f6b691abf Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Olav=20M=C3=B8yner?= <olav.moyner@gmail.com>
Date: Fri, 29 Dec 2023 18:33:17 +0100
Subject: [PATCH 09/16] Tags for coarse mesh + fix missing init

---
 src/meshes/coarse.jl | 7 +++++--
 src/meshes/mrst.jl   | 1 +
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/src/meshes/coarse.jl b/src/meshes/coarse.jl
index 398870bd..795bb4f0 100644
--- a/src/meshes/coarse.jl
+++ b/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)
+    initialize_entity_tags!(cg)
+    return cg
 end
 
 function Base.show(io::IO, t::MIME"text/plain", g::CoarseMesh)
diff --git a/src/meshes/mrst.jl b/src/meshes/mrst.jl
index 891db950..729ebc1b 100644
--- a/src/meshes/mrst.jl
+++ b/src/meshes/mrst.jl
@@ -20,6 +20,7 @@ function MRSTWrapMesh(G, N = nothing)
     end
     nf = size(N, 2)
     nc = G.cells.num
+    tags = MeshEntityTags()
     g = MRSTWrapMesh(G, N, nc, nf, tags)
     initialize_entity_tags!(g)
     return g

From b063e6c0b79461489f60f71b1cddcadf34f8e088 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Olav=20M=C3=B8yner?= <olav.moyner@gmail.com>
Date: Fri, 29 Dec 2023 19:45:03 +0100
Subject: [PATCH 10/16] Update coarse.jl

---
 src/meshes/coarse.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/meshes/coarse.jl b/src/meshes/coarse.jl
index 795bb4f0..987ff533 100644
--- a/src/meshes/coarse.jl
+++ b/src/meshes/coarse.jl
@@ -81,7 +81,7 @@ function CoarseMesh(G, p)
     end
     cfaces, coarse_neighbors = to_indir_facemap(coarse_faces)
     tags = MeshEntityTags()
-    cg = CoarseMesh(G, p, cells_p, cfaces, cboundary, coarse_neighbors, coarse_boundary_cells)
+    cg = CoarseMesh(G, p, cells_p, cfaces, cboundary, coarse_neighbors, coarse_boundary_cells, tags)
     initialize_entity_tags!(cg)
     return cg
 end

From 78e20092e132f59814eb3167bb9d59d3adc70ef4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Olav=20M=C3=B8yner?= <olav.moyner@gmail.com>
Date: Fri, 29 Dec 2023 20:52:05 +0100
Subject: [PATCH 11/16] Add utility to query tag

---
 src/core_types/core_types.jl | 12 +++++++++++-
 test/utils.jl                |  3 +++
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/src/core_types/core_types.jl b/src/core_types/core_types.jl
index d3e19086..3b41c95e 100644
--- a/src/core_types/core_types.jl
+++ b/src/core_types/core_types.jl
@@ -1056,7 +1056,7 @@ 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
+export set_mesh_entity_tag!, get_mesh_entity_tag, mesh_entity_has_tag
 struct MeshEntityTags{T}
     tags::Dict{JutulEntity, EntityTags{T}}
 end
@@ -1170,3 +1170,13 @@ function get_mesh_entity_tag(met::MeshEntityTags, entity::JutulEntity, tag_group
     end
     return out
 end
+
+function mesh_entity_has_tag(m::JutulMesh, arg...; kwarg...)
+    return mesh_entity_has_tag(mesh_entity_tags(m), arg...; kwarg...)
+end
+
+function mesh_entity_has_tag(met::MeshEntityTags, entity::JutulEntity, tag_group::Symbol, tag_value::Symbol, ix)
+    tag = get_mesh_entity_tag(met, entity, tag_group, tag_value)
+    ix = searchsortedfirst(tag, ix)
+    return ix < (length(tag)+1)
+end
diff --git a/test/utils.jl b/test/utils.jl
index a6f10e02..5fc5e961 100644
--- a/test/utils.jl
+++ b/test/utils.jl
@@ -252,6 +252,9 @@ end
         @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 mesh_entity_has_tag(g, Cells(), :group, :tag1, 1) == true
+        @test mesh_entity_has_tag(g, Cells(), :group, :tag1, 4) == false
+        @test mesh_entity_has_tag(g, Cells(), :group, :tag1, 3) == true
         @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

From 13a71b53bce8d34f54391fd6291e631a5a8b5a71 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Olav=20M=C3=B8yner?= <olav.moyner@gmail.com>
Date: Fri, 29 Dec 2023 21:17:36 +0100
Subject: [PATCH 12/16] Fix mesh_entity_has_tag logic

---
 src/core_types/core_types.jl | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/src/core_types/core_types.jl b/src/core_types/core_types.jl
index 3b41c95e..8e6e4608 100644
--- a/src/core_types/core_types.jl
+++ b/src/core_types/core_types.jl
@@ -1177,6 +1177,11 @@ end
 
 function mesh_entity_has_tag(met::MeshEntityTags, entity::JutulEntity, tag_group::Symbol, tag_value::Symbol, ix)
     tag = get_mesh_entity_tag(met, entity, tag_group, tag_value)
-    ix = searchsortedfirst(tag, ix)
-    return ix < (length(tag)+1)
+    pos = searchsortedfirst(tag, ix)
+    if pos > length(tag)
+        out = false
+    else
+        out = tag[pos] == ix
+    end
+    return out
 end

From ea5a464d720d6b85375e3f28bf69fe61852e0ed3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Olav=20M=C3=B8yner?= <olav.moyner@gmail.com>
Date: Tue, 2 Jan 2024 12:52:43 +0100
Subject: [PATCH 13/16] Use SMatrix for perm expansion

---
 src/discretization/finite-volume.jl | 81 ++++++++++++++++++++++-------
 1 file changed, 62 insertions(+), 19 deletions(-)

diff --git a/src/discretization/finite-volume.jl b/src/discretization/finite-volume.jl
index 957317d6..eedc984c 100644
--- a/src/discretization/finite-volume.jl
+++ b/src/discretization/finite-volume.jl
@@ -59,13 +59,14 @@ function compute_half_face_trans(cell_centroids, face_centroids, face_normals, f
     @assert(length(face_areas) == nf)
     # Check perm
     @assert(size(perm, 2) == nc)
+    vdim = Val(dim)
     for cell = 1:nc
         for fpos = facepos[cell]:(facepos[cell+1]-1)
             face = faces[fpos]
             cc = cell_centroids[:, cell]
             fc = face_centroids[:, face]
             A = face_areas[face]
-            K = expand_perm(perm[:, cell], dim)
+            K = expand_perm(perm[:, cell], vdim)
             C = fc - cc
             Nn = facesigns[fpos]*face_normals[:, face]
             T_hf[fpos] = compute_half_face_trans(A, K, C, Nn)
@@ -74,27 +75,68 @@ function compute_half_face_trans(cell_centroids, face_centroids, face_normals, f
     return T_hf
 end
 
-function expand_perm(K, dim; full = false)
+function expand_perm(K, dim)
+    return expand_perm(K, Val(dim))
+end
+
+function expand_perm(K, ::Val{1})
+    return only(K)
+end
+
+
+function expand_perm(K, ::Val{2})
+    T = eltype(K)
     n = length(K)
-    if n == dim
-        K_e = diagm(K)
-    elseif n == 1
-        K_e = first(K)
-        if full
-            # Expand to matrix
-            K_e = zeros(dim, dim) + I*K_e
-        end
+    if n == 1
+        K_xx = K_yy = only(K)
+        K_xy = zero(T)
+    elseif n == 2
+        K_xx = K[1]
+        K_yy = K[2]
+        K_xy = zero(T)
+    elseif n == 3
+        K_xx = K[1]
+        K_xy = K[2]
+        K_yy = K[2]
     else
-        if dim == 2
-            @assert n == 3 "Two-dimensional grids require 1/2/3 permeability entries per cell (was $n)"
-            K_e = [K[1] K[2]; K[2] K[3]]
-        else
-            @assert n == 6 "Three-dimensional grids require 1/3/6 permeability entries per cell (was $n)"
-            K_e =  [K[1] K[2] K[3];
-                    K[2] K[4] K[5];
-                    K[3] K[5] K[6]]
-        end
+        error("Permeability for two-dimensional grids must have 1/2/3 entries per cell, had $n")
+    end
+    K_e = @SMatrix [
+        K_xx K_xy;
+        K_xy K_yy
+        ]
+    return K_e
+end
+
+function expand_perm(K, ::Val{3})
+    T = eltype(K)
+    n = length(K)
+    K_xy = zero(T)
+    K_xz = zero(T)
+    K_yz = zero(T)
+    if n == 1
+        K_xx = K_yy = K_zz = only(K)
+    elseif n == 3
+        K_xx = K[1]
+        K_yy = K[2]
+        K_zz = K[3]
+    elseif n == 6
+        # First row
+        K_xx = K[1]
+        K_xy = K[2]
+        K_yz = K[3]
+        # Second row excluding symmetry
+        K_yy = K[4]
+        K_yz = K[5]
+        # Last entry
+        K_zz = K[6]
+    else
+        error("Permeability for three-dimensional meshes must have 1/3/6 entries per cell, had $n")
     end
+    K_e =  @SMatrix[
+        K_xx K_xy K_xz;
+        K_xy K_yy K_yz;
+        K_xz K_yz K_zz]
     return K_e
 end
 
@@ -104,6 +146,7 @@ end
 
 function compute_face_trans(T_hf, N)
     faces, facePos = get_facepos(N)
+    @assert length(T_hf) == length(faces)
     nf = size(N, 2)
     T = zeros(nf)
     for i in eachindex(faces)

From 806dc4b29469e25e1b13199e6c72f3566ff719c6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Olav=20M=C3=B8yner?= <olav.moyner@gmail.com>
Date: Tue, 2 Jan 2024 12:58:01 +0100
Subject: [PATCH 14/16] Fix for empty tags

---
 src/core_types/core_types.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/core_types/core_types.jl b/src/core_types/core_types.jl
index 8e6e4608..4a0ed656 100644
--- a/src/core_types/core_types.jl
+++ b/src/core_types/core_types.jl
@@ -1118,8 +1118,8 @@ function set_mesh_entity_tag!(met::MeshEntityTags{T}, entity::JutulEntity, tag_g
         @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."
+    @assert maximum(ix, init = one(T)) <= tags.count "Tag value must not exceed $(tags.count) for $entity"
+    @assert minimum(ix, init = tags.count) > 0 "Tags must have positive indices."
 
     tg = tags[tag_group]
     if haskey(tg, tag_value)

From 44d90601c36fd3ef088ba7f548c4a070b489b02d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Olav=20M=C3=B8yner?= <olav.moyner@gmail.com>
Date: Wed, 3 Jan 2024 12:15:44 +0100
Subject: [PATCH 15/16] Clean up of trans

---
 src/discretization/finite-volume.jl | 76 ++++++++++++++++++-----------
 1 file changed, 47 insertions(+), 29 deletions(-)

diff --git a/src/discretization/finite-volume.jl b/src/discretization/finite-volume.jl
index eedc984c..6faab4ea 100644
--- a/src/discretization/finite-volume.jl
+++ b/src/discretization/finite-volume.jl
@@ -1,8 +1,8 @@
 export compute_face_trans, compute_half_face_trans, compute_boundary_trans
 
-function compute_half_face_trans(d::DataDomain, k::Symbol = :permeability)
+function compute_half_face_trans(d::DataDomain, k::Symbol = :permeability; kwarg...)
     @assert haskey(d, k, Cells()) "$k must be defined in Cells."
-    return compute_half_face_trans(d, d[k])
+    return compute_half_face_trans(d, d[k]; kwarg...)
 end
 
 """
@@ -12,24 +12,32 @@ Compute half-face trans for the interior faces. The input `perm` can either be
 the symbol of some data defined on `Cells()`, a vector of numbers for each cell
 or a matrix with number of columns equal to the number of cells.
 """
-function compute_half_face_trans(g::DataDomain, perm)
-    return compute_half_face_trans(g[:cell_centroids], g[:face_centroids], g[:normals], g[:areas], perm, g[:neighbors])
+function compute_half_face_trans(g::DataDomain, perm; kwarg...)
+    return compute_half_face_trans(
+        g[:cell_centroids],
+        g[:face_centroids],
+        g[:normals],
+        g[:areas],
+        perm,
+        g[:neighbors]
+        ; kwarg...
+    )
 end
 
-function compute_half_face_trans(g::TwoPointFiniteVolumeGeometry, perm)
-    return compute_half_face_trans(g.cell_centroids, g.face_centroids, g.normals, g.areas, perm, g.neighbors)
+function compute_half_face_trans(g::TwoPointFiniteVolumeGeometry, perm; kwarg...)
+    return compute_half_face_trans(g.cell_centroids, g.face_centroids, g.normals, g.areas, perm, g.neighbors; kwarg...)
 end
 
-function compute_half_face_trans(cell_centroids, face_centroids, face_normals, face_areas, perm, N)
+function compute_half_face_trans(cell_centroids, face_centroids, face_normals, face_areas, perm, N; kwarg...)
     nc = size(cell_centroids, 2)
     @assert size(N) == (2, length(face_areas))
     faces, facepos = get_facepos(N, nc)
     facesigns = get_facesigns(N, faces, facepos, nc)
-    return compute_half_face_trans(cell_centroids, face_centroids, face_normals, face_areas, perm, faces, facepos, facesigns)
+    return compute_half_face_trans(cell_centroids, face_centroids, face_normals, face_areas, perm, faces, facepos, facesigns; kwarg...)
 end
 
 
-function compute_half_face_trans(cell_centroids, face_centroids, face_normals, face_areas, perm, faces, facepos, facesigns)
+function compute_half_face_trans(cell_centroids, face_centroids, face_normals, face_areas, perm, faces, facepos, facesigns; version = :xyz)
     nf = length(face_areas)
     dim = size(cell_centroids, 1)
 
@@ -45,31 +53,41 @@ function compute_half_face_trans(cell_centroids, face_centroids, face_normals, f
     end
 
     # Sanity check
-    @assert(dim == 2 || dim == 3 || dim == 1)
+    if !(dim == 2 || dim == 3 || dim == 1)
+        throw(ArgumentError("Dimension must 1, 2 or 3"))
+    end
     # Check cell centroids
-    @assert(size(cell_centroids, 1) == dim)
-    @assert(size(cell_centroids, 2) == nc)
+    cc_dim, cc_n = size(cell_centroids)
+    cc_dim == dim || throw(ArgumentError("Cell centroids had $cc_dim rows but grid had $dim dimension."))
+    cc_n == nc || throw(ArgumentError("Cell centroids had $cc_n columns but grid had $nc cells."))
     # Check face centroids
-    @assert(size(face_centroids, 1) == dim)
-    @assert(size(face_centroids, 2) == nf)
+    fc_dim, fc_n = size(face_centroids)
+    fc_dim == dim || throw(ArgumentError("Face centroids had $fc_dim rows but grid had $dim dimension."))
+    fc_n == nf || throw(ArgumentError("Face centroids had $fc_n columns but grid had $nf faces."))
     # Check normals
-    @assert(size(face_normals, 1) == dim)
-    @assert(size(face_normals, 2) == nf)
+    normal_dim, normal_n = size(face_normals)
+    normal_dim == dim || throw(ArgumentError("Face normals had $normal_dim rows but grid had $dim dimension."))
+    fc_n == nf || throw(ArgumentError("Face normals had $normal_n columns but grid had $nf faces."))
     # Check areas
-    @assert(length(face_areas) == nf)
+    # TODO: This isn't really checking anything since we get nf from areas...
+    length(face_areas) == nf || throw(ArgumentError("Face areas had $normal_n entries but grid had $nf faces."))
     # Check perm
-    @assert(size(perm, 2) == nc)
+    size(perm, 2) == nc || throw(ArgumentError("Permeability must have number of columns equal to number of cells (= $nc)."))
+    if !(version in (:xyz, :xyz_local))
+        throw(ArgumentError("version must be :xyz or :xyz_local"))
+    end
     vdim = Val(dim)
     for cell = 1:nc
         for fpos = facepos[cell]:(facepos[cell+1]-1)
             face = faces[fpos]
+            sgn = facesigns[fpos]
             cc = cell_centroids[:, cell]
             fc = face_centroids[:, face]
             A = face_areas[face]
-            K = expand_perm(perm[:, cell], vdim)
             C = fc - cc
-            Nn = facesigns[fpos]*face_normals[:, face]
-            T_hf[fpos] = compute_half_face_trans(A, K, C, Nn)
+            Nn = sgn*face_normals[:, face]
+            K = expand_perm(perm[:, cell], vdim)
+            T_hf[fpos] = half_face_trans(A, K, C, Nn)
         end
     end
     return T_hf
@@ -140,7 +158,7 @@ function expand_perm(K, ::Val{3})
     return K_e
 end
 
-function compute_half_face_trans(A, K, C, N)
+function half_face_trans(A, K, C, N)
     return A*(dot(K*C, N))/dot(C, C)
 end
 
@@ -152,7 +170,7 @@ function compute_face_trans(T_hf, N)
     for i in eachindex(faces)
         T[faces[i]] += 1.0/T_hf[i]
     end
-    @. T = 1.0 /T
+    @. T = 1.0 / T
     return T
 end
 
@@ -161,8 +179,8 @@ function compute_face_trans(g::JutulMesh, perm)
     return compute_face_trans(geo, perm)
 end
 
-function compute_face_trans(geometry::JutulGeometry, permeability)
-    T_hf = compute_half_face_trans(geometry, permeability)
+function compute_face_trans(geometry::JutulGeometry, permeability; kwarg...)
+    T_hf = compute_half_face_trans(geometry, permeability; kwarg...)
     return compute_face_trans(T_hf, geometry.neighbors)
 end
 
@@ -173,8 +191,8 @@ Compute face trans for the interior faces. The input `perm` can either be
 the symbol of some data defined on `Cells()`, a vector of numbers for each cell
 or a matrix with number of columns equal to the number of cells.
 """
-function compute_face_trans(d::DataDomain, arg...)
-    T_hf = compute_half_face_trans(d, arg...)
+function compute_face_trans(d::DataDomain, arg...; kwarg...)
+    T_hf = compute_half_face_trans(d, arg...; kwarg...)
     return compute_face_trans(T_hf, d[:neighbors])
 end
 
@@ -191,7 +209,7 @@ either be the symbol of some data defined on `Cells()`, a vector of numbers
 for each cell or a matrix with number of columns equal to the number of
 cells.
 """
-function compute_boundary_trans(d::DataDomain, perm)
+function compute_boundary_trans(d::DataDomain, perm; kwarg...)
     @assert hasentity(d, BoundaryFaces()) "Domain must have BoundaryFaces() to compute boundary transmissibilities"
     face_areas = d[:boundary_areas]
     cells = d[:boundary_neighbors]
@@ -209,7 +227,7 @@ function compute_boundary_trans(d::DataDomain, perm)
     faces = collect(1:nf)
     facepos = collect(1:(nc+1))
     facesigns = ones(nf)
-    return compute_half_face_trans(cell_centroids, face_centroids, face_normals, face_areas, perm, faces, facepos, facesigns)
+    return compute_half_face_trans(cell_centroids, face_centroids, face_normals, face_areas, perm, faces, facepos, facesigns; kwarg...)
 end
 
 export compute_face_gdz

From 65449034aa605493b4dfbd042b9597ea8df592a4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Olav=20M=C3=B8yner?= <olav.moyner@gmail.com>
Date: Wed, 3 Jan 2024 12:57:13 +0100
Subject: [PATCH 16/16] Add support for ijk trans

---
 src/discretization/finite-volume.jl | 33 ++++++++++++++++++++++++-----
 1 file changed, 28 insertions(+), 5 deletions(-)

diff --git a/src/discretization/finite-volume.jl b/src/discretization/finite-volume.jl
index 6faab4ea..fb2695a4 100644
--- a/src/discretization/finite-volume.jl
+++ b/src/discretization/finite-volume.jl
@@ -37,7 +37,7 @@ function compute_half_face_trans(cell_centroids, face_centroids, face_normals, f
 end
 
 
-function compute_half_face_trans(cell_centroids, face_centroids, face_normals, face_areas, perm, faces, facepos, facesigns; version = :xyz)
+function compute_half_face_trans(cell_centroids, face_centroids, face_normals, face_areas, perm, faces, facepos, facesigns; version = :xyz, face_dir = missing)
     nf = length(face_areas)
     dim = size(cell_centroids, 1)
 
@@ -73,8 +73,26 @@ function compute_half_face_trans(cell_centroids, face_centroids, face_normals, f
     length(face_areas) == nf || throw(ArgumentError("Face areas had $normal_n entries but grid had $nf faces."))
     # Check perm
     size(perm, 2) == nc || throw(ArgumentError("Permeability must have number of columns equal to number of cells (= $nc)."))
-    if !(version in (:xyz, :xyz_local))
-        throw(ArgumentError("version must be :xyz or :xyz_local"))
+    if !(version in (:xyz, :ijk))
+        throw(ArgumentError("version must be :xyz or :ijk"))
+    end
+    is_xyz = version == :xyz
+    if version == :ijk
+        if size(perm, 1) != dim
+            throw(ArgumentError("version = :ijk is only valid when perm is strictly diagonal."))
+        end
+        if ismissing(face_dir)
+            throw(ArgumentError("version = :ijk cannot be used without also passing face_dir."))
+        end
+        if length(face_dir) != nf
+            throw(ArgumentError("face_dir must have one entry per face."))
+        end
+        if maximum(face_dir) > dim
+            throw(ArgumentError("face_dir entry exceeds grid dim ($dim)."))
+        end
+        if minimum(face_dir) < 1
+            throw(ArgumentError("face_dir entry is less than 1."))
+        end
     end
     vdim = Val(dim)
     for cell = 1:nc
@@ -86,8 +104,13 @@ function compute_half_face_trans(cell_centroids, face_centroids, face_normals, f
             A = face_areas[face]
             C = fc - cc
             Nn = sgn*face_normals[:, face]
-            K = expand_perm(perm[:, cell], vdim)
-            T_hf[fpos] = half_face_trans(A, K, C, Nn)
+            if is_xyz
+                K = expand_perm(perm[:, cell], vdim)
+            else
+                K = perm[face_dir[face], cell]
+            end
+            T = half_face_trans(A, K, C, Nn)
+            T_hf[fpos] = T
         end
     end
     return T_hf