Merge pull request #114 from sintefmath/dev

Faster and more robust trajectory search, adjoint fix for multimodel, improved Gmsh support, litre fix and other minor improvements

Author: moyner

Project.toml

@@ -1,7 +1,7 @@
 name = "Jutul"
 uuid = "2b460a1a-8a2b-45b2-b125-b5c536396eb9"
 authors = ["Olav Møyner <[email protected]>"]
-version = "0.3.3"
+version = "0.3.4"
 
 [deps]
 AlgebraicMultigrid = "2169fc97-5a83-5252-b627-83903c6c433c"

ext/JutulGmshExt/interface.jl

@@ -1,16 +1,44 @@
 
-function Jutul.mesh_from_gmsh(pth; verbose = false, kwarg...)
-    Gmsh.initialize()
+function Jutul.mesh_from_gmsh(pth; manage_gmsh = true, kwarg...)
+    if manage_gmsh
+        Gmsh.initialize()
+    end
     ext = pth |> splitext |> last
     gmsh.open(pth)
     if lowercase(ext) == ".geo"
         gmsh.model.mesh.generate()
     end
+    g = missing
+    try
+        g = Jutul.mesh_from_gmsh(; kwarg...)
+    finally
+        if manage_gmsh
+            Gmsh.finalize()
+        end
+    end
+    if ismissing(g)
+        error("Failed to parse mesh")
+    end
+    return g
+end
 
+function Jutul.mesh_from_gmsh(; verbose = false, reverse_z = false, kwarg...)
     dim = gmsh.model.getDimension()
     dim == 3 || error("Only 3D models are supported")
 
     gmsh.model.mesh.removeDuplicateNodes()
+    if reverse_z
+        # Note: Gmsh API lets us send only the first 3 rows of the 4 by 4 matrix
+        # which is sufficient here.
+        gmsh.model.mesh.affineTransform(
+            [
+                1.0, 0.0, 0.0, 0.0,
+                0.0, 1.0, 0.0, 0.0,
+                0.0, 0.0, -1.0, 0.0
+            ]
+        )
+        gmsh.model.mesh.generate()
+    end
     node_tags, pts, = gmsh.model.mesh.getNodes()
     node_remap = Dict{UInt64, Int}()
     for (i, tag) in enumerate(node_tags)
@@ -30,9 +58,6 @@ function Jutul.mesh_from_gmsh(pth; verbose = false, kwarg...)
     face_lookup = generate_face_lookup(faces_to_nodes)
 
     cells_to_faces = parse_cells(remaps, faces_to_nodes, face_lookup, verbose = verbose)
-    # We are done with Gmsh.jl at this point, finalize it.
-    Gmsh.finalize()
-
     neighbors = build_neighbors(cells_to_faces, faces_to_nodes, face_lookup)
 
     # Make both of these in case we have rogue faces that are not connected to any cell.

src/ad/gradients.jl

@@ -85,7 +85,6 @@ function solve_adjoint_sensitivities(case::JutulCase, states::Vector, G; dt = ca
     )
 end
 
-
 function solve_adjoint_sensitivities(case::JutulCase, some_kind_of_result, G; kwarg...)
     if hasproperty(some_kind_of_result, :result)
         simresult = some_kind_of_result.result
@@ -317,6 +316,7 @@ function merge_sparsity!(sparsity::AbstractDict, s_new::AbstractDict)
         merge_sparsity!(sparsity[k], v)
     end
 end
+
 function merge_sparsity!(sparsity::AbstractVector, s_new::AbstractVector)
     for k in s_new
         push!(sparsity, k)
