Add ContinuousControl node (#1602)

Fixes #720.

---------

Co-authored-by: Martijn Visser <[email protected]>

core/src/callback.jl

@@ -233,12 +233,7 @@ function apply_discrete_control!(u, t, integrator)::Nothing
 
         # Loop over the variables listened to by this discrete control node
         for compound_variable in compound_variables
-
-            # Compute the value of the current variable
-            value = 0.0
-            for subvariable in compound_variable.subvariables
-                value += subvariable.weight * get_value(p, subvariable, t)
-            end
+            value = compound_variable_value(compound_variable, p, u, t)
 
             # The thresholds the value of this variable is being compared with
             greater_thans = compound_variable.greater_than
@@ -318,12 +313,12 @@ end
 Get a value for a condition. Currently supports getting levels from basins and flows
 from flow boundaries.
 """
-function get_value(p::Parameters, subvariable::NamedTuple, t::Float64)
+function get_value(subvariable::NamedTuple, p::Parameters, u::AbstractVector, t::Float64)
     (; flow_boundary, level_boundary) = p
     (; listen_node_id, look_ahead, variable, variable_ref) = subvariable
 
-    if !iszero(variable_ref.i)
-        return variable_ref[]
+    if !iszero(variable_ref.idx)
+        return get_value(variable_ref, u)
     end
 
     if variable == "level"
@@ -350,6 +345,14 @@ function get_value(p::Parameters, subvariable::NamedTuple, t::Float64)
     return value
 end
 
+function compound_variable_value(compound_variable::CompoundVariable, p, u, t)
+    value = zero(eltype(u))
+    for subvariable in compound_variable.subvariables
+        value += subvariable.weight * get_value(subvariable, p, u, t)
+    end
+    return value
+end
+
 function get_allocation_model(p::Parameters, subnetwork_id::Int32)::AllocationModel
     (; allocation) = p
     (; subnetwork_ids, allocation_models) = allocation

core/src/parameter.jl

@@ -1,6 +1,7 @@
 # EdgeType.flow and NodeType.FlowBoundary
 @enumx EdgeType flow control none
 @eval @enumx NodeType $(config.nodetypes...)
+@enumx ContinuousControlType None Continuous PID
 
 # Support creating a NodeType enum instance from a symbol or string
 function NodeType.T(s::Symbol)::NodeType.T
@@ -406,7 +407,7 @@ flow_rate: target flow rate
 min_flow_rate: The minimal flow rate of the pump
 max_flow_rate: The maximum flow rate of the pump
 control_mapping: dictionary from (node_id, control_state) to target flow rate
-is_pid_controlled: whether the flow rate of this pump is governed by PID control
+continuous_control_type: one of None, ContinuousControl, PidControl
 """
 @kwdef struct Pump{T} <: AbstractParameterNode
     node_id::Vector{NodeID}
@@ -417,7 +418,8 @@ is_pid_controlled: whether the flow rate of this pump is governed by PID control
     min_flow_rate::Vector{Float64} = zeros(length(node_id))
     max_flow_rate::Vector{Float64} = fill(Inf, length(node_id))
     control_mapping::Dict{Tuple{NodeID, String}, ControlStateUpdate}
-    is_pid_controlled::Vector{Bool} = fill(false, length(node_id))
+    continuous_control_type::Vector{ContinuousControlType.T} =
+        fill(ContinuousControlType.None, length(node_id))
 
     function Pump(
         node_id,
@@ -428,7 +430,7 @@ is_pid_controlled: whether the flow rate of this pump is governed by PID control
         min_flow_rate,
         max_flow_rate,
         control_mapping,
-        is_pid_controlled,
+        continuous_control_type,
     ) where {T}
         if valid_flow_rates(node_id, get_tmp(flow_rate, 0), control_mapping)
             return new{T}(
@@ -440,7 +442,7 @@ is_pid_controlled: whether the flow rate of this pump is governed by PID control
                 min_flow_rate,
                 max_flow_rate,
                 control_mapping,
-                is_pid_controlled,
+                continuous_control_type,
             )
         else
             error("Invalid Pump flow rate(s).")
@@ -459,7 +461,7 @@ flow_rate: target flow rate
 min_flow_rate: The minimal flow rate of the outlet
 max_flow_rate: The maximum flow rate of the outlet
 control_mapping: dictionary from (node_id, control_state) to target flow rate
-is_pid_controlled: whether the flow rate of this outlet is governed by PID control
+continuous_control_type: one of None, ContinuousControl, PidControl
 """
 @kwdef struct Outlet{T} <: AbstractParameterNode
     node_id::Vector{NodeID}
@@ -471,7 +473,8 @@ is_pid_controlled: whether the flow rate of this outlet is governed by PID contr
     max_flow_rate::Vector{Float64} = fill(Inf, length(node_id))
     min_crest_level::Vector{Float64} = fill(-Inf, length(node_id))
     control_mapping::Dict{Tuple{NodeID, String}, ControlStateUpdate} = Dict()
-    is_pid_controlled::Vector{Bool} = fill(false, length(node_id))
+    continuous_control_type::Vector{ContinuousControlType.T} =
+        fill(ContinuousControlType.None, length(node_id))
 
     function Outlet(
         node_id,
@@ -483,7 +486,7 @@ is_pid_controlled: whether the flow rate of this outlet is governed by PID contr
         max_flow_rate,
         min_crest_level,
         control_mapping,
-        is_pid_controlled,
+        continuous_control_type,
     ) where {T}
         if valid_flow_rates(node_id, get_tmp(flow_rate, 0), control_mapping)
             return new{T}(
@@ -496,7 +499,7 @@ is_pid_controlled: whether the flow rate of this outlet is governed by PID contr
                 max_flow_rate,
                 min_crest_level,
                 control_mapping,
-                is_pid_controlled,
+                continuous_control_type,
             )
         else
             error("Invalid Outlet flow rate(s).")
@@ -511,19 +514,35 @@ node_id: node ID of the Terminal node
     node_id::Vector{NodeID}
 end
 
+"""
+A variant on `Base.Ref` where the source array is a vector that is possibly wrapped in a ForwardDiff.DiffCache.
+Retrieve value with get_value(ref::PreallocationRef, val) where `val` determines the return type.
+"""
+struct PreallocationRef{T}
+    vector::T
+    idx::Int
+end
+
+get_value(ref::PreallocationRef, val) = get_tmp(ref.vector, val)[ref.idx]
+
+function set_value!(ref::PreallocationRef, value)::Nothing
+    get_tmp(ref.vector, value)[ref.idx] = value
+    return nothing
+end
+
 """
 The data for a single compound variable
 node_id:: The ID of the DiscreteControl that listens to this variable
 subvariables: data for one single subvariable
 greater_than: the thresholds this compound variable will be
-    compared against
+    compared against (in the case of DiscreteControl)
 """
-@kwdef struct CompoundVariable
+@kwdef struct CompoundVariable{T}
     node_id::NodeID
     subvariables::Vector{
         @NamedTuple{
             listen_node_id::NodeID,
-            variable_ref::Base.RefArray{Float64, Vector{Float64}, Nothing},
+            variable_ref::PreallocationRef{T},
             variable::String,
             weight::Float64,
             look_ahead::Float64,
@@ -543,16 +562,16 @@ logic_mapping: Dictionary: truth state => control state for the DiscreteControl
 control_mapping: dictionary node type => control mapping for that node type
 record: Namedtuple with discrete control information for results
 """
-@kwdef struct DiscreteControl <: AbstractParameterNode
+@kwdef struct DiscreteControl{T} <: AbstractParameterNode
     node_id::Vector{NodeID}
     controlled_nodes::Vector{Vector{NodeID}}
-    compound_variables::Vector{Vector{CompoundVariable}}
+    compound_variables::Vector{Vector{CompoundVariable{T}}}
     truth_state::Vector{Vector{Bool}}
     control_state::Vector{String} = fill("undefined_state", length(node_id))
     control_state_start::Vector{Float64} = zeros(length(node_id))
     logic_mapping::Vector{Dict{Vector{Bool}, String}}
     control_mappings::Dict{NodeType.T, Dict{Tuple{NodeID, String}, ControlStateUpdate}} =
-        Dict()
+        Dict{NodeType.T, Dict{Tuple{NodeID, String}, ControlStateUpdate}}()
     record::@NamedTuple{
         time::Vector{Float64},
         control_node_id::Vector{Int32},
@@ -566,26 +585,40 @@ record: Namedtuple with discrete control information for results
     )
 end
 
+@kwdef struct ContinuousControl{T} <: AbstractParameterNode
+    node_id::Vector{NodeID}
+    compound_variable::Vector{CompoundVariable{T}}
+    controlled_variable::Vector{String}
+    target_ref::Vector{PreallocationRef{T}}
+    func::Vector{ScalarInterpolation}
+end
+
 """
 PID control currently only supports regulating basin levels.
 
 node_id: node ID of the PidControl node
 active: whether this node is active and thus sets flow rates
+controlled_node_id: The node that is being controlled
 listen_node_id: the id of the basin being controlled
-pid_params: a vector interpolation for parameters changing over time.
-    The parameters are respectively target, proportional, integral, derivative,
-    where the last three are the coefficients for the PID equation.
+target: target level (possibly time dependent)
+target_ref: reference to the controlled flow_rate value
+proportional: proportionality coefficient error
+integral: proportionality coefficient error integral
+derivative: proportionality coefficient error derivative
 error: the current error; basin_target - current_level
+dictionary from (node_id, control_state) to target flow rate
 """
 @kwdef struct PidControl{T} <: AbstractParameterNode
     node_id::Vector{NodeID}
     active::Vector{Bool}
     listen_node_id::Vector{NodeID}
     target::Vector{ScalarInterpolation}
+    target_ref::Vector{PreallocationRef{T}}
     proportional::Vector{ScalarInterpolation}
     integral::Vector{ScalarInterpolation}
     derivative::Vector{ScalarInterpolation}
     error::T
+    controlled_basins::Vector{NodeID}
     control_mapping::Dict{Tuple{NodeID, String}, ControlStateUpdate}
 end
 
@@ -705,7 +738,8 @@ const ModelGraph{T} = MetaGraph{
     pump::Pump{T}
     outlet::Outlet{T}
     terminal::Terminal
-    discrete_control::DiscreteControl
+    discrete_control::DiscreteControl{T}
+    continuous_control::ContinuousControl{T}
     pid_control::PidControl{T}
     user_demand::UserDemand
     level_demand::LevelDemand