Read some interpolation tables into the Model struct

core/src/bmi.jl

@@ -70,6 +70,7 @@
         # use state
         state = load_structvector(db, config, BasinStateV1)
         n = length(get_ids(db, "Basin"))
+
     finally
         # always close the database, also in case of an error
         close(db)
@@ -104,7 +105,7 @@
     end
     @debug "Setup ODEProblem."
 
-    callback, saved_flow = create_callbacks(parameters, config; config.solver.saveat)
+    callback, saved = create_callbacks(parameters, config; config.solver.saveat)
     @debug "Created callbacks."
 
     # Initialize the integrator, providing all solver options as described in
@@ -138,7 +139,7 @@
 
     set_initial_discrete_controlled_parameters!(integrator, storage)
 
-    return Model(integrator, config, saved_flow)
+    return Model(integrator, config, saved)
 end
 
 """
@@ -171,6 +172,11 @@
     path = results_path(config, RESULTS_FILENAME.allocation)
     write_arrow(path, table, compress)
 
+    # exported levels
+    table = subgrid_level_table(model)
+    path = results_path(config, RESULTS_FILENAME.subgrid_levels)
+    write_arrow(path, table, compress)
+
     @debug "Wrote results."
     return model
 end
@@ -209,7 +215,7 @@
     parameters::Parameters,
     config::Config;
     saveat,
-)::Tuple{CallbackSet, SavedValues{Float64, Vector{Float64}}}
+)::Tuple{CallbackSet, SavedResults}
     (; starttime, basin, tabulated_rating_curve, discrete_control) = parameters
     callbacks = SciMLBase.DECallback[]
 
@@ -235,6 +241,20 @@
     save_flow_cb = SavingCallback(save_flow, saved_flow; saveat, save_start = false)
     push!(callbacks, save_flow_cb)
 
+    # interpolate the levels
+    saved_subgrid_level = SavedValues(Float64, Vector{Float64})
+    if config.results.subgrid
+        export_cb = SavingCallback(
+            save_subgrid_level,
+            saved_subgrid_level;
+            saveat,
+            save_start = false,
+        )
+        push!(callbacks, export_cb)
+    end
+
+    saved = SavedResults(saved_flow, saved_subgrid_level)
+
     n_conditions = length(discrete_control.node_id)
     if n_conditions > 0
         discrete_control_cb = VectorContinuousCallback(
@@ -247,7 +267,7 @@
     end
     callback = CallbackSet(callbacks...)
 
-    return callback, saved_flow
+    return callback, saved
 end
 
 """
@@ -486,6 +506,20 @@
     copy(nonzeros(get_tmp(integrator.p.connectivity.flow, u)))
 end
 
+function update_subgrid_level!(integrator)::Nothing
+    basin_level = get_tmp(integrator.p.basin.current_level, 0)
+    subgrid = integrator.p.subgrid
+    for (i, (index, interp)) in enumerate(zip(subgrid.basin_index, subgrid.interpolations))
+        subgrid.level[i] = interp(basin_level[index])
+    end
+end
+
+"Interpolate the levels and save them to SavedValues"
+function save_subgrid_level(u, t, integrator)
+    update_subgrid_level!(integrator)
+    return integrator.p.subgrid.level
+end
+
 "Load updates from 'Basin / time' into the parameters"
 function update_basin(integrator)::Nothing
     (; basin) = integrator.p
@@ -567,6 +601,8 @@
         model.integrator.p.basin.infiltration
     elseif name == "drainage"
         model.integrator.p.basin.drainage
+    elseif name == "subgrid_level"
+        model.integrator.p.subgrid.level
     else
         error("Unknown variable $name")
     end

core/src/config.jl

@@ -112,6 +112,7 @@ end
     outstate::Union{String, Nothing} = nothing
     compression::Compression = "zstd"
     compression_level::Int = 6
+    subgrid::Bool = false
 end
 
 @option struct Logging <: TableOption

core/src/consts.jl

@@ -3,4 +3,5 @@ const RESULTS_FILENAME = (
     flow = "flow.arrow",
     control = "control.arrow",
     allocation = "allocation.arrow",
+    subgrid_levels = "subgrid_levels.arrow",
 )

core/src/create.jl

@@ -785,6 +785,38 @@
     )
 end
 
+function Subgrid(db::DB, config::Config, basin::Basin)::Subgrid
+    node_to_basin = Dict(node_id => index for (index, node_id) in enumerate(basin.node_id))
+    tables = load_structvector(db, config, BasinSubgridV1)
+
+    basin_ids = Int[]
+    interpolations = ScalarInterpolation[]
+    has_error = false
+    for group in IterTools.groupby(row -> row.subgrid_id, tables)
+        subgrid_id = first(getproperty.(group, :subgrid_id))
+        node_id = NodeID(first(getproperty.(group, :node_id)))
+        basin_level = getproperty.(group, :basin_level)
+        subgrid_level = getproperty.(group, :subgrid_level)
+
+        is_valid =
+            valid_subgrid(subgrid_id, node_id, node_to_basin, basin_level, subgrid_level)
+
+        if !is_valid
+            has_error = true
+            # Ensure it doesn't extrapolate before the first value.
+            pushfirst!(subgrid_level, first(subgrid_level))
+            pushfirst!(basin_level, nextfloat(-Inf))
+            new_interp = LinearInterpolation(subgrid_level, basin_level; extrapolate = true)
+            push!(basin_ids, node_to_basin[node_id])
+            push!(interpolations, new_interp)
+        end
+    end
+
+    has_error && error("Invalid Basin / subgrid table.")
+
+    return Subgrid(basin_ids, interpolations, fill(NaN, length(basin_ids)))
+end
+
 function Parameters(db::DB, config::Config)::Parameters
     n_states = length(get_ids(db, "Basin")) + length(get_ids(db, "PidControl"))
     chunk_size = pickchunksize(n_states)
