Let integrator integrate flows for allocation

core/src/allocation_optim.jl

@@ -252,10 +252,11 @@ as the average flow over the last Δt_allocation of the source in the physical l
 function set_initial_capacities_source!(
     allocation_model::AllocationModel,
     p::Parameters,
+    u::ComponentVector,
 )::Nothing
     (; problem) = allocation_model
     (; graph, allocation) = p
-    (; mean_flows) = allocation
+    (; flow_dict) = allocation
     (; subnetwork_id) = allocation_model
     source_constraints = problem[:source]
     main_network_source_edges = get_main_network_connections(p, subnetwork_id)
@@ -266,7 +267,7 @@ function set_initial_capacities_source!(
             # If it is a source edge for this allocation problem
             if edge ∉ main_network_source_edges
                 # Reset the source to the averaged flow over the last allocation period
-                source_capacity = mean_flows[edge][]
+                source_capacity = u.flow_allocation_input[flow_dict[edge]]
                 JuMP.set_normalized_rhs(
                     source_constraints[edge],
                     # It is assumed that the allocation procedure does not have to be differentiated.
@@ -361,11 +362,11 @@ function get_basin_data(
     (; graph, basin, level_demand, allocation) = p
     (; vertical_flux) = basin
     (; Δt_allocation) = allocation_model
-    (; mean_flows) = allocation
+    (; flow_dict) = allocation
     @assert node_id.type == NodeType.Basin
     vertical_flux = get_tmp(vertical_flux, 0)
     _, basin_idx = id_index(basin.node_id, node_id)
-    influx = mean_flows[(node_id, node_id)][]
+    influx = u.flow_allocation_input[flow_dict[(node_id, node_id)]]
     _, basin_idx = id_index(basin.node_id, node_id)
     storage_basin = u.storage[basin_idx]
     control_inneighbors = inneighbor_labels_type(graph, node_id, EdgeType.control)
@@ -984,7 +985,7 @@ function set_initial_values!(
     u::ComponentVector,
     t::Float64,
 )::Nothing
-    set_initial_capacities_source!(allocation_model, p)
+    set_initial_capacities_source!(allocation_model, p, u)
     set_initial_capacities_edge!(allocation_model, p)
     set_initial_capacities_basin!(allocation_model, p, u, t)
     set_initial_capacities_buffer!(allocation_model)

core/src/callback.jl

@@ -405,7 +405,7 @@ end
 function update_allocation!(integrator)::Nothing
     (; p, t, u) = integrator
     (; allocation) = p
-    (; allocation_models, mean_flows) = allocation
+    (; allocation_models) = allocation
 
     # Don't run the allocation algorithm if allocation is not active
     # (Specifically for running Ribasim via the BMI)
@@ -417,9 +417,7 @@ function update_allocation!(integrator)::Nothing
 
     # Divide by the allocation Δt to obtain the mean flows
     # from the integrated flows
-    for value in values(mean_flows)
-        value[] /= Δt_allocation
-    end
+    u.flow_allocation_input /= Δt_allocation
 
     # If a main network is present, collect demands of subnetworks
     if has_main_network(allocation)
@@ -437,9 +435,9 @@ function update_allocation!(integrator)::Nothing
     end
 
     # Reset the mean source flows
-    for value in values(mean_flows)
-        value[] = 0.0
-    end
+    u.flow_allocation_input .= 0.0
+
+    return nothing
 end
 
 "Load updates from 'TabulatedRatingCurve / time' into the parameters"

core/src/model.jl

@@ -36,7 +36,7 @@ function Model(config_path::AbstractString)::Model
 end
 
 function get_u0(p::Parameters, state::StructVector)::ComponentVector
-    (; basin, pid_control, graph) = p
+    (; basin, pid_control, graph, allocation) = p
 
     storage = get_storages_from_levels(basin, state.level)
 
@@ -63,6 +63,10 @@ function get_u0(p::Parameters, state::StructVector)::ComponentVector
     drainage_bmi = zeros(n_basins)
     infiltration_bmi = zeros(n_basins)
 
+    # Flows for allocation
+    flow_allocation_input = zeros(length(allocation.flow_dict))
+
+    # NOTE: This is the source of truth for the state component names
     return ComponentVector{Float64}(;
         storage,
         integral,
@@ -75,6 +79,7 @@ function get_u0(p::Parameters, state::StructVector)::ComponentVector
         evaporation_bmi,
         drainage_bmi,
         infiltration_bmi,
+        flow_allocation_input,
     )
 end
 

core/src/parameter.jl

@@ -64,7 +64,7 @@ main_network_connections: (from_id, to_id) from the main network to the subnetwo
 priorities: All used priority values.
 subnetwork_demands: The demand of an edge from the main network to a subnetwork
 subnetwork_allocateds: The allocated flow of an edge from the main network to a subnetwork
-mean_flows: Flows averaged over Δt_allocation over edges that are allocation sources
+flow_dict: ...
 record_demand: A record of demands and allocated flows for nodes that have these
 record_flow: A record of all flows computed by allocation optimization, eventually saved to
     output file
@@ -76,7 +76,7 @@ struct Allocation
     priorities::Vector{Int32}
     subnetwork_demands::Dict{Tuple{NodeID, NodeID}, Vector{Float64}}
     subnetwork_allocateds::Dict{Tuple{NodeID, NodeID}, Vector{Float64}}
-    mean_flows::Dict{Tuple{NodeID, NodeID}, Base.RefValue{Float64}}
+    flow_dict::Dict{Tuple{NodeID, NodeID}, Int}
     record_demand::@NamedTuple{
         time::Vector{Float64},
         subnetwork_id::Vector{Int32},

core/src/read.jl

@@ -983,20 +983,24 @@ function Allocation(db::DB, config::Config, graph::MetaGraph)::Allocation
         optimization_type = String[],
     )
 
-    mean_flows = Dict{Tuple{NodeID, NodeID}, Base.RefValue{Float64}}()
+    flow_dict = Dict{Tuple{NodeID, NodeID}, Int}()
+    flow_counter = 0
 
     # Find edges which serve as sources in allocation
     for edge_metadata in values(graph.edge_data)
         (; subnetwork_id_source, edge) = edge_metadata
         if subnetwork_id_source != 0
-            mean_flows[edge] = Ref(0.0)
+            flow_counter += 1
+            flow_dict[edge] = flow_counter
         end
     end
 
     # Find basins with a level demand
     for node_id in values(graph.vertex_labels)
         if has_external_demand(graph, node_id, :level_demand)[1]
-            mean_flows[(node_id, node_id)] = Ref(0.0)
+            edge = (node_id, node_id)
+            flow_counter += 1
+            flow_dict[edge] = flow_counter
         end
     end
 
@@ -1007,7 +1011,7 @@ function Allocation(db::DB, config::Config, graph::MetaGraph)::Allocation
         get_all_priorities(db, config),
         Dict{Tuple{NodeID, NodeID}, Vector{Float64}}(),
         Dict{Tuple{NodeID, NodeID}, Vector{Float64}}(),
-        mean_flows,
+        flow_dict,
         record_demand,
         record_flow,
     )

core/src/solve.jl

@@ -25,7 +25,8 @@ function water_balance!(
     formulate_flows!(p, storage, t)
 
     # Now formulate du
-    formulate_du!(du, graph, basin, storage)
+    formulate_du_storage!(du, graph, basin, storage)
+    formulate_du_integration_flows!(du, graph, p, storage)
 
     # PID control (changes the du of PID controlled basins)
     continuous_control!(u, du, pid_control, p, integral, t)
@@ -596,12 +597,14 @@ function formulate_flow!(
     return nothing
 end
 
-function formulate_integration_flows!(
+function formulate_du_integration_flows!(
     du::ComponentVector,
     graph::MetaGraph,
-    basin::Basin,
+    p::Parameters,
     storage::AbstractVector,
 )::Nothing
+    (; basin, allocation) = p
+
     # Flows over edges
     flow = get_tmp(graph[].flow, storage)
     du.flow_integrated .= flow
@@ -611,10 +614,19 @@ function formulate_integration_flows!(
     forcings_integrated(du) .= vertical_flux
     forcings_bmi(du) .= vertical_flux
 
+    # Allocation_flows
+    for (edge, idx) in allocation.flow_dict
+        du.flow_allocation_input[idx] = if edge[1] == edge[2]
+            get_influx(basin, edge[1])
+        else
+            get_flow(graph, edge..., 0)
+        end
+    end
+
     return nothing
 end
 
-function formulate_du!(
+function formulate_du_storage!(
     du::ComponentVector,
     graph::MetaGraph,
     basin::Basin,
@@ -631,7 +643,6 @@ function formulate_du!(
             du.storage[i] -= get_flow(graph, basin_id, outflow_id, storage)
         end
     end
-    formulate_integration_flows!(du, graph, basin, storage)
     return nothing
 end
 

core/src/util.jl

@@ -749,15 +749,41 @@ function get_n_flows(db::DB)::Int
     return only(only(result))
 end
 
+function get_n_allocation_flow_inputs(db::DB)::Int
+    n_sources = only(
+        only(
+            execute(
+                columntable,
+                db,
+                "SELECT COUNT(*) From Edge where 'subnetwork_id' != 0",
+            ),
+        ),
+    )
+    n_level_demands = only(
+        only(
+            execute(
+                columntable,
+                db,
+                "SELECT COUNT(*) FROM Edge WHERE from_node_type = 'level_demand'",
+            ),
+        ),
+    )
+    return n_sources + n_level_demands
+end
+
 function get_n_states(db::DB)::Int
-    return 9 * get_n_node(db, "Basin") + get_n_node(db, "PidControl") + get_n_flows(db)
+    return 9 * get_n_node(db, "Basin") +
+           get_n_node(db, "PidControl") +
+           get_n_flows(db) +
+           get_n_allocation_flow_inputs(db)
 end
 
 function get_n_states(p::Parameters)::Int
-    (; basin, pid_control, graph) = p
+    (; basin, pid_control, graph, allocation) = p
     return 9 * length(basin.node_id) +
            length(pid_control.node_id) +
-           length(graph[].flow_dict)
+           length(graph[].flow_dict) +
+           length(allocation.flow_dict)
 end
 
 function forcings_integrated(u::ComponentVector)

core/test/allocation_test.jl

@@ -6,17 +6,12 @@
 
     toml_path = normpath(@__DIR__, "../../generated_testmodels/subnetwork/ribasim.toml")
     @test ispath(toml_path)
-    cfg = Ribasim.Config(toml_path)
-    db_path = Ribasim.input_path(cfg, cfg.database)
-    db = SQLite.DB(db_path)
-
-    p = Ribasim.Parameters(db, cfg)
-    (; graph, allocation) = p
-    close(db)
+    model = Ribasim.Model(toml_path)
+    (; u, p) = model.integrator
+    (; flow_dict) = p.allocation
 
-    allocation.mean_flows[(NodeID(:FlowBoundary, 1), NodeID(:Basin, 2))][] = 4.5
+    u.flow_allocation_input[flow_dict[(NodeID(:FlowBoundary, 1), NodeID(:Basin, 2))]] = 4.5
     allocation_model = p.allocation.allocation_models[1]
-    u = ComponentVector(; storage = zeros(length(p.basin.node_id)))
     Ribasim.allocate!(p, allocation_model, 0.0, u, OptimizationType.allocate)
 
     # Last priority (= 2) flows
@@ -203,24 +198,18 @@ end
         "../../generated_testmodels/main_network_with_subnetworks/ribasim.toml",
     )
     @test ispath(toml_path)
-    cfg = Ribasim.Config(toml_path)
-    db_path = Ribasim.input_path(cfg, cfg.database)
-    db = SQLite.DB(db_path)
-    p = Ribasim.Parameters(db, cfg)
-    close(db)
-
+    model = Ribasim.Model(toml_path)
+    (; u, p, t) = model.integrator
     (; allocation, user_demand, graph, basin) = p
     (;
         allocation_models,
         subnetwork_demands,
         subnetwork_allocateds,
         record_flow,
-        mean_flows,
+        flow_dict,
     ) = allocation
-    t = 0.0
 
     # Collecting demands
-    u = ComponentVector(; storage = zeros(length(basin.node_id)))
     for allocation_model in allocation_models[2:end]
         Ribasim.allocate!(p, allocation_model, t, u, OptimizationType.internal_sources)
         Ribasim.allocate!(p, allocation_model, t, u, OptimizationType.collect_demands)
@@ -247,8 +236,8 @@ end
 
     # Running full allocation algorithm
     (; Δt_allocation) = allocation_models[1]
-    mean_flows[(NodeID(:FlowBoundary, 1), NodeID(:Basin, 2))][] = 4.5 * Δt_allocation
-    u = ComponentVector(; storage = zeros(length(p.basin.node_id)))
+    u.flow_allocation_input[flow_dict[(NodeID(:FlowBoundary, 1), NodeID(:Basin, 2))]] =
+        4.5 * Δt_allocation
     Ribasim.update_allocation!((; p, t, u))
 
     @test subnetwork_allocateds[NodeID(:Basin, 2), NodeID(:Pump, 11)] ≈
@@ -282,23 +271,18 @@ end
         "../../generated_testmodels/subnetworks_with_sources/ribasim.toml",
     )
     @test ispath(toml_path)
-    cfg = Ribasim.Config(toml_path)
-    db_path = Ribasim.input_path(cfg, cfg.database)
-    db = SQLite.DB(db_path)
-    p = Ribasim.Parameters(db, cfg)
-    close(db)
-
+    model = Ribasim.Model(toml_path)
+    (; p, u, t) = model.integrator
     (; allocation, user_demand, graph, basin) = p
-    (; allocation_models, subnetwork_demands, subnetwork_allocateds, mean_flows) =
-        allocation
-    t = 0.0
+    (; allocation_models, subnetwork_demands, subnetwork_allocateds, flow_dict) = allocation
 
     # Set flows of sources in
-    mean_flows[(NodeID(:FlowBoundary, 58), NodeID(:Basin, 16))][] = 1.0
-    mean_flows[(NodeID(:FlowBoundary, 59), NodeID(:Basin, 44))][] = 1e-3
+    u.flow_allocation_input[flow_dict[(NodeID(:FlowBoundary, 58), NodeID(:Basin, 16))]] =
+        1.0
+    u.flow_allocation_input[flow_dict[(NodeID(:FlowBoundary, 59), NodeID(:Basin, 44))]] =
+        1e-3
 
     # Collecting demands
-    u = ComponentVector(; storage = zeros(length(basin.node_id)))
     for allocation_model in allocation_models[2:end]
         Ribasim.allocate!(p, allocation_model, t, u, OptimizationType.internal_sources)
         Ribasim.allocate!(p, allocation_model, t, u, OptimizationType.collect_demands)
@@ -396,7 +380,7 @@ end
     toml_path = normpath(@__DIR__, "../../generated_testmodels/flow_demand/ribasim.toml")
     @test ispath(toml_path)
     model = Ribasim.Model(toml_path)
-    (; p) = model.integrator
+    (; p, u, t) = model.integrator
     (; graph, allocation, flow_demand, user_demand, level_boundary) = p
 
     # Test has_external_demand
@@ -436,11 +420,9 @@ end
           F[(node_id_with_flow_demand, NodeID(:Basin, 3))] + 0.0
     @test constraint_flow_demand_outflow.set.upper == 0.0
 
-    t = 0.0
-    (; u) = model.integrator
     optimization_type = OptimizationType.internal_sources
-    for (edge, value) in allocation.mean_flows
-        value[] = Ribasim.get_flow(graph, edge..., 0)
+    for (edge, idx) in allocation.flow_dict
+        u.flow_allocation_input[idx] = Ribasim.get_flow(graph, edge..., 0)
     end
     Ribasim.set_initial_values!(allocation_model, p, u, t)
 

core/test/run_models_test.jl

@@ -235,7 +235,7 @@ end
     @test model.integrator.sol.u[end].storage ≈
           Float32[519.8817, 519.8798, 339.3959, 1418.4331] skip = Sys.isapple() atol = 1.5
 
-    @test length(logger.logs) == 14
+    @test length(logger.logs) == 11
     @test logger.logs[1].level == Debug
     @test logger.logs[1].message == "Read database into memory."