Control from allocation

core/src/allocation_optim.jl

@@ -1009,9 +1009,38 @@
         demand_priorities_all[demand_priority_idx],
         optimization_type,
     )
+
+    if optimization_type == OptimizationType.allocate
+        apply_control_from_allocation!(p, allocation_model)
+    end
     return nothing
 end
 
+function apply_control_from_allocation!(p::Parameters, allocation_model::AllocationModel)
+    (; pump, outlet, graph) = p
+
+    apply_control_from_allocation!(pump, allocation_model, graph)
+    apply_control_from_allocation!(outlet, allocation_model, graph)
+end
+
+function apply_control_from_allocation!(
+    node::Union{Pump, Outlet},
+    allocation_model::AllocationModel,
+    graph::MetaGraph,
+)
+    (; node_id, control_type, inflow_link) = node
+    (; flow, subnetwork_id) = allocation_model
+    cache = node.flow_rate_cache[Float64[]] # TODO: Make non-allocating
+
+    for (id, c_type, link_in) in zip(node_id, control_type, inflow_link)
+        in_subnetwork = (graph[id].subnetwork_id == subnetwork_id)
+        allocation_controlled = (c_type == ControlType.Allocation)
+        if in_subnetwork && allocation_controlled
+            cache[id.idx] = flow[link_in.link]
+        end
+    end
+end
+
 """
 Set the initial capacities and demands which are reduced by usage.
 """

core/src/callback.jl

@@ -651,7 +651,7 @@ end
 function apply_parameter_update!(parameter_update)::Nothing
     (; name, value, ref) = parameter_update
 
-    # Ignore this parameter update of the associated node does
+    # Ignore this parameter update if the associated node does
     # not have an 'active' field
     if name == :active && ref.i == 0
         return nothing

core/src/parameter.jl

@@ -16,7 +16,7 @@ const SolverStats = @NamedTuple{
 # LinkType.flow and NodeType.FlowBoundary
 @enumx LinkType flow control none
 @eval @enumx NodeType $(config.nodetypes...)
-@enumx ContinuousControlType None Continuous PID
+@enumx ControlType None Continuous PID Allocation
 @enumx Substance Continuity = 1 Initial = 2 LevelBoundary = 3 FlowBoundary = 4 UserDemand =
     5 Drainage = 6 Precipitation = 7
 Base.to_index(id::Substance.T) = Int(id)  # used to index into concentration matrices
@@ -650,7 +650,7 @@ max_flow_rate: The maximum flow rate of the pump
 min_upstream_level: The upstream level below which the Pump flow goes to zero
 max_downstream_level: The downstream level above which the Pump flow goes to zero
 control_mapping: dictionary from (node_id, control_state) to target flow rate
-continuous_control_type: one of None, ContinuousControl, PidControl
+control_type: one of None, ContinuousControl, PidControl, Allocation
 """
 @kwdef struct Pump <: AbstractParameterNode
     node_id::Vector{NodeID}
@@ -664,8 +664,7 @@ continuous_control_type: one of None, ContinuousControl, PidControl
     min_upstream_level::Vector{ScalarInterpolation} = ScalarInterpolation[]
     max_downstream_level::Vector{ScalarInterpolation} = ScalarInterpolation[]
     control_mapping::Dict{Tuple{NodeID, String}, ControlStateUpdate}
-    continuous_control_type::Vector{ContinuousControlType.T} =
-        fill(ContinuousControlType.None, length(node_id))
+    control_type::Vector{ControlType.T} = fill(ControlType.None, length(node_id))
 
     function Pump(
         node_id,
@@ -679,7 +678,7 @@ continuous_control_type: one of None, ContinuousControl, PidControl
         min_upstream_level,
         max_downstream_level,
         control_mapping,
-        continuous_control_type,
+        control_type,
     )
         if valid_flow_rates(node_id, flow_rate[Float64[]], control_mapping)
             return new(
@@ -694,7 +693,7 @@ continuous_control_type: one of None, ContinuousControl, PidControl
                 min_upstream_level,
                 max_downstream_level,
                 control_mapping,
-                continuous_control_type,
+                control_type,
             )
         else
             error("Invalid Pump flow rate(s).")
@@ -716,7 +715,7 @@ max_flow_rate: The maximum flow rate of the outlet
 min_upstream_level: The upstream level below which the Outlet flow goes to zero
 max_downstream_level: The downstream level above which the Outlet flow goes to zero
 control_mapping: dictionary from (node_id, control_state) to target flow rate
-continuous_control_type: one of None, ContinuousControl, PidControl
+control_type: one of None, ContinuousControl, PidControl, Allocation
 """
 @kwdef struct Outlet <: AbstractParameterNode
     node_id::Vector{NodeID}
@@ -730,8 +729,7 @@ continuous_control_type: one of None, ContinuousControl, PidControl
     min_upstream_level::Vector{ScalarInterpolation} = ScalarInterpolation[]
     max_downstream_level::Vector{ScalarInterpolation} = ScalarInterpolation[]
     control_mapping::Dict{Tuple{NodeID, String}, ControlStateUpdate} = Dict()
-    continuous_control_type::Vector{ContinuousControlType.T} =
-        fill(ContinuousControlType.None, length(node_id))
+    control_type::Vector{ControlType.T} = fill(ControlType.None, length(node_id))
 
     function Outlet(
         node_id,
@@ -745,7 +743,7 @@ continuous_control_type: one of None, ContinuousControl, PidControl
         min_upstream_level,
         max_downstream_level,
         control_mapping,
-        continuous_control_type,
+        control_type,
     )
         if valid_flow_rates(node_id, flow_rate[Float64[]], control_mapping)
             return new(
@@ -760,7 +758,7 @@ continuous_control_type: one of None, ContinuousControl, PidControl
                 min_upstream_level,
                 max_downstream_level,
                 control_mapping,
-                continuous_control_type,
+                control_type,
             )
         else
             error("Invalid Outlet flow rate(s).")

core/src/read.jl

@@ -742,7 +742,10 @@ function Pump(db::DB, config::Config, graph::MetaGraph)::Pump
     flow_rate_cache = cache(length(node_id))
     flow_rate_cache[Float64[]] .= [itp(0) for itp in parsed_parameters.flow_rate]
 
-    return Pump(;
+    control_type = fill(ControlType.None, length(node_id))
+    set_allocation_controlled!(control_type, node_id, parsed_parameters.control_mapping)
+
+    Pump(;
         node_id,
         inflow_link = inflow_link.(Ref(graph), node_id),
         outflow_link = outflow_link.(Ref(graph), node_id),
@@ -754,6 +757,7 @@ function Pump(db::DB, config::Config, graph::MetaGraph)::Pump
         parsed_parameters.min_upstream_level,
         parsed_parameters.max_downstream_level,
         parsed_parameters.control_mapping,
+        control_type,
     )
 end
 
@@ -801,6 +805,9 @@ function Outlet(db::DB, config::Config, graph::MetaGraph)::Outlet
     flow_rate_cache[Float64[], length(node_id)] .=
         [itp(0) for itp in parsed_parameters.flow_rate]
 
+    control_type = fill(ControlType.None, length(node_id))
+    set_allocation_controlled!(control_type, node_id, parsed_parameters.control_mapping)
+
     return Outlet(;
         node_id,
         inflow_link = inflow_link.(Ref(graph), node_id),
@@ -813,6 +820,7 @@ function Outlet(db::DB, config::Config, graph::MetaGraph)::Outlet
         parsed_parameters.control_mapping,
         parsed_parameters.min_upstream_level,
         parsed_parameters.max_downstream_level,
+        control_type,
     )
 end
 
@@ -1878,7 +1886,7 @@ function Parameters(db::DB, config::Config)::Parameters
     )
 
     collect_control_mappings!(p)
-    set_continuous_control_type!(p)
+    set_control_type!(p)
     set_listen_variable_refs!(p)
     set_discrete_controlled_variable_refs!(p)
     set_continuously_controlled_variable_refs!(p)

core/src/solve.jl

@@ -36,7 +36,7 @@ function water_balance!(du::Vector, u::Vector, p::Parameters, t::Number)::Nothin
         t,
         current_low_storage_factor,
         current_level;
-        continuous_control_type = ContinuousControlType.Continuous,
+        control_type = ControlType.Continuous,
     )
 
     # Compute PID control