@@ -355,7 +355,7 @@ function state_gradient_outer!(∂F∂x, F, model, state, extra_arg; sparsity =
     state_gradient_inner!(∂F∂x, F, model, state, nothing, extra_arg; sparsity = sparsity)
 end
 
-function state_gradient_inner!(∂F∂x, F, model, state, tag, extra_arg, eval_model = model; sparsity = nothing)
+function state_gradient_inner!(∂F∂x, F, model, state, tag, extra_arg, eval_model = model; sparsity = nothing, symbol = nothing)
     layout = matrix_layout(model.context)
     get_partial(x::AbstractFloat, i) = 0.0
     get_partial(x::ForwardDiff.Dual, i) = x.partials[i]
@@ -376,8 +376,12 @@ function state_gradient_inner!(∂F∂x, F, model, state, tag, extra_arg, eval_m
         end
     end
 
-    function diff_entity!(∂F∂x, state, i, S, ne, np, offset)
-        state_i = local_ad(state, i, S)
+    function diff_entity!(∂F∂x, state, i, S, ne, np, offset, symbol)
+        if !isnothing(symbol)
+            state_i = local_ad(state, i, S, symbol)
+        else
+            state_i = local_ad(state, i, S)
+        end
         v = F(eval_model, state_i, extra_arg...)
         store_partials!(∂F∂x, v, i, ne, np, offset)
     end
@@ -395,7 +399,7 @@ function state_gradient_inner!(∂F∂x, F, model, state, tag, extra_arg, eval_m
             ltag = get_entity_tag(tag, e)
             S = typeof(get_ad_entity_scalar(1.0, np, tag = ltag))
             for i in it_rng
-                diff_entity!(∂F∂x, state, i, S, ne, np, offset)
+                diff_entity!(∂F∂x, state, i, S, ne, np, offset, symbol)
             end
         end
         offset += ne*np
@@ -471,7 +475,7 @@ function next_lagrange_multiplier!(adjoint_storage, i, G, state, state0, state_n
     else
         dt_next = timesteps[i+1]
         forces_next = forces_for_timestep(backward_sim, all_forces, timesteps, i+1)
-        @tic "jacobian (with state0)" adjoint_reassemble!(backward_sim, state_next, state, dt_next, forces_next, t + dt_next)
+        @tic "jacobian (with state0)" adjoint_reassemble!(backward_sim, state_next, state, dt_next, forces_next, t + dt)
         lsys_next = backward_sim.storage.LinearizedSystem
         op = linear_operator(lsys_next, skip_red = true)
         # In-place version of
@@ -673,6 +677,9 @@ function evaluate_objective(G, model, states, timesteps, all_forces; large_value
     function convert_state_to_jutul_storage(model::MultiModel, x::AbstractDict)
         s = JutulStorage()
         for (k, v) in pairs(x)
+            if k == :substates
+                continue
+            end
             s[k] = JutulStorage(v)
         end
         return s
@@ -869,9 +876,10 @@ function adjoint_transfer_canonical_order_inner!(λ, dx, model, ::BlockMajorLayo
         if bz > 0
             error("Assumed that block major has a single entity group")
         end
-        bz = degrees_of_freedom_per_entity(model, e)
+        bz = degrees_of_freedom_per_entity(model, e)::Int
     end
     n = length(dx) ÷ bz
+    n::Int
     if to_canonical
         for b in 1:bz
             for i in 1:n

src/ad/local_ad.jl

@@ -6,14 +6,20 @@ struct LocalPerspectiveAD{T, N, A<:AbstractArray{T,N}, I} <: AbstractArray{T,N}
 end
 
 function LocalPerspectiveAD(a::A, index::I_t) where {A<:AbstractArray, I_t<:Integer}
-    LocalPerspectiveAD{eltype(a), ndims(A), A, I_t}(index, a)
+    return LocalPerspectiveAD{eltype(a), ndims(A), A, I_t}(index, a)
 end
 
 struct LocalStateAD{T, I, E} # Data type, index, entity tag
     index::I
     data::T
 end
 
+struct MultiModelLocalStateAD{T, I, E} # Data type, index, entity tag
+    symbol::Symbol
+    index::I
+    data::T
+end
+
 Base.keys(x::LocalStateAD) = keys(getfield(x, :data))
 
 struct ValueStateAD{T} # Data type
@@ -22,6 +28,12 @@ end
 
 Base.keys(x::ValueStateAD) = keys(getfield(x, :data))
 
+function convert_to_immutable_storage(x::ValueStateAD)
+    data = getfield(x, :data)
+    data = convert_to_immutable_storage(data)
+    return ValueStateAD(data)
+end
+
 const StateType = Union{NamedTuple,AbstractDict,JutulStorage}
 
 as_value(x::StateType) = ValueStateAD(x)
@@ -31,7 +43,6 @@ export local_ad
 @inline local_ad(v, ::Nothing) = as_value(v)
 @inline local_ad(v, i) = v
 
-
 @inline function new_entity_index(state::LocalStateAD{T, I, E}, index::I) where {T, I, E}
     return LocalStateAD{T, I, E}(index, getfield(state, :data))
 end
@@ -40,7 +51,6 @@ end
     return x
 end
 
-
 @inline local_entity(a::LocalPerspectiveAD) = a.index
 
 @inline function value_or_ad(A::LocalPerspectiveAD{T}, v::T, entity) where T
@@ -80,7 +90,6 @@ end
 end
 @inline Base.haskey(state::ValueStateAD, f::Symbol) = haskey(getfield(state, :data), f)
 
-
 # Match in type - pass index on
 @inline next_level_local_ad(x::AbstractArray{T}, ::Type{T}, index) where T = local_ad(x, index)
 
@@ -140,10 +149,26 @@ end
     return next_level_local_ad(val, E, index)
 end
 
+@inline function Base.getproperty(state::MultiModelLocalStateAD{T, I, E}, f::Symbol) where {T, I, E}
+    index = getfield(state, :index)
+    inner_state = getfield(state, :data)
+    val = getproperty(inner_state, f)
+    sym = getfield(state, :symbol)
+    if sym == f
+        return next_level_local_ad(val, E, index)
+    else
+        return as_value(val)
+    end
+end
+
 @inline function Base.getindex(state::LocalStateAD, s::Symbol)
     Base.getproperty(state, s)
 end
 
+@inline function Base.getindex(state::MultiModelLocalStateAD, s::Symbol)
+    Base.getproperty(state, s)
+end
+
 @inline function Base.getproperty(state::ValueStateAD{T}, f::Symbol) where {T}
     inner_state = getfield(state, :data)
     val = getproperty(inner_state, f)
@@ -163,10 +188,18 @@ Create local_ad for state for index I of AD tag of type ad_tag
     local_state_ad(state, index, ad_tag)
 end
 
+@inline function local_ad(state, index, ad_tag, symbol)
+    local_state_ad(state, index, ad_tag, symbol)
+end
+
 @inline function local_state_ad(state::T, index::I, ad_tag::∂T) where {T, I<:Integer, ∂T}
     return LocalStateAD{T, I, ad_tag}(index, state)
 end
 
+@inline function local_state_ad(state::T, index::I, ad_tag::∂T, symbol::Symbol) where {T, I<:Integer, ∂T}
+    return MultiModelLocalStateAD{T, I, ad_tag}(symbol, index, state)
+end
+
 function Base.show(io::IO, t::MIME"text/plain", x::LocalStateAD{T, I, E}) where {T, I, E}
     print(io, "Local state for $(unpack_tag(E)) -> $(getfield(x, :index)) with fields $(keys(getfield(x, :data)))")
 end

src/core_types/core_types.jl

@@ -523,6 +523,10 @@ function convert_to_immutable_storage(S::JutulStorage)
     return JutulStorage(tup)
 end
 
+function convert_to_immutable_storage(S::NamedTuple)
+    return S
+end
+
 function Base.getindex(S::JutulStorage, i::Int)
     d = data(S)
     if d isa OrderedDict

src/ext/gmsh_ext.jl

@@ -1,11 +1,15 @@
 export mesh_from_gmsh
 
 """
-    G = mesh_from_gmsh(pth; verbose = false)
+    G = mesh_from_gmsh(pth)
+    G = mesh_from_gmsh()
+    G = mesh_from_gmsh(pth; verbose = true)
 
 Parse a Gmsh file and return a Jutul `UnstructuredMesh` (in 3D only). Requires
-the Gmsh.jl package to be loaded. Please note that Gmsh is GPL licensed unless
-you have obtained another type of license from the authors.
+the Gmsh.jl package to be loaded. If no path is provided in `pth` it is assumed
+that you are managing the Gmsh state manually and it will use the current
+selected mesh inside Gmsh. Please note that Gmsh is GPL licensed unless you have
+obtained another type of license from the authors.
 """
 function mesh_from_gmsh
 

src/meshes/meshes.jl

@@ -279,3 +279,7 @@ function face_geometry_helper(G, D, nc, nf, e)
     end
     return (areas, normals, centroids)
 end
+
+
+import Base: promote_rule
+promote_rule(::Type{UnstructuredMesh}, ::Type{CartesianMesh}) = UnstructuredMesh