Complete allocation algorithm with main network and subnetworks (#964)

Fixes #922.

Project.toml

@@ -3,6 +3,7 @@ authors = ["Deltares and contributors <[email protected]>"]
 description = "Meta-project used to share the Manifest of Ribasim and its dependents"
 
 [deps]
+ArchGDAL = "c9ce4bd3-c3d5-55b8-8973-c0e20141b8c3"
 Arrow = "69666777-d1a9-59fb-9406-91d4454c9d45"
 BasicModelInterface = "59605e27-edc0-445a-b93d-c09a3a50b330"
 Dictionaries = "85a47980-9c8c-11e8-2b9f-f7ca1fa99fb4"

core/src/bmi.jl

@@ -481,12 +481,27 @@ function get_allocation_model(p::Parameters, allocation_network_id::Int)::Alloca
     (; 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.")
+        error("Invalid allocation network ID $allocation_network_id.")
     else
         return allocation_models[idx]
     end
 end
 
+function get_main_network_connections(
+    p::Parameters,
+    allocation_network_id::Int,
+)::Vector{Tuple{NodeID, NodeID}}
+    (; allocation) = p
+    (; allocation_network_ids, main_network_connections) = allocation
+    idx = findsorted(allocation_network_ids, allocation_network_id)
+    if isnothing(idx)
+        error("Invalid allocation network ID $allocation_network_id.")
+    else
+        return main_network_connections[idx]
+    end
+    return
+end
+
 """
 Update the fractional flow fractions in an allocation problem.
 """
@@ -588,7 +603,20 @@ 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.allocation_models
+    (; allocation) = p
+    (; allocation_models) = allocation
+
+    # If a main network is present, collect demands of subnetworks
+    if has_main_network(allocation)
+        for allocation_model in Iterators.drop(allocation_models, 1)
+            allocate!(p, allocation_model, t; collect_demands = true)
+        end
+    end
+
+    # Solve the allocation problems
+    # If a main network is present this is solved first,
+    # which provides allocation to the subnetworks
+    for allocation_model in allocation_models
         allocate!(p, allocation_model, t)
     end
 end

core/src/create.jl

@@ -639,7 +639,7 @@ function User(db::DB, config::Config)::User
         error("Problems encountered when parsing User static and time node IDs.")
     end
 
-    # The highest priority number given, which corresponds to the least important demands
+    # All provided priorities
     priorities = sort(unique(union(static.priority, time.priority)))
 
     active = BitVector()
@@ -808,6 +808,7 @@ function Parameters(db::DB, config::Config)::Parameters
         AllocationModel[],
         Vector{Tuple{NodeID, NodeID}}[],
         Dict{Tuple{NodeID, NodeID}, Float64}(),
+        Dict{Tuple{NodeID, NodeID}, Float64}(),
     )
 
     if !valid_edges(graph)

core/src/solve.jl

@@ -34,6 +34,7 @@ struct Allocation
     allocation_models::Vector{AllocationModel}
     main_network_connections::Vector{Vector{Tuple{NodeID, NodeID}}}
     subnetwork_demands::Dict{Tuple{NodeID, NodeID}, Vector{Float64}}
+    subnetwork_allocateds::Dict{Tuple{NodeID, NodeID}, Vector{Float64}}
 end
 
 @enumx EdgeType flow control none

core/src/utils.jl

@@ -1263,3 +1263,11 @@ function allocation_path_exists_in_graph(
     end
     return false
 end
+
+function has_main_network(allocation::Allocation)::Bool
+    return first(allocation.allocation_network_ids) == 1
+end
+
+function is_main_network(allocation_network_id::Int)::Bool
+    return allocation_network_id == 1
+end

core/test/allocation_test.jl

@@ -1,5 +1,4 @@
 @testitem "Allocation solve" begin
-    using PreallocationTools: get_tmp
     using Ribasim: NodeID
     import SQLite
     import JuMP
@@ -181,7 +180,9 @@ end
     p = Ribasim.Parameters(db, cfg)
     close(db)
     (; allocation, graph) = p
-    (; main_network_connections) = allocation
+    (; main_network_connections, allocation_network_ids) = allocation
+    @test Ribasim.has_main_network(allocation)
+    @test Ribasim.is_main_network(first(allocation_network_ids))
 
     # Connections from main network to subnetworks
     @test isempty(main_network_connections[1])
@@ -210,3 +211,58 @@ end
     @test Ribasim.get_allocation_model(p, 7).problem[:source].axes[1] ==
           Tuple{NodeID, NodeID}[(10, 38)]
 end
+
+@testitem "allocation with main network optimization problem" begin
+    using SQLite
+    using Ribasim: NodeID
+    using JuMP
+
+    toml_path = normpath(
+        @__DIR__,
+        "../../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)
+
+    (; allocation, user, graph) = p
+    (; allocation_models, subnetwork_demands, subnetwork_allocateds) = allocation
+    t = 0.0
+
+    # Collecting demands
+    for allocation_model in allocation_models[2:end]
+        Ribasim.allocate!(p, allocation_model, t; collect_demands = true)
+    end
+
+    @test subnetwork_demands[(NodeID(2), NodeID(11))] ≈ [4.0, 4.0, 0.0]
+    @test subnetwork_demands[(NodeID(6), NodeID(24))] ≈ [0.001333333333, 0.0, 0.0]
+    @test subnetwork_demands[(NodeID(10), NodeID(38))] ≈ [0.001, 0.002, 0.002]
+
+    # Solving for the main network,
+    # containing subnetworks as users
+    allocation_model = allocation_models[1]
+    (; problem) = allocation_model
+    Ribasim.allocate!(p, allocation_model, t)
+
+    # Main network objective function
+    objective = JuMP.objective_function(problem)
+    objective_variables = keys(objective.terms)
+    F_abs = problem[:F_abs]
+    @test F_abs[NodeID(11)] ∈ objective_variables
+    @test F_abs[NodeID(24)] ∈ objective_variables
+    @test F_abs[NodeID(38)] ∈ objective_variables
+
+    # Running full allocation algorithm
+    Ribasim.set_flow!(graph, NodeID(1), NodeID(2), 4.5)
+    Ribasim.update_allocation!((; p, t))
+
+    @test subnetwork_allocateds[NodeID(2), NodeID(11)] ≈ [4.0, 0.49766666, 0.0]
+    @test subnetwork_allocateds[NodeID(6), NodeID(24)] ≈ [0.00133333333, 0.0, 0.0]
+    @test subnetwork_allocateds[NodeID(10), NodeID(38)] ≈ [0.001, 0.0, 0.0]
+
+    @test user.allocated[2] ≈ [4.0, 0.0, 0.0]
+    @test user.allocated[7] ≈ [0.001, 0.0, 0.0]
+end

python/ribasim/tests/test_model.py

@@ -157,6 +157,12 @@ def test_write_adds_fid_in_tables(basic, tmp_path):
         fids = df["fid"]
         assert fids.equals(pd.Series(range(1, len(fids) + 1)))
 
+        query = "select fid from Node"
+        df = pd.read_sql_query(query, connection)
+        assert "fid" in df.columns
+        fids = df["fid"]
+        assert fids.equals(pd.Series(range(1, len(fids) + 1)))
+
         query = "select fid from Edge"
         df = pd.read_sql_query(query, connection)
         assert "fid" in df.columns