@@ -49,7 +49,7 @@ function water_balance!(du::Vector, u::Vector, p::Parameters, t::Number)::Nothin
         t,
         current_low_storage_factor,
         current_level;
-        continuous_control_type = ContinuousControlType.PID,
+        control_type = ControlType.PID,
     )
 
     return nothing
@@ -562,7 +562,7 @@ function formulate_flow!(
     t::Number,
     current_low_storage_factor::Vector,
     current_level::Vector,
-    continuous_control_type_::ContinuousControlType.T,
+    control_type_::ControlType.T,
 )::Nothing
     du_pump = view(du, p.state_ranges.pump)
     all_nodes_active = p.all_nodes_active[]
@@ -577,7 +577,7 @@ function formulate_flow!(
         max_flow_rate,
         min_upstream_level,
         max_downstream_level,
-        continuous_control_type,
+        control_type,
     ) in zip(
         pump.node_id,
         pump.inflow_link,
@@ -589,14 +589,13 @@ function formulate_flow!(
         pump.max_flow_rate,
         pump.min_upstream_level,
         pump.max_downstream_level,
-        pump.continuous_control_type,
+        pump.control_type,
     )
-        if !(active || all_nodes_active) ||
-           (continuous_control_type != continuous_control_type_)
+        if !(active || all_nodes_active) || (control_type != control_type_)
             continue
         end
 
-        flow_rate = if continuous_control_type == ContinuousControlType.None
+        flow_rate = if control_type == ControlType.None
             flow_rate_itp(t)
         else
             flow_rate_from_cache
@@ -626,7 +625,7 @@ function formulate_flow!(
     t::Number,
     current_low_storage_factor::Vector,
     current_level::Vector,
-    continuous_control_type_::ContinuousControlType.T,
+    control_type_::ControlType.T,
 )::Nothing
     du_outlet = view(du, p.state_ranges.outlet)
     all_nodes_active = p.all_nodes_active[]
@@ -639,7 +638,7 @@ function formulate_flow!(
         flow_rate_itp,
         min_flow_rate,
         max_flow_rate,
-        continuous_control_type,
+        control_type,
         min_upstream_level,
         max_downstream_level,
     ) in zip(
@@ -651,16 +650,15 @@ function formulate_flow!(
         outlet.flow_rate,
         outlet.min_flow_rate,
         outlet.max_flow_rate,
-        outlet.continuous_control_type,
+        outlet.control_type,
         outlet.min_upstream_level,
         outlet.max_downstream_level,
     )
-        if !(active || all_nodes_active) ||
-           (continuous_control_type != continuous_control_type_)
+        if !(active || all_nodes_active) || (control_type != control_type_)
             continue
         end
 
-        flow_rate = if continuous_control_type == ContinuousControlType.None
+        flow_rate = if control_type == ControlType.None
             flow_rate_itp(t)
         else
             flow_rate_from_cache
@@ -692,7 +690,7 @@ function formulate_flows!(
     t::Number,
     current_low_storage_factor::Vector,
     current_level::Vector;
-    continuous_control_type::ContinuousControlType.T = ContinuousControlType.None,
+    control_type::ControlType.T = ControlType.None,
 )::Nothing
     (;
         linear_resistance,
@@ -703,26 +701,18 @@ function formulate_flows!(
         user_demand,
     ) = p
 
-    formulate_flow!(
-        du,
-        pump,
-        p,
-        t,
-        current_low_storage_factor,
-        current_level,
-        continuous_control_type,
-    )
+    formulate_flow!(du, pump, p, t, current_low_storage_factor, current_level, control_type)
     formulate_flow!(
         du,
         outlet,
         p,
         t,
         current_low_storage_factor,
         current_level,
-        continuous_control_type,
+        control_type,
     )
 
-    if continuous_control_type == ContinuousControlType.None
+    if control_type == ControlType.None
         formulate_flow!(
             du,
             linear_resistance,

core/src/util.jl

@@ -400,41 +400,35 @@
 end
 
 """
-Set continuous_control_type for those pumps and outlets that are controlled by either
+Set control_type for those pumps and outlets that are controlled by either
 PidControl or ContinuousControl
 """
-function set_continuous_control_type!(p::Parameters)::Nothing
+function set_control_type!(p::Parameters)::Nothing
     (; continuous_control, pid_control) = p
-    errors = false
 
-    errors = set_continuous_control_type!(
-        p,
-        continuous_control.node_id,
-        ContinuousControlType.Continuous,
-    )
-    errors |=
-        set_continuous_control_type!(p, pid_control.node_id, ContinuousControlType.PID)
+    errors = set_control_type!(p, continuous_control.node_id, ControlType.Continuous)
+    errors |= set_control_type!(p, pid_control.node_id, ControlType.PID)
 
     if errors
         error("Errors occurred when parsing ContinuousControl and PidControl connectivity")
     end
     return nothing
 end
 
-function set_continuous_control_type!(
+function set_control_type!(
     p::Parameters,
     node_id::Vector{NodeID},
-    continuous_control_type::ContinuousControlType.T,
+    control_type::ControlType.T,
 )::Bool
     (; graph, pump, outlet) = p
     errors = false
 
     for id in node_id
         id_controlled = only(outneighbor_labels_type(graph, id, LinkType.control))
         if id_controlled.type == NodeType.Pump
-            pump.continuous_control_type[id_controlled.idx] = continuous_control_type
+            pump.control_type[id_controlled.idx] = control_type
         elseif id_controlled.type == NodeType.Outlet
-            outlet.continuous_control_type[id_controlled.idx] = continuous_control_type
+            outlet.control_type[id_controlled.idx] = control_type
         else
             errors = true
             @error "Only Pump and Outlet can be controlled by PidController, got $id_controlled"
@@ -1146,3 +1140,20 @@
 
     return tstops
 end
+
+"""
+If for a certain node the "Ribasim.allocation" control state is defined,
+make this the (initial) control type of that node
+"""
+function set_allocation_controlled!(
+    control_type::Vector{ControlType.T},
+    node_id::Vector{NodeID},
+    control_mapping::Dict{Tuple{NodeID, String}, ControlStateUpdate},
+)::Nothing
+    for id in node_id
+        if (id, "Ribasim.allocation") in keys(control_mapping)
+            control_type[id.idx] = ControlType.Allocation
+        end
+    end
+    return nothing
+end

python/ribasim_testmodels/ribasim_testmodels/__init__.py

@@ -6,6 +6,7 @@
 
 import ribasim_testmodels
 from ribasim_testmodels.allocation import (
+    allocation_control_model,
     allocation_example_model,
     allocation_training_model,
     bommelerwaard_model,
@@ -71,6 +72,7 @@
 from ribasim_testmodels.two_basin import two_basin_model
 
 __all__ = [
+    "allocation_control_model",
     "allocation_training_model",
     "allocation_example_model",
     "backwater_model",