From d18be21b26ece9f288cdd9a88832671643beb886 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart123-4@hotmail.com>
Date: Tue, 26 Nov 2024 15:36:19 +0100
Subject: [PATCH 01/11] Add deprecation warnings

---
 core/src/graph.jl               | 13 +++++++++++++
 python/ribasim/ribasim/model.py |  9 +++++++++
 2 files changed, 22 insertions(+)

diff --git a/core/src/graph.jl b/core/src/graph.jl
index d7e32d6a8..5edc7ae05 100644
--- a/core/src/graph.jl
+++ b/core/src/graph.jl
@@ -74,6 +74,13 @@ function create_graph(db::DB, config::Config)::MetaGraph
         end
         id_src = NodeID(from_node_type, from_node_id, db)
         id_dst = NodeID(to_node_type, to_node_id, db)
+        if !ismissing(subnetwork_id)
+            Base.depwarn(
+                "Sources for allocation are automatically inferred and no longer have to be specified in the `Edge` table.",
+                :create_graph;
+                force = true,
+            )
+        end
         if ismissing(subnetwork_id)
             subnetwork_id = 0
         end
@@ -101,6 +108,12 @@ function create_graph(db::DB, config::Config)::MetaGraph
     if errors
         error("Invalid edges found")
     end
+    # if edge_depwarn
+    #     Base.depwarn(
+    #         "Sources for allocation are automatically inferred and no longer have to be specified in the `Edge` table.",
+    #         :create_graph,
+    #     )
+    # end
 
     if incomplete_subnetwork(graph, node_ids)
         error("Incomplete connectivity in subnetwork")
diff --git a/python/ribasim/ribasim/model.py b/python/ribasim/ribasim/model.py
index d7274a67c..7457b4b45 100644
--- a/python/ribasim/ribasim/model.py
+++ b/python/ribasim/ribasim/model.py
@@ -1,6 +1,7 @@
 import datetime
 import logging
 import shutil
+import warnings
 from collections.abc import Generator
 from os import PathLike
 from pathlib import Path
@@ -70,6 +71,8 @@
 except ImportError:
     xugrid = MissingOptionalModule("xugrid")
 
+warnings.simplefilter("always", DeprecationWarning)
+
 
 class Model(FileModel):
     """A model of inland water resources systems."""
@@ -304,6 +307,12 @@ def write(self, filepath: str | PathLike[str]) -> Path:
         if self.use_validation:
             self._validate_model()
 
+        if self.edge.df.subnetwork_id.notna().any():
+            warnings.warn(
+                "Sources for allocation are automatically inferred and no longer have to be specified in the `Edge` table.",
+                DeprecationWarning,
+            )
+
         filepath = Path(filepath)
         self.filepath = filepath
         if not filepath.suffix == ".toml":

From 891aa8c7fa42b6f96ed252862352186e4055cbe7 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart123-4@hotmail.com>
Date: Tue, 26 Nov 2024 16:46:01 +0100
Subject: [PATCH 02/11] Towards automatically inferring source edges

---
 core/src/allocation_init.jl  | 13 ++++------
 core/src/allocation_optim.jl | 19 ++++++---------
 core/src/graph.jl            | 28 +++------------------
 core/src/parameter.jl        | 10 +++-----
 core/src/read.jl             | 13 +++++++---
 core/src/util.jl             |  8 ++----
 core/src/validation.jl       | 47 ------------------------------------
 7 files changed, 31 insertions(+), 107 deletions(-)

