Introduce flat demand vector and expose via BMI (#1081)

Fixes #893.

core/src/allocation.jl

@@ -1,7 +1,7 @@
 """Find the edges from the main network to a subnetwork."""
 function find_subnetwork_connections!(p::Parameters)::Nothing
     (; allocation, graph, user) = p
-    n_priorities = length(user.demand[1])
+    n_priorities = length(user.priorities)
     (; subnetwork_demands, subnetwork_allocateds) = allocation
     for node_id in graph[].node_ids[1]
         for outflow_id in outflow_ids(graph, node_id)
@@ -846,7 +846,7 @@ function set_objective_priority!(
 )::Nothing
     (; objective_type, problem, allocation_network_id) = allocation_model
     (; graph, user, allocation) = p
-    (; demand, node_id) = user
+    (; demand, demand_itp, demand_from_timeseries, node_id) = user
     (; main_network_connections, subnetwork_demands) = allocation
     edge_ids = graph[].edge_ids[allocation_network_id]
 
@@ -878,7 +878,14 @@ function set_objective_priority!(
         end
 
         user_idx = findsorted(node_id, node_id_user)
-        d = demand[user_idx][priority_idx](t)
+
+        if demand_from_timeseries[user_idx]
+            d = demand_itp[user_idx][priority_idx](t)
+            set_user_demand!(user, node_id_user, priority_idx, d)
+        else
+            d = get_user_demand(user, node_id_user, priority_idx)
+        end
+
         demand_max = max(demand_max, d)
         add_user_term!(ex, edge_id, objective_type, d, allocation_model)
     end
@@ -947,7 +954,7 @@ function assign_allocations!(
             push!(record.allocation_network_id, allocation_model.allocation_network_id)
             push!(record.user_node_id, Int(user_node_id))
             push!(record.priority, user.priorities[priority_idx])
-            push!(record.demand, user.demand[user_idx][priority_idx](t))
+            push!(record.demand, user.demand[user_idx])
             push!(record.allocated, allocated)
             # TODO: This is now the last abstraction before the allocation update,
             # should be the average abstraction since the last allocation solve

core/src/bmi.jl

@@ -681,6 +681,8 @@ function BMI.get_value_ptr(model::Model, name::AbstractString)
         model.integrator.p.basin.drainage
     elseif name == "subgrid_level"
         model.integrator.p.subgrid.level
+    elseif name == "demand"
+        model.integrator.p.demand
     else
         error("Unknown variable $name")
     end

core/src/create.jl

@@ -637,7 +637,7 @@ function User(db::DB, config::Config)::User
     static = load_structvector(db, config, UserStaticV1)
     time = load_structvector(db, config, UserTimeV1)
 
-    static_node_ids, time_node_ids, node_ids, node_names, valid =
+    static_node_ids, time_node_ids, node_ids, _, valid =
         static_and_time_node_ids(db, static, time, "User")
 
     if !valid
@@ -650,7 +650,7 @@ function User(db::DB, config::Config)::User
     active = BitVector()
     min_level = Float64[]
     return_factor = Float64[]
-    interpolations = Vector{ScalarInterpolation}[]
+    demand_itp = Vector{ScalarInterpolation}[]
 
     errors = false
     trivial_timespan = [nextfloat(-Inf), prevfloat(Inf)]
@@ -662,25 +662,34 @@ function User(db::DB, config::Config)::User
         group in IterTools.groupby(row -> row.priority, time)
     )
 
+    demand = Float64[]
+
+    # Whether the demand of a user node is given by a timeseries
+    demand_from_timeseries = BitVector()
+
     for node_id in node_ids
         first_row = nothing
-        demand = Vector{ScalarInterpolation}()
+        demand_itp_node_id = Vector{ScalarInterpolation}()
 
         if node_id in static_node_ids
+            push!(demand_from_timeseries, false)
             rows = searchsorted(static.node_id, node_id)
             static_id = view(static, rows)
             for p in priorities
                 idx = findsorted(static_id.priority, p)
                 demand_p = !isnothing(idx) ? static_id[idx].demand : 0.0
                 demand_p_itp = LinearInterpolation([demand_p, demand_p], trivial_timespan)
-                push!(demand, demand_p_itp)
+                push!(demand_itp_node_id, demand_p_itp)
+                push!(demand, demand_p)
             end
-            push!(interpolations, demand)
+            push!(demand_itp, demand_itp_node_id)
             first_row = first(static_id)
             is_active = coalesce(first_row.active, true)
 
         elseif node_id in time_node_ids
+            push!(demand_from_timeseries, true)
             for p in priorities
+                push!(demand, 0.0)
                 if p in keys(time_priority_dict)
                     demand_p_itp, is_valid = get_scalar_interpolation(
                         config.starttime,
@@ -691,17 +700,17 @@ function User(db::DB, config::Config)::User
                         default_value = 0.0,
                     )
                     if is_valid
-                        push!(demand, demand_p_itp)
+                        push!(demand_itp_node_id, demand_p_itp)
                     else
                         @error "The demand(t) relationship for User #$node_id of priority $p from the time table has repeated timestamps, this can not be interpolated."
                         errors = true
                     end
                 else
                     demand_p_itp = LinearInterpolation([0.0, 0.0], trivial_timespan)
-                    push!(demand, demand_p_itp)
+                    push!(demand_itp_node_id, demand_p_itp)
                 end
             end
-            push!(interpolations, demand)
+            push!(demand_itp, demand_itp_node_id)
 
             first_row_idx = searchsortedfirst(time.node_id, node_id)
             first_row = time[first_row_idx]
@@ -736,20 +745,14 @@ function User(db::DB, config::Config)::User
         abstracted = Float64[],
     )
 
-    record = (
-        time = Vector{Float64}(),
-        allocation_network_id = Vector{Int}(),
-        user_node_id = Vector{Int}(),
-        priority = Vector{Int}(),
-        demand = Vector{Float64}(),
-        allocated = Vector{Float64}(),
-        abstracted = Vector{Float64}(),
-    )
+    node_ids = NodeID.(node_ids)
 
     return User(
-        NodeID.(node_ids),
+        node_ids,
         active,
-        interpolations,
+        demand,
+        demand_itp,
+        demand_from_timeseries,
         allocated,
         return_factor,
         min_level,

core/src/solve.jl

@@ -432,7 +432,9 @@ record: Collected data of allocation optimizations for output file.
 struct User <: AbstractParameterNode
     node_id::Vector{NodeID}
     active::BitVector
-    demand::Vector{Vector{ScalarInterpolation}}
+    demand::Vector{Float64}
+    demand_itp::Vector{Vector{ScalarInterpolation}}
+    demand_from_timeseries::BitVector
     allocated::Vector{Vector{Float64}}
     return_factor::Vector{Float64}
     min_level::Vector{Float64}
@@ -451,17 +453,21 @@ struct User <: AbstractParameterNode
         node_id,
         active,
         demand,
+        demand_itp,
+        demand_from_timeseries,
         allocated,
         return_factor,
         min_level,
         priorities,
         record,
     )
-        if valid_demand(node_id, demand, priorities)
+        if valid_demand(node_id, demand_itp, priorities)
             return new(
                 node_id,
                 active,
                 demand,
+                demand_itp,
+                demand_from_timeseries,
                 allocated,
                 return_factor,
                 min_level,
@@ -805,7 +811,7 @@ function formulate_flow!(
     t::Number,
 )::Nothing
     (; graph, basin) = p
-    (; node_id, allocated, demand, active, return_factor, min_level) = user
+    (; node_id, allocated, active, demand_itp, return_factor, min_level) = user
 
     for (i, id) in enumerate(node_id)
         src_id = inflow_id(graph, id)
@@ -822,7 +828,7 @@ function formulate_flow!(
         # If allocation is not optimized then allocated = Inf, so the result is always
         # effectively allocated = demand.
         for priority_idx in eachindex(allocated[i])
-            alloc = min(allocated[i][priority_idx], demand[i][priority_idx](t))
+            alloc = min(allocated[i][priority_idx], demand_itp[i][priority_idx](t))
             q += alloc
         end
 

core/src/utils.jl

@@ -1282,3 +1282,23 @@ end
 function is_main_network(allocation_network_id::Int)::Bool
     return allocation_network_id == 1
 end
+
+function get_user_demand(user::User, node_id::NodeID, priority_idx::Int)::Float64
+    (; demand) = user
+    user_idx = findsorted(user.node_id, node_id)
+    n_priorities = length(user.priorities)
+    return demand[(user_idx - 1) * n_priorities + priority_idx]
+end
+
+function set_user_demand!(
+    user::User,
+    node_id::NodeID,
+    priority_idx::Int,
+    value::Float64,
+)::Nothing
+    (; demand) = user
+    user_idx = findsorted(user.node_id, node_id)
+    n_priorities = length(user.priorities)
+    demand[(user_idx - 1) * n_priorities + priority_idx] = value
+    return nothing
+end

core/src/validation.jl

@@ -639,14 +639,14 @@ end
 
 function valid_demand(
     node_id::Vector{NodeID},
-    demand::Vector{
+    demand_itp::Vector{
         Vector{LinearInterpolation{Vector{Float64}, Vector{Float64}, true, Float64}},
     },
     priorities::Vector{Int},
 )::Bool
     errors = false
 
-    for (col, id) in zip(demand, node_id)
+    for (col, id) in zip(demand_itp, node_id)
         for (demand_p_itp, p_itp) in zip(col, priorities)
             if any(demand_p_itp.u .< 0.0)
                 @error "Demand of user node $id with priority $p_itp should be non-negative"

core/test/allocation_test.jl

@@ -39,6 +39,12 @@
     @test allocated[1] ≈ [0.0, 0.5]
     @test allocated[2] ≈ [4.0, 0.0]
     @test allocated[3] ≈ [0.0, 0.0]
+
+    # Test getting and setting user demands
+    (; user) = p
+    Ribasim.set_user_demand!(user, NodeID(11), 2, Float64(π))
+    @test user.demand[4] ≈ π
+    @test Ribasim.get_user_demand(user, NodeID(11), 2) ≈ π
 end
 
 @testitem "Allocation objective types" begin

core/test/validation_test.jl

@@ -411,7 +411,9 @@ end
         @test_throws "Invalid demand" Ribasim.User(
             [Ribasim.NodeID(1)],
             [true],
+            [0.0],
             [[LinearInterpolation([-5.0, -5.0], [-1.8, 1.8])]],
+            [true],
             [0.0, -0.0],
             [0.9],
             [0.9],