add Allocation struct (#923)

Fixes #920.

core/src/bmi.jl

@@ -481,16 +481,15 @@ function discrete_control_affect!(
     return control_state_change
 end
 
-function get_allocation_model(
-    p::Parameters,
-    allocation_network_id::Int,
-)::Union{AllocationModel, Nothing}
-    for allocation_model in p.allocation_models
-        if allocation_model.allocation_network_id == allocation_network_id
-            return allocation_model
-        end
+function get_allocation_model(p::Parameters, allocation_network_id::Int)::AllocationModel
+    (; allocation) = p
+    (; allocation_network_ids, allocation_models) = allocation
+    idx = findsorted(allocation_network_ids, allocation_network_id)
+    if isnothing(idx)
+        error("Invalid allocation network id $allocation_network_id.")
+    else
+        return allocation_models[idx]
     end
-    return nothing
 end
 
 """
@@ -594,7 +593,7 @@ end
 "Solve the allocation problem for all users and assign allocated abstractions to user nodes."
 function update_allocation!(integrator)::Nothing
     (; p, t) = integrator
-    for allocation_model in integrator.p.allocation_models
+    for allocation_model in integrator.p.allocation.allocation_models
         allocate!(p, allocation_model, t)
     end
 end

core/src/create.jl

@@ -197,9 +197,11 @@ const nonconservative_nodetypes =
     Set{String}(["Basin", "LevelBoundary", "FlowBoundary", "Terminal", "User"])
 
 function generate_allocation_models!(p::Parameters, config::Config)::Nothing
-    (; graph, allocation_models) = p
+    (; graph, allocation) = p
+    (; allocation_network_ids, allocation_models) = allocation
 
-    for allocation_network_id in keys(graph[].node_ids)
+    for allocation_network_id in sort(collect(keys(graph[].node_ids)))
+        push!(allocation_network_ids, allocation_network_id)
         push!(
             allocation_models,
             AllocationModel(config, allocation_network_id, p, config.allocation.timestep),
@@ -792,7 +794,12 @@ function Parameters(db::DB, config::Config)::Parameters
     n_states = length(get_ids(db, "Basin")) + length(get_ids(db, "PidControl"))
     chunk_sizes = get_chunk_sizes(config, n_states)
     graph = create_graph(db, config, chunk_sizes)
-    allocation_models = Vector{AllocationModel}()
+    allocation = Allocation(
+        Int[],
+        AllocationModel[],
+        Vector{Tuple{NodeID, NodeID}}[],
+        Dict{Tuple{NodeID, NodeID}, Float64}(),
+    )
 
     if !valid_edges(graph)
         error("Invalid edge(s) found.")
@@ -829,7 +836,7 @@ function Parameters(db::DB, config::Config)::Parameters
     p = Parameters(
         config.starttime,
         graph,
-        allocation_models,
+        allocation,
         basin,
         linear_resistance,
         manning_resistance,

core/src/solve.jl

@@ -21,6 +21,21 @@ struct AllocationModel
     Δt_allocation::Float64
 end
 
+"""
+Object for all information about allocation
+allocation_network_ids: The unique sorted allocation network IDs
+allocation models: The allocation models for the main network and subnetworks corresponding to
+    allocation_network_ids
+main_network_connections: (from_id, to_id) from the main network to the subnetwork per subnetwork
+subnetwork_demands: The demand of an edge from the main network to a subnetwork
+"""
+struct Allocation
+    allocation_network_ids::Vector{Int}
+    allocation_models::Vector{AllocationModel}
+    main_network_connections::Vector{Vector{Tuple{NodeID, NodeID}}}
+    subnetwork_demands::Dict{Tuple{NodeID, NodeID}, Float64}
+end
+
 @enumx EdgeType flow control none
 
 """
@@ -445,7 +460,7 @@ struct Parameters{T, C1, C2}
         MetaGraphsNext.var"#11#13",
         Float64,
     }
-    allocation_models::Vector{AllocationModel}
+    allocation::Allocation
     basin::Basin{T, C1}
     linear_resistance::LinearResistance
     manning_resistance::ManningResistance

core/src/utils.jl

@@ -29,12 +29,19 @@ function create_graph(db::DB, config::Config, chunk_sizes::Vector{Int})::MetaGra
         db,
         "SELECT fid, from_node_id, to_node_id, edge_type, allocation_network_id FROM Edge ORDER BY fid",
     )