diff --git a/core/src/allocation_init.jl b/core/src/allocation_init.jl
index 90e033229..144976a0c 100644
--- a/core/src/allocation_init.jl
+++ b/core/src/allocation_init.jl
@@ -96,7 +96,7 @@ function get_subnetwork_capacity(
 end
 
 const allocation_source_nodetypes =
-    Set{NodeType.T}([NodeType.LevelBoundary, NodeType.FlowBoundary])
+    Set{NodeType.T}([NodeType.LevelBoundary, NodeType.FlowBoundary, NodeType.UserDemand])
 
 """
 Add the edges connecting the main network work to a subnetwork to both the main network
@@ -138,10 +138,6 @@ function get_capacity(
     capacity = get_subnetwork_capacity(p, subnetwork_id)
     add_subnetwork_connections!(capacity, p, subnetwork_id)
 
-    if !valid_sources(p, capacity, subnetwork_id)
-        error("Errors in sources in allocation network.")
-    end
-
     return capacity
 end
 
@@ -290,9 +286,10 @@ function add_constraints_source!(
 
     # Find the edges in the whole model which are a source for
     # this subnetwork
-    for edge_metadata in values(graph.edge_data)
-        (; edge) = edge_metadata
-        if graph[edge...].subnetwork_id_source == subnetwork_id
+    for edge in keys(p.allocation.mean_input_flows[subnetwork_id])
+        from_id, to_id = edge
+        if graph[from_id].subnetwork_id == subnetwork_id &&
+           graph[to_id].subnetwork_id == subnetwork_id
             push!(edges_source, edge)
         end
     end
diff --git a/core/src/allocation_optim.jl b/core/src/allocation_optim.jl
index 5c7bb3a81..5d8e40530 100644
--- a/core/src/allocation_optim.jl
+++ b/core/src/allocation_optim.jl
@@ -203,21 +203,18 @@ function set_initial_capacities_source!(
     p::Parameters,
 )::Nothing
     (; sources) = allocation_model
-    (; graph, allocation) = p
+    (; allocation) = p
     (; mean_input_flows) = allocation
     (; subnetwork_id) = allocation_model
     main_network_source_edges = get_main_network_connections(p, subnetwork_id)
 
-    for edge_metadata in values(graph.edge_data)
-        (; edge) = edge_metadata
-        if graph[edge...].subnetwork_id_source == subnetwork_id
-            # 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 = sources[edge]
-                @assert source.type == AllocationSourceType.edge
-                source.capacity[] = mean_input_flows[edge][]
-            end
+    for edge in keys(p.allocation.mean_input_flows[subnetwork_id])
+        # 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 = sources[edge]
+            @assert source.type == AllocationSourceType.edge
+            source.capacity[] = mean_input_flows[edge][]
         end
     end
     return nothing
diff --git a/core/src/graph.jl b/core/src/graph.jl
index 5edc7ae05..ef2f30b54 100644
--- a/core/src/graph.jl
+++ b/core/src/graph.jl
@@ -28,8 +28,7 @@ function create_graph(db::DB, config::Config)::MetaGraph
     )
     # Node IDs per subnetwork
     node_ids = Dict{Int32, Set{NodeID}}()
-    # Source edges per subnetwork
-    edges_source = Dict{Int32, Set{EdgeMetadata}}()
+
     # The metadata of the flow edges in the order in which they are in the input
     # and will be in the output
     flow_edges = EdgeMetadata[]
@@ -81,15 +80,8 @@ function create_graph(db::DB, config::Config)::MetaGraph
                 force = true,
             )
         end
-        if ismissing(subnetwork_id)
-            subnetwork_id = 0
-        end
-        edge_metadata = EdgeMetadata(;
-            id = edge_id,
-            type = edge_type,
-            subnetwork_id_source = subnetwork_id,
-            edge = (id_src, id_dst),
-        )
+        edge_metadata =
+            EdgeMetadata(; id = edge_id, type = edge_type, edge = (id_src, id_dst))
         if edge_type == EdgeType.flow
             push!(flow_edges, edge_metadata)
         end
@@ -98,28 +90,16 @@ function create_graph(db::DB, config::Config)::MetaGraph
             @error "Duplicate edge" id_src id_dst
         end
         graph[id_src, id_dst] = edge_metadata
-        if subnetwork_id != 0
-            if !haskey(edges_source, subnetwork_id)
-                edges_source[subnetwork_id] = Set{EdgeMetadata}()
-            end
-            push!(edges_source[subnetwork_id], edge_metadata)
-        end
     end
     if errors
         error("Invalid edges found")
     end
-    # if edge_depwarn
-    #     Base.depwarn(
-    #         "Sources for allocation are automatically inferred and no longer have to be specified in the `Edge` table.",
-    #         :create_graph,
-    #     )
-    # end
 
     if incomplete_subnetwork(graph, node_ids)
         error("Incomplete connectivity in subnetwork")
     end
 
-    graph_data = (; node_ids, edges_source, flow_edges, config.solver.saveat)
+    graph_data = (; node_ids, flow_edges, config.solver.saveat)
     @reset graph.graph_data = graph_data
 
     return graph
diff --git a/core/src/parameter.jl b/core/src/parameter.jl
index d4f32cde8..9f360bde5 100644
--- a/core/src/parameter.jl
+++ b/core/src/parameter.jl
@@ -128,7 +128,7 @@ The caches are always initialized with zeros
 """
 cache(len::Int)::Cache = LazyBufferCache(Returns(len); initializer! = set_zero!)
 
-@enumx AllocationSourceType edge basin main_to_sub user_return buffer
+@enumx AllocationSourceType boundary_node basin main_to_sub user_return buffer
 
 """
 Data structure for a single source within an allocation subnetwork.
@@ -177,7 +177,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_input_flows: Flows averaged over Δt_allocation over edges that are allocation sources
+mean_input_flows: Per subnetwork, flows averaged over Δt_allocation over edges that are allocation sources
 mean_realized_flows: Flows averaged over Δt_allocation over edges that realize a demand
 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
@@ -190,7 +190,7 @@ record_flow: A record of all flows computed by allocation optimization, eventual
     priorities::Vector{Int32}
     subnetwork_demands::Dict{Tuple{NodeID, NodeID}, Vector{Float64}} = Dict()
     subnetwork_allocateds::Dict{Tuple{NodeID, NodeID}, Vector{Float64}} = Dict()
-    mean_input_flows::Dict{Tuple{NodeID, NodeID}, Float64}
+    mean_input_flows::Dict{Int32, Dict{Tuple{NodeID, NodeID}, Float64}}
     mean_realized_flows::Dict{Tuple{NodeID, NodeID}, Float64}
     record_demand::@NamedTuple{
         time::Vector{Float64},
@@ -252,14 +252,11 @@ end
 Type for storing metadata of edges in the graph:
 id: ID of the edge (only used for labeling flow output)
 type: type of the edge
-subnetwork_id_source: ID of subnetwork where this edge is a source
-  (0 if not a source)
 edge: (from node ID, to node ID)
 """
 @kwdef struct EdgeMetadata
     id::Int32
     type::EdgeType.T
-    subnetwork_id_source::Int32
     edge::Tuple{NodeID, NodeID}
 end
 
@@ -899,7 +896,6 @@ const ModelGraph = MetaGraph{
     EdgeMetadata,
     @NamedTuple{
         node_ids::Dict{Int32, Set{NodeID}},
-        edges_source::Dict{Int32, Set{EdgeMetadata}},
         flow_edges::Vector{EdgeMetadata},
         saveat::Float64,
     },
diff --git a/core/src/read.jl b/core/src/read.jl
index edcce108c..0c565126d 100644
--- a/core/src/read.jl
+++ b/core/src/read.jl
@@ -1268,13 +1268,18 @@ function Subgrid(db::DB, config::Config, basin::Basin)::Subgrid
 end
 
 function Allocation(db::DB, config::Config, graph::MetaGraph)::Allocation
-    mean_input_flows = Dict{Tuple{NodeID, NodeID}, Float64}()
+    mean_input_flows = Dict{Int32, Dict{Tuple{NodeID, NodeID}, Float64}}()
 
     # 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_input_flows[edge] = 0.0
+        (; edge) = edge_metadata
+        id_source, _ = edge
+        if id_source in allocation_source_nodetypes
+            (; subnetwork_id) = graph[id_source]
+            if !haskey(mean_input_flows, subnetwork_id)
+                mean_input_flows[subnetwork_id] = Dict{Tuple{NodeID, NodeID}, Float64}()
+            end
+            mean_input_flows[subnetwork_id][edge] = 0.0
         end
     end
 
diff --git a/core/src/util.jl b/core/src/util.jl
index c05057911..517dd080e 100644
--- a/core/src/util.jl
+++ b/core/src/util.jl
@@ -1004,12 +1004,8 @@ function set_state_flow_edges(p::Parameters, u0::ComponentVector)::Parameters
     state_inflow_edges = Vector{EdgeMetadata}[]
     state_outflow_edges = Vector{EdgeMetadata}[]
 
-    placeholder_edge = EdgeMetadata(
-        0,
-        EdgeType.flow,
-        0,
-        (NodeID(:Terminal, 0, 0), NodeID(:Terminal, 0, 0)),
-    )
+    placeholder_edge =
+        EdgeMetadata(0, EdgeType.flow, (NodeID(:Terminal, 0, 0), NodeID(:Terminal, 0, 0)))
 
     for node_name in keys(u0)
         if hasfield(Parameters, node_name)
diff --git a/core/src/validation.jl b/core/src/validation.jl
index cfed36c6c..3eba4339a 100644
--- a/core/src/validation.jl
+++ b/core/src/validation.jl
@@ -626,53 +626,6 @@ function valid_discrete_control(p::Parameters, config::Config)::Bool
     return !errors
 end
 
-"""
-An allocation source edge is valid if either:
-    - The edge connects the main network to a subnetwork
-    - The edge comes from a source node
-"""
-function valid_sources(
-    p::Parameters,
-    capacity::JuMP.Containers.SparseAxisArray{Float64, 2, Tuple{NodeID, NodeID}},
-    subnetwork_id::Int32,
-)::Bool
-    (; graph) = p
-
-    errors = false
-
-    # Loop over edges that were assigned a capacity
-    for edge in keys(capacity.data)
-
-        # For an edge (id_a, id_b) in the physical model
-        # the reverse (id_b, id_a) can exist in the allocation subnetwork
-        if !haskey(graph, edge...)
-            edge = reverse(edge)
-        end
-
-        (id_source, id_dst) = edge
-
-        # Whether the current edge is a source for the current subnetwork
-        if graph[edge...].subnetwork_id_source == subnetwork_id
-            from_source_node = id_source.type in allocation_source_nodetypes
-
-            if is_main_network(subnetwork_id)
-                if !from_source_node
-                    errors = true
-                    @error "The source node of source edge $edge in the main network must be one of $allocation_source_nodetypes."
-                end
-            else
-                from_main_network = is_main_network(graph[id_source].subnetwork_id)
-
-                if !from_source_node && !from_main_network
-                    errors = true
-                    @error "The source node of source edge $edge for subnetwork $subnetwork_id is neither a source node nor is it coming from the main network."
-                end
-            end
-        end
-    end
-    return !errors
-end
-
 function valid_priorities(priorities::Vector{Int32}, use_allocation::Bool)::Bool
     if use_allocation && any(iszero, priorities)
         return false

From eeba58a7f5f7b7bddb8700c82ededd10cfe3e42f Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart123-4@hotmail.com>
Date: Wed, 27 Nov 2024 12:20:47 +0100
Subject: [PATCH 03/11] Finish refactor of handling boundary sources

---
 core/src/allocation_init.jl  | 57 +++++++++++++++++++++---------------
 core/src/allocation_optim.jl | 55 ++++++++++++++++------------------
 core/src/callback.jl         | 15 ++++++----
 core/src/parameter.jl        |  9 +++---
 core/src/read.jl             | 20 ++++++++-----
 core/src/util.jl             | 27 ++++++++++++-----
 core/test/allocation_test.jl | 57 ++++++++++++++++++------------------
 7 files changed, 135 insertions(+), 105 deletions(-)

diff --git a/core/src/allocation_init.jl b/core/src/allocation_init.jl
index 144976a0c..8d63b8d94 100644
--- a/core/src/allocation_init.jl
+++ b/core/src/allocation_init.jl
@@ -95,8 +95,8 @@ function get_subnetwork_capacity(
     return capacity
 end
 
-const allocation_source_nodetypes =
-    Set{NodeType.T}([NodeType.LevelBoundary, NodeType.FlowBoundary, NodeType.UserDemand])
+const boundary_source_nodetypes =
+    Set{NodeType.T}([NodeType.LevelBoundary, NodeType.FlowBoundary])
 
 """
 Add the edges connecting the main network work to a subnetwork to both the main network
@@ -275,30 +275,39 @@ The constraint indices are (edge_source_id, edge_dst_id).
 Constraint:
 flow over source edge <= source flow in subnetwork
 """
-function add_constraints_source!(
+function add_constraints_boundary_source!(
     problem::JuMP.Model,
     p::Parameters,
     subnetwork_id::Int32,
 )::Nothing
-    (; graph) = p
-    edges_source = Tuple{NodeID, NodeID}[]
+    edges_source =
+        [edge for edge in source_edges_subnetwork(p, subnetwork_id) if edge[1] != edge[2]]
     F = problem[:F]
 
-    # Find the edges in the whole model which are a source for
-    # this subnetwork
-    for edge in keys(p.allocation.mean_input_flows[subnetwork_id])
-        from_id, to_id = edge
-        if graph[from_id].subnetwork_id == subnetwork_id &&
-           graph[to_id].subnetwork_id == subnetwork_id
-            push!(edges_source, edge)
-        end
-    end
+    problem[:source_boundary] = JuMP.@constraint(
+        problem,
+        [edge_id = edges_source],
+        F[edge_id] <= 0.0,
+        base_name = "source_boundary"
+    )
+    return nothing
+end
+
+function add_constraints_main_network_source!(
+    problem::JuMP.Model,
+    p::Parameters,
+    subnetwork_id::Int32,
+)::Nothing
+    F = problem[:F]
+    (; main_network_connections, subnetwork_ids) = p.allocation
+    subnetwork_id = searchsortedfirst(subnetwork_ids, subnetwork_id)
+    edges_source = main_network_connections[subnetwork_id]
 
-    problem[:source] = JuMP.@constraint(
+    problem[:source_main_network] = JuMP.@constraint(
         problem,
         [edge_id = edges_source],
         F[edge_id] <= 0.0,
-        base_name = "source"
+        base_name = "source_main_network"
     )
     return nothing
 end
@@ -448,9 +457,9 @@ function allocation_problem(
 
     # Add constraints to problem
     add_constraints_conservation_node!(problem, p, subnetwork_id)
-
     add_constraints_capacity!(problem, capacity, p, subnetwork_id)
-    add_constraints_source!(problem, p, subnetwork_id)
+    add_constraints_boundary_source!(problem, p, subnetwork_id)
+    add_constraints_main_network_source!(problem, p, subnetwork_id)
     add_constraints_user_source!(problem, p, subnetwork_id)
     add_constraints_basin_flow!(problem)
     add_constraints_buffer!(problem)
@@ -481,12 +490,12 @@ function get_sources_in_order(
         sources[edge] = AllocationSource(; edge, type = AllocationSourceType.user_return)
     end
 
-    # Source edges (within subnetwork)
-    for edge in
-        sort(only(problem[:source].axes); by = edge -> (edge[1].value, edge[2].value))
-        if graph[edge[1]].subnetwork_id == graph[edge[2]].subnetwork_id
-            sources[edge] = AllocationSource(; edge, type = AllocationSourceType.edge)
-        end
+    # Boundary node sources
+    for edge in sort(
+        only(problem[:source_boundary].axes);
+        by = edge -> (edge[1].value, edge[2].value),
+    )
+        sources[edge] = AllocationSource(; edge, type = AllocationSourceType.boundary_node)
     end
 
     # Basins with level demand
diff --git a/core/src/allocation_optim.jl b/core/src/allocation_optim.jl
index 20f47fdfc..9b2a0eac7 100644
--- a/core/src/allocation_optim.jl
+++ b/core/src/allocation_optim.jl
@@ -127,8 +127,7 @@ function assign_allocations!(
         # If this edge is a source edge from the main network to a subnetwork,
         # and demands are being collected, add its flow to the demand of this edge
         if optimization_type == OptimizationType.collect_demands
-            if graph[edge...].subnetwork_id_source == subnetwork_id &&
-               edge ∈ main_network_source_edges
+            if edge in main_network_source_edges
                 allocated = flow[edge]
                 subnetwork_demands[edge][priority_idx] += allocated
             end
@@ -206,16 +205,12 @@ function set_initial_capacities_source!(
     (; allocation) = p
     (; mean_input_flows) = allocation
     (; subnetwork_id) = allocation_model
-    main_network_source_edges = get_main_network_connections(p, subnetwork_id)
 
-    for edge in keys(p.allocation.mean_input_flows[subnetwork_id])
-        # 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 = sources[edge]
-            @assert source.type == AllocationSourceType.edge
-            source.capacity = mean_input_flows[subnetwork_id][edge]
-        end
+    mean_input_flows_subnetwork_ = mean_input_flows_subnetwork(p, subnetwork_id)
+
+    for edge in keys(mean_input_flows_subnetwork_)
+        source = sources[edge]
+        source.capacity = mean_input_flows_subnetwork_[edge]
     end
     return nothing
 end
@@ -228,7 +223,7 @@ function reduce_source_capacity!(problem::JuMP.Model, source::AllocationSource):
 
     used_capacity =
         if source.type in (
-            AllocationSourceType.edge,
+            AllocationSourceType.boundary_node,
             AllocationSourceType.main_to_sub,
             AllocationSourceType.user_return,
         )
@@ -340,11 +335,10 @@ function get_basin_data(
     u::ComponentVector,
     node_id::NodeID,
 )
-    (; graph, allocation, basin) = p
-    (; Δt_allocation) = allocation_model
-    (; mean_input_flows) = allocation
+    (; graph, basin) = p
+    (; Δt_allocation, subnetwork_id) = allocation_model
     @assert node_id.type == NodeType.Basin
-    influx = mean_input_flows[(node_id, node_id)][]
+    influx = mean_input_flows_subnetwork(p, subnetwork_id)[(node_id, node_id)]
     storage_basin = basin.current_properties.current_storage[parent(u)][node_id.idx]
     control_inneighbors = inneighbor_labels_type(graph, node_id, EdgeType.control)
     if isempty(control_inneighbors)
@@ -840,9 +834,10 @@ function set_source_capacity!(
     optimization_type::OptimizationType.T,
 )::Nothing
     (; problem, sources) = allocation_model
-    constraints_source_edge = problem[:source]
-    constraints_source_basin = problem[:basin_outflow]
+    constraints_source_boundary = problem[:source_boundary]
     constraints_source_user_out = problem[:source_user]
+    constraints_source_main_network = problem[:source_main_network]
+    constraints_source_basin = problem[:basin_outflow]
     constraints_source_buffer = problem[:flow_buffer_outflow]
 
     for source in values(sources)
@@ -859,16 +854,17 @@ function set_source_capacity!(
             0.0
         end
 
-        constraint =
-            if source.type in (AllocationSourceType.edge, AllocationSourceType.main_to_sub)
-                constraints_source_edge[edge]
-            elseif source.type == AllocationSourceType.basin
-                constraints_source_basin[edge[1]]
-            elseif source.type == AllocationSourceType.user_return
-                constraints_source_user_out[edge[1]]
-            elseif source.type == AllocationSourceType.buffer
-                constraints_source_buffer[edge[1]]
-            end
+        constraint = if source.type == AllocationSourceType.boundary_node
+            constraints_source_boundary[edge]
+        elseif source.type == AllocationSourceType.main_to_sub
+            constraints_source_main_network[edge]
+        elseif source.type == AllocationSourceType.basin
+            constraints_source_basin[edge[1]]
+        elseif source.type == AllocationSourceType.user_return
+            constraints_source_user_out[edge[1]]
+        elseif source.type == AllocationSourceType.buffer
+            constraints_source_buffer[edge[1]]
+        end
 
         JuMP.set_normalized_rhs(constraint, capacity_effective)
     end
@@ -1043,7 +1039,8 @@ function empty_sources!(allocation_model::AllocationModel, allocation::Allocatio
     (; problem) = allocation_model
     (; subnetwork_demands) = allocation
 
-    for constraint_set_name in [:source, :source_user, :basin_outflow, :flow_buffer_outflow]
+    for constraint_set_name in
+        [:source_boundary, :source_user, :basin_outflow, :flow_buffer_outflow]
         constraint_set = problem[constraint_set_name]
         for key in only(constraint_set.axes)
             # Do not set the capacity to 0.0 if the edge
diff --git a/core/src/callback.jl b/core/src/callback.jl
index 213e97284..cda123cff 100644
--- a/core/src/callback.jl
+++ b/core/src/callback.jl
@@ -153,8 +153,11 @@ function update_cumulative_flows!(u, t, integrator)::Nothing
     end
 
     # Update realized flows for allocation input
-    for edge in keys(allocation.mean_input_flows)
-        allocation.mean_input_flows[edge] += flow_update_on_edge(integrator, edge)
+    for subnetwork_id in allocation.subnetwork_ids
+        mean_input_flows_subnetwork_ = mean_input_flows_subnetwork(p, subnetwork_id)
+        for edge in keys(mean_input_flows_subnetwork_)
+            mean_input_flows_subnetwork_[edge] += flow_update_on_edge(integrator, edge)
+        end
     end
 
     # Update realized flows for allocation output
@@ -773,8 +776,10 @@ function update_allocation!(integrator)::Nothing
 
     # Divide by the allocation Δt to get the mean input flows from the cumulative flows
     (; Δt_allocation) = allocation_models[1]
-    for edge in keys(mean_input_flows)
-        mean_input_flows[edge] /= Δt_allocation
+    for mean_input_flows_subnetwork in values(mean_input_flows)
+        for edge in keys(mean_input_flows_subnetwork)
+            mean_input_flows_subnetwork[edge] /= Δt_allocation
+        end
     end
 
     # Divide by the allocation Δt to get the mean realized flows from the cumulative flows
@@ -797,7 +802,7 @@ function update_allocation!(integrator)::Nothing
     end
 
     # Reset the mean flows
-    for mean_flows in (mean_input_flows, mean_realized_flows)
+    for mean_flows in (mean_input_flows..., mean_realized_flows)
         for edge in keys(mean_flows)
             mean_flows[edge] = 0.0
         end
diff --git a/core/src/parameter.jl b/core/src/parameter.jl
index 292be86e7..62e3f76e8 100644
--- a/core/src/parameter.jl
+++ b/core/src/parameter.jl
@@ -184,13 +184,14 @@ record_flow: A record of all flows computed by allocation optimization, eventual
     output file
 """
 @kwdef struct Allocation
-    subnetwork_ids::Vector{Int32} = []
-    allocation_models::Vector{AllocationModel} = []
-    main_network_connections::Vector{Vector{Tuple{NodeID, NodeID}}} = []
+    subnetwork_ids::Vector{Int32} = Int32[]
+    allocation_models::Vector{AllocationModel} = AllocationModel[]
+    main_network_connections::Vector{Vector{Tuple{NodeID, NodeID}}} =
+        Vector{Tuple{NodeID, NodeID}}[]
     priorities::Vector{Int32}
     subnetwork_demands::Dict{Tuple{NodeID, NodeID}, Vector{Float64}} = Dict()
     subnetwork_allocateds::Dict{Tuple{NodeID, NodeID}, Vector{Float64}} = Dict()
-    mean_input_flows::Dict{Int32, Dict{Tuple{NodeID, NodeID}, Float64}}
+    mean_input_flows::Vector{Dict{Tuple{NodeID, NodeID}, Float64}}
     mean_realized_flows::Dict{Tuple{NodeID, NodeID}, Float64}
     record_demand::@NamedTuple{
         time::Vector{Float64},
diff --git a/core/src/read.jl b/core/src/read.jl
index 0c565126d..d8e43e3eb 100644
--- a/core/src/read.jl
+++ b/core/src/read.jl
@@ -1268,25 +1268,31 @@ function Subgrid(db::DB, config::Config, basin::Basin)::Subgrid
 end
 
 function Allocation(db::DB, config::Config, graph::MetaGraph)::Allocation
-    mean_input_flows = Dict{Int32, Dict{Tuple{NodeID, NodeID}, Float64}}()
+    mean_input_flows = Dict{Tuple{NodeID, NodeID}, Float64}[]
+
+    subnetwork_ids = sort(collect(keys(graph[].node_ids)))
+
+    for subnetwork_id in subnetwork_ids
+        push!(mean_input_flows, Dict{Tuple{NodeID, NodeID}, Float64}())
+    end
 
     # Find edges which serve as sources in allocation
     for edge_metadata in values(graph.edge_data)
         (; edge) = edge_metadata
         id_source, _ = edge
-        if id_source in allocation_source_nodetypes
+        if id_source.type in boundary_source_nodetypes
             (; subnetwork_id) = graph[id_source]
-            if !haskey(mean_input_flows, subnetwork_id)
-                mean_input_flows[subnetwork_id] = Dict{Tuple{NodeID, NodeID}, Float64}()
-            end
-            mean_input_flows[subnetwork_id][edge] = 0.0
+            subnetwork_idx = searchsortedfirst(subnetwork_ids, subnetwork_id)
+            mean_input_flows[subnetwork_idx][edge] = 0.0
         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_input_flows[(node_id, node_id)] = 0.0
+            subnetwork_id = graph[node_id].subnetwork_id
+            subnetwork_idx = searchsortedfirst(subnetwork_ids, subnetwork_id)
+            mean_input_flows[subnetwork_idx][(node_id, node_id)] = 0.0
         end
     end
 
diff --git a/core/src/util.jl b/core/src/util.jl
index 852595d92..2eac3f267 100644
--- a/core/src/util.jl
+++ b/core/src/util.jl
@@ -587,15 +587,17 @@ function set_initial_allocation_mean_flows!(integrator)::Nothing
     du = get_du(integrator)
     water_balance!(du, u, p, t)
 
-    for edge in keys(mean_input_flows)
-        if edge[1] == edge[2]
-            q = get_influx(du, edge[1], p)
-        else
-            q = get_flow(du, p, t, edge)
+    for mean_input_flows_subnetwork in values(mean_input_flows)
+        for edge in keys(mean_input_flows_subnetwork)
+            if edge[1] == edge[2]
+                q = get_influx(du, edge[1], p)
+            else
+                q = get_flow(du, p, t, edge)
+            end
+            # Multiply by Δt_allocation as averaging divides by this factor
+            # in update_allocation!
+            mean_input_flows_subnetwork[edge] = q * Δt_allocation
         end
-        # Multiply by Δt_allocation as averaging divides by this factor
-        # in update_allocation!
-        mean_input_flows[edge] = q * Δt_allocation
     end
 
     # Mean realized demands for basins are calculated as Δstorage/Δt
@@ -1168,3 +1170,12 @@ function min_low_user_demand_level_factor(
         one(T)
     end
 end
+
+function mean_input_flows_subnetwork(p::Parameters, subnetwork_id::Int32)
+    (; mean_input_flows, subnetwork_ids) = p.allocation
+    subnetwork_idx = searchsortedfirst(subnetwork_ids, subnetwork_id)
+    return mean_input_flows[subnetwork_idx]
+end
+
+source_edges_subnetwork(p::Parameters, subnetwork_id::Int32) =
+    keys(mean_input_flows_subnetwork(p, subnetwork_id))
diff --git a/core/test/allocation_test.jl b/core/test/allocation_test.jl
index 6da2fec36..3dd6c7ba0 100644
--- a/core/test/allocation_test.jl
+++ b/core/test/allocation_test.jl
@@ -14,7 +14,8 @@
 
     (; graph, allocation) = p
 
-    allocation.mean_input_flows[(NodeID(:FlowBoundary, 1, p), NodeID(:Basin, 2, p))] = 4.5
+    allocation.mean_input_flows[1][(NodeID(:FlowBoundary, 1, p), NodeID(:Basin, 2, p))] =
+        4.5
     allocation_model = p.allocation.allocation_models[1]
     (; flow) = allocation_model
     u = ComponentVector(; storage = zeros(length(p.basin.node_id)))
@@ -105,11 +106,11 @@ end
 
     # In each subnetwork, the connection from the main network to the subnetwork is
     # interpreted as a source
-    @test Ribasim.get_allocation_model(p, Int32(3)).problem[:source].axes[1] ==
+    @test Ribasim.get_allocation_model(p, Int32(3)).problem[:source_main_network].axes[1] ==
           [(NodeID(:Basin, 2, p), NodeID(:Pump, 11, p))]
-    @test Ribasim.get_allocation_model(p, Int32(5)).problem[:source].axes[1] ==
+    @test Ribasim.get_allocation_model(p, Int32(5)).problem[:source_main_network].axes[1] ==
           [(NodeID(:Basin, 6, p), NodeID(:Pump, 24, p))]
-    @test Ribasim.get_allocation_model(p, Int32(7)).problem[:source].axes[1] ==
+    @test Ribasim.get_allocation_model(p, Int32(7)).problem[:source_main_network].axes[1] ==
           [(NodeID(:Basin, 10, p), NodeID(:Pump, 38, p))]
 end
 
@@ -174,7 +175,7 @@ end
 
     # Running full allocation algorithm
     (; Δt_allocation) = allocation_models[1]
-    mean_input_flows[(NodeID(:FlowBoundary, 1, p), NodeID(:Basin, 2, p))] =
+    mean_input_flows[1][(NodeID(:FlowBoundary, 1, p), NodeID(:Basin, 2, p))] =
         4.5 * Δt_allocation
     Ribasim.update_allocation!(model.integrator)
 
@@ -224,9 +225,9 @@ end
         allocation
     t = 0.0
 
-    # Set flows of sources in
-    mean_input_flows[(NodeID(:FlowBoundary, 58, p), NodeID(:Basin, 16, p))] = 1.0
-    mean_input_flows[(NodeID(:FlowBoundary, 59, p), NodeID(:Basin, 44, p))] = 1e-3
+    # Set flows of sources in subnetworks
+    mean_input_flows[2][(NodeID(:FlowBoundary, 58, p), NodeID(:Basin, 16, p))] = 1.0
+    mean_input_flows[4][(NodeID(:FlowBoundary, 59, p), NodeID(:Basin, 44, p))] = 1e-3
 
     # Collecting demands
     u = ComponentVector(; storage = zeros(length(basin.node_id)))
@@ -251,24 +252,24 @@ end
     du = get_du(model.integrator)
     Ribasim.formulate_storages!(current_storage, du, u, p, t)
 
-    current_storage ≈ Float32[
-        1.0346e6,
-        1.0346e6,
-        1.0346e6,
-        1.0346e6,
-        1.0346e6,
-        13.83,
-        40.10,
-        6.029e5,
-        4641,
-        2402,
-        6.03995,
-        928.8,
-        8.017,
-        10417,
-        5.619,
-        10417,
-        4.057,
+    @test current_storage ≈ Float32[
+        1.0346908f6,
+        1.03469f6,
+        1.0346894f6,
+        1.034689f6,
+        1.0346888f6,
+        13.833241,
+        40.109993,
+        187761.73,
+        4641.365,
+        2402.6687,
+        6.039952,
+        928.84283,
+        8.0175905,
+        10419.247,
+        5.619053,
+        10419.156,
+        4.057502,
     ]
 end
 
@@ -287,7 +288,7 @@ end
     (; user_demand, graph, allocation, basin, level_demand) = p
 
     # Initial "integrated" vertical flux
-    @test allocation.mean_input_flows[(NodeID(:Basin, 2, p), NodeID(:Basin, 2, p))] ≈ 1e2
+    @test allocation.mean_input_flows[1][(NodeID(:Basin, 2, p), NodeID(:Basin, 2, p))] ≈ 1e2
 
     Ribasim.solve!(model)
 
@@ -614,5 +615,5 @@ end
     # Given a max_level of Inf, the basin capacity is 0.0 because it is not possible for the basin level to be > Inf
     @test Ribasim.get_basin_capacity(allocation_models[1], u, p, t, basin.node_id[1]) == 0.0
     @test Ribasim.get_basin_capacity(allocation_models[1], u, p, t, basin.node_id[2]) == 0.0
-    @test Ribasim.get_basin_capacity(allocation_models[1], u, p, t, basin.node_id[3]) == 0.0
+    @test Ribasim.get_basin_capacity(allocation_models[2], u, p, t, basin.node_id[3]) == 0.0
 end

From 4344c12a8bc4b6cb344f48ed1354318c276ea1ca Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart123-4@hotmail.com>
Date: Wed, 27 Nov 2024 13:30:33 +0100
Subject: [PATCH 04/11] Pass tests

---
 core/src/read.jl                              | 76 ++++++++++---------
 core/test/validation_test.jl                  | 14 +---
 docs/guide/examples.ipynb                     |  4 +-
 .../ribasim_testmodels/allocation.py          | 27 ++++---
 .../ribasim_testmodels/invalid.py             |  1 -
 5 files changed, 56 insertions(+), 66 deletions(-)

diff --git a/core/src/read.jl b/core/src/read.jl
index d8e43e3eb..17e020b78 100644
--- a/core/src/read.jl
+++ b/core/src/read.jl
@@ -1269,50 +1269,52 @@ end
 
 function Allocation(db::DB, config::Config, graph::MetaGraph)::Allocation
     mean_input_flows = Dict{Tuple{NodeID, NodeID}, Float64}[]
-
+    mean_realized_flows = Dict{Tuple{NodeID, NodeID}, Float64}()
     subnetwork_ids = sort(collect(keys(graph[].node_ids)))
 
-    for subnetwork_id in subnetwork_ids
-        push!(mean_input_flows, Dict{Tuple{NodeID, NodeID}, Float64}())
-    end
-
-    # Find edges which serve as sources in allocation
-    for edge_metadata in values(graph.edge_data)
-        (; edge) = edge_metadata
-        id_source, _ = edge
-        if id_source.type in boundary_source_nodetypes
-            (; subnetwork_id) = graph[id_source]
-            subnetwork_idx = searchsortedfirst(subnetwork_ids, subnetwork_id)
-            mean_input_flows[subnetwork_idx][edge] = 0.0
+    if config.allocation.use_allocation
+        for subnetwork_id in subnetwork_ids
+            push!(mean_input_flows, Dict{Tuple{NodeID, NodeID}, Float64}())
         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]
-            subnetwork_id = graph[node_id].subnetwork_id
-            subnetwork_idx = searchsortedfirst(subnetwork_ids, subnetwork_id)
-            mean_input_flows[subnetwork_idx][(node_id, node_id)] = 0.0
+        # Find edges which serve as sources in allocation
+        for edge_metadata in values(graph.edge_data)
+            (; edge) = edge_metadata
+            id_source, _ = edge
+            if id_source.type in boundary_source_nodetypes
+                (; subnetwork_id) = graph[id_source]
+                subnetwork_idx = searchsortedfirst(subnetwork_ids, subnetwork_id)
+                mean_input_flows[subnetwork_idx][edge] = 0.0
+            end
         end
-    end
 
-    mean_realized_flows = Dict{Tuple{NodeID, NodeID}, Float64}()
+        # Find basins with a level demand
+        for node_id in values(graph.vertex_labels)
+            if has_external_demand(graph, node_id, :level_demand)[1]
+                subnetwork_id = graph[node_id].subnetwork_id
+                subnetwork_idx = searchsortedfirst(subnetwork_ids, subnetwork_id)
+                mean_input_flows[subnetwork_idx][(node_id, node_id)] = 0.0
+            end
+        end
 
-    # Find edges that realize a demand
-    for edge_metadata in values(graph.edge_data)
-        (; type, edge) = edge_metadata
-
-        src_id, dst_id = edge
-        user_demand_inflow = (type == EdgeType.flow) && (dst_id.type == NodeType.UserDemand)
-        level_demand_inflow =
-            (type == EdgeType.control) && (src_id.type == NodeType.LevelDemand)
-        flow_demand_inflow =
-            (type == EdgeType.flow) && has_external_demand(graph, dst_id, :flow_demand)[1]
-
-        if user_demand_inflow || flow_demand_inflow
-            mean_realized_flows[edge] = 0.0
-        elseif level_demand_inflow
-            mean_realized_flows[(dst_id, dst_id)] = 0.0
+        # Find edges that realize a demand
+        for edge_metadata in values(graph.edge_data)
+            (; type, edge) = edge_metadata
+
+            src_id, dst_id = edge
+            user_demand_inflow =
+                (type == EdgeType.flow) && (dst_id.type == NodeType.UserDemand)
+            level_demand_inflow =
+                (type == EdgeType.control) && (src_id.type == NodeType.LevelDemand)
+            flow_demand_inflow =
+                (type == EdgeType.flow) &&
+                has_external_demand(graph, dst_id, :flow_demand)[1]
+
+            if user_demand_inflow || flow_demand_inflow
+                mean_realized_flows[edge] = 0.0
+            elseif level_demand_inflow
+                mean_realized_flows[(dst_id, dst_id)] = 0.0
+            end
         end
     end
 
diff --git a/core/test/validation_test.jl b/core/test/validation_test.jl
index 547a11ff4..e46c0ef74 100644
--- a/core/test/validation_test.jl
+++ b/core/test/validation_test.jl
@@ -86,12 +86,7 @@ end
     graph[NodeID(:Pump, 6, 1)] = NodeMetadata(:pump, 9)
 
     function set_edge_metadata!(id_1, id_2, edge_type)
-        graph[id_1, id_2] = EdgeMetadata(;
-            id = 0,
-            type = edge_type,
-            subnetwork_id_source = 0,
-            edge = (id_1, id_2),
-        )
+        graph[id_1, id_2] = EdgeMetadata(; id = 0, type = edge_type, edge = (id_1, id_2))
         return nothing
     end
 
@@ -148,12 +143,7 @@ end
     graph[NodeID(:Basin, 7, 1)] = NodeMetadata(:basin, 0)
 
     function set_edge_metadata!(id_1, id_2, edge_type)
-        graph[id_1, id_2] = EdgeMetadata(;
-            id = 0,
-            type = edge_type,
-            subnetwork_id_source = 0,
-            edge = (id_1, id_2),
-        )
+        graph[id_1, id_2] = EdgeMetadata(; id = 0, type = edge_type, edge = (id_1, id_2))
         return nothing
     end
 
diff --git a/docs/guide/examples.ipynb b/docs/guide/examples.ipynb
index 31e899dd6..3e09e5986 100644
--- a/docs/guide/examples.ipynb
+++ b/docs/guide/examples.ipynb
@@ -1203,7 +1203,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "model.edge.add(model.flow_boundary[1], model.basin[2], subnetwork_id=1)\n",
+    "model.edge.add(model.flow_boundary[1], model.basin[2])\n",
     "model.edge.add(model.basin[2], model.user_demand[3])\n",
     "model.edge.add(model.basin[2], model.linear_resistance[4])\n",
     "model.edge.add(model.linear_resistance[4], model.basin[5])\n",
@@ -1499,7 +1499,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "model.edge.add(model.flow_boundary[1], model.basin[2], subnetwork_id=2)\n",
+    "model.edge.add(model.flow_boundary[1], model.basin[2])\n",
     "model.edge.add(model.basin[2], model.user_demand[3])\n",
     "model.edge.add(model.level_demand[4], model.basin[2])\n",
     "model.edge.add(model.user_demand[3], model.basin[5])\n",
diff --git a/python/ribasim_testmodels/ribasim_testmodels/allocation.py b/python/ribasim_testmodels/ribasim_testmodels/allocation.py
index e51fbbe40..da31a3ecc 100644
--- a/python/ribasim_testmodels/ribasim_testmodels/allocation.py
+++ b/python/ribasim_testmodels/ribasim_testmodels/allocation.py
@@ -154,7 +154,7 @@ def subnetwork_model() -> Model:
     )
     model.outlet.add(Node(13, Point(2, 4), subnetwork_id=2), outlet_data)
 
-    model.edge.add(model.flow_boundary[1], model.basin[2], subnetwork_id=2)
+    model.edge.add(model.flow_boundary[1], model.basin[2])
     model.edge.add(model.basin[2], model.outlet[3])
     model.edge.add(model.outlet[3], model.terminal[4])
     model.edge.add(model.basin[2], model.user_demand[10])
@@ -269,7 +269,7 @@ def looped_subnetwork_model() -> Model:
         ],
     )
 
-    model.edge.add(model.flow_boundary[5], model.basin[2], subnetwork_id=2)
+    model.edge.add(model.flow_boundary[5], model.basin[2])
     model.edge.add(model.basin[2], model.outlet[3])
     model.edge.add(model.outlet[3], model.terminal[4])
     model.edge.add(model.basin[2], model.user_demand[1])
@@ -348,7 +348,7 @@ def minimal_subnetwork_model() -> Model:
         ],
     )
 
-    model.edge.add(model.flow_boundary[1], model.basin[2], subnetwork_id=2)
+    model.edge.add(model.flow_boundary[1], model.basin[2])
     model.edge.add(model.basin[2], model.pump[3])
     model.edge.add(model.pump[3], model.basin[4])
     model.edge.add(model.basin[4], model.user_demand[5])
@@ -409,7 +409,7 @@ def allocation_example_model() -> Model:
     )
     model.terminal.add(Node(8, Point(5, 0), subnetwork_id=2))
 
-    model.edge.add(model.flow_boundary[1], model.basin[2], subnetwork_id=2)
+    model.edge.add(model.flow_boundary[1], model.basin[2])
     model.edge.add(model.basin[2], model.user_demand[3])
     model.edge.add(model.basin[2], model.linear_resistance[4])
     model.edge.add(model.linear_resistance[4], model.basin[5])
@@ -632,7 +632,7 @@ def main_network_with_subnetworks_model() -> Model:
         [user_demand.Static(return_factor=[0.9], priority=2, min_level=0.0)],
     )
 
-    model.edge.add(model.flow_boundary[1], model.basin[2], subnetwork_id=1)
+    model.edge.add(model.flow_boundary[1], model.basin[2])
     model.edge.add(model.basin[2], model.linear_resistance[3])
     model.edge.add(model.linear_resistance[3], model.basin[4])
     model.edge.add(model.basin[4], model.linear_resistance[5])
@@ -690,9 +690,9 @@ def main_network_with_subnetworks_model() -> Model:
     model.edge.add(model.user_demand[53], model.basin[50])
     model.edge.add(model.basin[44], model.user_demand[57])
     model.edge.add(model.user_demand[57], model.basin[44])
-    model.edge.add(model.basin[2], model.pump[11], subnetwork_id=3)
-    model.edge.add(model.basin[6], model.pump[24], subnetwork_id=5)
-    model.edge.add(model.basin[10], model.pump[38], subnetwork_id=7)
+    model.edge.add(model.basin[2], model.pump[11])
+    model.edge.add(model.basin[6], model.pump[24])
+    model.edge.add(model.basin[10], model.pump[38])
     model.edge.add(model.basin[8], model.user_demand[60])
     model.edge.add(model.user_demand[60], model.basin[8])
 
@@ -713,8 +713,8 @@ def subnetworks_with_sources_model() -> Model:
         [flow_boundary.Static(flow_rate=[0.003])],
     )
 
-    model.edge.add(model.flow_boundary[58], model.basin[16], subnetwork_id=3)
-    model.edge.add(model.flow_boundary[59], model.basin[44], subnetwork_id=7)
+    model.edge.add(model.flow_boundary[58], model.basin[16])
+    model.edge.add(model.flow_boundary[59], model.basin[44])
 
     return model
 
@@ -769,7 +769,7 @@ def level_demand_model() -> Model:
         [basin.Profile(area=1000.0, level=[0.0, 1.0]), basin.State(level=[2.0])],
     )
 
-    model.edge.add(model.flow_boundary[1], model.basin[2], subnetwork_id=2)
+    model.edge.add(model.flow_boundary[1], model.basin[2])
     model.edge.add(model.basin[2], model.user_demand[3])
     model.edge.add(model.level_demand[4], model.basin[2])
     model.edge.add(model.user_demand[3], model.basin[5])
@@ -842,7 +842,6 @@ def flow_demand_model() -> Model:
     model.edge.add(
         model.level_boundary[1],
         model.tabulated_rating_curve[2],
-        subnetwork_id=2,
     )
     model.edge.add(model.tabulated_rating_curve[2], model.basin[3])
     model.edge.add(model.basin[3], model.user_demand[4])
@@ -885,7 +884,7 @@ def linear_resistance_demand_model():
         [flow_demand.Static(priority=[1], demand=2.0)],
     )
 
-    model.edge.add(model.basin[1], model.linear_resistance[2], subnetwork_id=1)
+    model.edge.add(model.basin[1], model.linear_resistance[2])
     model.edge.add(model.linear_resistance[2], model.basin[3])
     model.edge.add(model.flow_demand[4], model.linear_resistance[2])
 
@@ -978,7 +977,7 @@ def fair_distribution_model():
         ],
     )
 
-    model.edge.add(model.level_boundary[1], model.pump[2], subnetwork_id=1)
+    model.edge.add(model.level_boundary[1], model.pump[2])
     model.edge.add(model.pump[2], model.basin[3])
     model.edge.add(model.basin[3], model.linear_resistance[4])
     model.edge.add(model.linear_resistance[4], model.basin[5])
diff --git a/python/ribasim_testmodels/ribasim_testmodels/invalid.py b/python/ribasim_testmodels/ribasim_testmodels/invalid.py
index 995d30993..46473b461 100644
--- a/python/ribasim_testmodels/ribasim_testmodels/invalid.py
+++ b/python/ribasim_testmodels/ribasim_testmodels/invalid.py
@@ -216,7 +216,6 @@ def invalid_priorities_model() -> Model:
     model.edge.add(
         model.level_boundary[1],
         model.tabulated_rating_curve[2],
-        subnetwork_id=2,
     )
     model.edge.add(model.tabulated_rating_curve[2], model.basin[3])
     model.edge.add(model.basin[3], model.user_demand[4])

From 9356cedf71e8fe8ca3c6707042d0d3ad48fb4989 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart123-4@hotmail.com>
Date: Wed, 27 Nov 2024 13:41:55 +0100
Subject: [PATCH 05/11] Make mypy happy

---
 python/ribasim/ribasim/model.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/python/ribasim/ribasim/model.py b/python/ribasim/ribasim/model.py
index 7457b4b45..d46b648b0 100644
--- a/python/ribasim/ribasim/model.py
+++ b/python/ribasim/ribasim/model.py
@@ -307,7 +307,7 @@ def write(self, filepath: str | PathLike[str]) -> Path:
         if self.use_validation:
             self._validate_model()
 
-        if self.edge.df.subnetwork_id.notna().any():
+        if self.edge.df is not None and self.edge.df["subnetwork_id"].notna().any():
             warnings.warn(
                 "Sources for allocation are automatically inferred and no longer have to be specified in the `Edge` table.",
                 DeprecationWarning,

From cfe67e70002095ab6ca5b187061a753c4486e273 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart123-4@hotmail.com>
Date: Wed, 27 Nov 2024 17:28:02 +0100
Subject: [PATCH 06/11] Update docs

---
 core/src/allocation_init.jl             | 13 +++++++++++--
 core/src/allocation_optim.jl            |  2 --
 docs/concept/allocation.qmd             |  4 ++--
 docs/reference/usage.qmd                |  2 +-
 python/ribasim/ribasim/geometry/edge.py |  3 +--
 5 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/core/src/allocation_init.jl b/core/src/allocation_init.jl
index 8d63b8d94..d80660861 100644
--- a/core/src/allocation_init.jl
+++ b/core/src/allocation_init.jl
@@ -268,18 +268,19 @@ function add_constraints_user_source!(
 end
 
 """
-Add the source constraints to the allocation problem.
+Add the boundary source constraints to the allocation problem.
 The actual threshold values will be set before each allocation solve.
 The constraint indices are (edge_source_id, edge_dst_id).
 
 Constraint:
-flow over source edge <= source flow in subnetwork
+flow over source edge <= source flow in physical layer
 """
 function add_constraints_boundary_source!(
     problem::JuMP.Model,
     p::Parameters,
     subnetwork_id::Int32,
 )::Nothing
+    # Source edges (without the basins)
     edges_source =
         [edge for edge in source_edges_subnetwork(p, subnetwork_id) if edge[1] != edge[2]]
     F = problem[:F]
@@ -293,6 +294,14 @@ function add_constraints_boundary_source!(
     return nothing
 end
 
+"""
+Add main network source constraints to the allocation problem.
+The actual threshold values will be set before each allocation solve.
+The constraint indices are (edge_source_id, edge_dst_id).
+
+Constraint:
+flow over main network to subnetwork connection edge <= either 0 or allocated amount from the main network
+"""
 function add_constraints_main_network_source!(
     problem::JuMP.Model,
     p::Parameters,
diff --git a/core/src/allocation_optim.jl b/core/src/allocation_optim.jl
index 9b2a0eac7..7ecb3b885 100644
--- a/core/src/allocation_optim.jl
+++ b/core/src/allocation_optim.jl
@@ -202,8 +202,6 @@ function set_initial_capacities_source!(
     p::Parameters,
 )::Nothing
     (; sources) = allocation_model
-    (; allocation) = p
-    (; mean_input_flows) = allocation
     (; subnetwork_id) = allocation_model
 
     mean_input_flows_subnetwork_ = mean_input_flows_subnetwork(p, subnetwork_id)
diff --git a/docs/concept/allocation.qmd b/docs/concept/allocation.qmd
index f90a83cec..91318fd84 100644
--- a/docs/concept/allocation.qmd
+++ b/docs/concept/allocation.qmd
@@ -42,9 +42,9 @@ The allocation problem is solved per subnetwork, which is given by a subset $S \
 
 Sources are indicated by a set of edges in the subnetwork
 $$
-E_S^\text{source} \subset E.
+E_S^\text{source} \subset E,
 $$
-That is, if $(i,j) \in E_S^\text{source}$, then the average over the last allocation interval $\Delta t_{\text{alloc}}$ of the of the flow over this edge
+which are automatically inferred as all edges that point out of LevelBoundary or FlowBoundary nodes. That is, if $(i,j) \in E_S^\text{source}$, then the average over the last allocation interval $\Delta t_{\text{alloc}}$ of the of the flow over this edge
 $$
     \frac{1}{\Delta t_{\text{alloc}}}\int_{t - \Delta t_{\text{alloc}}}^tQ_{ij}(t') dt'
 $$
diff --git a/docs/reference/usage.qmd b/docs/reference/usage.qmd
index 63eeedd99..cd209650d 100644
--- a/docs/reference/usage.qmd
+++ b/docs/reference/usage.qmd
@@ -199,7 +199,7 @@ to_node_type  | String                        | -
 to_node_id    | Int32                         | -
 edge_type     | String                        | must be "flow" or "control"
 geom          | LineString or MultiLineString | (optional)
-name          | String                        | (optional, does not have to be unique)
+name          | String                        | (deprecated)
 subnetwork_id | Int32                         | (optional, denotes source in allocation network)
 
 Similarly to the node table, you can use a geometry to visualize the connections between the
diff --git a/python/ribasim/ribasim/geometry/edge.py b/python/ribasim/ribasim/geometry/edge.py
index 9142c6817..f89e6f6ee 100644
--- a/python/ribasim/ribasim/geometry/edge.py
+++ b/python/ribasim/ribasim/geometry/edge.py
@@ -91,8 +91,7 @@ def add(
         name : str
             An optional name for the edge.
         subnetwork_id : int | None
-            An optional subnetwork id for the edge. This edge indicates a source for
-            the allocation algorithm, and should thus not be set for every edge in a subnetwork.
+            Deprecated: Sources for allocation are automatically inferred and no longer have to be specified in the `Edge` table.
         **kwargs : Dict
         """
         if not can_connect(from_node.node_type, to_node.node_type):

From b215c61fa7b8a7ed34df06bacbfa8763af3d4a47 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart123-4@hotmail.com>
Date: Thu, 28 Nov 2024 16:56:56 +0100
Subject: [PATCH 07/11] docs fix

---
 docs/dev/callstacks.qmd | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/docs/dev/callstacks.qmd b/docs/dev/callstacks.qmd
index 5bd4fe372..a36c75cdf 100644
--- a/docs/dev/callstacks.qmd
+++ b/docs/dev/callstacks.qmd
@@ -74,10 +74,12 @@ toml_path = normpath(
     @__DIR__,
     "../../generated_testmodels/main_network_with_subnetworks/ribasim.toml",
 )
-config = Ribasim.Config(toml_path; allocation_use_allocation = false)
+config = Ribasim.Config(toml_path)
 db_path = Ribasim.database_path(config)
 db = SQLite.DB(db_path)
 p = Ribasim.Parameters(db, config)
+empty!(p.allocation.subnetwork_ids)
+empty!(p.allocation.main_network_connections)
 graph, verts = tracecall((Ribasim,), Ribasim.initialize_allocation!, (p, config))
 plot_graph(graph)
 ```

From d9284f520ea28f1b576644b5273bc80b414ef441 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart123-4@hotmail.com>
Date: Thu, 28 Nov 2024 17:10:08 +0100
Subject: [PATCH 08/11] Remove subnetwork_id column from Edge table

---
 core/src/graph.jl                       | 18 ++----------------
 docs/reference/usage.qmd                |  3 +--
 python/ribasim/ribasim/geometry/edge.py |  2 --
 python/ribasim/ribasim/migrations.py    |  3 +++
 python/ribasim/ribasim/model.py         |  8 --------
 5 files changed, 6 insertions(+), 28 deletions(-)

diff --git a/core/src/graph.jl b/core/src/graph.jl
index ef2f30b54..666591e8e 100644
--- a/core/src/graph.jl
+++ b/core/src/graph.jl
@@ -56,15 +56,8 @@ function create_graph(db::DB, config::Config)::MetaGraph
     end
 
     errors = false
-    for (;
-        edge_id,
-        from_node_type,
-        from_node_id,
-        to_node_type,
-        to_node_id,
-        edge_type,
-        subnetwork_id,
-    ) in edge_rows
+    for (; edge_id, from_node_type, from_node_id, to_node_type, to_node_id, edge_type) in
+        edge_rows
         try
             # hasfield does not work
             edge_type = getfield(EdgeType, Symbol(edge_type))
@@ -73,13 +66,6 @@ function create_graph(db::DB, config::Config)::MetaGraph
         end
         id_src = NodeID(from_node_type, from_node_id, db)
         id_dst = NodeID(to_node_type, to_node_id, db)
-        if !ismissing(subnetwork_id)
-            Base.depwarn(
-                "Sources for allocation are automatically inferred and no longer have to be specified in the `Edge` table.",
-                :create_graph;
-                force = true,
-            )
-        end
         edge_metadata =
             EdgeMetadata(; id = edge_id, type = edge_type, edge = (id_src, id_dst))
         if edge_type == EdgeType.flow
diff --git a/docs/reference/usage.qmd b/docs/reference/usage.qmd
index cd209650d..afe2f1529 100644
--- a/docs/reference/usage.qmd
+++ b/docs/reference/usage.qmd
@@ -199,8 +199,7 @@ to_node_type  | String                        | -
 to_node_id    | Int32                         | -
 edge_type     | String                        | must be "flow" or "control"
 geom          | LineString or MultiLineString | (optional)
-name          | String                        | (deprecated)
-subnetwork_id | Int32                         | (optional, denotes source in allocation network)
+name          | String                        | (optional, does not have to be unique)
 
 Similarly to the node table, you can use a geometry to visualize the connections between the
 nodes in QGIS. For instance, you can draw a line connecting the two node coordinates.
diff --git a/python/ribasim/ribasim/geometry/edge.py b/python/ribasim/ribasim/geometry/edge.py
index f89e6f6ee..62a73cb28 100644
--- a/python/ribasim/ribasim/geometry/edge.py
+++ b/python/ribasim/ribasim/geometry/edge.py
@@ -90,8 +90,6 @@ def add(
             The geometry of a line. If not supplied, it creates a straight line between the nodes.
         name : str
             An optional name for the edge.
-        subnetwork_id : int | None
-            Deprecated: Sources for allocation are automatically inferred and no longer have to be specified in the `Edge` table.
         **kwargs : Dict
         """
         if not can_connect(from_node.node_type, to_node.node_type):
diff --git a/python/ribasim/ribasim/migrations.py b/python/ribasim/ribasim/migrations.py
index fc51eff52..239effaa8 100644
--- a/python/ribasim/ribasim/migrations.py
+++ b/python/ribasim/ribasim/migrations.py
@@ -27,6 +27,9 @@ def edgeschema_migration(gdf: GeoDataFrame, schema_version: int) -> GeoDataFrame
         warnings.warn("Migrating outdated Edge table.", UserWarning)
         assert gdf["edge_id"].is_unique, "Edge IDs have to be unique."
         gdf.set_index("edge_id", inplace=True)
+    if "subnetwork_id" in gdf.columns and schema_version == 0:
+        warnings.warn("Migrating outdated Edge table.", UserWarning)
+        gdf.drop(columns="subnetwork_id", inplace = True, errors = "ignore")
 
     return gdf
 
diff --git a/python/ribasim/ribasim/model.py b/python/ribasim/ribasim/model.py
index d46b648b0..d50c95d58 100644
--- a/python/ribasim/ribasim/model.py
+++ b/python/ribasim/ribasim/model.py
@@ -71,8 +71,6 @@
 except ImportError:
     xugrid = MissingOptionalModule("xugrid")
 
-warnings.simplefilter("always", DeprecationWarning)
-
 
 class Model(FileModel):
     """A model of inland water resources systems."""
@@ -307,12 +305,6 @@ def write(self, filepath: str | PathLike[str]) -> Path:
         if self.use_validation:
             self._validate_model()
 
-        if self.edge.df is not None and self.edge.df["subnetwork_id"].notna().any():
-            warnings.warn(
-                "Sources for allocation are automatically inferred and no longer have to be specified in the `Edge` table.",
-                DeprecationWarning,
-            )
-
         filepath = Path(filepath)
         self.filepath = filepath
         if not filepath.suffix == ".toml":

From 57cbe51a5ced22e97d6a4bdf35a40430b7b2234b Mon Sep 17 00:00:00 2001
From: Martijn Visser <mgvisser@gmail.com>
Date: Mon, 2 Dec 2024 11:29:14 +0100
Subject: [PATCH 09/11] Fix linting job

---
 python/ribasim/ribasim/migrations.py | 2 +-
 python/ribasim/ribasim/model.py      | 1 -
 2 files changed, 1 insertion(+), 2 deletions(-)

diff --git a/python/ribasim/ribasim/migrations.py b/python/ribasim/ribasim/migrations.py
index 239effaa8..5f3e09d3c 100644
--- a/python/ribasim/ribasim/migrations.py
+++ b/python/ribasim/ribasim/migrations.py
@@ -29,7 +29,7 @@ def edgeschema_migration(gdf: GeoDataFrame, schema_version: int) -> GeoDataFrame
         gdf.set_index("edge_id", inplace=True)
     if "subnetwork_id" in gdf.columns and schema_version == 0:
         warnings.warn("Migrating outdated Edge table.", UserWarning)
-        gdf.drop(columns="subnetwork_id", inplace = True, errors = "ignore")
+        gdf.drop(columns="subnetwork_id", inplace=True, errors="ignore")
 
     return gdf
 
diff --git a/python/ribasim/ribasim/model.py b/python/ribasim/ribasim/model.py
index d50c95d58..d7274a67c 100644
--- a/python/ribasim/ribasim/model.py
+++ b/python/ribasim/ribasim/model.py
@@ -1,7 +1,6 @@
 import datetime
 import logging
 import shutil
-import warnings
 from collections.abc import Generator
 from os import PathLike
 from pathlib import Path

From de9ec5b06cd7a6d017eebb9372419206f639fa9c Mon Sep 17 00:00:00 2001
From: Martijn Visser <mgvisser@gmail.com>
Date: Mon, 2 Dec 2024 13:59:16 +0100
Subject: [PATCH 10/11] Remove `Edge.subnetwork_id` from the data model

---
 core/src/graph.jl                       | 3 +--
 core/src/model.jl                       | 1 +
 core/src/validation.jl                  | 8 ++++++++
 docs/changelog.qmd                      | 3 +++
 python/ribasim/ribasim/geometry/edge.py | 5 ++---
 python/ribasim/ribasim/migrations.py    | 2 +-
 python/ribasim/tests/test_model.py      | 1 -
 ribasim_qgis/core/nodes.py              | 1 -
 8 files changed, 16 insertions(+), 8 deletions(-)

diff --git a/core/src/graph.jl b/core/src/graph.jl
index 666591e8e..e2a3273e0 100644
--- a/core/src/graph.jl
+++ b/core/src/graph.jl
@@ -19,8 +19,7 @@ function create_graph(db::DB, config::Config)::MetaGraph
             FromNode.node_type AS from_node_type,
             ToNode.node_id AS to_node_id,
             ToNode.node_type AS to_node_type,
-            Edge.edge_type,
-            Edge.subnetwork_id
+            Edge.edge_type
         FROM Edge
         LEFT JOIN Node AS FromNode ON FromNode.node_id = Edge.from_node_id
         LEFT JOIN Node AS ToNode ON ToNode.node_id = Edge.to_node_id
diff --git a/core/src/model.jl b/core/src/model.jl
index e0047ca0d..2197d72b7 100644
--- a/core/src/model.jl
+++ b/core/src/model.jl
@@ -47,6 +47,7 @@ function Model(config::Config)::Model
     # so we can directly close it again.
     db = SQLite.DB(db_path)
 
+    database_warning(db)
     if !valid_nodes(db)
         error("Invalid nodes found.")
     end
diff --git a/core/src/validation.jl b/core/src/validation.jl
index 3eba4339a..81a372421 100644
--- a/core/src/validation.jl
+++ b/core/src/validation.jl
@@ -182,6 +182,14 @@ function valid_nodes(db::DB)::Bool
     return !errors
 end
 
+function database_warning(db::DB)::Nothing
+    cols = SQLite.columns(db, "Edge")
+    if "subnetwork_id" in cols.name
+        @warn "The 'subnetwork_id' column in the 'Edge' table is deprecated since ribasim v2024.12."
+    end
+    return nothing
+end
+
 """
 Test for each node given its node type whether the nodes that
 # are downstream ('down-edge') of this node are of an allowed type
diff --git a/docs/changelog.qmd b/docs/changelog.qmd
index d9c04350e..373ab6cc8 100644
--- a/docs/changelog.qmd
+++ b/docs/changelog.qmd
@@ -5,6 +5,9 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 
 ## [Unreleased]
 
+### Changed
+- The Edge table no longer supports `subnetwork_id`; this is automatically inferred. [#1956](https://github.com/Deltares/Ribasim/pull/1956)
+
 ## [v2024.11.0] - 2024-10-08
 
 This major new release contains many improvements.
diff --git a/python/ribasim/ribasim/geometry/edge.py b/python/ribasim/ribasim/geometry/edge.py
index 62a73cb28..faa2e92fa 100644
--- a/python/ribasim/ribasim/geometry/edge.py
+++ b/python/ribasim/ribasim/geometry/edge.py
@@ -47,7 +47,6 @@ class EdgeSchema(_GeoBaseSchema):
     from_node_id: Series[Int32] = pa.Field(default=0)
     to_node_id: Series[Int32] = pa.Field(default=0)
     edge_type: Series[str] = pa.Field(default="flow")
-    subnetwork_id: Series[pd.Int32Dtype] = pa.Field(default=pd.NA, nullable=True)
     geometry: GeoSeries[LineString] = pa.Field(default=None, nullable=True)
 
     @classmethod
@@ -73,7 +72,6 @@ def add(
         to_node: NodeData,
         geometry: LineString | MultiLineString | None = None,
         name: str = "",
-        subnetwork_id: int | None = None,
         edge_id: Optional[NonNegativeInt] = None,
         **kwargs,
     ):
@@ -90,6 +88,8 @@ def add(
             The geometry of a line. If not supplied, it creates a straight line between the nodes.
         name : str
             An optional name for the edge.
+        id : int
+            An optional non-negative edge ID. If not supplied, it will be automatically generated.
         **kwargs : Dict
         """
         if not can_connect(from_node.node_type, to_node.node_type):
@@ -120,7 +120,6 @@ def add(
                 "to_node_id": [to_node.node_id],
                 "edge_type": [edge_type],
                 "name": [name],
-                "subnetwork_id": [subnetwork_id],
                 **kwargs,
             },
             geometry=geometry_to_append,
diff --git a/python/ribasim/ribasim/migrations.py b/python/ribasim/ribasim/migrations.py
index 5f3e09d3c..31a326e83 100644
--- a/python/ribasim/ribasim/migrations.py
+++ b/python/ribasim/ribasim/migrations.py
@@ -27,7 +27,7 @@ def edgeschema_migration(gdf: GeoDataFrame, schema_version: int) -> GeoDataFrame
         warnings.warn("Migrating outdated Edge table.", UserWarning)
         assert gdf["edge_id"].is_unique, "Edge IDs have to be unique."
         gdf.set_index("edge_id", inplace=True)
-    if "subnetwork_id" in gdf.columns and schema_version == 0:
+    if "subnetwork_id" in gdf.columns:
         warnings.warn("Migrating outdated Edge table.", UserWarning)
         gdf.drop(columns="subnetwork_id", inplace=True, errors="ignore")
 
diff --git a/python/ribasim/tests/test_model.py b/python/ribasim/tests/test_model.py
index ba703d5fa..f5e341561 100644
--- a/python/ribasim/tests/test_model.py
+++ b/python/ribasim/tests/test_model.py
@@ -145,7 +145,6 @@ def test_edge_table(basic):
     df = model.edge.df
     assert df.geometry.is_unique
     assert df.from_node_id.dtype == np.int32
-    assert df.subnetwork_id.dtype == pd.Int32Dtype()
     assert df.crs == CRS.from_epsg(28992)
 
 
diff --git a/ribasim_qgis/core/nodes.py b/ribasim_qgis/core/nodes.py
index 9a1859c77..c638a58e0 100644
--- a/ribasim_qgis/core/nodes.py
+++ b/ribasim_qgis/core/nodes.py
@@ -253,7 +253,6 @@ def attributes(cls) -> list[QgsField]:
             QgsField("from_node_id", QVariant.Int),
             QgsField("to_node_id", QVariant.Int),
             QgsField("edge_type", QVariant.String),
-            QgsField("subnetwork_id", QVariant.Int),
         ]
 
     @classmethod

From bc9eeaec4e5f4d0b54d981a1b9e48acfbd98e800 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart123-4@hotmail.com>
Date: Tue, 3 Dec 2024 14:59:01 +0100
Subject: [PATCH 11/11] Comments adressed

---
 core/src/callback.jl | 5 ++++-
 core/src/read.jl     | 7 +++++--
 2 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/core/src/callback.jl b/core/src/callback.jl
index cda123cff..467e71060 100644
--- a/core/src/callback.jl
+++ b/core/src/callback.jl
@@ -802,11 +802,14 @@ function update_allocation!(integrator)::Nothing
     end
 
     # Reset the mean flows
-    for mean_flows in (mean_input_flows..., mean_realized_flows)
+    for mean_flows in mean_input_flows
         for edge in keys(mean_flows)
             mean_flows[edge] = 0.0
         end
     end
+    for edge in keys(mean_realized_flows)
+        mean_realized_flows[edge] = 0.0
+    end
 end
 
 "Load updates from 'TabulatedRatingCurve / time' into the parameters"
diff --git a/core/src/read.jl b/core/src/read.jl
index 17e020b78..e99d207e0 100644
--- a/core/src/read.jl
+++ b/core/src/read.jl
@@ -1283,8 +1283,11 @@ function Allocation(db::DB, config::Config, graph::MetaGraph)::Allocation
             id_source, _ = edge
             if id_source.type in boundary_source_nodetypes
                 (; subnetwork_id) = graph[id_source]
-                subnetwork_idx = searchsortedfirst(subnetwork_ids, subnetwork_id)
-                mean_input_flows[subnetwork_idx][edge] = 0.0
+                # Check whether the source node is part of a subnetwork
+                if subnetwork_id ≠ 0
+                    subnetwork_idx = searchsortedfirst(subnetwork_ids, subnetwork_id)
+                    mean_input_flows[subnetwork_idx][edge] = 0.0
+                end
             end
         end