State gradients and minor fixes

Project.toml

@@ -1,7 +1,7 @@
 name = "Jutul"
 uuid = "2b460a1a-8a2b-45b2-b125-b5c536396eb9"
 authors = ["Olav Møyner <olav.moyner@gmail.com>"]
-version = "0.2.37"
+version = "0.2.38"
 
 [deps]
 AlgebraicMultigrid = "2169fc97-5a83-5252-b627-83903c6c433c"

ext/JutulMakieExt/interactive_3d.jl

@@ -68,7 +68,7 @@ function plot_interactive_impl(grid, states;
     )
     has_primitives = !isnothing(primitives)
     active_filters = []
-    if states isa AbstractDict
+    if states isa AbstractDict || states isa DataDomain
         states = [states]
     end
     if states isa AbstractVecOrMat && eltype(states)<:AbstractFloat

ext/JutulMakieExt/mesh_plots.jl

@@ -66,7 +66,11 @@ function Jutul.plot_mesh_impl!(ax, m;
         tri = tri[keep, :]
         tri, pts = remove_unused_points(tri, pts)
     end
-    f = mesh!(ax, pts, tri; color = color, backlight = 1, kwarg...)
+    if length(pts) > 0
+        f = mesh!(ax, pts, tri; color = color, backlight = 1, kwarg...)
+    else
+        f = nothing
+    end
     return f
 end
 

src/ad/gradients.jl

@@ -63,6 +63,10 @@ function solve_adjoint_sensitivities(model, states, reports_or_timesteps, G;
             jutul_message("Adjoints", "Storing sensitivities.", color = :blue)
         end
         out = store_sensitivities(parameter_model, ∇G, storage.parameter_map)
+        s0_map = storage.state0_map
+        if !ismissing(s0_map)
+            store_sensitivities!(out, storage.backward.model, ∇G, s0_map)
+        end
     end
     if extra_output
         return (out, storage)
@@ -107,24 +111,27 @@ function setup_adjoint_storage(model;
         parameters = setup_parameters(model),
         n_objective = nothing,
         targets = parameter_targets(model),
+        include_state0 = false,
         use_sparsity = true,
         linear_solver = select_linear_solver(model, mode = :adjoint, rtol = 1e-6),
         param_obj = true,
         info_level = 0,
         kwarg...
     )
-    # Set up the generic adjoint storage
-    storage = setup_adjoint_storage_base(
-            model, state0, parameters,
-            use_sparsity = use_sparsity,
-            linear_solver = linear_solver,
-            n_objective = n_objective,
-            info_level = info_level
-        )
     # Create parameter model for ∂Fₙ / ∂p
     parameter_model = adjoint_parameter_model(model, targets)
+    n_prm = number_of_degrees_of_freedom(parameter_model)
     # Note that primary is here because the target parameters are now the primaries for the parameter_model
     parameter_map, = variable_mapper(parameter_model, :primary, targets = targets; kwarg...)
+    if include_state0
+        state0_map, = variable_mapper(model, :primary)
+        n_state0 = number_of_degrees_of_freedom(model)
+        state0_vec = zeros(n_state0)
+    else
+        state0_map = missing
+        state0_vec = missing
+        n_state0 = 0
+    end
     # Transfer over parameters and state0 variables since many parameters are now variables
     state0_p = swap_variables(state0, parameters, parameter_model, variables = true)
     parameters_p = swap_variables(state0, parameters, parameter_model, variables = false)
@@ -137,11 +144,21 @@ function setup_adjoint_storage(model;
         dobj_dparam = nothing
         param_buf = nothing
     end
+    # Set up the generic adjoint storage
+    storage = setup_adjoint_storage_base(
+            model, state0, parameters,
+            use_sparsity = use_sparsity,
+            linear_solver = linear_solver,
+            n_objective = n_objective,
+            info_level = info_level,
+    )
     storage[:dparam] = dobj_dparam
     storage[:param_buf] = param_buf
     storage[:parameter] = parameter_sim
     storage[:parameter_map] = parameter_map
-    storage[:n] = number_of_degrees_of_freedom(parameter_model)
+    storage[:state0_map] = state0_map
+    storage[:dstate0] = state0_vec
+    storage[:n] = n_prm
 
     return storage
 end
@@ -151,7 +168,7 @@ function setup_adjoint_storage_base(model, state0, parameters;
         linear_solver = select_linear_solver(model, mode = :adjoint, rtol = 1e-8),
         n_objective = nothing,
         info_level = 0
-        )
+    )
     primary_model = adjoint_model_copy(model)
     # Standard model for: ∂Fₙᵀ / ∂xₙ
     forward_sim = Simulator(primary_model, state0 = deepcopy(state0), parameters = deepcopy(parameters), mode = :forward, extra_timing = nothing)
@@ -236,6 +253,8 @@ function solve_adjoint_sensitivities!(∇G, storage, states, state0, timesteps,
     end
     rescale_sensitivities!(∇G, storage.parameter.model, storage.parameter_map)
     @assert all(isfinite, ∇G)
+    # Finally deal with initial state gradients
+    update_state0_sensitivities!(storage)
     return ∇G
 end
 
@@ -720,6 +739,28 @@ function store_sensitivities!(out, model, variables, result, prm_map, ::Equation
     return out
 end
 
+function store_sensitivities!(out, model, variables, result, prm_map, ::Union{EquationMajorLayout, BlockMajorLayout})
+    scalar_valued_objective = result isa AbstractVector
+    @assert scalar_valued_objective "Only supported for scalar objective"
+
+    us = get_primary_variable_ordered_entities(model)
+    @assert length(us) == 1 "This function is not implemented for more than one entity type for primary variables"
+    u = only(us)
+    bz = degrees_of_freedom_per_entity(model, u)
+    ne = count_active_entities(model.domain, u)
+
+    offset = 1
+    for (k, var) in pairs(variables)
+        m = degrees_of_freedom_per_entity(model, var)
+        var::ScalarVariable
+        pos = offset:bz:(bz*(ne-1)+offset)
+        @assert length(pos) == ne "$(length(pos))"
+        out[k] = result[pos]
+        offset += 1
+    end
+    return out
+end
+
 function extract_sensitivity_subset(r, var, n, m, offset)
     if var isa ScalarVariable
         v = r
@@ -845,3 +886,37 @@ function adjoint_transfer_canonical_order_inner!(λ, dx, model, ::BlockMajorLayo
         end
     end
 end
+
+function update_state0_sensitivities!(storage)
+    state0_map = storage.state0_map
+    if !ismissing(state0_map)
+        sim = storage.backward
+        model = sim.model
+        if model isa MultiModel
+            for (k, v) in pairs(model.models)
+                @assert matrix_layout(v.context) isa EquationMajorLayout
+            end
+        else
+            @assert matrix_layout(model.context) isa EquationMajorLayout
+        end
+        # Assume that this gets called at the end when everything has been set
+        # up in terms of the simulators
+        λ = storage.lagrange
+        ∇x = storage.dstate0
+        @. ∇x = 0.0
+        # order = collect(eachindex(λ))
+        # renum = similar(order)
+        # TODO: Finish this part and remove the assertions above
+        # Get order to put values into canonical order
+        # adjoint_transfer_canonical_order!(renum, order, model)
+        # λ_renum = similar(λ)
+        # @. λ_renum[renum] = λ
+        # model_prm = storage.parameter.model
+        lsys_b = sim.storage.LinearizedSystem
+        op_b = linear_operator(lsys_b, skip_red = true)
+        # tmp = zeros(size(∇x))
+        sens_add_mult!(∇x, op_b, λ)
+        # adjoint_transfer_canonical_order!(∇x, tmp, model)
+        rescale_sensitivities!(∇x, sim.model, storage.state0_map)
+    end
+end

src/core_types/core_types.jl

@@ -884,6 +884,15 @@ function Base.getindex(case::JutulCase, ix::Int)
     return case[ix:ix]
 end
 
+function Base.lastindex(case::JutulCase)
+    return length(case.dt)
+end
+
+function Base.lastindex(case::JutulCase, d::Int)
+    d == 1 || throw(ArgumentError("JutulCase is 1D."))
+    return Base.lastindex(case)
+end
+
 function Base.getindex(case::JutulCase, ix)
     (; model, dt, forces, state0, parameters, input_data) = case
     f = deepcopy(forces)
@@ -1159,9 +1168,9 @@ function Base.show(io::IO, t::MIME"text/plain", options::MeshEntityTags{T}) wher
             kv = "<no tags>"
         else
             s = map(x -> "$x $(keys(v[x]))", collect(kv))
-            kv = join(s, ",")
+            kv = join(s, ",\n\t")
         end
-        println(io, "    $k: $(kv)")
+        println(io, "    $k:\n\t$(kv)")
     end
 end
 

src/core_types/domains.jl

@@ -99,7 +99,11 @@ function Base.show(io::IO, t::MIME"text/plain", d::DataDomain)
             for (k, v) in data
                 vals, e = v
                 if e == u
-                    sz = join(map(x -> "$x", size(vals)), "×")
+                    if vals isa AbstractVecOrMat
+                        sz = join(map(x -> "$x", size(vals)), "×")
+                    else
+                        sz = "$(typeof(vals))"
+                    end
                     print(io, "    :$k => $sz $(typeof(vals))\n")
                 end
             end

src/domains.jl

@@ -56,14 +56,30 @@ count_entities(D::Union{DataDomain, DiscretizedDomain}, entity) = D.entities[ent
 count_active_entities(D, entity; kwarg...) = count_entities(D, entity)
 count_active_entities(D::DiscretizedDomain, entity; kwarg...) = count_active_entities(D, D.global_map, entity; kwarg...)
 
+
+"""
+    number_of_cells(D::Union{DataDomain, DiscretizedDomain})
+
+Get the number of cells in a `DataDomain` or `DiscretizedDomain`.
+"""
 function number_of_cells(D::Union{DataDomain, DiscretizedDomain})
     return count_entities(D, Cells())
 end
 
+"""
+    number_of_faces(D::Union{DataDomain, DiscretizedDomain})
+
+Get the number of faces in a `DataDomain` or `DiscretizedDomain`.
+"""
 function number_of_faces(D::Union{DataDomain, DiscretizedDomain})
     return count_entities(D, Faces())
 end
 
+"""
+    number_of_half_faces(D::Union{DataDomain, DiscretizedDomain})
+
+Get the number of half-faces in a `DataDomain` or `DiscretizedDomain`.
+"""
 function number_of_half_faces(D::Union{DataDomain, DiscretizedDomain})
     return 2*number_of_faces(D)
 end

src/ext/makie_ext.jl

@@ -23,7 +23,20 @@ function plot_multimodel_interactive_impl
 
 end
 
-
+"""
+    plot_mesh(mesh)
+    plot_mesh(mesh;
+        cells = nothing,
+        faces = nothing,
+        boundaryfaces = nothing,
+        outer = false,
+        color = :lightblue,
+    )
+
+Plot a `mesh` with uniform colors. Optionally, indices `cells`, `faces` or
+`boundaryfaces` can be passed to limit the plotting to a specific selection of
+entities.
+"""
 function plot_mesh(arg...; kwarg...)
     check_plotting_availability()
     plot_mesh_impl(arg...; kwarg...)
@@ -33,6 +46,13 @@ function plot_mesh_impl
 
 end
 
+
+"""
+    plot_mesh!(ax, mesh)
+
+Mutating version of `plot_mesh` that plots into an existing Makie `Axis`
+instance.
+"""
 function plot_mesh!(arg...; kwarg...)
     check_plotting_availability()
     plot_mesh_impl!(arg...; kwarg...)
@@ -42,6 +62,11 @@ function plot_mesh_impl!
 
 end
 
+"""
+    plot_mesh_edges(mesh; kwarg...)
+
+Plot the edges of all cells on the exterior of a mesh.
+"""
 function plot_mesh_edges(arg...; kwarg...)
     check_plotting_availability()
     plot_mesh_edges_impl(arg...; kwarg...)
@@ -51,6 +76,12 @@ function plot_mesh_edges_impl
 
 end
 
+"""
+    plot_mesh_edges!(ax, mesh; kwarg...)
+
+Plot the edges of all cells on the exterior of a mesh into existing Makie
+`Axis` `ax`.
+"""
 function plot_mesh_edges!(arg...; kwarg...)
     check_plotting_availability()
     plot_mesh_edges_impl!(arg...; kwarg...)
@@ -82,6 +113,18 @@ function plotting_check_interactive
 
 end
 
+"""
+    check_plotting_availability(; throw = true, interactive = false)
+
+Check if plotting through at least one `Makie` backend is available in the Julia
+session (after package has been loaded by for example `using GLMakie`). The
+argument `throw` can be used to control if this function acts as a programmatic
+check (`throw=false`) there the return value indicates availability, or if an
+error message is to be printed telling the user how to get plotting working
+(`throw=true`)
+
+An additional check for specifically `interactive` plots can also be added.
+"""
 function check_plotting_availability(; throw = true, interactive = false)
     ok = true
     try

src/interpolation.jl

@@ -74,6 +74,11 @@ end
 
 function LinearInterpolant(X::V, F::T; static = false, constant_dx = missing) where {T<:AbstractVector, V<:AbstractVector}
     length(X) == length(F) || throw(ArgumentError("X and F values must have equal length."))
+    if length(X) == 1
+        # Handle single inputs by constant extrapolation
+        push!(X, only(X) + one(eltype(X)))
+        push!(F, only(F))
+    end
     if !issorted(X)
         ix = sortperm(X)
         X = X[ix]

src/linsolve/default.jl

@@ -90,7 +90,7 @@ function MultiLinearizedSystem(subsystems, context, layout; r = nothing, dx = no
     for i in urng
         J = subsystems[i, i].jac
         ni, mi = size(J)
-        @assert ni == mi
+        @assert ni == mi "Mismatch in block size: $ni != $mi"
         e = eltype(J)
         if e <: Real
             bz = 1

src/meshes/cart.jl

@@ -7,9 +7,11 @@ Create a Cartesian mesh with dimensions specified by the `Tuple` `dims`.
 
 # Arguments
 - `dims::Tuple`: Number of grid cells in each direction. For example, `(nx, ny)` will give a 2D grids with `nx` cells in the x-direction.
-- `Δ::Tuple=ones(length(dims))`: Equal length to `dims`. First option: A `Tuple` of scalars where each entry is the length of each cell in that direction. For
-example, specifying `(Δx, Δy) for a uniform grid with each grid cell having area of `Δx*Δy`. Second option: A `Tuple` of vectors where each entry
-contains the cell sizes in the direction.
+- `Δ::Tuple=Tuple(ones(length(dims)))`: Equal length to `dims`. First option: A
+`Tuple` of scalars where each entry is the length of each cell in that
+direction. For example, specifying `(Δx, Δy) for a uniform grid with each grid
+cell having area of `Δx*Δy`. Second option: A `Tuple` of vectors where each
+entry contains the cell sizes in the direction.
 - `origin=zeros(length(dims))`: The origin of the first corner in the grid.
 
 # Examples
@@ -22,7 +24,7 @@ CartesianMesh (3D) with 3x5x2=30 cells
 Generate a non-uniform 2D mesh:
 ```julia-repl
 julia> CartesianMesh((2, 3), ([1.0, 2.0], [0.1, 3.0, 2.5]))
-CartesianMesh (3D) with 3x5x2=30 cells
+CartesianMesh (2D) with 2x3x1=6 cells
 ```
 """
 struct CartesianMesh{D, Δ, O, T} <: FiniteVolumeMesh