@@ -805,6 +837,7 @@
     user = User(db, config)
 
     basin = Basin(db, config, chunk_size)
+    subgrid_level = Subgrid(db, config, basin)
 
     # Set is_pid_controlled to true for those pumps and outlets that are PID controlled
     for id in pid_control.node_id
@@ -835,6 +868,8 @@
         discrete_control,
         pid_control,
         user,
+        Dict{Int, Symbol}(),
+        subgrid_level,
     )
 
     # Allocation data structures

core/src/io.jl

@@ -179,8 +179,8 @@ function flow_table(
     to_node_id::Vector{Int},
     flow::FlatVector{Float64},
 }
-    (; config, saved_flow, integrator) = model
-    (; t, saveval) = saved_flow
+    (; config, saved, integrator) = model
+    (; t, saveval) = saved.flow
     (; connectivity) = integrator.p
     (; graph) = connectivity
 
@@ -252,6 +252,27 @@ function allocation_table(
     )
 end
 
+function subgrid_level_table(
+    model::Model,
+)::@NamedTuple{
+    time::Vector{DateTime},
+    subgrid_id::Vector{Int},
+    subgrid_level::Vector{Float64},
+}
+    (; config, saved, integrator) = model
+    (; t, saveval) = saved.subgrid_level
+    subgrid = integrator.p.subgrid
+
+    nelem = length(subgrid.basin_index)
+    ntsteps = length(t)
+    unique_elem_id = collect(1:nelem)
+
+    time = repeat(datetime_since.(t, config.starttime); inner = nelem)
+    subgrid_id = repeat(unique_elem_id; outer = ntsteps)
+    subgrid_level = FlatVector(saveval)
+    return (; time, subgrid_id, subgrid_level)
+end
+
 "Write a result table to disk as an Arrow file"
 function write_arrow(
     path::AbstractString,

core/src/lib.jl

@@ -1,3 +1,8 @@
+struct SavedResults
+    flow::SavedValues{Float64, Vector{Float64}}
+    subgrid_level::SavedValues{Float64, Vector{Float64}}
+end
+
 """
     Model(config_path::AbstractString)
     Model(config::Config)
@@ -11,13 +16,13 @@ A Model can be created from the path to a TOML configuration file, or a Config o
 struct Model{T}
     integrator::T
     config::Config
-    saved_flow::SavedValues{Float64, Vector{Float64}}
+    saved::SavedResults
     function Model(
         integrator::T,
         config,
-        saved_flow,
+        saved,
     ) where {T <: SciMLBase.AbstractODEIntegrator}
-        new{T}(integrator, config, saved_flow)
+        new{T}(integrator, config, saved)
     end
 end
 

core/src/solve.jl

@@ -462,6 +462,13 @@ struct User <: AbstractParameterNode
     }
 end
 
+"Subgrid linearly interpolates basin levels."
+struct Subgrid
+    basin_index::Vector{Int}
+    interpolations::Vector{ScalarInterpolation}
+    level::Vector{Float64}
+end
+
 # TODO Automatically add all nodetypes here
 struct Parameters{T, TSparse, C1, C2}
     starttime::DateTime
@@ -479,6 +486,8 @@ struct Parameters{T, TSparse, C1, C2}
     discrete_control::DiscreteControl
     pid_control::PidControl{T}
     user::User
+    lookup::Dict{Int, Symbol}
+    subgrid::Subgrid
 end
 
 """

core/src/validation.jl

@@ -8,6 +8,7 @@
 @schema "ribasim.basin.time" BasinTime
 @schema "ribasim.basin.profile" BasinProfile
 @schema "ribasim.basin.state" BasinState
+@schema "ribasim.basin.subgrid" BasinSubgrid
 @schema "ribasim.terminal.static" TerminalStatic
 @schema "ribasim.fractionalflow.static" FractionalFlowStatic
 @schema "ribasim.flowboundary.static" FlowBoundaryStatic
@@ -30,7 +31,7 @@ tablename(sv::Type{SchemaVersion{T, N}}) where {T, N} = tablename(sv())
 tablename(sv::SchemaVersion{T, N}) where {T, N} =
     join(filter(!isnothing, nodetype(sv)), delimiter)
 isnode(sv::Type{SchemaVersion{T, N}}) where {T, N} = isnode(sv())
-isnode(::SchemaVersion{T, N}) where {T, N} = length(split(string(T), ".")) == 3
+isnode(::SchemaVersion{T, N}) where {T, N} = length(split(string(T), '.'; limit = 3)) == 3
 nodetype(sv::Type{SchemaVersion{T, N}}) where {T, N} = nodetype(sv())
 
 """
@@ -43,9 +44,9 @@ function nodetype(
     # so we parse the related record Ribasim.BasinTimeV1
     # to derive BasinTime from it.
     record = Legolas.record_type(sv)
-    node = last(split(string(Symbol(record)), "."))
+    node = last(split(string(Symbol(record)), '.'; limit = 3))
 
-    elements = split(string(T), ".")
+    elements = split(string(T), '.'; limit = 3)
     if isnode(sv)
         n = elements[2]
         k = Symbol(elements[3])
@@ -202,6 +203,13 @@ end
     level::Float64
 end
 
+@version BasinSubgridV1 begin
+    subgrid_id::Int
+    node_id::Int
+    basin_level::Float64
+    subgrid_level::Float64
+end
+
 @version FractionalFlowStaticV1 begin
     node_id::Int
     fraction::Float64
@@ -362,6 +370,7 @@ sort_by_id_level(row) = (row.node_id, row.level)
 sort_by_id_state_level(row) = (row.node_id, row.control_state, row.level)
 sort_by_priority(row) = (row.node_id, row.priority)
 sort_by_priority_time(row) = (row.node_id, row.priority, row.time)
+sort_by_subgrid_level(row) = (row.subgrid_id, row.basin_level)
 
 # get the right sort by function given the Schema, with sort_by_id as the default
 sort_by_function(table::StructVector{<:Legolas.AbstractRecord}) = sort_by_id
@@ -371,6 +380,7 @@ sort_by_function(table::StructVector{TabulatedRatingCurveStaticV1}) = sort_by_id
 sort_by_function(table::StructVector{BasinProfileV1}) = sort_by_id_level
 sort_by_function(table::StructVector{UserStaticV1}) = sort_by_priority
 sort_by_function(table::StructVector{UserTimeV1}) = sort_by_priority_time
+sort_by_function(table::StructVector{BasinSubgridV1}) = sort_by_subgrid_level
 
 const TimeSchemas = Union{
     BasinTimeV1,
@@ -619,3 +629,33 @@ function valid_fractional_flow(
     end
     return !errors
 end
+
+"""
+Validate the entries for a single subgrid element.
+"""
+function valid_subgrid(
+    subgrid_id::Int,
+    node_id::NodeID,
+    node_to_basin::Dict{NodeID, Int},
+    basin_level::Vector{Float64},
+    subgrid_level::Vector{Float64},
+)::Bool
+    errors = false
+
+    if !(node_id in keys(node_to_basin))
+        errors = true
+        @error "The node_id of the Basin / subgrid_level does not refer to a basin." node_id subgrid_id
+    end
+
+    if !allunique(basin_level)
+        errors = true
+        @error "Basin / subgrid_level subgrid_id $(subgrid_id) has repeated basin levels, this cannot be interpolated."
+    end
+
+    if !allunique(subgrid_level)
+        errors = true
+        @error "Basin / subgrid_level subgrid_id $(subgrid_id) has repeated element levels, this cannot be interpolated."
+    end
+
+    return !errors
+end

core/test/run_models_test.jl

@@ -6,6 +6,10 @@
     model = Ribasim.run(toml_path)
     @test model isa Ribasim.Model
     @test successful_retcode(model)
+
+    # The exporter interpolates 1:1 for three subgrid elements, but shifted by 1.0 meter.
+    subgrid = model.integrator.p.subgrid
+    @test all(diff(subgrid.level) .≈ 1.0)
 end
 
 @testitem "bucket model" begin
@@ -191,7 +195,7 @@ end
     (; level) = level_boundary
     level = level[1]
 
-    timesteps = model.saved_flow.t
+    timesteps = model.saved.flow.t
     flow = DataFrame(Ribasim.flow_table(model))
     outlet_flow =
         filter([:from_node_id, :to_node_id] => (from, to) -> from === 2 && to === 3, flow)
@@ -301,6 +305,6 @@ end
     # https://www.hec.usace.army.mil/confluence/rasdocs/ras1dtechref/latest/theoretical-basis-for-one-dimensional-and-two-dimensional-hydrodynamic-calculations/1d-steady-flow-water-surface-profiles/friction-loss-evaluation
     @test all(isapprox.(h_expected, h_actual; atol = 0.02))
     # Test for conservation of mass, flow at the beginning == flow at the end
-    @test model.saved_flow.saveval[end][2] ≈ 5.0 atol = 0.001 skip = Sys.isapple()
-    @test model.saved_flow.saveval[end][end - 1] ≈ 5.0 atol = 0.001 skip = Sys.isapple()
+    @test model.saved.flow.saveval[end][2] ≈ 5.0 atol = 0.001 skip = Sys.isapple()
+    @test model.saved.flow.saveval[end][end - 1] ≈ 5.0 atol = 0.001 skip = Sys.isapple()
 end