Get rid of specifying sources in allocation by edges (#1956)

Fixes #1953. To do: come up with a unified way to handle deprecation, and incorporate this in the developer docs and release process. Maybe also use a 'deprecation' label instead of a 'breaking' label? --------- Co-authored-by: Martijn Visser <[email protected]>
Deltares · Dec 3, 2024 · 5b5e99f · 5b5e99f
1 parent 9311290
commit 5b5e99f
Show file tree

Hide file tree

Showing 22 changed files with 229 additions and 264 deletions.
diff --git a/core/src/allocation_init.jl b/core/src/allocation_init.jl
@@ -95,7 +95,7 @@ function get_subnetwork_capacity(
     return capacity
 end
 
-const allocation_source_nodetypes =
+const boundary_source_nodetypes =
     Set{NodeType.T}([NodeType.LevelBoundary, NodeType.FlowBoundary])
 
 """
@@ -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
 
@@ -272,36 +268,55 @@ 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_source!(
+function add_constraints_boundary_source!(
     problem::JuMP.Model,
     p::Parameters,
     subnetwork_id::Int32,
 )::Nothing
-    (; graph) = p
-    edges_source = Tuple{NodeID, NodeID}[]
+    # 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]
 
-    # 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
-            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
 
-    problem[:source] = JuMP.@constraint(
+"""
+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,
+    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_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
@@ -451,9 +466,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)
@@ -484,12 +499,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
@@ -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
@@ -203,22 +202,13 @@ function set_initial_capacities_source!(
     p::Parameters,
 )::Nothing
     (; sources) = allocation_model
-    (; graph, 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
-        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
@@ -231,7 +221,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,
         )
@@ -343,11 +333,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)
@@ -841,9 +830,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)
@@ -860,16 +850,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
@@ -1062,7 +1053,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
@@ -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,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/graph.jl b/core/src/graph.jl
@@ -19,17 +19,15 @@ 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
         """,
     )
     # 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[]
@@ -57,15 +55,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))
@@ -74,15 +65,8 @@ 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)
-            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
@@ -91,12 +75,6 @@ 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")
@@ -106,7 +84,7 @@ function create_graph(db::DB, config::Config)::MetaGraph
         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/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