Fix the Julia tests

core/src/allocation_init.jl

@@ -1,18 +1,3 @@
-function get_main_network_connections(
-    p::Parameters,
-    subnetwork_id::Int32,
-)::Vector{Tuple{NodeID, NodeID}}
-    (; allocation) = p
-    (; subnetwork_ids, main_network_connections) = allocation
-    idx = findsorted(subnetwork_ids, subnetwork_id)
-    if isnothing(idx)
-        error("Invalid allocation network ID $subnetwork_id.")
-    else
-        return main_network_connections[idx]
-    end
-    return
-end
-
 """
 Get the fixed capacity (∈[0,∞]) of the links in the subnetwork in a JuMP.Containers.SparseAxisArray,
 which is a type of sparse arrays that in this case takes NodeID in stead of Int as indices.
@@ -88,14 +73,14 @@ function add_subnetwork_connections!(
 
     # Add the connections to the main network
     if is_main_network(subnetwork_id)
-        for connections in main_network_connections
+        for connections in values(main_network_connections)
             for connection in connections
                 capacity[connection...] = Inf
             end
         end
     else
         # Add the connections to this subnetwork
-        for connection in get_main_network_connections(p, subnetwork_id)
+        for connection in main_network_connections[subnetwork_id]
             capacity[connection...] = Inf
         end
     end
@@ -187,7 +172,8 @@ function add_constraints_capacity!(
     p::Parameters,
     subnetwork_id::Int32,
 )::Nothing
-    main_network_source_links = get_main_network_connections(p, subnetwork_id)
+    (; main_network_connections) = p.allocation
+    main_network_source_links = main_network_connections[subnetwork_id]
     F = problem[:F]
 
     # Find the links within the subnetwork with finite capacity

core/src/allocation_optim.jl

@@ -47,7 +47,10 @@ function set_objective_demand_priority!(
     # Terms for subnetworks acting as UserDemand on the main network
     if is_main_network(subnetwork_id)
         # Loop over the connections between main and subnetwork
-        for connections_subnetwork in main_network_connections[2:end]
+        for (subnetwork_id, connections_subnetwork) in main_network_connections
+            if is_main_network(subnetwork_id)
+                continue
+            end
             for connection in connections_subnetwork
                 d = subnetwork_demands[connection][demand_priority_idx]
                 F_inlet = F[connection]
@@ -111,13 +114,8 @@ function assign_allocations!(
 )::Nothing
     (; subnetwork_id, capacity, flow) = allocation_model
     (; graph, user_demand, allocation) = p
-    (;
-        subnetwork_demands,
-        subnetwork_allocateds,
-        subnetwork_ids,
-        main_network_connections,
-    ) = allocation
-    main_network_source_links = get_main_network_connections(p, subnetwork_id)
+    (; subnetwork_demands, subnetwork_allocateds, main_network_connections) = allocation
+    main_network_source_links = main_network_connections[subnetwork_id]
     for link in keys(capacity.data)
         # If this link does not exist in the physical model then it comes from a
         # bidirectional link, and thus does not have directly allocating flow
@@ -144,11 +142,11 @@ function assign_allocations!(
     # Write the flows to the subnetworks as allocated flows
     # in the allocation object
     if is_main_network(subnetwork_id)
-        for (subnetwork_id, main_network_source_links) in
-            zip(subnetwork_ids, main_network_connections)
+        for (subnetwork_id, main_network_source_links) in main_network_connections
             if is_main_network(subnetwork_id)
                 continue
             end
+            @show main_network_source_links
             for link_id in main_network_source_links
                 subnetwork_allocateds[link_id][demand_priority_idx] = flow[link_id]
             end
@@ -172,9 +170,9 @@ function set_initial_capacities_inlet!(
     (; sources) = allocation_model
     (; allocation) = p
     (; subnetwork_id) = allocation_model
-    (; subnetwork_allocateds) = allocation
+    (; subnetwork_allocateds, main_network_connections) = allocation
 
-    main_network_source_links = get_main_network_connections(p, subnetwork_id)
+    main_network_source_links = main_network_connections[subnetwork_id]
 
     for link_id in main_network_source_links
         source_capacity = if optimization_type == OptimizationType.internal_sources
@@ -279,9 +277,10 @@ function set_initial_capacities_link!(
     allocation_model::AllocationModel,
     p::Parameters,
 )::Nothing
+    (; main_network_connections) = p.allocation
     (; problem, capacity, subnetwork_id) = allocation_model
     constraints_capacity = problem[:capacity]
-    main_network_source_links = get_main_network_connections(p, subnetwork_id)
+    main_network_source_links = main_network_connections[subnetwork_id]
 
     for (link_id, c) in capacity.data
 
@@ -1070,7 +1069,7 @@ function collect_demands!(
 )::Nothing
     (; allocation) = p
     (; subnetwork_id) = allocation_model
-    (; demand_priorities_all, subnetwork_demands) = allocation
+    (; demand_priorities_all, subnetwork_demands, main_network_connections) = allocation
 
     ## Find internal sources
     optimization_type = OptimizationType.internal_sources
@@ -1092,7 +1091,7 @@ function collect_demands!(
     ## Collect demand
     optimization_type = OptimizationType.collect_demands
 
-    main_network_source_links = get_main_network_connections(p, subnetwork_id)
+    main_network_source_links = main_network_connections[subnetwork_id]
 
     # Reset the subnetwork demands to 0.0
     for main_network_connection in keys(subnetwork_demands)

core/src/parameter.jl

@@ -218,8 +218,8 @@ record_flow: A record of all flows computed by allocation optimization, eventual
 @kwdef struct Allocation
     subnetwork_ids::Vector{Int32} = Int32[]
     allocation_models::Vector{AllocationModel} = AllocationModel[]
-    main_network_connections::Vector{Vector{Tuple{NodeID, NodeID}}} =
-        Vector{Tuple{NodeID, NodeID}}[]
+    main_network_connections::Dict{Int32, Vector{Tuple{NodeID, NodeID}}} =
+        Dict{Int, Vector{Tuple{NodeID, NodeID}}}()
     demand_priorities_all::Vector{Int32}
     subnetwork_demands::Dict{Tuple{NodeID, NodeID}, Vector{Float64}} = Dict()
     subnetwork_allocateds::Dict{Tuple{NodeID, NodeID}, Vector{Float64}} = Dict()

core/src/read.jl

@@ -252,7 +252,12 @@ function get_allocation_sources_in_order!(
     config::Config,
 )::OrderedDict{Tuple{NodeID, NodeID}, AllocationSource}
     (; graph, user_demand, allocation) = p
-    (; subnetwork_demands, subnetwork_allocateds, demand_priorities_all) = allocation
+    (;
+        subnetwork_demands,
+        subnetwork_allocateds,
+        demand_priorities_all,
+        main_network_connections,
+    ) = allocation
     n_demand_priorities = length(demand_priorities_all)
 
     default_source_priority = config.allocation.source_priority
@@ -306,9 +311,8 @@ function get_allocation_sources_in_order!(
                     end
                 elseif source_type == AllocationSourceType.flow_demand
                     # If the row is for a flow demand, make a source for the connected connector node
-                    id_with_demand = only(
-                        outneighbor_labels_type(graph, only(node_ids), LinkType.control),
-                    )
+                    id_with_demand =
+                        only(outneighbor_labels_type(graph, node_id, LinkType.control))
                     push!(links, (id_with_demand, id_with_demand))
                 else # if source_type == AllocationSourceType.user_demand
                     push!(links, user_demand.outflow_link[node_id.idx].link)
@@ -334,6 +338,7 @@ function get_allocation_sources_in_order!(
                     source_type = AllocationSourceType.subnetwork_inlet
                     link = (main_network_id, node_id)
                     push!(links, link)
+                    push!(main_network_connections[row.subnetwork_id], link)
                     # Allocate memory for the demands and demand priorities
                     # from the subnetwork via this link
                     subnetwork_demands[link] = zeros(n_demand_priorities)
@@ -386,7 +391,7 @@ function initialize_allocation!(p::Parameters, db::DB, config::Config)::Nothing
 
     for subnetwork_id in subnetwork_ids_
         push!(subnetwork_ids, subnetwork_id)
-        push!(main_network_connections, Tuple{NodeID, NodeID}[])
+        main_network_connections[subnetwork_id] = Tuple{NodeID, NodeID}[]
     end
 
     sources = get_allocation_sources_in_order!(p, db, config)

core/test/allocation_test.jl

@@ -85,9 +85,9 @@ end
 
     # Connections from main network to subnetworks
     @test isempty(main_network_connections[1])
-    @test only(main_network_connections[2]) == (NodeID(:Basin, 2, p), NodeID(:Pump, 11, p))
-    @test only(main_network_connections[3]) == (NodeID(:Basin, 6, p), NodeID(:Pump, 24, p))
-    @test only(main_network_connections[4]) == (NodeID(:Basin, 10, p), NodeID(:Pump, 38, p))
+    @test only(main_network_connections[3]) == (NodeID(:Basin, 2, p), NodeID(:Pump, 11, p))
+    @test only(main_network_connections[5]) == (NodeID(:Basin, 6, p), NodeID(:Pump, 24, p))
+    @test only(main_network_connections[7]) == (NodeID(:Basin, 10, p), NodeID(:Pump, 38, p))
 
     # main-sub connections are part of main network allocation network
     allocation_model_main_network = Ribasim.get_allocation_model(p, Int32(1))

docs/dev/allocation.qmd

@@ -95,7 +95,7 @@ The function `get_capacity` obtains the capacities of the links within a subnetw
 
 ### Handling the connection between the main network and subnetworks
 The function `find_subnetwork_connetions` finds the links that connected the main network to a subnetwork. `subnetwork_demands` and `subnetwork_allocateds` will be created, which stores demands and allocated values for subnetworks as a whole.
-`main_network_connections` is a vector of links that connect a subnetwork with the main network.
+`main_network_connections` is a dictionary of links that connect a subnetwork with the main network.
 
 ## The optimization problem {#sec-optimization-problem}
 ### Setting up the optimization variables