Support conditions on linear combinations of variables for DiscreteControl

core/src/callback.jl

@@ -194,15 +194,20 @@ function discrete_control_condition(out, u, t, integrator)
     (; p) = integrator
     (; discrete_control) = p
 
-    for (i, (listen_node_id, variable, greater_than, look_ahead)) in enumerate(
+    for (i, (listen_node_ids, variables, weights, greater_than, look_aheads)) in enumerate(
         zip(
             discrete_control.listen_node_id,
             discrete_control.variable,
+            discrete_control.weight,
             discrete_control.greater_than,
             discrete_control.look_ahead,
         ),
     )
-        value = get_value(p, listen_node_id, variable, look_ahead, u, t)
+        value = 0.0
+        for (listen_node_id, variable, weight, look_ahead) in
+            zip(listen_node_ids, variables, weights, look_aheads)
+            value += weight * get_value(p, listen_node_id, variable, look_ahead, u, t)
+        end
         diff = value - greater_than
         out[i] = diff
     end
@@ -272,16 +277,20 @@ function discrete_control_affect_upcrossing!(integrator, condition_idx)
     # only possibly the du. Parameter changes can change the flow on an edge discontinuously,
     # giving the possibility of logical paradoxes where certain parameter changes immediately
     # undo the truth state that caused that parameter change.
-    is_basin = id_index(basin.node_id, discrete_control.listen_node_id[condition_idx])[1]
+    listen_node_ids = discrete_control.listen_node_id[condition_idx]
+    is_basin =
+        length(listen_node_ids) == 1 ? id_index(basin.node_id, only(listen_node_ids))[1] :
+        false
+
     # NOTE: The above no longer works when listen feature ids can be something other than node ids
     # I think the more durable option is to give all possible condition types a different variable string,
     # e.g. basin.level and level_boundary.level
-    if variable[condition_idx] == "level" && control_state_change && is_basin
+    if variable[condition_idx][1] == "level" && control_state_change && is_basin
         # Calling water_balance is expensive, but it is a sure way of getting
         # du for the basin of this level condition
         du = zero(u)
         water_balance!(du, u, p, t)
-        _, condition_basin_idx = id_index(basin.node_id, listen_node_id[condition_idx])
+        _, condition_basin_idx = id_index(basin.node_id, listen_node_id[condition_idx][1])
 
         if du[condition_basin_idx] < 0.0
             condition_value[condition_idx] = false
@@ -308,13 +317,18 @@ function discrete_control_affect_downcrossing!(integrator, condition_idx)
     # only possibly the du. Parameter changes can change the flow on an edge discontinuously,
     # giving the possibility of logical paradoxes where certain parameter changes immediately
     # undo the truth state that caused that parameter change.
-    if variable[condition_idx] == "level" && control_state_change
+    listen_node_ids = discrete_control.listen_node_id[condition_idx]
+    is_basin =
+        length(listen_node_ids) == 1 ? id_index(basin.node_id, only(listen_node_ids))[1] :
+        false
+
+    if variable[condition_idx][1] == "level" && control_state_change && is_basin
         # Calling water_balance is expensive, but it is a sure way of getting
         # du for the basin of this level condition
         du = zero(u)
         water_balance!(du, u, p, t)
         has_index, condition_basin_idx =
-            id_index(basin.node_id, listen_node_id[condition_idx])
+            id_index(basin.node_id, listen_node_id[condition_idx][1])
 
         if has_index && du[condition_basin_idx] > 0.0
             condition_value[condition_idx] = true
@@ -369,7 +383,7 @@ function discrete_control_affect!(
             discrete_control.logic_mapping[(discrete_control_node_id, truth_state)]
         else
             error(
-                "Control state specified for neither $truth_state_crossing_specific nor $truth_state for DiscreteControl node $discrete_control_node_id.",
+                "Control state specified for neither $truth_state_crossing_specific nor $truth_state for $discrete_control_node_id.",
             )
         end
 

core/src/parameter.jl

@@ -445,8 +445,10 @@ end
 """
 node_id: node ID of the DiscreteControl node; these are not unique but repeated
     by the amount of conditions of this DiscreteControl node
-listen_node_id: the ID of the node being condition on
-variable: the name of the variable in the condition
+listen_node_id: the IDs of the nodes being condition on
+variable: the names of the variables in the condition
+weight: the weight of the variables in the condition
+look_ahead: the look ahead of variables  in the condition in seconds
 greater_than: The threshold value in the condition
 condition_value: The current value of each condition
 control_state: Dictionary: node ID => (control state, control state start)
@@ -455,9 +457,10 @@ record: Namedtuple with discrete control information for results
 """
 struct DiscreteControl <: AbstractParameterNode
     node_id::Vector{NodeID}
-    listen_node_id::Vector{NodeID}
-    variable::Vector{String}
-    look_ahead::Vector{Float64}
+    listen_node_id::Vector{Vector{NodeID}}
+    variable::Vector{Vector{String}}
+    weight::Vector{Vector{Float64}}
+    look_ahead::Vector{Vector{Float64}}
     greater_than::Vector{Float64}
     condition_value::Vector{Bool}
     control_state::Dict{NodeID, Tuple{String, Float64}}

core/src/read.jl

@@ -542,9 +542,53 @@ function Basin(db::DB, config::Config, chunk_sizes::Vector{Int})::Basin
     )
 end
 
+"""
+Parse compound variables from the compound variable table, and turn
+variables in the condition table into trivial compound variables
+"""
+function get_compound_variables(compound_variable, condition)
+    listen_node_id = Vector{NodeID}[]
+    variable = Vector{String}[]
+    weight = Vector{Float64}[]
+    look_ahead = Vector{Float64}[]
+
+    for cond in condition
+        if cond.listen_node_type == "compound"
+            compound_variable_data = filter(
+                row -> (row.node_id, row.name) == (cond.node_id, cond.variable),
+                compound_variable,
+            )
+            listen_node_id_data =
+                NodeID.(
+                    compound_variable_data.listen_node_type,
+                    compound_variable_data.listen_node_id,
+                )
+            @assert !isempty(listen_node_id_data) "No compound variable data found for name $(cond.variable)."
+            variable_data = compound_variable_data.variable
+            weight_data = compound_variable_data.weight
+            look_ahead_data = coalesce.(compound_variable_data.look_ahead, 0.0)
+        else
+            listen_node_id_data = [NodeID(cond.listen_node_type, cond.listen_node_id)]
+            variable_data = [cond.variable]
+            weight_data = [1.0]
+            look_ahead_data = [coalesce(cond.look_ahead, 0.0)]
+        end
+
+        push!(listen_node_id, listen_node_id_data)
+        push!(variable, variable_data)
+        push!(weight, weight_data)
+        push!(look_ahead, look_ahead_data)
+    end
+    return listen_node_id, variable, weight, look_ahead
+end
+
 function DiscreteControl(db::DB, config::Config)::DiscreteControl
+    compound_variable = load_structvector(db, config, DiscreteControlCompoundvariableV1)
     condition = load_structvector(db, config, DiscreteControlConditionV1)
 
+    listen_node_id, variable, weight, look_ahead =
+        get_compound_variables(compound_variable, condition)
+
     condition_value = fill(false, length(condition.node_id))
     control_state::Dict{NodeID, Tuple{String, Float64}} = Dict()
 
@@ -557,7 +601,6 @@ function DiscreteControl(db::DB, config::Config)::DiscreteControl
     end
 
     logic = load_structvector(db, config, DiscreteControlLogicV1)
-
     logic_mapping = Dict{Tuple{NodeID, String}, String}()
 
     for (node_id, truth_state, control_state_) in
@@ -567,7 +610,6 @@ function DiscreteControl(db::DB, config::Config)::DiscreteControl
     end
 
     logic_mapping = expand_logic_mapping(logic_mapping)
-    look_ahead = coalesce.(condition.look_ahead, 0.0)
 
     record = (
         time = Float64[],
@@ -578,8 +620,9 @@ function DiscreteControl(db::DB, config::Config)::DiscreteControl
 
     return DiscreteControl(
         NodeID.(NodeType.DiscreteControl, condition.node_id), # Not unique
-        NodeID.(condition.listen_node_type, condition.listen_node_id),
-        condition.variable,
+        listen_node_id,
+        variable,
+        weight,
         look_ahead,
         condition.greater_than,
         condition_value,

core/src/schema.jl

@@ -2,6 +2,7 @@
 # The identifier is parsed as ribasim.nodetype.kind, no capitals or underscores are allowed.
 @schema "ribasim.discretecontrol.condition" DiscreteControlCondition
 @schema "ribasim.discretecontrol.logic" DiscreteControlLogic
+@schema "ribasim.discretecontrol.compoundvariable" DiscreteControlCompoundvariable
 @schema "ribasim.basin.static" BasinStatic
 @schema "ribasim.basin.time" BasinTime
 @schema "ribasim.basin.profile" BasinProfile
@@ -183,6 +184,16 @@ end
     node_id::Int32
 end
 
+@version DiscreteControlCompoundvariableV1 begin
+    node_id::Int32
+    name::String
+    listen_node_type::String
+    listen_node_id::Int
+    variable::String
+    weight::Float64
+    look_ahead::Union{Missing, Float64}
+end
+
 @version DiscreteControlConditionV1 begin
     node_id::Int32
     listen_node_type::String

core/src/validation.jl

@@ -102,6 +102,10 @@
 sort_by_priority(row) = (row.node_id, row.priority)
 sort_by_priority_time(row) = (row.node_id, row.priority, row.time)
 sort_by_subgrid_level(row) = (row.subgrid_id, row.basin_level)
+sort_by_condition(row) =
+    (row.node_id, row.listen_node_type, row.listen_node_id, row.variable, row.greater_than)
+sort_by_id_name(row) = (row.node_id, row.name)
+sort_by_truth_state(row) = (row.node_id, row.truth_state)
 
 # get the right sort by function given the Schema, with sort_by_id as the default
 sort_by_function(table::StructVector{<:Legolas.AbstractRecord}) = sort_by_id
@@ -110,6 +114,9 @@
 sort_by_function(table::StructVector{UserDemandStaticV1}) = sort_by_priority
 sort_by_function(table::StructVector{UserDemandTimeV1}) = sort_by_priority_time
 sort_by_function(table::StructVector{BasinSubgridV1}) = sort_by_subgrid_level
+sort_by_function(table::StructVector{DiscreteControlCompoundvariableV1}) = sort_by_id_name
+sort_by_function(table::StructVector{DiscreteControlConditionV1}) = sort_by_condition
+sort_by_function(table::StructVector{DiscreteControlLogicV1}) = sort_by_truth_state
 
 const TimeSchemas = Union{
     BasinTimeV1,
@@ -552,29 +559,32 @@
             end
         end
     end
-    for (Δt, var, node_id) in zip(look_ahead, variable, listen_node_id)
-        if !iszero(Δt)
-            node_type = node_id.type
-            # TODO: If more transient listen variables must be supported, this validation must be more specific
-            # (e.g. for some node some variables are transient, some not).
-            if node_type ∉ [NodeType.FlowBoundary, NodeType.LevelBoundary]
-                errors = true
-                @error "Look ahead supplied for non-timeseries listen variable '$var' from listen node $node_id."
-            else
-                if Δt < 0
+    for (look_aheads, variables, listen_node_ids) in
+        zip(look_ahead, variable, listen_node_id)
+        for (Δt, var, node_id) in zip(look_aheads, variables, listen_node_ids)
+            if !iszero(Δt)
+                node_type = node_id.type
+                # TODO: If more transient listen variables must be supported, this validation must be more specific
+                # (e.g. for some node some variables are transient, some not).
+                if node_type ∉ [NodeType.FlowBoundary, NodeType.LevelBoundary]
                     errors = true
-                    @error "Negative look ahead supplied for listen variable '$var' from listen node $node_id."
+                    @error "Look ahead supplied for non-timeseries listen variable '$var' from listen node $node_id."
                 else
-                    node = getfield(p, graph[node_id].type)
-                    idx = if node_type == NodeType.Basin
-                        id_index(node.node_id, node_id)
-                    else
-                        searchsortedfirst(node.node_id, node_id)
-                    end
-                    interpolation = getfield(node, Symbol(var))[idx]
-                    if t_end + Δt > interpolation.t[end]
+                    if Δt < 0
                         errors = true
-                        @error "Look ahead for listen variable '$var' from listen node $node_id goes past timeseries end during simulation."
+                        @error "Negative look ahead supplied for listen variable '$var' from listen node $node_id."
+                    else
+                        node = getfield(p, graph[node_id].type)
+                        idx = if node_type == NodeType.Basin
+                            id_index(node.node_id, node_id)
+                        else
+                            searchsortedfirst(node.node_id, node_id)
+                        end
+                        interpolation = getfield(node, Symbol(var))[idx]
+                        if t_end + Δt > interpolation.t[end]
+                            errors = true
+                            @error "Look ahead for listen variable '$var' from listen node $node_id goes past timeseries end during simulation."
+                        end
                     end
                 end
             end

core/test/control_test.jl

@@ -54,7 +54,7 @@ end
     p = model.integrator.p
     (; discrete_control, flow_boundary) = p
 
-    Δt = discrete_control.look_ahead[1]
+    Δt = discrete_control.look_ahead[1][1]
 
     t = Ribasim.tsaves(model)
     t_control = discrete_control.record.time[2]
@@ -78,7 +78,7 @@ end
     p = model.integrator.p
     (; discrete_control, level_boundary) = p
 
-    Δt = discrete_control.look_ahead[1]
+    Δt = discrete_control.look_ahead[1][1]
 
     t = Ribasim.tsaves(model)
     t_control = discrete_control.record.time[2]
@@ -206,3 +206,23 @@ end
     @test isapprox(level[t_idx_target_jump], target_high, atol = 1e-1)
     @test isapprox(level[end], target_low, atol = 1e-1)
 end
+
+@testitem "Compound condition" begin
+    using Ribasim: NodeID
+
+    toml_path = normpath(
+        @__DIR__,
+        "../../generated_testmodels/compound_variable_condition/ribasim.toml",
+    )
+    @test ispath(toml_path)
+    model = Ribasim.run(toml_path)
+    (; discrete_control) = model.integrator.p
+    (; listen_node_id, variable, weight, record) = discrete_control
+
+    @test listen_node_id == [[NodeID(:FlowBoundary, 2), NodeID(:FlowBoundary, 3)]]
+    @test variable == [["flow_rate", "flow_rate"]]
+    @test weight == [[0.5, 0.5]]
+    @test record.time ≈ [0.0, model.integrator.sol.t[end] / 2]
+    @test record.truth_state == ["F", "T"]
+    @test record.control_state == ["Off", "On"]
+end

docs/core/usage.qmd

@@ -204,6 +204,7 @@ name it must have in the database if it is stored there.
 - Terminal: Water sink without state or properties
   - `Terminal / static`: - (only node IDs)
 - DiscreteControl: Set parameters of other nodes based on model state conditions (e.g. basin level)
+  - `DiscreteControl / compoundvariable` : Linear combinations of variables, e.g. averages or differences of levels
   - `DisceteControl / condition`: Conditions of the form 'the level in the basin with node id `n` is bigger than 2.0 m'
   - `DisceteControl / logic`: Translates the truth value of a set of conditions to parameter values for a controlled node
 - PidControl: Controls the level in a basin by continuously controlling the flow rate of a connected pump or outlet. See also [Wikipedia](https://en.wikipedia.org/wiki/PID_controller) and [PID controller in equations](equations.qmd#sec-PID).
@@ -625,16 +626,37 @@ node_id               | Int32   | -            | sorted
 
 DiscreteControl is implemented based on [VectorContinuousCallback](https://docs.sciml.ai/DiffEqDocs/stable/features/callback_functions/#VectorContinuousCallback).
 
-## DiscreteControl / condition
+## DiscreteControl / compoundvariable
+
+The compound variable schema defines linear combinations of variables which can be used in conditions. This means that
+this schema defines new variables with the given name that look like
+$$
+\text{weight}_1 * \text{variable}_1 + \text{weight}_2 * \text{variable}_2 + \ldots,
+$$
 
-The condition schema defines conditions of the form 'the discrete_control node with this node id listens to whether the given variable of the node with the given listen feature id is grater than the given value'. If the condition variable comes from a time-series, a look ahead $\Delta t$ can be supplied.
+which can be for instance an average or a difference of variables. If a variable comes from a time-series, a look ahead $\Delta t$ can be supplied.
 
 column            | type     | unit    | restriction
 ----------------- | -------- | ------- | -----------
 node_id           | Int32    | -       | sorted
+name              | String   | -       | sorted per node_id
 listen_node_type  | String   | -       | known node type
 listen_node_id    | Int32    | -       | sorted per node_id
 variable          | String   | -       | must be "level" or "flow_rate", sorted per listen_node_id
+weight            | Float64  | -       | -
+look_ahead        | Float64  | $s$     | Only on transient boundary conditions, non-negative (optional, default 0).
+
+## DiscreteControl / condition
+
+The condition schema defines conditions of the form 'the discrete_control node with this node id listens to whether the given variable of the node with the given listen node id is greater than the given value'. If the condition variable comes from a time-series, a look ahead $\Delta t$ can be supplied.
+Compound variables from the schema above can also be used, in which case the listen node ID and look ahead in this schema are ignored.
+
+column            | type     | unit    | restriction
+----------------- | -------- | ------- | -----------
+node_id           | Int32    | -       | sorted
+listen_node_type  | String   | -       | known node type or "compound"
+listen_node_id    | Int32    | -       | sorted per node_id
+variable          | String   | -       | must be "level", "flow_rate" or a compound variable name, sorted per listen_node_id
 greater_than      | Float64  | various | sorted per variable
 look_ahead        | Float64  | $s$     | Only on transient boundary conditions, non-negative (optional, default 0)