A Boolean for allocation controlled (#2912)

fixes #2889

example of an outlet that switches between distinct allocation states

<img width="596" height="396" alt="Screenshot 2026-02-26 at 09 50 06"
src="https://github.com/user-attachments/assets/3e1e19fd-56ac-4f25-88b8-dea4064ec606"
/>

---------

Co-authored-by: Martijn Visser <mgvisser@gmail.com>

Author: verheem

.gitignore

@@ -172,3 +172,5 @@ debug/*
 
 .DS_Store
 lock.pid
+
+CLAUDE.md

core/src/allocation_init.jl

@@ -47,21 +47,9 @@ function add_flow!(
         allocation_model::AllocationModel,
         p_independent::ParametersIndependent,
     )::Nothing
-    (; problem, subnetwork_id, scaling) = allocation_model
+    (; problem, subnetwork_id, scaling, flow_links_subnetwork) = allocation_model
     (; graph) = p_independent
 
-    node_ids_subnetwork = graph[].node_ids[subnetwork_id]
-    flow_links_subnetwork = Vector{Tuple{NodeID, NodeID}}()
-
-    # Sort link metadata for deterministic problem generation
-    for link_metadata in sort!(collect(values(graph.edge_data)))
-        (; type, link) = link_metadata
-        if (type == LinkType.flow) &&
-                ((link[1] ∈ node_ids_subnetwork) || (link[2] ∈ node_ids_subnetwork))
-            push!(flow_links_subnetwork, link)
-        end
-    end
-
     # Define decision variables: flow over flow links (scaling.flow * m^3/s)
     problem[:flow] = JuMP.@variable(
         problem,
@@ -73,6 +61,7 @@ function add_flow!(
     return nothing
 end
 
+
 """
 Add flow conservation constraints for conservative nodes.
 Ensures that inflow equals outflow for nodes that conserve mass (pumps, outlets,
@@ -1004,6 +993,20 @@ function AllocationModel(
         end
     end
 
+    (; graph) = p_independent
+
+    node_ids_subnetwork = graph[].node_ids[subnetwork_id]
+    flow_links_subnetwork = Vector{Tuple{NodeID, NodeID}}()
+
+    # Sort link metadata for deterministic problem generation
+    for link_metadata in sort!(collect(values(graph.edge_data)))
+        (; type, link) = link_metadata
+        if (type == LinkType.flow) &&
+                ((link[1] ∈ node_ids_subnetwork) || (link[2] ∈ node_ids_subnetwork))
+            push!(flow_links_subnetwork, link)
+        end
+    end
+
     allocation_model = AllocationModel(;
         subnetwork_id,
         node_ids_in_subnetwork,
@@ -1012,6 +1015,7 @@ function AllocationModel(
         scaling,
         has_demand_priority,
         secondary_network_demand,
+        flow_links_subnetwork,
     )
 
     # Volume and flow

core/src/allocation_optim.jl

@@ -1039,7 +1039,7 @@ function save_flows!(
     return nothing
 end
 
-# Set the flow rate of allocation controlled pumps and outlets to
+# Set the flow rate of allocation controlled pumps, outlets and rating curves to
 # their flow determined by allocation
 function apply_control_from_allocation!(
         node::Union{Pump, Outlet, TabulatedRatingCurve},
@@ -1100,14 +1100,38 @@ function delete_control_constraints!(
     return nothing
 end
 
+function update_flow_variable_bounds!(
+        allocation_model::AllocationModel,
+        p_independent::ParametersIndependent,
+    )::Nothing
+    (; problem, scaling, flow_links_subnetwork) = allocation_model
+    flow = problem[:flow]
+    for flow_link in flow_links_subnetwork
+        flow_var = flow[flow_link]
+        JuMP.is_fixed(flow_var) && continue
+        JuMP.set_lower_bound(
+            flow_var,
+            flow_capacity_lower_bound(flow_link, p_independent) / scaling.flow,
+        )
+        JuMP.set_upper_bound(
+            flow_var,
+            flow_capacity_upper_bound(flow_link, p_independent) / scaling.flow,
+        )
+    end
+    return nothing
+end
+
 function update_control_states!(
         allocation_model::AllocationModel,
         p_independent::ParametersIndependent,
     )::Nothing
     delete_control_constraints!(allocation_model, :pump)
     delete_control_constraints!(allocation_model, :outlet)
+    delete_control_constraints!(allocation_model, :tabulated_rating_curve_constraint)
     add_pump!(allocation_model, p_independent)
     add_outlet!(allocation_model, p_independent)
+    add_tabulated_rating_curve!(allocation_model, p_independent)
+    update_flow_variable_bounds!(allocation_model, p_independent)
     return nothing
 end
 

core/src/allocation_util.jl

@@ -502,3 +502,10 @@ function get_primary_network(allocation_models::Vector{AllocationModel})::Alloca
     end
     error("Queries primary network while no primary network found in allocation models.")
 end
+
+function delete_flow!(
+        allocation_model::AllocationModel
+    )::Nothing
+    (; problem) = allocation_model
+    return JuMP.delete(problem, problem[:flow])
+end

core/src/callback.jl

@@ -685,22 +685,19 @@ function set_new_control_state!(
         for target_node_id in discrete_control.controlled_nodes[discrete_control_id.idx]
             set_control_params!(p, target_node_id, control_state_new)
 
-            if control_state_new == "Ribasim.allocation"
-                if target_node_id.type == NodeType.Pump
-                    pump.allocation_controlled[target_node_id.idx] = true
-                elseif target_node_id.type == NodeType.Outlet
-                    outlet.allocation_controlled[target_node_id.idx] = true
-                elseif target_node_id.type == NodeType.TabulatedRatingCurve
-                    tabulated_rating_curve.allocation_controlled[target_node_id.idx] = true
-                end
-            else
-                if target_node_id.type == NodeType.Pump
-                    pump.allocation_controlled[target_node_id.idx] = false
-                elseif target_node_id.type == NodeType.Outlet
-                    outlet.allocation_controlled[target_node_id.idx] = false
-                elseif target_node_id.type == NodeType.TabulatedRatingCurve
-                    tabulated_rating_curve.allocation_controlled[target_node_id.idx] = false
-                end
+            # Update allocation_controlled based on the new control state
+            if target_node_id.type == NodeType.Pump
+                control_state_update = pump.control_mapping[(target_node_id, control_state_new)]
+                pump.allocation_controlled[target_node_id.idx] =
+                    control_state_update.allocation_controlled
+            elseif target_node_id.type == NodeType.Outlet
+                control_state_update = outlet.control_mapping[(target_node_id, control_state_new)]
+                outlet.allocation_controlled[target_node_id.idx] =
+                    control_state_update.allocation_controlled
+            elseif target_node_id.type == NodeType.TabulatedRatingCurve
+                control_state_update = tabulated_rating_curve.control_mapping[(target_node_id, control_state_new)]
+                tabulated_rating_curve.allocation_controlled[target_node_id.idx] =
+                    control_state_update.allocation_controlled
             end
         end
 

core/src/parameter.jl

@@ -259,6 +259,7 @@ scaling: The flow and storage scaling factors to make the optimization problem m
     sources::OrderedDict{Int32, NodeID} = OrderedDict()
     secondary_network_demand::OrderedDict{Tuple{NodeID, NodeID}, Vector{Float64}} =
         OrderedDict()
+    flow_links_subnetwork::Vector{Tuple{NodeID, NodeID}} = Vector{Tuple{NodeID, NodeID}}()
     scaling::ScalingFactors = ScalingFactors()
     temporary_constraints::Vector{JuMP.ConstraintRef} = JuMP.ConstraintRef[]
     route_priority_expression::JuMP.AffExpr = JuMP.AffExpr()
@@ -370,6 +371,7 @@ The parameter update associated with a certain control state for discrete contro
     itp_update_linear::Vector{ParameterUpdate{ScalarLinearInterpolation}} =
         ParameterUpdate{ScalarLinearInterpolation}[]
     itp_update_lookup::Vector{ParameterUpdate{IndexLookup}} = ParameterUpdate{IndexLookup}[]
+    allocation_controlled::Bool = false
 end
 
 """

core/src/read.jl

@@ -329,6 +329,8 @@ parse_control_states!(::BasinForcing, args...) = nothing
 
 function initialize_control_mapping!(node::AbstractParameterNode, static::StructVector)
     isempty(static) && return
+    has_alloc_controlled = hasproperty(first(static), :allocation_controlled)
+    node_has_alloc_controlled = hasfield(typeof(node), :allocation_controlled)
     static_groups = IterTools.groupby(row -> row.node_id, static)
     static_group, static_idx = iterate(static_groups)
     for node_id in node.node_id
@@ -338,8 +340,23 @@ function initialize_control_mapping!(node::AbstractParameterNode, static::Struct
                 IterTools.groupby(row -> coalesce(row.control_state, nothing), static_group)
                 first_row = first(control_state_group)
                 control_state = first_row.control_state
+
+                # Read allocation_controlled from static data if available
+                alloc_controlled = if has_alloc_controlled
+                    coalesce(first_row.allocation_controlled, false)
+                else
+                    false
+                end
+
+                # Set the node's allocation_controlled flag if any row has it set,
+                # so we always start with allocation enabled for all allocation controlled nodes
+                if alloc_controlled && node_has_alloc_controlled
+                    node.allocation_controlled[node_id.idx] = true
+                end
+
                 ismissing(control_state) && continue
-                node.control_mapping[(node_id, control_state)] = ControlStateUpdate()
+                node.control_mapping[(node_id, control_state)] =
+                    ControlStateUpdate(; allocation_controlled = alloc_controlled)
             end
             if static_idx[1]
                 static_group, static_idx = iterate(static_groups, static_idx)

core/src/schema.jl

@@ -89,6 +89,7 @@ module Schema
             min_upstream_level::Union{Missing, Float64}
             max_downstream_level::Union{Missing, Float64}
             control_state::Union{Missing, String}
+            allocation_controlled::Union{Missing, Bool}
         end
 
         struct Time <: Table
@@ -115,6 +116,7 @@ module Schema
             min_upstream_level::Union{Missing, Float64}
             max_downstream_level::Union{Missing, Float64}
             control_state::Union{Missing, String}
+            allocation_controlled::Union{Missing, Bool}
         end
 
         struct Time <: Table
@@ -215,6 +217,7 @@ module Schema
             flow_rate::Float64
             max_downstream_level::Union{Missing, Float64}
             control_state::Union{Missing, String}
+            allocation_controlled::Union{Missing, Bool}
         end
 
         struct Time <: Table

core/src/util.jl

@@ -411,9 +411,6 @@ function set_control_type!(node::AbstractParameterNode, graph::MetaGraph)::Nothi
         else
             ContinuousControlType.None
         end
-
-        node.allocation_controlled[node_id.idx] =
-            (node_id, "Ribasim.allocation") in keys(control_mapping)
     end
 
     errors && error("Errors encountered when parsing control type of $(typeof(node)).")

core/src/validation.jl

@@ -1,7 +1,3 @@
-# Reserved control states that can be used by Control nodes
-# These are control states that start with "Ribasim." and are reserved for special behavior
-const RESERVED_CONTROL_STATES = Set(["Ribasim.allocation"])
-
 # Allowed types for downstream (to_node_id) nodes given the type of the upstream (from_node_id) node
 neighbortypes(nodetype::Symbol) = neighbortypes(Val(config.snake_case(nodetype)))
 neighbortypes(::Val{:pump}) = OrderedSet((:basin, :terminal, :level_boundary, :junction))
@@ -553,25 +549,6 @@ function valid_link_types(db::DB)::Bool
     return !errors
 end
 
-"""
-Validate that all control states starting with "Ribasim." are reserved and supported.
-"""
-function valid_reserved_control_states(control_states::Set{String})::Bool
-    errors = false
-
-    for control_state in control_states
-        # Check if control state starts with "Ribasim." prefix
-        if startswith(control_state, "Ribasim.")
-            # If it does, it must be in the reserved set
-            if control_state ∉ RESERVED_CONTROL_STATES
-                errors = true
-                @error "Unknown reserved control state '$control_state'. Control states starting with 'Ribasim.' must be one of: $(join(RESERVED_CONTROL_STATES, ", "))"
-            end
-        end
-    end
-
-    return !errors
-end
 
 """
 Check:
@@ -586,9 +563,6 @@ function valid_discrete_control(p::ParametersIndependent, config::Config)::Bool
     t_end = seconds_since(config.endtime, config.starttime)
     errors = false
 
-    # Collect all control states to validate reserved ones
-    all_control_states = Set{String}()
-
     for (id, compound_variables) in zip(node_id, discrete_control.compound_variables)
 
         # The number of conditions of this DiscreteControl node
@@ -605,7 +579,6 @@ function valid_discrete_control(p::ParametersIndependent, config::Config)::Bool
 
         for (truth_state, control_state) in logic_mapping[id.idx]
             push!(control_states_discrete_control, control_state)
-            push!(all_control_states, control_state)
 
             if length(truth_state) != n_conditions
                 push!(truth_states_wrong_length, truth_state)
@@ -683,11 +656,6 @@ function valid_discrete_control(p::ParametersIndependent, config::Config)::Bool
         end
     end
 
-    # Validate reserved control states
-    if !valid_reserved_control_states(all_control_states)
-        errors = true
-    end
-
     return !errors
 end