Set capacity of outflow of node with flow demand

core/src/allocation_init.jl

@@ -251,7 +251,7 @@ function process_allocation_graph_edges!(
 
             # Find flow constraints
             if is_flow_constraining(node)
-                problem_node_idx = Ribasim.findsorted(node.node_id, node_id_2)
+                problem_node_idx = findsorted(node.node_id, node_id_2)
                 edge_capacity = min(edge_capacity, node.max_flow_rate[problem_node_idx])
             end
 
@@ -893,6 +893,25 @@ function add_constraints_buffer!(problem::JuMP.Model)::Nothing
     return nothing
 end
 
+function add_constraints_flow_demand_outflow!(
+    problem::JuMP.Model,
+    p::Parameters,
+    allocation_network_id::Int,
+)::Nothing
+    (; graph) = p
+    F = problem[:F]
+    node_ids = graph[].node_ids[allocation_network_id]
+    node_ids_flow_demand = [
+        node_id for node_id in node_ids if has_external_demand(graph, node_id, :flow_demand)
+    ]
+    problem[:flow_demand_outflow] = JuMP.@constraint(
+        problem,
+        [node_id = node_ids_flow_demand],
+        F[(node_id, only(outflow_ids_allocation(graph, node_id)))] <= 0.0
+    )
+    return nothing
+end
+
 """
 Construct the allocation problem for the current subnetwork as a JuMP.jl model.
 """
@@ -920,6 +939,7 @@ function allocation_problem(
     add_constraints_absolute_value_basin!(problem, config)
     add_constraints_fractional_flow!(problem, p, allocation_network_id)
     add_constraints_basin_flow!(problem)
+    add_constraints_flow_demand_outflow!(problem, p, allocation_network_id)
     add_constraints_buffer!(problem)
 
     return problem

core/src/allocation_optim.jl

@@ -497,6 +497,33 @@ function adjust_buffers!(allocation_model::AllocationModel, priority_idx::Int)::
     return nothing
 end
 
+function set_capacity_flow_demand_outflow!(
+    allocation_model::AllocationModel,
+    p::Parameters,
+    priority_idx::Int,
+)::Nothing
+    (; graph, allocation, flow_demand) = p
+    (; priorities) = allocation
+    (; problem) = allocation_model
+    priority = priorities[priority_idx]
+    constraints = problem[:flow_demand_outflow]
+
+    for node_id in only(constraints.axes)
+        constraint = constraints[node_id]
+        node_id_flow_demand = only(inneighbor_labels_type(graph, node_id, EdgeType.control))
+        node_idx = findsorted(flow_demand.node_id, node_id_flow_demand)
+        priority_flow_demand = flow_demand.priority[node_idx]
+
+        capacity = if priority == priority_flow_demand
+            0.0
+        else
+            Inf
+        end
+
+        JuMP.set_normalized_rhs(constraint, capacity)
+    end
+end
+
 """
 Save the demands and allocated flows for UserDemand and Basin.
 Note: Basin supply (negative demand) is only saved for the first priority.
@@ -664,6 +691,8 @@ function allocate!(
         # https://jump.dev/JuMP.jl/v1.16/manual/objective/#Modify-an-objective-coefficient
         set_objective_priority!(allocation_model, p, u, t, priority_idx)
 
+        set_capacity_flow_demand_outflow!(allocation_model, p, priority_idx)
+
         # Solve the allocation problem for this priority
         JuMP.optimize!(problem)
         @debug JuMP.solution_summary(problem)

core/test/allocation_test.jl

@@ -368,6 +368,7 @@ end
 @testitem "flow_demand" begin
     using JuMP
     using Ribasim: NodeID, NodeType
+    using MathOptInterface
 
     toml_path = normpath(@__DIR__, "../../generated_testmodels/flow_demand/ribasim.toml")
     @test ispath(toml_path)
@@ -395,14 +396,18 @@ end
 
     # Test flow conservation constraint containing flow buffer
     constraint_with_flow_buffer = JuMP.constraint_object(
-        allocation_model.problem[:flow_conservation_flow_demand][NodeID(
-            NodeType.TabulatedRatingCurve,
-            2,
-        )],
+        problem[:flow_conservation_flow_demand][NodeID(NodeType.TabulatedRatingCurve, 2)],
     )
     @test constraint_with_flow_buffer.func ==
           F[(NodeID(NodeType.TabulatedRatingCurve, 2), NodeID(NodeType.Basin, 3))] -
           F[(NodeID(NodeType.LevelBoundary, 1), NodeID(NodeType.TabulatedRatingCurve, 2))] +
           F_flow_buffer_in[NodeID(NodeType.TabulatedRatingCurve, 2)] -
           F_flow_buffer_out[NodeID(NodeType.TabulatedRatingCurve, 2)]
+
+    constraint_flow_demand_outflow = JuMP.constraint_object(
+        problem[:flow_demand_outflow][NodeID(NodeType.TabulatedRatingCurve, 2)],
+    )
+    @test constraint_flow_demand_outflow.func ==
+          F[(NodeID(NodeType.TabulatedRatingCurve, 2), NodeID(NodeType.Basin, 3))] + 0.0
+    @test constraint_flow_demand_outflow.set == MathOptInterface.LessThan{Float64}(0.0)
 end