Allocation fractional flow constraints (#854)

Fixes #727.

core/src/allocation.jl

@@ -3,15 +3,18 @@ Find all nodes in the subnetwork which will be used in the allocation network.
 Some nodes are skipped to optimize allocation optimization.
 """
 function allocation_graph_used_nodes!(p::Parameters, allocation_network_id::Int)::Nothing
-    (; graph) = p
+    (; graph, basin, fractional_flow) = p
 
     node_ids = graph[].node_ids[allocation_network_id]
     used_nodes = Set{NodeID}()
-
     for node_id in node_ids
+        has_fractional_flow_outneighbors =
+            get_fractional_flow_connected_basins(node_id, basin, fractional_flow, graph)[3]
         node_type = graph[node_id].type
         if node_type in [:user, :basin]
             push!(used_nodes, node_id)
+        elseif has_fractional_flow_outneighbors
+            push!(used_nodes, node_id)
         elseif count(x -> true, inoutflow_ids(graph, node_id)) > 2
             # use count since the length of the iterator is unknown
             push!(used_nodes, node_id)
@@ -378,7 +381,7 @@ function add_constraints_capacity!(
     F = problem[:F]
     edge_ids = graph[].edge_ids[allocation_network_id]
     edge_ids_finite_capacity = Tuple{NodeID, NodeID}[]
-    for (i, edge) in enumerate(edge_ids)
+    for edge in edge_ids
         if !isinf(capacity[edge...])
             push!(edge_ids_finite_capacity, edge)
         end
@@ -577,6 +580,57 @@ function add_constraints_absolute_value!(
     return nothing
 end
 
+"""
+Add the fractional flow constraints to the allocation problem.
+The constraint indices are allocation edges over a fractional flow node.
+
+Constraint:
+flow after fractional_flow node <= fraction * inflow
+"""
+function add_constraints_fractional_flow!(
+    problem::JuMP.Model,
+    p::Parameters,
+    allocation_network_id::Int,
+)::Nothing
+    (; graph, fractional_flow) = p
+    F = problem[:F]
+    node_ids = graph[].node_ids[allocation_network_id]
+
+    edges_to_fractional_flow = Tuple{NodeID, NodeID}[]
+    fractions = Dict{Tuple{NodeID, NodeID}, Float64}()
+    inflows = Dict{NodeID, JuMP.AffExpr}()
+    for node_id in node_ids
+        for outflow_id_ in outflow_ids(graph, node_id)
+            if graph[outflow_id_].type == :fractional_flow
+                # The fractional flow nodes themselves are not represented in
+                # the allocation graph
+                dst_id = outflow_id(graph, outflow_id_)
+                # For now only consider fractional flow nodes which end in a basin
+                if haskey(graph, node_id, dst_id) && graph[dst_id].type == :basin
+                    edge = (node_id, dst_id)
+                    push!(edges_to_fractional_flow, edge)
+                    node_idx = findsorted(fractional_flow.node_id, outflow_id_)
+                    fractions[edge] = fractional_flow.fraction[node_idx]
+                    inflows[node_id] = sum([
+                        F[(inflow_id_, node_id)] for
+                        inflow_id_ in inflow_ids(graph, node_id)
+                    ])
+                end
+            end
+        end
+    end
+
+    if !isempty(edges_to_fractional_flow)
+        problem[:fractional_flow] = JuMP.@constraint(
+            problem,
+            [edge = edges_to_fractional_flow],
+            F[edge] <= fractions[edge] * inflows[edge[1]],
+            base_name = "fractional_flow"
+        )
+    end
+    return nothing
+end
+
 """
 Construct the allocation problem for the current subnetwork as a JuMP.jl model.
 """
@@ -600,7 +654,7 @@ function allocation_problem(
     add_constraints_flow_conservation!(problem, p, allocation_network_id)
     add_constraints_user_returnflow!(problem, p, allocation_network_id)
     add_constraints_absolute_value!(problem, p, allocation_network_id, config)
-    # TODO: The fractional flow constraints
+    add_constraints_fractional_flow!(problem, p, allocation_network_id)
 
     return problem
 end

core/src/bmi.jl

@@ -32,10 +32,6 @@ function BMI.initialize(T::Type{Model}, config::Config)::Model
     try
         parameters = Parameters(db, config)
 
-        if !valid_n_neighbors(parameters)
-            error("Invalid number of connections for certain node types.")
-        end
-
         if !valid_discrete_control(parameters, config)
             error("Invalid discrete control state definition(s).")
         end
@@ -485,6 +481,48 @@ function discrete_control_affect!(
     return control_state_change
 end
 
+function get_allocation_model(
+    p::Parameters,
+    allocation_network_id::Int,
+)::Union{AllocationModel, Nothing}
+    for allocation_model in p.allocation_models
+        if allocation_model.allocation_network_id == allocation_network_id
+            return allocation_model
+        end
+    end
+    return nothing
+end
+
+"""
+Update the fractional flow fractions in an allocation problem.
+"""
+function set_fractional_flow_in_allocation!(
+    p::Parameters,
+    node_id::NodeID,
+    fraction::Number,
+)::Nothing
+    (; graph) = p
+
+    allocation_network_id = graph[node_id].allocation_network_id
+    # Get the allocation model this fractional flow node is in
+    allocation_model = get_allocation_model(p, allocation_network_id)
+    if !isnothing(allocation_model)
+        problem = allocation_model.problem
+        # The allocation edge which jumps over the fractional flow node
+        edge = (inflow_id(graph, node_id), outflow_id(graph, node_id))
+        if haskey(graph, edge...)
+            # The constraint for this fractional flow node
+            constraint = problem[:fractional_flow][edge]
+            # Set the new fraction on all inflow terms in the constraint
+            for inflow_id in inflow_ids_allocation(graph, edge[1])
+                flow = problem[:F][(inflow_id, edge[1])]
+                JuMP.set_normalized_coefficient(constraint, flow, -fraction)
+            end
+        end
+    end
+    return nothing
+end
+
 function set_control_params!(p::Parameters, node_id::NodeID, control_state::String)
     node = getfield(p, p.graph[node_id].type)
     idx = searchsortedfirst(node.node_id, node_id)
@@ -495,6 +533,11 @@ function set_control_params!(p::Parameters, node_id::NodeID, control_state::Stri
             vec = get_tmp(getfield(node, field), 0)
             vec[idx] = value
         end
+
+        # Set new fractional flow fractions in allocation problem
+        if node isa FractionalFlow && field == :fraction
+            set_fractional_flow_in_allocation!(p, node_id, value)
+        end
     end
 end
 

core/src/create.jl

@@ -846,6 +846,11 @@ function Parameters(db::DB, config::Config)::Parameters
         Dict{Int, Symbol}(),
         subgrid_level,
     )
+
+    if !valid_n_neighbors(p)
+        error("Invalid number of connections for certain node types.")
+    end
+
     # Allocation data structures
     if config.allocation.use_allocation
         generate_allocation_models!(p, config)

core/src/utils.jl

@@ -1206,7 +1206,7 @@ is_flow_constraining(node::AbstractParameterNode) = hasfield(typeof(node), :max_
 
 """Whether the given node is flow direction constraining (only in direction of edges)."""
 is_flow_direction_constraining(node::AbstractParameterNode) =
-    (nameof(typeof(node)) ∈ [:Pump, :Outlet, :TabulatedRatingCurve])
+    (nameof(typeof(node)) ∈ [:Pump, :Outlet, :TabulatedRatingCurve, :FractionalFlow])
 
 """Find out whether a path exists between a start node and end node in the given allocation graph."""
 function allocation_path_exists_in_graph(

core/test/allocation_test.jl

@@ -109,3 +109,59 @@ end
     @test objective.terms[F[(NodeID(4), NodeID(5))]] ≈ 62.585499316005475
     @test objective.terms[F[(NodeID(2), NodeID(4))]] ≈ 62.585499316005475
 end
+
+@testitem "Allocation with controlled fractional flow" begin
+    using DataFrames
+    using Ribasim: NodeID
+    using OrdinaryDiffEq: solve!
+    using JuMP
+
+    toml_path = normpath(
+        @__DIR__,
+        "../../generated_testmodels/fractional_flow_subnetwork/ribasim.toml",
+    )
+    model = Ribasim.BMI.initialize(Ribasim.Model, toml_path)
+    problem = model.integrator.p.allocation_models[1].problem
+    F = problem[:F]
+    @test JuMP.normalized_coefficient(
+        problem[:fractional_flow][(NodeID(3), NodeID(5))],
+        F[(NodeID(2), NodeID(3))],
+    ) ≈ -0.25
+    @test JuMP.normalized_coefficient(
+        problem[:fractional_flow][(NodeID(3), NodeID(8))],
+        F[(NodeID(2), NodeID(3))],
+    ) ≈ -0.75
+
+    solve!(model)
+    record_allocation = DataFrame(model.integrator.p.user.record)
+    record_control = model.integrator.p.discrete_control.record
+    groups = groupby(record_allocation, [:user_node_id, :priority])
+    fractional_flow = model.integrator.p.fractional_flow
+    (; control_mapping) = fractional_flow
+    t_control = record_control.time[2]
+
+    allocated_6_before = groups[(6, 1)][groups[(6, 1)].time .< t_control, :].allocated
+    allocated_9_before = groups[(9, 1)][groups[(9, 1)].time .< t_control, :].allocated
+    allocated_6_after = groups[(6, 1)][groups[(6, 1)].time .> t_control, :].allocated
+    allocated_9_after = groups[(9, 1)][groups[(9, 1)].time .> t_control, :].allocated
+    @test all(
+        allocated_9_before ./ allocated_6_before .<=
+        control_mapping[(NodeID(7), "A")].fraction /
+        control_mapping[(NodeID(4), "A")].fraction,
+    )
+    @test all(allocated_9_after ./ allocated_6_after .<= 1.0)
+
+    @test record_control.truth_state == ["F", "T"]
+    @test record_control.control_state == ["A", "B"]
+
+    fractional_flow_constraints =
+        model.integrator.p.allocation_models[1].problem[:fractional_flow]
+    @test JuMP.normalized_coefficient(
+        problem[:fractional_flow][(NodeID(3), NodeID(5))],
+        F[(NodeID(2), NodeID(3))],
+    ) ≈ -0.75
+    @test JuMP.normalized_coefficient(
+        problem[:fractional_flow][(NodeID(3), NodeID(8))],
+        F[(NodeID(2), NodeID(3))],
+    ) ≈ -0.25
+end

docs/core/allocation.qmd

@@ -166,10 +166,6 @@ $$ {#eq-fractionalflowconstraint}
 
 - Flow sign: Furthermore there are the non-negativity constraints for the flows and allocations, see [The optimization variables](allocation.qmd#the-optimization-variables).
 
-:::{.callout-note}
-Currently the fractional flow constraints are not taken into account in the implementation.
-:::
-
 ## Final notes on the allocation problem
 
 ### Users using their own return flow

python/ribasim_testmodels/ribasim_testmodels/__init__.py

@@ -6,8 +6,9 @@
 
 import ribasim_testmodels
 from ribasim_testmodels.allocation import (
-    minimal_subnetwork_model,
     # looped_subnetwork_model,
+    fractional_flow_subnetwork_model,
+    minimal_subnetwork_model,
     subnetwork_model,
     user_model,
 )
@@ -78,6 +79,7 @@
     "user_model",
     "subnetwork_model",
     "minimal_subnetwork_model",
+    "fractional_flow_subnetwork_model",
     # Disable until this issue is resolved:
     # https://github.com/Deltares/Ribasim/issues/692
     # "looped_subnetwork_model",