+    # Node IDs per subnetwork
     node_ids = Dict{Int, Set{NodeID}}()
+    # Allocation edges per subnetwork
     edge_ids = Dict{Int, Set{Tuple{NodeID, NodeID}}}()
+    # Source edges per subnetwork
     edges_source = Dict{Int, Set{EdgeMetadata}}()
+    # The number of flow edges
     flow_counter = 0
+    # Dictionary from flow edge to index in flow vector
     flow_dict = Dict{Tuple{NodeID, NodeID}, Int}()
+    # The number of nodes with vertical flow (interaction with outside of model)
     flow_vertical_counter = 0
+    # Dictionary from node ID to index in vertical flow vector
     flow_vertical_dict = Dict{NodeID, Int}()
     graph = MetaGraph(
         DiGraph();

core/test/allocation_test.jl

@@ -15,7 +15,7 @@
 
     graph = p.graph
     Ribasim.set_flow!(graph, NodeID(1), NodeID(2), 4.5) # Source flow
-    allocation_model = p.allocation_models[1]
+    allocation_model = p.allocation.allocation_models[1]
     Ribasim.allocate!(p, allocation_model, 0.0)
 
     F = allocation_model.problem[:F]
@@ -45,7 +45,7 @@ end
     config = Ribasim.Config(toml_path; allocation_objective_type = "quadratic_absolute")
     model = Ribasim.run(config)
     @test successful_retcode(model)
-    problem = model.integrator.p.allocation_models[1].problem
+    problem = model.integrator.p.allocation.allocation_models[1].problem
     objective = JuMP.objective_function(problem)
     @test objective isa JuMP.QuadExpr # Quadratic expression
     F = problem[:F]
@@ -61,7 +61,7 @@ end
     config = Ribasim.Config(toml_path; allocation_objective_type = "quadratic_relative")
     model = Ribasim.run(config)
     @test successful_retcode(model)
-    problem = model.integrator.p.allocation_models[1].problem
+    problem = model.integrator.p.allocation.allocation_models[1].problem
     objective = JuMP.objective_function(problem)
     @test objective isa JuMP.QuadExpr # Quadratic expression
     @test objective.aff.constant == 2.0
@@ -78,7 +78,7 @@ end
     config = Ribasim.Config(toml_path; allocation_objective_type = "linear_absolute")
     model = Ribasim.run(config)
     @test successful_retcode(model)
-    problem = model.integrator.p.allocation_models[1].problem
+    problem = model.integrator.p.allocation.allocation_models[1].problem
     objective = JuMP.objective_function(problem)
     @test objective isa JuMP.AffExpr # Affine expression
     @test :F_abs in keys(problem.obj_dict)
@@ -95,7 +95,7 @@ end
     config = Ribasim.Config(toml_path; allocation_objective_type = "linear_relative")
     model = Ribasim.run(config)
     @test successful_retcode(model)
-    problem = model.integrator.p.allocation_models[1].problem
+    problem = model.integrator.p.allocation.allocation_models[1].problem
     objective = JuMP.objective_function(problem)
     @test objective isa JuMP.AffExpr # Affine expression
     @test :F_abs in keys(problem.obj_dict)
@@ -121,7 +121,7 @@ end
         "../../generated_testmodels/fractional_flow_subnetwork/ribasim.toml",
     )
     model = Ribasim.BMI.initialize(Ribasim.Model, toml_path)
-    problem = model.integrator.p.allocation_models[1].problem
+    problem = model.integrator.p.allocation.allocation_models[1].problem
     F = problem[:F]
     @test JuMP.normalized_coefficient(
         problem[:fractional_flow][(NodeID(3), NodeID(5))],
@@ -155,7 +155,7 @@ end
     @test record_control.control_state == ["A", "B"]
 
     fractional_flow_constraints =
-        model.integrator.p.allocation_models[1].problem[:fractional_flow]
+        model.integrator.p.allocation.allocation_models[1].problem[:fractional_flow]
     @test JuMP.normalized_coefficient(
         problem[:fractional_flow][(NodeID(3), NodeID(5))],
         F[(NodeID(2), NodeID(3))],