From 4812bcac34862f12bcabf3865b32acd99b323bcd Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Wed, 10 Apr 2024 13:42:18 +0200
Subject: [PATCH 01/25] Python side

---
 core/src/schema.jl                            | 10 ++++
 python/ribasim/ribasim/config.py              |  5 ++
 python/ribasim/ribasim/model.py               | 12 +++-
 .../ribasim/ribasim/nodes/discrete_control.py | 13 +++-
 python/ribasim/ribasim/schemas.py             |  9 +++
 .../ribasim_testmodels/__init__.py            |  2 +
 .../ribasim_testmodels/discrete_control.py    | 60 +++++++++++++++++++
 7 files changed, 106 insertions(+), 5 deletions(-)

diff --git a/core/src/schema.jl b/core/src/schema.jl
index d59190c3b..c8dd7094d 100644
--- a/core/src/schema.jl
+++ b/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,15 @@ 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
+end
+
 @version DiscreteControlConditionV1 begin
     node_id::Int32
     listen_node_type::String
diff --git a/python/ribasim/ribasim/config.py b/python/ribasim/ribasim/config.py
index f9605111d..ecc49badc 100644
--- a/python/ribasim/ribasim/config.py
+++ b/python/ribasim/ribasim/config.py
@@ -21,6 +21,7 @@
     BasinStaticSchema,
     BasinSubgridSchema,
     BasinTimeSchema,
+    DiscreteControlCompoundVariableSchema,
     DiscreteControlConditionSchema,
     DiscreteControlLogicSchema,
     FlowBoundaryStaticSchema,
@@ -287,6 +288,10 @@ class ManningResistance(MultiNodeModel):
 
 
 class DiscreteControl(MultiNodeModel):
+    compound_variable: TableModel[DiscreteControlCompoundVariableSchema] = Field(
+        default_factory=TableModel[DiscreteControlCompoundVariableSchema],
+        json_schema_extra={"sort_keys": ["name"]},
+    )
     condition: TableModel[DiscreteControlConditionSchema] = Field(
         default_factory=TableModel[DiscreteControlConditionSchema],
         json_schema_extra={
diff --git a/python/ribasim/ribasim/model.py b/python/ribasim/ribasim/model.py
index 34c808f73..472f7d4cc 100644
--- a/python/ribasim/ribasim/model.py
+++ b/python/ribasim/ribasim/model.py
@@ -321,11 +321,17 @@ def plot_control_listen(self, ax):
             df_listen_edge = pd.concat([df_listen_edge, to_add])
 
         # Listen edges from DiscreteControl
-        condition = self.discrete_control.condition.df
-        if condition is not None:
-            to_add = condition[
+        for table in (
+            self.discrete_control.condition.df,
+            self.discrete_control.compound_variable.df,
+        ):
+            if table is None:
+                continue
+
+            to_add = table[
                 ["node_id", "listen_node_id", "listen_node_type"]
             ].drop_duplicates()
+            to_add = to_add[to_add["listen_node_type"] != "compound"]
             to_add.columns = ["control_node_id", "listen_node_id", "listen_node_type"]
             to_add["control_node_type"] = "DiscreteControl"
             df_listen_edge = pd.concat([df_listen_edge, to_add])
diff --git a/python/ribasim/ribasim/nodes/discrete_control.py b/python/ribasim/ribasim/nodes/discrete_control.py
index eeb0b28ca..02b22840c 100644
--- a/python/ribasim/ribasim/nodes/discrete_control.py
+++ b/python/ribasim/ribasim/nodes/discrete_control.py
@@ -1,9 +1,18 @@
 from pandas import DataFrame
 
 from ribasim.input_base import TableModel
-from ribasim.schemas import DiscreteControlConditionSchema, DiscreteControlLogicSchema
+from ribasim.schemas import (
+    DiscreteControlCompoundVariableSchema,
+    DiscreteControlConditionSchema,
+    DiscreteControlLogicSchema,
+)
 
-__all__ = ["Condition", "Logic"]
+__all__ = ["Condition", "Logic", "CompoundVariable"]
+
+
+class CompoundVariable(TableModel[DiscreteControlCompoundVariableSchema]):
+    def __init__(self, **kwargs):
+        super().__init__(df=DataFrame(dict(**kwargs)))
 
 
 class Condition(TableModel[DiscreteControlConditionSchema]):
diff --git a/python/ribasim/ribasim/schemas.py b/python/ribasim/ribasim/schemas.py
index 0cf033671..aa4c738c2 100644
--- a/python/ribasim/ribasim/schemas.py
+++ b/python/ribasim/ribasim/schemas.py
@@ -48,6 +48,15 @@ class BasinTimeSchema(_BaseSchema):
     urban_runoff: Series[float] = pa.Field(nullable=True)
 
 
+class DiscreteControlCompoundVariableSchema(_BaseSchema):
+    node_id: Series[Int32] = pa.Field(nullable=False, default=0)
+    name: Series[str] = pa.Field(nullable=False)
+    listen_node_type: Series[str] = pa.Field(nullable=False)
+    listen_node_id: Series[Int32] = pa.Field(nullable=False, default=0)
+    variable: Series[str] = pa.Field(nullable=False)
+    weight: Series[float] = pa.Field(nullable=False)
+
+
 class DiscreteControlConditionSchema(_BaseSchema):
     node_id: Series[Int32] = pa.Field(nullable=False, default=0)
     listen_node_type: Series[str] = pa.Field(nullable=False)
diff --git a/python/ribasim_testmodels/ribasim_testmodels/__init__.py b/python/ribasim_testmodels/ribasim_testmodels/__init__.py
index ceeb1fad2..4f9ce86c3 100644
--- a/python/ribasim_testmodels/ribasim_testmodels/__init__.py
+++ b/python/ribasim_testmodels/ribasim_testmodels/__init__.py
@@ -28,6 +28,7 @@
 )
 from ribasim_testmodels.bucket import bucket_model, leaky_bucket_model
 from ribasim_testmodels.discrete_control import (
+    compound_variable_condition_model,
     flow_condition_model,
     level_boundary_condition_model,
     level_setpoint_with_minmax_model,
@@ -64,6 +65,7 @@
     "basic_model",
     "basic_transient_model",
     "bucket_model",
+    "compound_variable_condition_model",
     "discrete_control_of_pid_control_model",
     "dutch_waterways_model",
     "flow_boundary_time_model",
diff --git a/python/ribasim_testmodels/ribasim_testmodels/discrete_control.py b/python/ribasim_testmodels/ribasim_testmodels/discrete_control.py
index aae7674c9..2f14e8e9a 100644
--- a/python/ribasim_testmodels/ribasim_testmodels/discrete_control.py
+++ b/python/ribasim_testmodels/ribasim_testmodels/discrete_control.py
@@ -390,3 +390,63 @@ def level_setpoint_with_minmax_model() -> Model:
     )
 
     return model
+
+
+def compound_variable_condition_model() -> Model:
+    """
+    Set up a minimal model containing a condition on a compound variable
+    for discrete control.
+    """
+
+    model = Model(
+        starttime="2020-01-01",
+        endtime="2021-01-01",
+        crs="EPSG:28992",
+    )
+
+    model.basin.add(
+        Node(1, Point(1, 0)),
+        [
+            basin.Profile(area=1000.0, level=[0.0, 1.0]),
+            basin.State(level=[1.0]),
+        ],
+    )
+    model.flow_boundary.add(
+        Node(2, Point(0, 0)), [flow_boundary.Static(flow_rate=[0.0])]
+    )
+    model.flow_boundary.add(
+        Node(3, Point(0, 1)),
+        [flow_boundary.Time(time=["2020-01-01", "2022-01-01"], flow_rate=[0.0, 2.0])],
+    )
+    model.pump.add(
+        Node(4, Point(2, 0)),
+        [pump.Static(control_state=["Off", "On"], flow_rate=[0.0, 1.0])],
+    )
+    model.terminal.add(Node(5, Point(3, 0)))
+    model.discrete_control.add(
+        Node(6, Point(1, 1)),
+        [
+            discrete_control.CompoundVariable(
+                name="flow_mean",
+                listen_node_type="FlowBoundary",
+                listen_node_id=[2, 3],
+                variable="flow",
+                weight=0.5,
+            ),
+            discrete_control.Condition(
+                listen_node_type="compound",
+                listen_node_id=[0],  # Irrelevant
+                variable="flow_mean",
+                greater_than=1.0,
+            ),
+            discrete_control.Logic(truth_state=["T", "F"], control_state=["On", "Off"]),
+        ],
+    )
+
+    model.edge.add(model.flow_boundary[2], model.basin[1])
+    model.edge.add(model.flow_boundary[3], model.basin[1])
+    model.edge.add(model.basin[1], model.pump[4])
+    model.edge.add(model.pump[4], model.terminal[5])
+    model.edge.add(model.discrete_control[6], model.pump[4])
+
+    return model

From 44fcaaea381fbe62f95758447748094b2502085c Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Wed, 10 Apr 2024 16:05:31 +0200
Subject: [PATCH 02/25] Change name of new schema

---
 core/src/schema.jl                            |  4 ++--
 python/ribasim/ribasim/config.py              |  6 +++---
 .../ribasim/ribasim/nodes/discrete_control.py |  6 +++---
 python/ribasim/ribasim/schemas.py             |  2 +-
 .../ribasim_testmodels/discrete_control.py    |  2 +-
 ribasim_qgis/core/nodes.py                    | 21 +++++++++++++++++++
 6 files changed, 31 insertions(+), 10 deletions(-)

diff --git a/core/src/schema.jl b/core/src/schema.jl
index c8dd7094d..2d3896e82 100644
--- a/core/src/schema.jl
+++ b/core/src/schema.jl
@@ -2,7 +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.discretecontrol.compoundvariable" DiscreteControlCompoundvariable
 @schema "ribasim.basin.static" BasinStatic
 @schema "ribasim.basin.time" BasinTime
 @schema "ribasim.basin.profile" BasinProfile
@@ -184,7 +184,7 @@ end
     node_id::Int32
 end
 
-@version DiscreteControlCompoundVariableV1 begin
+@version DiscreteControlCompoundvariableV1 begin
     node_id::Int32
     name::String
     listen_node_type::String
diff --git a/python/ribasim/ribasim/config.py b/python/ribasim/ribasim/config.py
index ecc49badc..64a5019e9 100644
--- a/python/ribasim/ribasim/config.py
+++ b/python/ribasim/ribasim/config.py
@@ -21,7 +21,7 @@
     BasinStaticSchema,
     BasinSubgridSchema,
     BasinTimeSchema,
-    DiscreteControlCompoundVariableSchema,
+    DiscreteControlCompoundvariableSchema,
     DiscreteControlConditionSchema,
     DiscreteControlLogicSchema,
     FlowBoundaryStaticSchema,
@@ -288,8 +288,8 @@ class ManningResistance(MultiNodeModel):
 
 
 class DiscreteControl(MultiNodeModel):
-    compound_variable: TableModel[DiscreteControlCompoundVariableSchema] = Field(
-        default_factory=TableModel[DiscreteControlCompoundVariableSchema],
+    compoundvariable: TableModel[DiscreteControlCompoundvariableSchema] = Field(
+        default_factory=TableModel[DiscreteControlCompoundvariableSchema],
         json_schema_extra={"sort_keys": ["name"]},
     )
     condition: TableModel[DiscreteControlConditionSchema] = Field(
diff --git a/python/ribasim/ribasim/nodes/discrete_control.py b/python/ribasim/ribasim/nodes/discrete_control.py
index 02b22840c..cc3cfcaf9 100644
--- a/python/ribasim/ribasim/nodes/discrete_control.py
+++ b/python/ribasim/ribasim/nodes/discrete_control.py
@@ -2,15 +2,15 @@
 
 from ribasim.input_base import TableModel
 from ribasim.schemas import (
-    DiscreteControlCompoundVariableSchema,
+    DiscreteControlCompoundvariableSchema,
     DiscreteControlConditionSchema,
     DiscreteControlLogicSchema,
 )
 
-__all__ = ["Condition", "Logic", "CompoundVariable"]
+__all__ = ["Condition", "Logic", "Compoundvariable"]
 
 
-class CompoundVariable(TableModel[DiscreteControlCompoundVariableSchema]):
+class Compoundvariable(TableModel[DiscreteControlCompoundvariableSchema]):
     def __init__(self, **kwargs):
         super().__init__(df=DataFrame(dict(**kwargs)))
 
diff --git a/python/ribasim/ribasim/schemas.py b/python/ribasim/ribasim/schemas.py
index aa4c738c2..213d97667 100644
--- a/python/ribasim/ribasim/schemas.py
+++ b/python/ribasim/ribasim/schemas.py
@@ -48,7 +48,7 @@ class BasinTimeSchema(_BaseSchema):
     urban_runoff: Series[float] = pa.Field(nullable=True)
 
 
-class DiscreteControlCompoundVariableSchema(_BaseSchema):
+class DiscreteControlCompoundvariableSchema(_BaseSchema):
     node_id: Series[Int32] = pa.Field(nullable=False, default=0)
     name: Series[str] = pa.Field(nullable=False)
     listen_node_type: Series[str] = pa.Field(nullable=False)
diff --git a/python/ribasim_testmodels/ribasim_testmodels/discrete_control.py b/python/ribasim_testmodels/ribasim_testmodels/discrete_control.py
index 2f14e8e9a..ea9eb3899 100644
--- a/python/ribasim_testmodels/ribasim_testmodels/discrete_control.py
+++ b/python/ribasim_testmodels/ribasim_testmodels/discrete_control.py
@@ -426,7 +426,7 @@ def compound_variable_condition_model() -> Model:
     model.discrete_control.add(
         Node(6, Point(1, 1)),
         [
-            discrete_control.CompoundVariable(
+            discrete_control.Compoundvariable(
                 name="flow_mean",
                 listen_node_type="FlowBoundary",
                 listen_node_id=[2, 3],
diff --git a/ribasim_qgis/core/nodes.py b/ribasim_qgis/core/nodes.py
index 95edb13bd..70750e419 100644
--- a/ribasim_qgis/core/nodes.py
+++ b/ribasim_qgis/core/nodes.py
@@ -691,6 +691,27 @@ def attributes(cls) -> list[QgsField]:
         ]
 
 
+class DiscreteControlCompoundvariable(Input):
+    @classmethod
+    def input_type(cls) -> str:
+        return "DiscreteControl / compoundvariable"
+
+    @classmethod
+    def geometry_type(cls) -> str:
+        return "No Geometry"
+
+    @classmethod
+    def attributes(cls) -> list[QgsField]:
+        return [
+            QgsField("node_id", QVariant.Int),
+            QgsField("name", QVariant.String),
+            QgsField("listen_node_type", QVariant.String),
+            QgsField("listen_node_id", QVariant.Int),
+            QgsField("variable", QVariant.String),
+            QgsField("weight", QVariant.Double),
+        ]
+
+
 class DiscreteControlCondition(Input):
     @classmethod
     def input_type(cls) -> str:

From 1bc18e7bc62ff2dffd48e30aa7137df30fd77c7a Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Wed, 10 Apr 2024 16:49:39 +0200
Subject: [PATCH 03/25] Read compound variable data in core

---
 core/src/parameter.jl             |  7 +++--
 core/src/read.jl                  | 47 ++++++++++++++++++++++++++++---
 core/src/schema.jl                |  1 +
 python/ribasim/ribasim/schemas.py |  1 +
 4 files changed, 49 insertions(+), 7 deletions(-)

diff --git a/core/src/parameter.jl b/core/src/parameter.jl
index 16774df21..2d8c5bca0 100644
--- a/core/src/parameter.jl
+++ b/core/src/parameter.jl
@@ -455,9 +455,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}}
diff --git a/core/src/read.jl b/core/src/read.jl
index 20643c943..af4e9345b 100644
--- a/core/src/read.jl
+++ b/core/src/read.jl
@@ -542,9 +542,49 @@ function Basin(db::DB, config::Config, chunk_sizes::Vector{Int})::Basin
     )
 end
 
+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 +597,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 +606,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 +616,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,
diff --git a/core/src/schema.jl b/core/src/schema.jl
index 2d3896e82..95d92d38b 100644
--- a/core/src/schema.jl
+++ b/core/src/schema.jl
@@ -191,6 +191,7 @@ end
     listen_node_id::Int
     variable::String
     weight::Float64
+    look_ahead::Union{Missing, Float64}
 end
 
 @version DiscreteControlConditionV1 begin
diff --git a/python/ribasim/ribasim/schemas.py b/python/ribasim/ribasim/schemas.py
index 213d97667..24dd99f76 100644
--- a/python/ribasim/ribasim/schemas.py
+++ b/python/ribasim/ribasim/schemas.py
@@ -55,6 +55,7 @@ class DiscreteControlCompoundvariableSchema(_BaseSchema):
     listen_node_id: Series[Int32] = pa.Field(nullable=False, default=0)
     variable: Series[str] = pa.Field(nullable=False)
     weight: Series[float] = pa.Field(nullable=False)
+    look_ahead: Series[float] = pa.Field(nullable=True)
 
 
 class DiscreteControlConditionSchema(_BaseSchema):

From 0cec53786665261ab9771ebb93a1b3a87e07d3be Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Wed, 10 Apr 2024 17:04:53 +0200
Subject: [PATCH 04/25] Handle compound variables in discrete_control callback

---
 core/src/callback.jl                                     | 9 +++++++--
 .../ribasim_testmodels/discrete_control.py               | 2 +-
 2 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/core/src/callback.jl b/core/src/callback.jl
index bd94043a5..ee49c6086 100644
--- a/core/src/callback.jl
+++ b/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
diff --git a/python/ribasim_testmodels/ribasim_testmodels/discrete_control.py b/python/ribasim_testmodels/ribasim_testmodels/discrete_control.py
index ea9eb3899..55aa1849b 100644
--- a/python/ribasim_testmodels/ribasim_testmodels/discrete_control.py
+++ b/python/ribasim_testmodels/ribasim_testmodels/discrete_control.py
@@ -430,7 +430,7 @@ def compound_variable_condition_model() -> Model:
                 name="flow_mean",
                 listen_node_type="FlowBoundary",
                 listen_node_id=[2, 3],
-                variable="flow",
+                variable="flow_rate",
                 weight=0.5,
             ),
             discrete_control.Condition(

From 78e8364036f1ac720cd03bb13ffe1b678e8ea368 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Wed, 10 Apr 2024 18:24:48 +0200
Subject: [PATCH 05/25] Fix part of tests

---
 core/src/callback.jl      |  7 +++++--
 core/src/validation.jl    | 43 +++++++++++++++++++++------------------
 core/test/control_test.jl |  4 ++--
 3 files changed, 30 insertions(+), 24 deletions(-)

diff --git a/core/src/callback.jl b/core/src/callback.jl
index ee49c6086..0ac0701d4 100644
--- a/core/src/callback.jl
+++ b/core/src/callback.jl
@@ -277,7 +277,10 @@ 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
@@ -374,7 +377,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
 
diff --git a/core/src/validation.jl b/core/src/validation.jl
index 202eec6a8..478fca586 100644
--- a/core/src/validation.jl
+++ b/core/src/validation.jl
@@ -552,29 +552,32 @@ function valid_discrete_control(p::Parameters, config::Config)::Bool
             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
diff --git a/core/test/control_test.jl b/core/test/control_test.jl
index fdb2e80f8..55d934a10 100644
--- a/core/test/control_test.jl
+++ b/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]

From ff54bb34d22c2167661222a7ff8156a3642ea387 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Wed, 10 Apr 2024 22:28:46 +0200
Subject: [PATCH 06/25] Fix the last test in a hacky way

---
 core/src/callback.jl | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

diff --git a/core/src/callback.jl b/core/src/callback.jl
index 0ac0701d4..7e6daec53 100644
--- a/core/src/callback.jl
+++ b/core/src/callback.jl
@@ -281,15 +281,16 @@ function discrete_control_affect_upcrossing!(integrator, 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
@@ -316,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

From acca9496859abb49abc1b67ad178bad07350f0e6 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Thu, 11 Apr 2024 07:09:16 +0200
Subject: [PATCH 07/25] Add tests

---
 core/test/control_test.jl                     | 20 +++++++++++++++++++
 .../ribasim_testmodels/discrete_control.py    |  4 ++--
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/core/test/control_test.jl b/core/test/control_test.jl
index 55d934a10..e0913ba1d 100644
--- a/core/test/control_test.jl
+++ b/core/test/control_test.jl
@@ -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
diff --git a/python/ribasim_testmodels/ribasim_testmodels/discrete_control.py b/python/ribasim_testmodels/ribasim_testmodels/discrete_control.py
index 55aa1849b..a81b3e64c 100644
--- a/python/ribasim_testmodels/ribasim_testmodels/discrete_control.py
+++ b/python/ribasim_testmodels/ribasim_testmodels/discrete_control.py
@@ -416,7 +416,7 @@ def compound_variable_condition_model() -> Model:
     )
     model.flow_boundary.add(
         Node(3, Point(0, 1)),
-        [flow_boundary.Time(time=["2020-01-01", "2022-01-01"], flow_rate=[0.0, 2.0])],
+        [flow_boundary.Time(time=["2020-01-01", "2021-01-01"], flow_rate=[0.0, 2.0])],
     )
     model.pump.add(
         Node(4, Point(2, 0)),
@@ -437,7 +437,7 @@ def compound_variable_condition_model() -> Model:
                 listen_node_type="compound",
                 listen_node_id=[0],  # Irrelevant
                 variable="flow_mean",
-                greater_than=1.0,
+                greater_than=0.5,
             ),
             discrete_control.Logic(truth_state=["T", "F"], control_state=["On", "Off"]),
         ],

From 72e65b57b07385670aeb086ace947f6722fc51d5 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Thu, 11 Apr 2024 07:15:24 +0200
Subject: [PATCH 08/25] small plotting fix

---
 python/ribasim/ribasim/model.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/python/ribasim/ribasim/model.py b/python/ribasim/ribasim/model.py
index cbe9798e9..bb0b74f2c 100644
--- a/python/ribasim/ribasim/model.py
+++ b/python/ribasim/ribasim/model.py
@@ -323,7 +323,7 @@ def plot_control_listen(self, ax):
         # Listen edges from DiscreteControl
         for table in (
             self.discrete_control.condition.df,
-            self.discrete_control.compound_variable.df,
+            self.discrete_control.compoundvariable.df,
         ):
             if table is None:
                 continue

From 44eab71344d1f39ac54689446009b02cca674f07 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Thu, 11 Apr 2024 07:21:08 +0200
Subject: [PATCH 09/25] small qgis fix

---
 ribasim_qgis/core/nodes.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/ribasim_qgis/core/nodes.py b/ribasim_qgis/core/nodes.py
index 70750e419..7fe9390f3 100644
--- a/ribasim_qgis/core/nodes.py
+++ b/ribasim_qgis/core/nodes.py
@@ -709,6 +709,7 @@ def attributes(cls) -> list[QgsField]:
             QgsField("listen_node_id", QVariant.Int),
             QgsField("variable", QVariant.String),
             QgsField("weight", QVariant.Double),
+            QgsField("look_ahead", QVariant.Double),
         ]
 
 
@@ -729,6 +730,7 @@ def attributes(cls) -> list[QgsField]:
             QgsField("listen_node_id", QVariant.Int),
             QgsField("variable", QVariant.String),
             QgsField("greater_than", QVariant.Double),
+            QgsField("look_ahead", QVariant.Double),
         ]
 
 

From 25997fea2bb007cc77798e1e18617f3604aaa0a3 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Thu, 11 Apr 2024 07:35:17 +0200
Subject: [PATCH 10/25] Update schemas in usage.qmd

---
 docs/core/usage.qmd | 17 +++++++++++++++--
 1 file changed, 15 insertions(+), 2 deletions(-)

diff --git a/docs/core/usage.qmd b/docs/core/usage.qmd
index 3613fa409..4cd4b3788 100644
--- a/docs/core/usage.qmd
+++ b/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,6 +626,18 @@ node_id               | Int32   | -            | sorted
 
 DiscreteControl is implemented based on [VectorContinuousCallback](https://docs.sciml.ai/DiffEqDocs/stable/features/callback_functions/#VectorContinuousCallback).
 
+## DiscreteControl / compoundvariable
+
+column            | type     | unit    | restriction
+----------------- | -------- | ------- | -----------
+node_id           | Int32    | -       | sorted
+name              | String   | -       | -
+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 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.
@@ -632,9 +645,9 @@ The condition schema defines conditions of the form 'the discrete_control node w
 column            | type     | unit    | restriction
 ----------------- | -------- | ------- | -----------
 node_id           | Int32    | -       | sorted
-listen_node_type  | String   | -       | known node type
+listen_node_type  | String   | -       | known node type or "compound"
 listen_node_id    | Int32    | -       | sorted per node_id
-variable          | String   | -       | must be "level" or "flow_rate", sorted per listen_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)
 

From 3144ef1098b6528137492c6a6597354db88c5b74 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Thu, 11 Apr 2024 07:42:23 +0200
Subject: [PATCH 11/25] Update docstrings

---
 core/src/parameter.jl | 6 ++++--
 core/src/read.jl      | 4 ++++
 2 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/core/src/parameter.jl b/core/src/parameter.jl
index 2d8c5bca0..231f0a9c4 100644
--- a/core/src/parameter.jl
+++ b/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)
diff --git a/core/src/read.jl b/core/src/read.jl
index af4e9345b..bed7abb0c 100644
--- a/core/src/read.jl
+++ b/core/src/read.jl
@@ -542,6 +542,10 @@ 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}[]

From e842bdfd9ea586f7acc8f9858cc7a23bc572c401 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Thu, 11 Apr 2024 09:00:07 +0200
Subject: [PATCH 12/25] update sorting and schema explanations

---
 core/src/validation.jl           |  2 ++
 docs/core/usage.qmd              | 13 +++++++++++--
 python/ribasim/ribasim/config.py |  6 ++----
 3 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/core/src/validation.jl b/core/src/validation.jl
index 478fca586..1280049dc 100644
--- a/core/src/validation.jl
+++ b/core/src/validation.jl
@@ -102,6 +102,7 @@ sort_by_id_state_level(row) = (row.node_id, row.control_state, row.level)
 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_id_name(row) = (row.node_id, row.name)
 
 # 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 +111,7 @@ sort_by_function(table::StructVector{BasinProfileV1}) = sort_by_id_level
 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
 
 const TimeSchemas = Union{
     BasinTimeV1,
diff --git a/docs/core/usage.qmd b/docs/core/usage.qmd
index 4cd4b3788..aa7c5f73a 100644
--- a/docs/core/usage.qmd
+++ b/docs/core/usage.qmd
@@ -628,10 +628,18 @@ DiscreteControl is implemented based on [VectorContinuousCallback](https://docs.
 
 ## 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,
+$$
+
+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   | -       | -
+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
@@ -640,7 +648,8 @@ look_ahead        | Float64  | $s$     | Only on transient boundary conditions,
 
 ## 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 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.
+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
 ----------------- | -------- | ------- | -----------
diff --git a/python/ribasim/ribasim/config.py b/python/ribasim/ribasim/config.py
index 64a5019e9..3fed1f8a7 100644
--- a/python/ribasim/ribasim/config.py
+++ b/python/ribasim/ribasim/config.py
@@ -290,13 +290,11 @@ class ManningResistance(MultiNodeModel):
 class DiscreteControl(MultiNodeModel):
     compoundvariable: TableModel[DiscreteControlCompoundvariableSchema] = Field(
         default_factory=TableModel[DiscreteControlCompoundvariableSchema],
-        json_schema_extra={"sort_keys": ["name"]},
+        json_schema_extra={"sort_keys": ["node_id", "name"]},
     )
     condition: TableModel[DiscreteControlConditionSchema] = Field(
         default_factory=TableModel[DiscreteControlConditionSchema],
-        json_schema_extra={
-            "sort_keys": ["node_id", "listen_node_id", "variable", "greater_than"]
-        },
+        json_schema_extra={"sort_keys": ["node_id"]},
     )
     logic: TableModel[DiscreteControlLogicSchema] = Field(
         default_factory=TableModel[DiscreteControlLogicSchema],

From c975b0c0bac347fdf66398740fa7b5bc11950ddc Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Thu, 11 Apr 2024 11:19:54 +0200
Subject: [PATCH 13/25] Fix table sorting

---
 core/src/validation.jl           |  5 +++++
 python/ribasim/ribasim/config.py | 10 +++++++++-
 python/ribasim/tests/test_io.py  |  1 +
 3 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/core/src/validation.jl b/core/src/validation.jl
index 1280049dc..d4b46b203 100644
--- a/core/src/validation.jl
+++ b/core/src/validation.jl
@@ -102,7 +102,10 @@ sort_by_id_state_level(row) = (row.node_id, row.control_state, row.level)
 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
@@ -112,6 +115,8 @@ 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,
diff --git a/python/ribasim/ribasim/config.py b/python/ribasim/ribasim/config.py
index 3fed1f8a7..e00f20ed2 100644
--- a/python/ribasim/ribasim/config.py
+++ b/python/ribasim/ribasim/config.py
@@ -294,7 +294,15 @@ class DiscreteControl(MultiNodeModel):
     )
     condition: TableModel[DiscreteControlConditionSchema] = Field(
         default_factory=TableModel[DiscreteControlConditionSchema],
-        json_schema_extra={"sort_keys": ["node_id"]},
+        json_schema_extra={
+            "sort_keys": [
+                "node_id",
+                "listen_node_type",
+                "listen_node_id",
+                "variable",
+                "greater_than",
+            ]
+        },
     )
     logic: TableModel[DiscreteControlLogicSchema] = Field(
         default_factory=TableModel[DiscreteControlLogicSchema],
diff --git a/python/ribasim/tests/test_io.py b/python/ribasim/tests/test_io.py
index bd6fcc7a8..e9b724eaa 100644
--- a/python/ribasim/tests/test_io.py
+++ b/python/ribasim/tests/test_io.py
@@ -102,6 +102,7 @@ def test_sort(level_setpoint_with_minmax, tmp_path):
     assert table.df.iloc[0]["greater_than"] == 15.0
     assert table._sort_keys == [
         "node_id",
+        "listen_node_type",
         "listen_node_id",
         "variable",
         "greater_than",

From 0ba1d6fdfb66884bb0b60afe95d7d85fc056b453 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Mon, 15 Apr 2024 15:34:40 +0200
Subject: [PATCH 14/25] Refactor variable input

---
 core/src/read.jl                              | 50 ++++---------------
 core/src/schema.jl                            | 11 ++--
 core/src/validation.jl                        |  8 +--
 python/ribasim/ribasim/config.py              | 18 ++++---
 python/ribasim/ribasim/model.py               |  2 +-
 .../ribasim/ribasim/nodes/discrete_control.py |  6 +--
 python/ribasim/ribasim/schemas.py             | 23 ++++-----
 .../ribasim_testmodels/allocation.py          | 18 ++++---
 .../ribasim_testmodels/discrete_control.py    | 44 +++++++++-------
 .../ribasim_testmodels/dutch_waterways.py     |  6 ++-
 .../ribasim_testmodels/invalid.py             |  6 ++-
 .../ribasim_testmodels/pid_control.py         |  6 ++-
 ribasim_qgis/core/nodes.py                    |  9 +---
 13 files changed, 94 insertions(+), 113 deletions(-)

diff --git a/core/src/read.jl b/core/src/read.jl
index bed7abb0c..b9bf0c8ec 100644
--- a/core/src/read.jl
+++ b/core/src/read.jl
@@ -542,52 +542,22 @@ 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)
+function DiscreteControl(db::DB, config::Config)::DiscreteControl
+    compound_variable = load_structvector(db, config, DiscreteControlVariableV1)
+    condition = load_structvector(db, config, DiscreteControlConditionV1)
+
     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)
+    for id in unique(condition.node_id)
+        group = filter(row -> row.node_id == id, compound_variable)
+        push!(listen_node_id, NodeID.(group.listen_node_type, group.listen_node_id))
+        push!(variable, group.variable)
+        push!(weight, coalesce.(group.weight, 1.0))
+        push!(look_ahead, coalesce.(group.look_ahead, 0.0))
     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()
diff --git a/core/src/schema.jl b/core/src/schema.jl
index 95d92d38b..60147c38a 100644
--- a/core/src/schema.jl
+++ b/core/src/schema.jl
@@ -1,8 +1,8 @@
 # These schemas define the name of database tables and the configuration file structure
 # The identifier is parsed as ribasim.nodetype.kind, no capitals or underscores are allowed.
+@schema "ribasim.discretecontrol.variable" DiscreteControlVariable
 @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
@@ -184,23 +184,18 @@ end
     node_id::Int32
 end
 
-@version DiscreteControlCompoundvariableV1 begin
+@version DiscreteControlVariableV1 begin
     node_id::Int32
-    name::String
     listen_node_type::String
     listen_node_id::Int
     variable::String
-    weight::Float64
+    weight::Union{Missing, Float64}
     look_ahead::Union{Missing, Float64}
 end
 
 @version DiscreteControlConditionV1 begin
     node_id::Int32
-    listen_node_type::String
-    listen_node_id::Int32
-    variable::String
     greater_than::Float64
-    look_ahead::Union{Missing, Float64}
 end
 
 @version DiscreteControlLogicV1 begin
diff --git a/core/src/validation.jl b/core/src/validation.jl
index d4b46b203..846cca5b4 100644
--- a/core/src/validation.jl
+++ b/core/src/validation.jl
@@ -104,7 +104,9 @@ 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_variable(row) =
+    (row.node_id, row.listen_node_type, row.listen_node_id, row.variable)
+sort_by_id_greater_than(row) = (row.node_id, row.greater_than)
 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
@@ -114,8 +116,8 @@ sort_by_function(table::StructVector{BasinProfileV1}) = sort_by_id_level
 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{DiscreteControlVariableV1}) = sort_by_variable
+sort_by_function(table::StructVector{DiscreteControlConditionV1}) = sort_by_id_greater_than
 sort_by_function(table::StructVector{DiscreteControlLogicV1}) = sort_by_truth_state
 
 const TimeSchemas = Union{
diff --git a/python/ribasim/ribasim/config.py b/python/ribasim/ribasim/config.py
index e00f20ed2..dd7fbbf0d 100644
--- a/python/ribasim/ribasim/config.py
+++ b/python/ribasim/ribasim/config.py
@@ -21,9 +21,9 @@
     BasinStaticSchema,
     BasinSubgridSchema,
     BasinTimeSchema,
-    DiscreteControlCompoundvariableSchema,
     DiscreteControlConditionSchema,
     DiscreteControlLogicSchema,
+    DiscreteControlVariableSchema,
     FlowBoundaryStaticSchema,
     FlowBoundaryTimeSchema,
     FlowDemandStaticSchema,
@@ -288,18 +288,22 @@ class ManningResistance(MultiNodeModel):
 
 
 class DiscreteControl(MultiNodeModel):
-    compoundvariable: TableModel[DiscreteControlCompoundvariableSchema] = Field(
-        default_factory=TableModel[DiscreteControlCompoundvariableSchema],
-        json_schema_extra={"sort_keys": ["node_id", "name"]},
-    )
-    condition: TableModel[DiscreteControlConditionSchema] = Field(
-        default_factory=TableModel[DiscreteControlConditionSchema],
+    variable: TableModel[DiscreteControlVariableSchema] = Field(
+        default_factory=TableModel[DiscreteControlVariableSchema],
         json_schema_extra={
             "sort_keys": [
                 "node_id",
                 "listen_node_type",
                 "listen_node_id",
                 "variable",
+            ]
+        },
+    )
+    condition: TableModel[DiscreteControlConditionSchema] = Field(
+        default_factory=TableModel[DiscreteControlConditionSchema],
+        json_schema_extra={
+            "sort_keys": [
+                "node_id",
                 "greater_than",
             ]
         },
diff --git a/python/ribasim/ribasim/model.py b/python/ribasim/ribasim/model.py
index d8d6b08e0..8bbb1666a 100644
--- a/python/ribasim/ribasim/model.py
+++ b/python/ribasim/ribasim/model.py
@@ -323,7 +323,7 @@ def plot_control_listen(self, ax):
         # Listen edges from DiscreteControl
         for table in (
             self.discrete_control.condition.df,
-            self.discrete_control.compoundvariable.df,
+            self.discrete_control.variable.df,
         ):
             if table is None:
                 continue
diff --git a/python/ribasim/ribasim/nodes/discrete_control.py b/python/ribasim/ribasim/nodes/discrete_control.py
index cc3cfcaf9..df7b9f42b 100644
--- a/python/ribasim/ribasim/nodes/discrete_control.py
+++ b/python/ribasim/ribasim/nodes/discrete_control.py
@@ -2,15 +2,15 @@
 
 from ribasim.input_base import TableModel
 from ribasim.schemas import (
-    DiscreteControlCompoundvariableSchema,
     DiscreteControlConditionSchema,
     DiscreteControlLogicSchema,
+    DiscreteControlVariableSchema,
 )
 
-__all__ = ["Condition", "Logic", "Compoundvariable"]
+__all__ = ["Condition", "Logic", "Variable"]
 
 
-class Compoundvariable(TableModel[DiscreteControlCompoundvariableSchema]):
+class Variable(TableModel[DiscreteControlVariableSchema]):
     def __init__(self, **kwargs):
         super().__init__(df=DataFrame(dict(**kwargs)))
 
diff --git a/python/ribasim/ribasim/schemas.py b/python/ribasim/ribasim/schemas.py
index 24dd99f76..260831282 100644
--- a/python/ribasim/ribasim/schemas.py
+++ b/python/ribasim/ribasim/schemas.py
@@ -48,23 +48,9 @@ class BasinTimeSchema(_BaseSchema):
     urban_runoff: Series[float] = pa.Field(nullable=True)
 
 
-class DiscreteControlCompoundvariableSchema(_BaseSchema):
-    node_id: Series[Int32] = pa.Field(nullable=False, default=0)
-    name: Series[str] = pa.Field(nullable=False)
-    listen_node_type: Series[str] = pa.Field(nullable=False)
-    listen_node_id: Series[Int32] = pa.Field(nullable=False, default=0)
-    variable: Series[str] = pa.Field(nullable=False)
-    weight: Series[float] = pa.Field(nullable=False)
-    look_ahead: Series[float] = pa.Field(nullable=True)
-
-
 class DiscreteControlConditionSchema(_BaseSchema):
     node_id: Series[Int32] = pa.Field(nullable=False, default=0)
-    listen_node_type: Series[str] = pa.Field(nullable=False)
-    listen_node_id: Series[Int32] = pa.Field(nullable=False, default=0)
-    variable: Series[str] = pa.Field(nullable=False)
     greater_than: Series[float] = pa.Field(nullable=False)
-    look_ahead: Series[float] = pa.Field(nullable=True)
 
 
 class DiscreteControlLogicSchema(_BaseSchema):
@@ -73,6 +59,15 @@ class DiscreteControlLogicSchema(_BaseSchema):
     control_state: Series[str] = pa.Field(nullable=False)
 
 
+class DiscreteControlVariableSchema(_BaseSchema):
+    node_id: Series[Int32] = pa.Field(nullable=False, default=0)
+    listen_node_type: Series[str] = pa.Field(nullable=False)
+    listen_node_id: Series[Int32] = pa.Field(nullable=False, default=0)
+    variable: Series[str] = pa.Field(nullable=False)
+    weight: Series[float] = pa.Field(nullable=True)
+    look_ahead: Series[float] = pa.Field(nullable=True)
+
+
 class FlowBoundaryStaticSchema(_BaseSchema):
     node_id: Series[Int32] = pa.Field(nullable=False, default=0)
     active: Series[pa.BOOL] = pa.Field(nullable=True)
diff --git a/python/ribasim_testmodels/ribasim_testmodels/allocation.py b/python/ribasim_testmodels/ribasim_testmodels/allocation.py
index 1e102e8c0..71353cf93 100644
--- a/python/ribasim_testmodels/ribasim_testmodels/allocation.py
+++ b/python/ribasim_testmodels/ribasim_testmodels/allocation.py
@@ -406,11 +406,13 @@ def fractional_flow_subnetwork_model() -> Model:
     model.discrete_control.add(
         Node(10, Point(-1, 2), subnetwork_id=2),
         [
-            discrete_control.Condition(
+            discrete_control.Variable(
                 listen_node_type="FlowBoundary",
                 listen_node_id=[1],
                 variable="flow_rate",
-                greater_than=3e-3,
+            ),
+            discrete_control.Condition(
+                greater_than=[3e-3],
             ),
             discrete_control.Logic(truth_state=["F", "T"], control_state=["A", "B"]),
         ],
@@ -500,11 +502,13 @@ def allocation_example_model() -> Model:
     model.discrete_control.add(
         Node(11, Point(4.5, 0.25), subnetwork_id=2),
         [
-            discrete_control.Condition(
+            discrete_control.Variable(
                 listen_node_type="Basin",
                 listen_node_id=[5],
                 variable="level",
-                greater_than=0.52,
+            ),
+            discrete_control.Condition(
+                greater_than=[0.52],
             ),
             discrete_control.Logic(
                 truth_state=["T", "F"], control_state=["divert", "close"]
@@ -679,11 +683,13 @@ def main_network_with_subnetworks_model() -> Model:
     model.discrete_control.add(
         Node(33, Point(13, 5), subnetwork_id=5),
         [
-            discrete_control.Condition(
+            discrete_control.Variable(
                 listen_node_type="Basin",
                 listen_node_id=[25],
                 variable="level",
-                greater_than=0.003,
+            ),
+            discrete_control.Condition(
+                greater_than=[0.003],
             ),
             discrete_control.Logic(truth_state=["F", "T"], control_state=["A", "B"]),
         ],
diff --git a/python/ribasim_testmodels/ribasim_testmodels/discrete_control.py b/python/ribasim_testmodels/ribasim_testmodels/discrete_control.py
index a81b3e64c..1a6037d8a 100644
--- a/python/ribasim_testmodels/ribasim_testmodels/discrete_control.py
+++ b/python/ribasim_testmodels/ribasim_testmodels/discrete_control.py
@@ -54,10 +54,12 @@ def pump_discrete_control_model() -> Model:
     model.discrete_control.add(
         Node(5, Point(1, 1)),
         [
-            discrete_control.Condition(
+            discrete_control.Variable(
                 listen_node_type="Basin",
                 listen_node_id=[1, 3],
                 variable="level",
+            ),
+            discrete_control.Condition(
                 greater_than=[0.8, 0.4],
             ),
             discrete_control.Logic(
@@ -69,10 +71,12 @@ def pump_discrete_control_model() -> Model:
     model.discrete_control.add(
         Node(6, Point(2, -1)),
         [
-            discrete_control.Condition(
+            discrete_control.Variable(
                 listen_node_type="Basin",
-                listen_node_id=3,
+                listen_node_id=[3],
                 variable="level",
+            ),
+            discrete_control.Condition(
                 greater_than=[0.45],
             ),
             discrete_control.Logic(
@@ -140,13 +144,15 @@ def flow_condition_model() -> Model:
     model.discrete_control.add(
         Node(5, Point(1, 1)),
         [
-            discrete_control.Condition(
+            discrete_control.Variable(
                 listen_node_type="FlowBoundary",
-                listen_node_id=1,
+                listen_node_id=[1],
                 variable="flow_rate",
-                greater_than=[20 / (86400)],
                 look_ahead=60 * 86400,
             ),
+            discrete_control.Condition(
+                greater_than=[20 / (86400)],
+            ),
             discrete_control.Logic(truth_state=["T", "F"], control_state=["off", "on"]),
         ],
     )
@@ -203,13 +209,15 @@ def level_boundary_condition_model() -> Model:
     model.discrete_control.add(
         Node(6, Point(1.5, 1)),
         [
-            discrete_control.Condition(
+            discrete_control.Variable(
                 listen_node_type="LevelBoundary",
                 listen_node_id=[1],
                 variable="level",
-                greater_than=6.0,
                 look_ahead=60 * 86400,
             ),
+            discrete_control.Condition(
+                greater_than=[6.0],
+            ),
             discrete_control.Logic(truth_state=["T", "F"], control_state=["on", "off"]),
         ],
     )
@@ -275,11 +283,13 @@ def tabulated_rating_curve_control_model() -> Model:
     model.discrete_control.add(
         Node(4, Point(1, 1)),
         [
-            discrete_control.Condition(
+            discrete_control.Variable(
                 listen_node_type="Basin",
                 listen_node_id=[1],
                 variable="level",
-                greater_than=0.5,
+            ),
+            discrete_control.Condition(
+                greater_than=[0.5],
             ),
             discrete_control.Logic(
                 truth_state=["T", "F"], control_state=["low", "high"]
@@ -342,10 +352,12 @@ def level_setpoint_with_minmax_model() -> Model:
     model.discrete_control.add(
         Node(7, Point(1, 0)),
         [
-            discrete_control.Condition(
+            discrete_control.Variable(
                 listen_node_type="Basin",
-                listen_node_id=1,
+                listen_node_id=[1],
                 variable="level",
+            ),
+            discrete_control.Condition(
                 # min, setpoint, max
                 greater_than=[5.0, 10.0, 15.0],
             ),
@@ -426,18 +438,14 @@ def compound_variable_condition_model() -> Model:
     model.discrete_control.add(
         Node(6, Point(1, 1)),
         [
-            discrete_control.Compoundvariable(
-                name="flow_mean",
+            discrete_control.Variable(
                 listen_node_type="FlowBoundary",
                 listen_node_id=[2, 3],
                 variable="flow_rate",
                 weight=0.5,
             ),
             discrete_control.Condition(
-                listen_node_type="compound",
-                listen_node_id=[0],  # Irrelevant
-                variable="flow_mean",
-                greater_than=0.5,
+                greater_than=[0.5],
             ),
             discrete_control.Logic(truth_state=["T", "F"], control_state=["On", "Off"]),
         ],
diff --git a/python/ribasim_testmodels/ribasim_testmodels/dutch_waterways.py b/python/ribasim_testmodels/ribasim_testmodels/dutch_waterways.py
index fd990957b..b13bb4bcb 100644
--- a/python/ribasim_testmodels/ribasim_testmodels/dutch_waterways.py
+++ b/python/ribasim_testmodels/ribasim_testmodels/dutch_waterways.py
@@ -121,10 +121,12 @@ def dutch_waterways_model() -> Model:
     model.discrete_control.add(
         Node(17, Point(183612, 441833), name="Controller Driel"),
         [
-            discrete_control.Condition(
+            discrete_control.Variable(
                 listen_node_type="FlowBoundary",
-                listen_node_id=1,
+                listen_node_id=[1],
                 variable="flow_rate",
+            ),
+            discrete_control.Condition(
                 greater_than=[250, 275, 750, 800],
             ),
             discrete_control.Logic(
diff --git a/python/ribasim_testmodels/ribasim_testmodels/invalid.py b/python/ribasim_testmodels/ribasim_testmodels/invalid.py
index 906a7b485..29c9a8405 100644
--- a/python/ribasim_testmodels/ribasim_testmodels/invalid.py
+++ b/python/ribasim_testmodels/ribasim_testmodels/invalid.py
@@ -148,16 +148,18 @@ def invalid_discrete_control_model() -> Model:
     model.discrete_control.add(
         Node(5, Point(1, 1)),
         [
-            discrete_control.Condition(
+            discrete_control.Variable(
                 listen_node_type=["Basin", "FlowBoundary", "FlowBoundary"],
                 listen_node_id=[1, 4, 4],
                 variable=["level", "flow_rate", "flow_rate"],
-                greater_than=[0.5, 1.5, 1.5],
                 # Invalid: look_ahead can only be specified for timeseries variables.
                 # Invalid: this look_ahead will go past the provided timeseries during simulation.
                 # Invalid: look_ahead must be non-negative.
                 look_ahead=[100.0, 40 * 24 * 60 * 60, -10.0],
             ),
+            discrete_control.Condition(
+                greater_than=[0.5, 1.5, 1.5],
+            ),
             # Invalid: DiscreteControl node #4 has 2 conditions so
             # truth states have to be of length 2
             discrete_control.Logic(truth_state=["FFFF"], control_state=["foo"]),
diff --git a/python/ribasim_testmodels/ribasim_testmodels/pid_control.py b/python/ribasim_testmodels/ribasim_testmodels/pid_control.py
index d8c591c09..05ebc31b7 100644
--- a/python/ribasim_testmodels/ribasim_testmodels/pid_control.py
+++ b/python/ribasim_testmodels/ribasim_testmodels/pid_control.py
@@ -129,11 +129,13 @@ def discrete_control_of_pid_control_model() -> Model:
     model.discrete_control.add(
         Node(7, Point(0, 1)),
         [
-            discrete_control.Condition(
+            discrete_control.Variable(
                 listen_node_type="LevelBoundary",
                 listen_node_id=[1],
                 variable="level",
-                greater_than=5.0,
+            ),
+            discrete_control.Condition(
+                greater_than=[5.0],
             ),
             discrete_control.Logic(
                 truth_state=["T", "F"], control_state=["target_high", "target_low"]
diff --git a/ribasim_qgis/core/nodes.py b/ribasim_qgis/core/nodes.py
index 7fe9390f3..5c0393a81 100644
--- a/ribasim_qgis/core/nodes.py
+++ b/ribasim_qgis/core/nodes.py
@@ -691,10 +691,10 @@ def attributes(cls) -> list[QgsField]:
         ]
 
 
-class DiscreteControlCompoundvariable(Input):
+class DiscreteControlVariable(Input):
     @classmethod
     def input_type(cls) -> str:
-        return "DiscreteControl / compoundvariable"
+        return "DiscreteControl / variable"
 
     @classmethod
     def geometry_type(cls) -> str:
@@ -704,7 +704,6 @@ def geometry_type(cls) -> str:
     def attributes(cls) -> list[QgsField]:
         return [
             QgsField("node_id", QVariant.Int),
-            QgsField("name", QVariant.String),
             QgsField("listen_node_type", QVariant.String),
             QgsField("listen_node_id", QVariant.Int),
             QgsField("variable", QVariant.String),
@@ -726,11 +725,7 @@ def geometry_type(cls) -> str:
     def attributes(cls) -> list[QgsField]:
         return [
             QgsField("node_id", QVariant.Int),
-            QgsField("listen_node_type", QVariant.String),
-            QgsField("listen_node_id", QVariant.Int),
-            QgsField("variable", QVariant.String),
             QgsField("greater_than", QVariant.Double),
-            QgsField("look_ahead", QVariant.Double),
         ]
 
 

From 17d2a1531e3ff970938e1d91e17b6e8f383d56d1 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Mon, 15 Apr 2024 17:00:33 +0200
Subject: [PATCH 15/25] Add compound_variable_id

---
 core/src/read.jl                              | 21 +++++++++++++------
 core/src/schema.jl                            |  2 ++
 core/src/validation.jl                        |  1 -
 docs/python/examples.ipynb                    |  8 +++----
 python/ribasim/ribasim/config.py              |  2 +-
 python/ribasim/ribasim/model.py               | 11 +++-------
 python/ribasim/ribasim/schemas.py             |  2 ++
 .../ribasim_testmodels/allocation.py          |  6 ++++++
 .../ribasim_testmodels/discrete_control.py    | 13 ++++++++++++
 .../ribasim_testmodels/dutch_waterways.py     |  2 ++
 .../ribasim_testmodels/invalid.py             |  2 ++
 .../ribasim_testmodels/pid_control.py         |  2 ++
 12 files changed, 52 insertions(+), 20 deletions(-)

diff --git a/core/src/read.jl b/core/src/read.jl
index b9bf0c8ec..a0aec5fda 100644
--- a/core/src/read.jl
+++ b/core/src/read.jl
@@ -546,17 +546,26 @@ function DiscreteControl(db::DB, config::Config)::DiscreteControl
     compound_variable = load_structvector(db, config, DiscreteControlVariableV1)
     condition = load_structvector(db, config, DiscreteControlConditionV1)
 
+    node_id = NodeID[]
     listen_node_id = Vector{NodeID}[]
     variable = Vector{String}[]
     weight = Vector{Float64}[]
     look_ahead = Vector{Float64}[]
 
     for id in unique(condition.node_id)
-        group = filter(row -> row.node_id == id, compound_variable)
-        push!(listen_node_id, NodeID.(group.listen_node_type, group.listen_node_id))
-        push!(variable, group.variable)
-        push!(weight, coalesce.(group.weight, 1.0))
-        push!(look_ahead, coalesce.(group.look_ahead, 0.0))
+        group_id = filter(row -> row.node_id == id, compound_variable)
+        for group_variable in
+            StructVector.(IterTools.groupby(row -> row.compound_variable_id, group_id))
+            first_row = first(group_variable)
+            push!(node_id, NodeID(NodeType.DiscreteControl, first_row.node_id))
+            push!(
+                listen_node_id,
+                NodeID.(group_variable.listen_node_type, group_variable.listen_node_id),
+            )
+            push!(variable, group_variable.variable)
+            push!(weight, coalesce.(group_variable.weight, 1.0))
+            push!(look_ahead, coalesce.(group_variable.look_ahead, 0.0))
+        end
     end
 
     condition_value = fill(false, length(condition.node_id))
@@ -589,7 +598,7 @@ function DiscreteControl(db::DB, config::Config)::DiscreteControl
     )
 
     return DiscreteControl(
-        NodeID.(NodeType.DiscreteControl, condition.node_id), # Not unique
+        node_id, # Not unique
         listen_node_id,
         variable,
         weight,
diff --git a/core/src/schema.jl b/core/src/schema.jl
index 60147c38a..60eb87ac8 100644
--- a/core/src/schema.jl
+++ b/core/src/schema.jl
@@ -186,6 +186,7 @@ end
 
 @version DiscreteControlVariableV1 begin
     node_id::Int32
+    compound_variable_id::Int
     listen_node_type::String
     listen_node_id::Int
     variable::String
@@ -195,6 +196,7 @@ end
 
 @version DiscreteControlConditionV1 begin
     node_id::Int32
+    compound_variable_id::Int
     greater_than::Float64
 end
 
diff --git a/core/src/validation.jl b/core/src/validation.jl
index 846cca5b4..580b4d5ba 100644
--- a/core/src/validation.jl
+++ b/core/src/validation.jl
@@ -118,7 +118,6 @@ 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{DiscreteControlVariableV1}) = sort_by_variable
 sort_by_function(table::StructVector{DiscreteControlConditionV1}) = sort_by_id_greater_than
-sort_by_function(table::StructVector{DiscreteControlLogicV1}) = sort_by_truth_state
 
 const TimeSchemas = Union{
     BasinTimeV1,
diff --git a/docs/python/examples.ipynb b/docs/python/examples.ipynb
index 5a76e3ae7..f78357420 100644
--- a/docs/python/examples.ipynb
+++ b/docs/python/examples.ipynb
@@ -492,10 +492,10 @@
     "model.discrete_control.add(\n",
     "    Node(7, Point(1.0, 0.0)),\n",
     "    [\n",
+    "        discrete_control.Variable(\n",
+    "            listen_node_id=[1], listen_node_type=\"Basin\", variable=\"level\"\n",
+    "        ),\n",
     "        discrete_control.Condition(\n",
-    "            listen_node_id=[1, 1, 1],\n",
-    "            listen_node_type=[\"Basin\", \"Basin\", \"Basin\"],\n",
-    "            variable=[\"level\", \"level\", \"level\"],\n",
     "            greater_than=[5.0, 10.0, 15.0],\n",
     "        ),\n",
     "        discrete_control.Logic(\n",
@@ -691,7 +691,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Now run the model with `level_setpoint_with_minmax/ribasim.toml`.\n",
+    "Now run the model with `ribasim level_setpoint_with_minmax/ribasim.toml`.\n",
     "After running the model, read back the results:\n"
    ]
   },
diff --git a/python/ribasim/ribasim/config.py b/python/ribasim/ribasim/config.py
index dd7fbbf0d..d837fc8fc 100644
--- a/python/ribasim/ribasim/config.py
+++ b/python/ribasim/ribasim/config.py
@@ -310,7 +310,7 @@ class DiscreteControl(MultiNodeModel):
     )
     logic: TableModel[DiscreteControlLogicSchema] = Field(
         default_factory=TableModel[DiscreteControlLogicSchema],
-        json_schema_extra={"sort_keys": ["node_id", "truth_state"]},
+        json_schema_extra={"sort_keys": ["node_id"]},
     )
 
 
diff --git a/python/ribasim/ribasim/model.py b/python/ribasim/ribasim/model.py
index 8bbb1666a..966a65888 100644
--- a/python/ribasim/ribasim/model.py
+++ b/python/ribasim/ribasim/model.py
@@ -321,14 +321,9 @@ def plot_control_listen(self, ax):
             df_listen_edge = pd.concat([df_listen_edge, to_add])
 
         # Listen edges from DiscreteControl
-        for table in (
-            self.discrete_control.condition.df,
-            self.discrete_control.variable.df,
-        ):
-            if table is None:
-                continue
-
-            to_add = table[
+        df_variable = self.discrete_control.variable.df
+        if df_variable is not None:
+            to_add = df_variable[
                 ["node_id", "listen_node_id", "listen_node_type"]
             ].drop_duplicates()
             to_add = to_add[to_add["listen_node_type"] != "compound"]
diff --git a/python/ribasim/ribasim/schemas.py b/python/ribasim/ribasim/schemas.py
index 260831282..2a8b5eab5 100644
--- a/python/ribasim/ribasim/schemas.py
+++ b/python/ribasim/ribasim/schemas.py
@@ -50,6 +50,7 @@ class BasinTimeSchema(_BaseSchema):
 
 class DiscreteControlConditionSchema(_BaseSchema):
     node_id: Series[Int32] = pa.Field(nullable=False, default=0)
+    compound_variable_id: Series[Int32] = pa.Field(nullable=False, default=0)
     greater_than: Series[float] = pa.Field(nullable=False)
 
 
@@ -61,6 +62,7 @@ class DiscreteControlLogicSchema(_BaseSchema):
 
 class DiscreteControlVariableSchema(_BaseSchema):
     node_id: Series[Int32] = pa.Field(nullable=False, default=0)
+    compound_variable_id: Series[Int32] = pa.Field(nullable=False, default=0)
     listen_node_type: Series[str] = pa.Field(nullable=False)
     listen_node_id: Series[Int32] = pa.Field(nullable=False, default=0)
     variable: Series[str] = pa.Field(nullable=False)
diff --git a/python/ribasim_testmodels/ribasim_testmodels/allocation.py b/python/ribasim_testmodels/ribasim_testmodels/allocation.py
index 71353cf93..51980abc5 100644
--- a/python/ribasim_testmodels/ribasim_testmodels/allocation.py
+++ b/python/ribasim_testmodels/ribasim_testmodels/allocation.py
@@ -410,9 +410,11 @@ def fractional_flow_subnetwork_model() -> Model:
                 listen_node_type="FlowBoundary",
                 listen_node_id=[1],
                 variable="flow_rate",
+                compound_variable_id=1,
             ),
             discrete_control.Condition(
                 greater_than=[3e-3],
+                compound_variable_id=1,
             ),
             discrete_control.Logic(truth_state=["F", "T"], control_state=["A", "B"]),
         ],
@@ -506,9 +508,11 @@ def allocation_example_model() -> Model:
                 listen_node_type="Basin",
                 listen_node_id=[5],
                 variable="level",
+                compound_variable_id=1,
             ),
             discrete_control.Condition(
                 greater_than=[0.52],
+                compound_variable_id=1,
             ),
             discrete_control.Logic(
                 truth_state=["T", "F"], control_state=["divert", "close"]
@@ -687,9 +691,11 @@ def main_network_with_subnetworks_model() -> Model:
                 listen_node_type="Basin",
                 listen_node_id=[25],
                 variable="level",
+                compound_variable_id=1,
             ),
             discrete_control.Condition(
                 greater_than=[0.003],
+                compound_variable_id=1,
             ),
             discrete_control.Logic(truth_state=["F", "T"], control_state=["A", "B"]),
         ],
diff --git a/python/ribasim_testmodels/ribasim_testmodels/discrete_control.py b/python/ribasim_testmodels/ribasim_testmodels/discrete_control.py
index 1a6037d8a..d8d911674 100644
--- a/python/ribasim_testmodels/ribasim_testmodels/discrete_control.py
+++ b/python/ribasim_testmodels/ribasim_testmodels/discrete_control.py
@@ -58,9 +58,11 @@ def pump_discrete_control_model() -> Model:
                 listen_node_type="Basin",
                 listen_node_id=[1, 3],
                 variable="level",
+                compound_variable_id=[1, 2],
             ),
             discrete_control.Condition(
                 greater_than=[0.8, 0.4],
+                compound_variable_id=[1, 2],
             ),
             discrete_control.Logic(
                 truth_state=["FF", "TF", "FT", "TT"],
@@ -75,9 +77,11 @@ def pump_discrete_control_model() -> Model:
                 listen_node_type="Basin",
                 listen_node_id=[3],
                 variable="level",
+                compound_variable_id=1,
             ),
             discrete_control.Condition(
                 greater_than=[0.45],
+                compound_variable_id=1,
             ),
             discrete_control.Logic(
                 truth_state=["T", "F"],
@@ -149,9 +153,11 @@ def flow_condition_model() -> Model:
                 listen_node_id=[1],
                 variable="flow_rate",
                 look_ahead=60 * 86400,
+                compound_variable_id=1,
             ),
             discrete_control.Condition(
                 greater_than=[20 / (86400)],
+                compound_variable_id=1,
             ),
             discrete_control.Logic(truth_state=["T", "F"], control_state=["off", "on"]),
         ],
@@ -214,9 +220,11 @@ def level_boundary_condition_model() -> Model:
                 listen_node_id=[1],
                 variable="level",
                 look_ahead=60 * 86400,
+                compound_variable_id=1,
             ),
             discrete_control.Condition(
                 greater_than=[6.0],
+                compound_variable_id=1,
             ),
             discrete_control.Logic(truth_state=["T", "F"], control_state=["on", "off"]),
         ],
@@ -287,6 +295,7 @@ def tabulated_rating_curve_control_model() -> Model:
                 listen_node_type="Basin",
                 listen_node_id=[1],
                 variable="level",
+                compound_variable_id=1,
             ),
             discrete_control.Condition(
                 greater_than=[0.5],
@@ -356,10 +365,12 @@ def level_setpoint_with_minmax_model() -> Model:
                 listen_node_type="Basin",
                 listen_node_id=[1],
                 variable="level",
+                compound_variable_id=1,
             ),
             discrete_control.Condition(
                 # min, setpoint, max
                 greater_than=[5.0, 10.0, 15.0],
+                compound_variable_id=1,
             ),
             discrete_control.Logic(
                 truth_state=["FFF", "U**", "T*F", "**D", "TTT"],
@@ -443,9 +454,11 @@ def compound_variable_condition_model() -> Model:
                 listen_node_id=[2, 3],
                 variable="flow_rate",
                 weight=0.5,
+                compound_variable_id=1,
             ),
             discrete_control.Condition(
                 greater_than=[0.5],
+                compound_variable_id=1,
             ),
             discrete_control.Logic(truth_state=["T", "F"], control_state=["On", "Off"]),
         ],
diff --git a/python/ribasim_testmodels/ribasim_testmodels/dutch_waterways.py b/python/ribasim_testmodels/ribasim_testmodels/dutch_waterways.py
index b13bb4bcb..9fb315a7f 100644
--- a/python/ribasim_testmodels/ribasim_testmodels/dutch_waterways.py
+++ b/python/ribasim_testmodels/ribasim_testmodels/dutch_waterways.py
@@ -125,9 +125,11 @@ def dutch_waterways_model() -> Model:
                 listen_node_type="FlowBoundary",
                 listen_node_id=[1],
                 variable="flow_rate",
+                compound_variable_id=1,
             ),
             discrete_control.Condition(
                 greater_than=[250, 275, 750, 800],
+                compound_variable_id=1,
             ),
             discrete_control.Logic(
                 truth_state=["FFFF", "U***", "T**F", "***D", "TTTT"],
diff --git a/python/ribasim_testmodels/ribasim_testmodels/invalid.py b/python/ribasim_testmodels/ribasim_testmodels/invalid.py
index 29c9a8405..463603bf6 100644
--- a/python/ribasim_testmodels/ribasim_testmodels/invalid.py
+++ b/python/ribasim_testmodels/ribasim_testmodels/invalid.py
@@ -156,9 +156,11 @@ def invalid_discrete_control_model() -> Model:
                 # Invalid: this look_ahead will go past the provided timeseries during simulation.
                 # Invalid: look_ahead must be non-negative.
                 look_ahead=[100.0, 40 * 24 * 60 * 60, -10.0],
+                compound_variable_id=[1, 2, 3],
             ),
             discrete_control.Condition(
                 greater_than=[0.5, 1.5, 1.5],
+                compound_variable_id=[1, 2, 3],
             ),
             # Invalid: DiscreteControl node #4 has 2 conditions so
             # truth states have to be of length 2
diff --git a/python/ribasim_testmodels/ribasim_testmodels/pid_control.py b/python/ribasim_testmodels/ribasim_testmodels/pid_control.py
index 05ebc31b7..4a543acc3 100644
--- a/python/ribasim_testmodels/ribasim_testmodels/pid_control.py
+++ b/python/ribasim_testmodels/ribasim_testmodels/pid_control.py
@@ -133,9 +133,11 @@ def discrete_control_of_pid_control_model() -> Model:
                 listen_node_type="LevelBoundary",
                 listen_node_id=[1],
                 variable="level",
+                compound_variable_id=1,
             ),
             discrete_control.Condition(
                 greater_than=[5.0],
+                compound_variable_id=1,
             ),
             discrete_control.Logic(
                 truth_state=["T", "F"], control_state=["target_high", "target_low"]

From 75b56a077f379e1510fbf9223f1ffa80e6f717e0 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Tue, 16 Apr 2024 11:49:38 +0200
Subject: [PATCH 16/25] Fix bugs

---
 core/src/callback.jl                          | 106 +++++++++++++-----
 core/src/parameter.jl                         |   7 +-
 core/src/read.jl                              |  80 ++++++++++---
 core/src/validation.jl                        |  12 +-
 core/test/control_test.jl                     |  12 +-
 python/ribasim/ribasim/config.py              |   3 +-
 .../ribasim_testmodels/discrete_control.py    |   1 +
 7 files changed, 158 insertions(+), 63 deletions(-)

diff --git a/core/src/callback.jl b/core/src/callback.jl
index 7e6daec53..9dbb48acd 100644
--- a/core/src/callback.jl
+++ b/core/src/callback.jl
@@ -9,10 +9,17 @@ function set_initial_discrete_controlled_parameters!(
     (; p) = integrator
     (; discrete_control) = p
 
-    n_conditions = length(discrete_control.condition_value)
+    n_conditions = sum(length(vec) for vec in discrete_control.condition_value)
     condition_diffs = zeros(Float64, n_conditions)
     discrete_control_condition(condition_diffs, storage0, integrator.t, integrator)
-    discrete_control.condition_value .= (condition_diffs .> 0.0)
+
+    idx_start = 1
+    for (i, vec) in enumerate(discrete_control.condition_value)
+        l = length(vec)
+        idx_end = idx_start + l - 1
+        discrete_control.condition_value[i] .= (condition_diffs[idx_start:idx_end] .> 0)
+        idx_start += l
+    end
 
     # For every discrete_control node find a condition_idx it listens to
     for discrete_control_node_id in unique(discrete_control.node_id)
@@ -88,7 +95,7 @@ function create_callbacks(
 
     saved = SavedResults(saved_flow, saved_vertical_flux, saved_subgrid_level)
 
-    n_conditions = length(discrete_control.node_id)
+    n_conditions = sum(length(vec) for vec in discrete_control.greater_than; init = 0)
     if n_conditions > 0
         discrete_control_cb = VectorContinuousCallback(
             discrete_control_condition,
@@ -193,23 +200,25 @@ Listens for changes in condition truths.
 function discrete_control_condition(out, u, t, integrator)
     (; p) = integrator
     (; discrete_control) = p
-
-    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,
-        ),
+    condition_idx = 0
+
+    for (listen_node_ids, variables, weights, greater_thans, look_aheads) in zip(
+        discrete_control.listen_node_id,
+        discrete_control.variable,
+        discrete_control.weight,
+        discrete_control.greater_than,
+        discrete_control.look_ahead,
     )
         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
+        for greater_than in greater_thans
+            condition_idx += 1
+            diff = value - greater_than
+            out[condition_idx] = diff
+        end
     end
 end
 
@@ -259,6 +268,22 @@ function get_value(
     return value
 end
 
+function get_discrete_control_indices(discrete_control::DiscreteControl, condition_idx::Int)
+    (; greater_than) = discrete_control
+    condition_idx_now = 1
+
+    for (compound_variable_idx, vec) in enumerate(greater_than)
+        l = length(vec)
+
+        if condition_idx_now + l > condition_idx
+            greater_than_idx = condition_idx - condition_idx_now + 1
+            return compound_variable_idx, greater_than_idx
+        end
+
+        condition_idx_now += l
+    end
+end
+
 """
 An upcrossing means that a condition (always greater than) becomes true.
 """
@@ -267,7 +292,9 @@ function discrete_control_affect_upcrossing!(integrator, condition_idx)
     (; discrete_control, basin) = p
     (; variable, condition_value, listen_node_id) = discrete_control
 
-    condition_value[condition_idx] = true
+    compound_variable_idx, greater_than_idx =
+        get_discrete_control_indices(discrete_control, condition_idx)
+    condition_value[compound_variable_idx][greater_than_idx] = true
 
     control_state_change = discrete_control_affect!(integrator, condition_idx, true)
 
@@ -277,7 +304,7 @@ 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.
-    listen_node_ids = discrete_control.listen_node_id[condition_idx]
+    listen_node_ids = discrete_control.listen_node_id[compound_variable_idx]
     is_basin =
         length(listen_node_ids) == 1 ? id_index(basin.node_id, only(listen_node_ids))[1] :
         false
@@ -285,15 +312,16 @@ function discrete_control_affect_upcrossing!(integrator, condition_idx)
     # 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][1] == "level" && control_state_change && is_basin
+    if variable[compound_variable_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][1])
+        _, condition_basin_idx =
+            id_index(basin.node_id, listen_node_id[compound_variable_idx][1])
 
         if du[condition_basin_idx] < 0.0
-            condition_value[condition_idx] = false
+            condition_value[compound_variable_idx][greater_than_idx] = false
             discrete_control_affect!(integrator, condition_idx, false)
         end
     end
@@ -307,7 +335,9 @@ function discrete_control_affect_downcrossing!(integrator, condition_idx)
     (; discrete_control, basin) = p
     (; variable, condition_value, listen_node_id) = discrete_control
 
-    condition_value[condition_idx] = false
+    compound_variable_idx, greater_than_idx =
+        get_discrete_control_indices(discrete_control, condition_idx)
+    condition_value[compound_variable_idx][greater_than_idx] = false
 
     control_state_change = discrete_control_affect!(integrator, condition_idx, false)
 
@@ -317,21 +347,23 @@ 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.
-    listen_node_ids = discrete_control.listen_node_id[condition_idx]
+    compound_variable_idx, greater_than_idx =
+        get_discrete_control_indices(discrete_control, condition_idx)
+    listen_node_ids = discrete_control.listen_node_id[compound_variable_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
+    if variable[compound_variable_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][1])
+            id_index(basin.node_id, listen_node_id[compound_variable_idx][1])
 
         if has_index && du[condition_basin_idx] > 0.0
-            condition_value[condition_idx] = true
+            condition_value[compound_variable_idx][greater_than_idx] = true
             discrete_control_affect!(integrator, condition_idx, true)
         end
     end
@@ -349,20 +381,34 @@ function discrete_control_affect!(
     (; discrete_control, graph) = p
 
     # Get the discrete_control node that listens to this condition
-    discrete_control_node_id = discrete_control.node_id[condition_idx]
+
+    compound_variable_idx, _ = get_discrete_control_indices(discrete_control, condition_idx)
+    discrete_control_node_id = discrete_control.node_id[compound_variable_idx]
 
     # Get the indices of all conditions that this control node listens to
-    condition_ids = discrete_control.node_id .== discrete_control_node_id
+    where_node_id = searchsorted(discrete_control.node_id, discrete_control_node_id)
 
     # Get the truth state for this discrete_control node
-    truth_values = [ifelse(b, "T", "F") for b in discrete_control.condition_value]
-    truth_state = join(truth_values[condition_ids], "")
+    truth_values = cat(
+        [
+            [ifelse(b, "T", "F") for b in discrete_control.condition_value[i]] for
+            i in where_node_id
+        ]...;
+        dims = 1,
+    )
+    truth_state = join(truth_values, "")
 
     # Get the truth specific about the latest crossing
     if !ismissing(upcrossing)
-        truth_values[condition_idx] = upcrossing ? "U" : "D"
+        truth_value_idx =
+            condition_idx - sum(
+                length(vec) for
+                vec in discrete_control.condition_value[1:(where_node_id.start - 1)];
+                init = 0,
+            )
+        truth_values[truth_value_idx] = upcrossing ? "U" : "D"
     end
-    truth_state_crossing_specific = join(truth_values[condition_ids], "")
+    truth_state_crossing_specific = join(truth_values, "")
 
     # What the local control state should be
     control_state_new =
diff --git a/core/src/parameter.jl b/core/src/parameter.jl
index 231f0a9c4..e898280ff 100644
--- a/core/src/parameter.jl
+++ b/core/src/parameter.jl
@@ -457,12 +457,15 @@ record: Namedtuple with discrete control information for results
 """
 struct DiscreteControl <: AbstractParameterNode
     node_id::Vector{NodeID}
+    # Definition of compound variables
     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}
+    # Definition of conditions (one or more greater_than per compound variable)
+    greater_than::Vector{Vector{Float64}}
+    condition_value::Vector{BitVector}
+    # Definition of logic
     control_state::Dict{NodeID, Tuple{String, Float64}}
     logic_mapping::Dict{Tuple{NodeID, String}, String}
     record::@NamedTuple{
diff --git a/core/src/read.jl b/core/src/read.jl
index a0aec5fda..62a7ddc89 100644
--- a/core/src/read.jl
+++ b/core/src/read.jl
@@ -542,33 +542,79 @@ function Basin(db::DB, config::Config, chunk_sizes::Vector{Int})::Basin
     )
 end
 
-function DiscreteControl(db::DB, config::Config)::DiscreteControl
-    compound_variable = load_structvector(db, config, DiscreteControlVariableV1)
-    condition = load_structvector(db, config, DiscreteControlConditionV1)
-
+function parse_variables_and_conditions(compound_variable, condition)
     node_id = NodeID[]
     listen_node_id = Vector{NodeID}[]
     variable = Vector{String}[]
     weight = Vector{Float64}[]
     look_ahead = Vector{Float64}[]
+    greater_than = Vector{Float64}[]
+    condition_value = BitVector[]
+    errors = false
 
     for id in unique(condition.node_id)
-        group_id = filter(row -> row.node_id == id, compound_variable)
-        for group_variable in
-            StructVector.(IterTools.groupby(row -> row.compound_variable_id, group_id))
-            first_row = first(group_variable)
-            push!(node_id, NodeID(NodeType.DiscreteControl, first_row.node_id))
-            push!(
-                listen_node_id,
-                NodeID.(group_variable.listen_node_type, group_variable.listen_node_id),
+        condition_group_id = filter(row -> row.node_id == id, condition)
+        variable_group_id = filter(row -> row.node_id == id, compound_variable)
+        for compound_variable_id in unique(condition_group_id.compound_variable_id)
+            condition_group_variable = filter(
+                row -> row.compound_variable_id == compound_variable_id,
+                condition_group_id,
             )
-            push!(variable, group_variable.variable)
-            push!(weight, coalesce.(group_variable.weight, 1.0))
-            push!(look_ahead, coalesce.(group_variable.look_ahead, 0.0))
+            variable_group_variable = filter(
+                row -> row.compound_variable_id == compound_variable_id,
+                variable_group_id,
+            )
+            discrete_control_id = NodeID(NodeType.DiscreteControl, id)
+            if isempty(variable_group_variable)
+                errors = true
+                @error "compound_variable_id $compound_variable_id for $discrete_control_id in condition table but not in variable table"
+            else
+                push!(node_id, discrete_control_id)
+                push!(
+                    listen_node_id,
+                    NodeID.(
+                        variable_group_variable.listen_node_type,
+                        variable_group_variable.listen_node_id,
+                    ),
+                )
+                push!(variable, variable_group_variable.variable)
+                push!(weight, coalesce.(variable_group_variable.weight, 1.0))
+                push!(look_ahead, coalesce.(variable_group_variable.look_ahead, 0.0))
+                push!(greater_than, condition_group_variable.greater_than)
+                push!(
+                    condition_value,
+                    BitVector(zeros(length(condition_group_variable.greater_than))),
+                )
+            end
         end
     end
+    return node_id,
+    listen_node_id,
+    variable,
+    weight,
+    look_ahead,
+    greater_than,
+    condition_value,
+    !errors
+end
+
+function DiscreteControl(db::DB, config::Config)::DiscreteControl
+    condition = load_structvector(db, config, DiscreteControlConditionV1)
+    compound_variable = load_structvector(db, config, DiscreteControlVariableV1)
+
+    node_id,
+    listen_node_id,
+    variable,
+    weight,
+    look_ahead,
+    greater_than,
+    condition_value,
+    valid = parse_variables_and_conditions(compound_variable, condition)
+
+    if !valid
+        error("Problems encountered when parsing DiscreteControl variables and conditions.")
+    end
 
-    condition_value = fill(false, length(condition.node_id))
     control_state::Dict{NodeID, Tuple{String, Float64}} = Dict()
 
     rows = execute(db, "SELECT from_node_id, edge_type FROM Edge ORDER BY fid")
@@ -603,7 +649,7 @@ function DiscreteControl(db::DB, config::Config)::DiscreteControl
         variable,
         weight,
         look_ahead,
-        condition.greater_than,
+        greater_than,
         condition_value,
         control_state,
         logic_mapping,
diff --git a/core/src/validation.jl b/core/src/validation.jl
index 580b4d5ba..9d6a621d2 100644
--- a/core/src/validation.jl
+++ b/core/src/validation.jl
@@ -102,12 +102,9 @@ sort_by_id_state_level(row) = (row.node_id, row.control_state, row.level)
 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_variable(row) =
     (row.node_id, row.listen_node_type, row.listen_node_id, row.variable)
-sort_by_id_greater_than(row) = (row.node_id, row.greater_than)
-sort_by_truth_state(row) = (row.node_id, row.truth_state)
+sort_by_condition(row) = (row.node_id, row.compound_variable_id, row.greater_than)
 
 # 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
@@ -117,7 +114,7 @@ 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{DiscreteControlVariableV1}) = sort_by_variable
-sort_by_function(table::StructVector{DiscreteControlConditionV1}) = sort_by_id_greater_than
+sort_by_function(table::StructVector{DiscreteControlConditionV1}) = sort_by_condition
 
 const TimeSchemas = Union{
     BasinTimeV1,
@@ -499,7 +496,8 @@ Check:
 """
 function valid_discrete_control(p::Parameters, config::Config)::Bool
     (; discrete_control, graph) = p
-    (; node_id, logic_mapping, look_ahead, variable, listen_node_id) = discrete_control
+    (; node_id, logic_mapping, look_ahead, variable, listen_node_id, greater_than) =
+        discrete_control
 
     t_end = seconds_since(config.endtime, config.starttime)
     errors = false
@@ -512,7 +510,7 @@ function valid_discrete_control(p::Parameters, config::Config)::Bool
         truth_states_wrong_length = String[]
 
         # The number of conditions of this DiscreteControl node
-        n_conditions = length(searchsorted(node_id, id))
+        n_conditions = sum(length(greater_than[i]) for i in searchsorted(node_id, id))
 
         for (key, control_state) in logic_mapping
             id_, truth_state = key
diff --git a/core/test/control_test.jl b/core/test/control_test.jl
index e0913ba1d..21bcc653c 100644
--- a/core/test/control_test.jl
+++ b/core/test/control_test.jl
@@ -37,11 +37,11 @@
     # Control times
     t_1 = discrete_control.record.time[3]
     t_1_index = findfirst(>=(t_1), t)
-    @test level[1, t_1_index] <= discrete_control.greater_than[1]
+    @test level[1, t_1_index] <= discrete_control.greater_than[1][1]
 
     t_2 = discrete_control.record.time[4]
     t_2_index = findfirst(>=(t_2), t)
-    @test level[2, t_2_index] >= discrete_control.greater_than[2]
+    @test level[2, t_2_index] >= discrete_control.greater_than[2][1]
 
     flow = get_tmp(graph[].flow, 0)
     @test all(iszero, flow)
@@ -60,7 +60,7 @@ end
     t_control = discrete_control.record.time[2]
     t_control_index = searchsortedfirst(t, t_control)
 
-    greater_than = discrete_control.greater_than[1]
+    greater_than = discrete_control.greater_than[1][1]
     flow_t_control = flow_boundary.flow_rate[1](t_control)
     flow_t_control_ahead = flow_boundary.flow_rate[1](t_control + Δt)
 
@@ -84,7 +84,7 @@ end
     t_control = discrete_control.record.time[2]
     t_control_index = searchsortedfirst(t, t_control)
 
-    greater_than = discrete_control.greater_than[1]
+    greater_than = discrete_control.greater_than[1][1]
     level_t_control = level_boundary.level[1](t_control)
     level_t_control_ahead = level_boundary.level[1](t_control + Δt)
 
@@ -167,8 +167,8 @@ end
     t_in = discrete_control.record.time[3]
     t_none_2 = discrete_control.record.time[4]
 
-    level_min = greater_than[1]
-    setpoint = greater_than[2]
+    level_min = greater_than[1][1]
+    setpoint = greater_than[1][2]
 
     t_1_none_index = findfirst(>=(t_none_1), t)
     t_in_index = findfirst(>=(t_in), t)
diff --git a/python/ribasim/ribasim/config.py b/python/ribasim/ribasim/config.py
index d837fc8fc..30408fe05 100644
--- a/python/ribasim/ribasim/config.py
+++ b/python/ribasim/ribasim/config.py
@@ -304,13 +304,14 @@ class DiscreteControl(MultiNodeModel):
         json_schema_extra={
             "sort_keys": [
                 "node_id",
+                "compound_variable_id",
                 "greater_than",
             ]
         },
     )
     logic: TableModel[DiscreteControlLogicSchema] = Field(
         default_factory=TableModel[DiscreteControlLogicSchema],
-        json_schema_extra={"sort_keys": ["node_id"]},
+        json_schema_extra={"sort_keys": ["node_id", "truth_state"]},
     )
 
 
diff --git a/python/ribasim_testmodels/ribasim_testmodels/discrete_control.py b/python/ribasim_testmodels/ribasim_testmodels/discrete_control.py
index d8d911674..650e73876 100644
--- a/python/ribasim_testmodels/ribasim_testmodels/discrete_control.py
+++ b/python/ribasim_testmodels/ribasim_testmodels/discrete_control.py
@@ -299,6 +299,7 @@ def tabulated_rating_curve_control_model() -> Model:
             ),
             discrete_control.Condition(
                 greater_than=[0.5],
+                compound_variable_id=1,
             ),
             discrete_control.Logic(
                 truth_state=["T", "F"], control_state=["low", "high"]

From a65f6f69266caa7749715d5d495898b5bb701b57 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Tue, 16 Apr 2024 11:57:32 +0200
Subject: [PATCH 17/25] Small fixes

---
 core/src/callback.jl            | 2 +-
 python/ribasim/tests/test_io.py | 4 +---
 2 files changed, 2 insertions(+), 4 deletions(-)

diff --git a/core/src/callback.jl b/core/src/callback.jl
index 9dbb48acd..05b49dda4 100644
--- a/core/src/callback.jl
+++ b/core/src/callback.jl
@@ -9,7 +9,7 @@ function set_initial_discrete_controlled_parameters!(
     (; p) = integrator
     (; discrete_control) = p
 
-    n_conditions = sum(length(vec) for vec in discrete_control.condition_value)
+    n_conditions = sum(length(vec) for vec in discrete_control.condition_value; init = 0)
     condition_diffs = zeros(Float64, n_conditions)
     discrete_control_condition(condition_diffs, storage0, integrator.t, integrator)
 
diff --git a/python/ribasim/tests/test_io.py b/python/ribasim/tests/test_io.py
index e9b724eaa..e6507d77e 100644
--- a/python/ribasim/tests/test_io.py
+++ b/python/ribasim/tests/test_io.py
@@ -102,9 +102,7 @@ def test_sort(level_setpoint_with_minmax, tmp_path):
     assert table.df.iloc[0]["greater_than"] == 15.0
     assert table._sort_keys == [
         "node_id",
-        "listen_node_type",
-        "listen_node_id",
-        "variable",
+        "compound_variable_id",
         "greater_than",
     ]
     table.sort()

From 9765dc2d977947b0c3104eed8a1969a4ffe7ddbd Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Tue, 16 Apr 2024 12:47:46 +0200
Subject: [PATCH 18/25] Fix examples notebook

---
 docs/python/examples.ipynb | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/docs/python/examples.ipynb b/docs/python/examples.ipynb
index f78357420..d225579d3 100644
--- a/docs/python/examples.ipynb
+++ b/docs/python/examples.ipynb
@@ -493,10 +493,13 @@
     "    Node(7, Point(1.0, 0.0)),\n",
     "    [\n",
     "        discrete_control.Variable(\n",
-    "            listen_node_id=[1], listen_node_type=\"Basin\", variable=\"level\"\n",
+    "            listen_node_id=[1],\n",
+    "            listen_node_type=\"Basin\",\n",
+    "            variable=\"level\",\n",
+    "            compound_variable_id=1,\n",
     "        ),\n",
     "        discrete_control.Condition(\n",
-    "            greater_than=[5.0, 10.0, 15.0],\n",
+    "            greater_than=[5.0, 10.0, 15.0], compound_variable_id=1\n",
     "        ),\n",
     "        discrete_control.Logic(\n",
     "            truth_state=[\"FFF\", \"U**\", \"T*F\", \"**D\", \"TTT\"],\n",
@@ -1183,10 +1186,12 @@
     "model.discrete_control.add(\n",
     "    Node(11, Point(4.5, 0.25), subnetwork_id=1),\n",
     "    [\n",
-    "        discrete_control.Condition(\n",
+    "        discrete_control.Variable(\n",
     "            listen_node_id=[5],\n",
     "            listen_node_type=[\"Basin\"],\n",
     "            variable=[\"level\"],\n",
+    "        ),\n",
+    "        discrete_control.Condition(\n",
     "            greater_than=[0.52],\n",
     "        ),\n",
     "        discrete_control.Logic(\n",

From 267d412a48eb9aab3ee5887c6213be5858941ad8 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Tue, 16 Apr 2024 13:12:26 +0200
Subject: [PATCH 19/25] Update docstrings

---
 core/src/callback.jl  | 24 ++++++------------------
 core/src/parameter.jl | 15 +++++++--------
 core/src/read.jl      |  2 ++
 core/src/util.jl      | 16 ++++++++++++++++
 4 files changed, 31 insertions(+), 26 deletions(-)

diff --git a/core/src/callback.jl b/core/src/callback.jl
index 05b49dda4..fb1918831 100644
--- a/core/src/callback.jl
+++ b/core/src/callback.jl
@@ -13,11 +13,13 @@ function set_initial_discrete_controlled_parameters!(
     condition_diffs = zeros(Float64, n_conditions)
     discrete_control_condition(condition_diffs, storage0, integrator.t, integrator)
 
+    # Set the discrete control value (bool) per compound variable
     idx_start = 1
-    for (i, vec) in enumerate(discrete_control.condition_value)
+    for (compound_variable_idx, vec) in enumerate(discrete_control.condition_value)
         l = length(vec)
         idx_end = idx_start + l - 1
-        discrete_control.condition_value[i] .= (condition_diffs[idx_start:idx_end] .> 0)
+        discrete_control.condition_value[compound_variable_idx] .=
+            (condition_diffs[idx_start:idx_end] .> 0)
         idx_start += l
     end
 
@@ -202,6 +204,7 @@ function discrete_control_condition(out, u, t, integrator)
     (; discrete_control) = p
     condition_idx = 0
 
+    # Loop over compound variables
     for (listen_node_ids, variables, weights, greater_thans, look_aheads) in zip(
         discrete_control.listen_node_id,
         discrete_control.variable,
@@ -214,6 +217,7 @@ function discrete_control_condition(out, u, t, integrator)
             zip(listen_node_ids, variables, weights, look_aheads)
             value += weight * get_value(p, listen_node_id, variable, look_ahead, u, t)
         end
+        # Loop over greater_than values for this compound_variable
         for greater_than in greater_thans
             condition_idx += 1
             diff = value - greater_than
@@ -268,22 +272,6 @@ function get_value(
     return value
 end
 
-function get_discrete_control_indices(discrete_control::DiscreteControl, condition_idx::Int)
-    (; greater_than) = discrete_control
-    condition_idx_now = 1
-
-    for (compound_variable_idx, vec) in enumerate(greater_than)
-        l = length(vec)
-
-        if condition_idx_now + l > condition_idx
-            greater_than_idx = condition_idx - condition_idx_now + 1
-            return compound_variable_idx, greater_than_idx
-        end
-
-        condition_idx_now += l
-    end
-end
-
 """
 An upcrossing means that a condition (always greater than) becomes true.
 """
diff --git a/core/src/parameter.jl b/core/src/parameter.jl
index e898280ff..9ea874d74 100644
--- a/core/src/parameter.jl
+++ b/core/src/parameter.jl
@@ -443,14 +443,13 @@ struct Terminal <: AbstractParameterNode
 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 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
+node_id: node ID of the DiscreteControl node per compound variable (can contain repeats)
+listen_node_id: the IDs of the nodes being condition on per compound variable
+variable: the names of the variables in the condition per compound variable
+weight: the weight of the variables in the condition per compound variable
+look_ahead: the look ahead of variables in the condition in seconds per compound_variable
+greater_than: The threshold values per compound variable
+condition_value: The current truth value of each condition per compound_variable per greater_than
 control_state: Dictionary: node ID => (control state, control state start)
 logic_mapping: Dictionary: (control node ID, truth state) => control state
 record: Namedtuple with discrete control information for results
diff --git a/core/src/read.jl b/core/src/read.jl
index 62a7ddc89..bd07fd17d 100644
--- a/core/src/read.jl
+++ b/core/src/read.jl
@@ -552,9 +552,11 @@ function parse_variables_and_conditions(compound_variable, condition)
     condition_value = BitVector[]
     errors = false
 
+    # Loop over unique discrete_control node IDs (on which at least one condition is defined)
     for id in unique(condition.node_id)
         condition_group_id = filter(row -> row.node_id == id, condition)
         variable_group_id = filter(row -> row.node_id == id, compound_variable)
+        # Loop over compound variables for this node ID
         for compound_variable_id in unique(condition_group_id.compound_variable_id)
             condition_group_variable = filter(
                 row -> row.compound_variable_id == compound_variable_id,
diff --git a/core/src/util.jl b/core/src/util.jl
index 9f28670cf..d84567632 100644
--- a/core/src/util.jl
+++ b/core/src/util.jl
@@ -709,3 +709,19 @@ function get_influx(basin::Basin, basin_idx::Int)::Float64
     return precipitation[basin_idx] - evaporation[basin_idx] + drainage[basin_idx] -
            infiltration[basin_idx]
 end
+
+function get_discrete_control_indices(discrete_control::DiscreteControl, condition_idx::Int)
+    (; greater_than) = discrete_control
+    condition_idx_now = 1
+
+    for (compound_variable_idx, vec) in enumerate(greater_than)
+        l = length(vec)
+
+        if condition_idx_now + l > condition_idx
+            greater_than_idx = condition_idx - condition_idx_now + 1
+            return compound_variable_idx, greater_than_idx
+        end
+
+        condition_idx_now += l
+    end
+end

From c662fe7f37ea30b32f17fee2f143250ca34c1819 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Tue, 16 Apr 2024 13:28:39 +0200
Subject: [PATCH 20/25] Update usage.qmd

---
 core/src/schema.jl         |  6 +++---
 docs/core/usage.qmd        | 44 +++++++++++++++++---------------------
 ribasim_qgis/core/nodes.py |  2 ++
 3 files changed, 25 insertions(+), 27 deletions(-)

diff --git a/core/src/schema.jl b/core/src/schema.jl
index 60eb87ac8..0875e48c5 100644
--- a/core/src/schema.jl
+++ b/core/src/schema.jl
@@ -186,9 +186,9 @@ end
 
 @version DiscreteControlVariableV1 begin
     node_id::Int32
-    compound_variable_id::Int
+    compound_variable_id::Int32
     listen_node_type::String
-    listen_node_id::Int
+    listen_node_id::Int32
     variable::String
     weight::Union{Missing, Float64}
     look_ahead::Union{Missing, Float64}
@@ -196,7 +196,7 @@ end
 
 @version DiscreteControlConditionV1 begin
     node_id::Int32
-    compound_variable_id::Int
+    compound_variable_id::Int32
     greater_than::Float64
 end
 
diff --git a/docs/core/usage.qmd b/docs/core/usage.qmd
index aa7c5f73a..a736c41cd 100644
--- a/docs/core/usage.qmd
+++ b/docs/core/usage.qmd
@@ -204,8 +204,8 @@ 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'
+  - `DiscreteControl / variable` : Linear combinations of variables, e.g. averages or differences of levels
+  - `DisceteControl / condition`: Conditions of the form 'the compound variable with ID `n` is bigger than 2.0'
   - `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).
   - `PidControl / static`: The proportional, integral, and derivative parameters, the target value and which basin should be controlled
@@ -626,39 +626,35 @@ node_id               | Int32   | -            | sorted
 
 DiscreteControl is implemented based on [VectorContinuousCallback](https://docs.sciml.ai/DiffEqDocs/stable/features/callback_functions/#VectorContinuousCallback).
 
-## DiscreteControl / compoundvariable
+## DiscreteControl / variable
 
 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
+this schema defines new variables with the given `compound_variable_id` that look like
 $$
 \text{weight}_1 * \text{variable}_1 + \text{weight}_2 * \text{variable}_2 + \ldots,
 $$
 
 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).
+column               | type     | unit    | restriction
+-------------------- | -------- | ------- | -----------
+node_id              | Int32    | -       | sorted
+compound_variable_id | 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  | -       | (optional, default 1.0)
+look_ahead           | Float64  | $s$     | Only on transient boundary conditions, non-negative (optional, default 0.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)
+The condition schema defines conditions of the form 'the `discrete_control` node with this `node_id`listens to whether the variable given by the `node_id` and `compound_variable_id` is greater than `greater_than`'. If the condition variable comes from a time-series, a look ahead $\Delta t$ can be supplied. Multiple conditions with different `greater_than` values can be defined on the same compound_variable.
+
+column               | type     | unit    | restriction
+-------------------- | -------- | ------- | -----------
+node_id              | Int32    | -       | sorted
+compound_variable_id | Int32    | -       | -
+greater_than         | Float64  | various | sorted per variable
 
 ## DiscreteControl / logic
 
diff --git a/ribasim_qgis/core/nodes.py b/ribasim_qgis/core/nodes.py
index 5c0393a81..97682e780 100644
--- a/ribasim_qgis/core/nodes.py
+++ b/ribasim_qgis/core/nodes.py
@@ -704,6 +704,7 @@ def geometry_type(cls) -> str:
     def attributes(cls) -> list[QgsField]:
         return [
             QgsField("node_id", QVariant.Int),
+            QgsField("compound_variable_id", QVariant.Int),
             QgsField("listen_node_type", QVariant.String),
             QgsField("listen_node_id", QVariant.Int),
             QgsField("variable", QVariant.String),
@@ -725,6 +726,7 @@ def geometry_type(cls) -> str:
     def attributes(cls) -> list[QgsField]:
         return [
             QgsField("node_id", QVariant.Int),
+            QgsField("compound_variable_id", QVariant.Int),
             QgsField("greater_than", QVariant.Double),
         ]
 

From c6627a172dbf1aec10ef207b71db1824782d2af5 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Tue, 16 Apr 2024 15:07:07 +0200
Subject: [PATCH 21/25] VectorContinuousCallback -> FunctionCallingCallback and
 remove level_set_point_with_minmax model and associated tests

---
 core/src/callback.jl                          | 211 ++----------
 core/src/model.jl                             |   2 -
 core/src/util.jl                              |  18 +-
 core/src/validation.jl                        |   1 -
 core/test/control_test.jl                     |  32 +-
 docs/python/examples.ipynb                    | 321 ------------------
 python/ribasim/tests/conftest.py              |   5 -
 python/ribasim/tests/test_io.py               |  35 --
 .../ribasim_testmodels/__init__.py            |   2 -
 .../ribasim_testmodels/discrete_control.py    |  93 -----
 10 files changed, 35 insertions(+), 685 deletions(-)

diff --git a/core/src/callback.jl b/core/src/callback.jl
index fb1918831..d5025340c 100644
--- a/core/src/callback.jl
+++ b/core/src/callback.jl
@@ -1,35 +1,3 @@
-"""
-Set parameters of nodes that are controlled by DiscreteControl to the
-values corresponding to the initial state of the model.
-"""
-function set_initial_discrete_controlled_parameters!(
-    integrator,
-    storage0::Vector{Float64},
-)::Nothing
-    (; p) = integrator
-    (; discrete_control) = p
-
-    n_conditions = sum(length(vec) for vec in discrete_control.condition_value; init = 0)
-    condition_diffs = zeros(Float64, n_conditions)
-    discrete_control_condition(condition_diffs, storage0, integrator.t, integrator)
-
-    # Set the discrete control value (bool) per compound variable
-    idx_start = 1
-    for (compound_variable_idx, vec) in enumerate(discrete_control.condition_value)
-        l = length(vec)
-        idx_end = idx_start + l - 1
-        discrete_control.condition_value[compound_variable_idx] .=
-            (condition_diffs[idx_start:idx_end] .> 0)
-        idx_start += l
-    end
-
-    # For every discrete_control node find a condition_idx it listens to
-    for discrete_control_node_id in unique(discrete_control.node_id)
-        condition_idx =
-            searchsortedfirst(discrete_control.node_id, discrete_control_node_id)
-        discrete_control_affect!(integrator, condition_idx, missing)
-    end
-end
 
 """
 Create the different callbacks that are used to store results
@@ -99,13 +67,7 @@ function create_callbacks(
 
     n_conditions = sum(length(vec) for vec in discrete_control.greater_than; init = 0)
     if n_conditions > 0
-        discrete_control_cb = VectorContinuousCallback(
-            discrete_control_condition,
-            discrete_control_affect_upcrossing!,
-            discrete_control_affect_downcrossing!,
-            n_conditions;
-            save_positions = (false, false),
-        )
+        discrete_control_cb = FunctionCallingCallback(apply_discrete_control!)
         push!(callbacks, discrete_control_cb)
     end
     callback = CallbackSet(callbacks...)
@@ -196,33 +158,49 @@ function save_vertical_flux(u, t, integrator)
     return vertical_flux_mean
 end
 
+function apply_discrete_control!(u, t, integrator)::Nothing
+    (; p) = integrator
+    (; discrete_control) = p
+
+    discrete_control_condition!(u, t, integrator)
+
+    # For every compound variable see whether it changes a control state
+    for compound_variable_idx in eachindex(discrete_control.node_id)
+        discrete_control_affect!(integrator, compound_variable_idx)
+    end
+end
+
 """
-Listens for changes in condition truths.
+Update discrete control condition truths.
 """
-function discrete_control_condition(out, u, t, integrator)
+function discrete_control_condition!(u, t, integrator)
     (; p) = integrator
     (; discrete_control) = p
-    condition_idx = 0
 
     # Loop over compound variables
-    for (listen_node_ids, variables, weights, greater_thans, look_aheads) in zip(
+    for (
+        listen_node_ids,
+        variables,
+        weights,
+        greater_thans,
+        look_aheads,
+        condition_values,
+    ) in zip(
         discrete_control.listen_node_id,
         discrete_control.variable,
         discrete_control.weight,
         discrete_control.greater_than,
         discrete_control.look_ahead,
+        discrete_control.condition_value,
     )
         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
-        # Loop over greater_than values for this compound_variable
-        for greater_than in greater_thans
-            condition_idx += 1
-            diff = value - greater_than
-            out[condition_idx] = diff
-        end
+
+        condition_values .= false
+        condition_values[1:searchsortedlast(greater_thans, value)] .= true
     end
 end
 
@@ -272,105 +250,14 @@ function get_value(
     return value
 end
 
-"""
-An upcrossing means that a condition (always greater than) becomes true.
-"""
-function discrete_control_affect_upcrossing!(integrator, condition_idx)
-    (; p, u, t) = integrator
-    (; discrete_control, basin) = p
-    (; variable, condition_value, listen_node_id) = discrete_control
-
-    compound_variable_idx, greater_than_idx =
-        get_discrete_control_indices(discrete_control, condition_idx)
-    condition_value[compound_variable_idx][greater_than_idx] = true
-
-    control_state_change = discrete_control_affect!(integrator, condition_idx, true)
-
-    # Check whether the control state change changed the direction of the crossing
-    # NOTE: This works for level conditions, but not for flow conditions on an
-    # arbitrary edge. That is because parameter changes do not change the instantaneous level,
-    # 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.
-    listen_node_ids = discrete_control.listen_node_id[compound_variable_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[compound_variable_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[compound_variable_idx][1])
-
-        if du[condition_basin_idx] < 0.0
-            condition_value[compound_variable_idx][greater_than_idx] = false
-            discrete_control_affect!(integrator, condition_idx, false)
-        end
-    end
-end
-
-"""
-An downcrossing means that a condition (always greater than) becomes false.
-"""
-function discrete_control_affect_downcrossing!(integrator, condition_idx)
-    (; p, u, t) = integrator
-    (; discrete_control, basin) = p
-    (; variable, condition_value, listen_node_id) = discrete_control
-
-    compound_variable_idx, greater_than_idx =
-        get_discrete_control_indices(discrete_control, condition_idx)
-    condition_value[compound_variable_idx][greater_than_idx] = false
-
-    control_state_change = discrete_control_affect!(integrator, condition_idx, false)
-
-    # Check whether the control state change changed the direction of the crossing
-    # NOTE: This works for level conditions, but not for flow conditions on an
-    # arbitrary edge. That is because parameter changes do not change the instantaneous level,
-    # 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.
-    compound_variable_idx, greater_than_idx =
-        get_discrete_control_indices(discrete_control, condition_idx)
-    listen_node_ids = discrete_control.listen_node_id[compound_variable_idx]
-    is_basin =
-        length(listen_node_ids) == 1 ? id_index(basin.node_id, only(listen_node_ids))[1] :
-        false
-
-    if variable[compound_variable_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[compound_variable_idx][1])
-
-        if has_index && du[condition_basin_idx] > 0.0
-            condition_value[compound_variable_idx][greater_than_idx] = true
-            discrete_control_affect!(integrator, condition_idx, true)
-        end
-    end
-end
-
 """
 Change parameters based on the control logic.
 """
-function discrete_control_affect!(
-    integrator,
-    condition_idx::Int,
-    upcrossing::Union{Bool, Missing},
-)::Bool
+function discrete_control_affect!(integrator, compound_variable_idx)
     p = integrator.p
     (; discrete_control, graph) = p
 
-    # Get the discrete_control node that listens to this condition
-
-    compound_variable_idx, _ = get_discrete_control_indices(discrete_control, condition_idx)
+    # Get the discrete_control node to which this compound variable belongs
     discrete_control_node_id = discrete_control.node_id[compound_variable_idx]
 
     # Get the indices of all conditions that this control node listens to
@@ -386,56 +273,24 @@ function discrete_control_affect!(
     )
     truth_state = join(truth_values, "")
 
-    # Get the truth specific about the latest crossing
-    if !ismissing(upcrossing)
-        truth_value_idx =
-            condition_idx - sum(
-                length(vec) for
-                vec in discrete_control.condition_value[1:(where_node_id.start - 1)];
-                init = 0,
-            )
-        truth_values[truth_value_idx] = upcrossing ? "U" : "D"
-    end
-    truth_state_crossing_specific = join(truth_values, "")
-
     # What the local control state should be
     control_state_new =
-        if haskey(
-            discrete_control.logic_mapping,
-            (discrete_control_node_id, truth_state_crossing_specific),
-        )
-            truth_state_used = truth_state_crossing_specific
-            discrete_control.logic_mapping[(
-                discrete_control_node_id,
-                truth_state_crossing_specific,
-            )]
-        elseif haskey(
-            discrete_control.logic_mapping,
-            (discrete_control_node_id, truth_state),
-        )
-            truth_state_used = truth_state
+        if haskey(discrete_control.logic_mapping, (discrete_control_node_id, truth_state))
             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 $discrete_control_node_id.",
+                "No control state specified for $discrete_control_node_id for truth state $truth_state.",
             )
         end
 
-    # What the local control state is
-    # TODO: Check time elapsed since control change
     control_state_now, _ = discrete_control.control_state[discrete_control_node_id]
-
-    control_state_change = false
-
     if control_state_now != control_state_new
-        control_state_change = true
-
         # Store control action in record
         record = discrete_control.record
 
         push!(record.time, integrator.t)
         push!(record.control_node_id, Int32(discrete_control_node_id))
-        push!(record.truth_state, truth_state_used)
+        push!(record.truth_state, truth_state)
         push!(record.control_state, control_state_new)
 
         # Loop over nodes which are under control of this control node
@@ -447,7 +302,7 @@ function discrete_control_affect!(
         discrete_control.control_state[discrete_control_node_id] =
             (control_state_new, integrator.t)
     end
-    return control_state_change
+    return nothing
 end
 
 function get_allocation_model(p::Parameters, allocation_network_id::Int32)::AllocationModel
diff --git a/core/src/model.jl b/core/src/model.jl
index 241ba2ec6..0eca2bc73 100644
--- a/core/src/model.jl
+++ b/core/src/model.jl
@@ -163,8 +163,6 @@ function Model(config::Config)::Model
         @show Ribasim.to
     end
 
-    set_initial_discrete_controlled_parameters!(integrator, storage)
-
     return Model(integrator, config, saved)
 end
 
diff --git a/core/src/util.jl b/core/src/util.jl
index d84567632..636227673 100644
--- a/core/src/util.jl
+++ b/core/src/util.jl
@@ -421,7 +421,7 @@ function expand_logic_mapping(
     logic_mapping_expanded = Dict{Tuple{NodeID, String}, String}()
 
     for (node_id, truth_state) in keys(logic_mapping)
-        pattern = r"^[TFUD\*]+$"
+        pattern = r"^[TF\*]+$"
         if !occursin(pattern, truth_state)
             error("Truth state \'$truth_state\' contains illegal characters or is empty.")
         end
@@ -709,19 +709,3 @@ function get_influx(basin::Basin, basin_idx::Int)::Float64
     return precipitation[basin_idx] - evaporation[basin_idx] + drainage[basin_idx] -
            infiltration[basin_idx]
 end
-
-function get_discrete_control_indices(discrete_control::DiscreteControl, condition_idx::Int)
-    (; greater_than) = discrete_control
-    condition_idx_now = 1
-
-    for (compound_variable_idx, vec) in enumerate(greater_than)
-        l = length(vec)
-
-        if condition_idx_now + l > condition_idx
-            greater_than_idx = condition_idx - condition_idx_now + 1
-            return compound_variable_idx, greater_than_idx
-        end
-
-        condition_idx_now += l
-    end
-end
diff --git a/core/src/validation.jl b/core/src/validation.jl
index 9d6a621d2..6551690ad 100644
--- a/core/src/validation.jl
+++ b/core/src/validation.jl
@@ -552,7 +552,6 @@ function valid_discrete_control(p::Parameters, config::Config)::Bool
 
             if !isempty(undefined_control_states)
                 undefined_list = collect(undefined_control_states)
-                node_type = typeof(node).name.name
                 @error "These control states from $id are not defined for controlled $id_outneighbor: $undefined_list."
                 errors = true
             end
diff --git a/core/test/control_test.jl b/core/test/control_test.jl
index 21bcc653c..1bee012a9 100644
--- a/core/test/control_test.jl
+++ b/core/test/control_test.jl
@@ -145,41 +145,11 @@ end
     @test discrete_control.record.control_state == ["high", "low"]
     @test discrete_control.record.time[1] == 0.0
     t = Ribasim.datetime_since(discrete_control.record.time[2], model.config.starttime)
-    @test Date(t) == Date("2020-03-15")
+    @test Date(t) == Date("2020-03-16")
     # then the rating curve is updated to the "low" control_state
     @test only(p.tabulated_rating_curve.tables).t[2] == 1.2
 end
 
-@testitem "Setpoint with bounds control" begin
-    toml_path = normpath(
-        @__DIR__,
-        "../../generated_testmodels/level_setpoint_with_minmax/ribasim.toml",
-    )
-    @test ispath(toml_path)
-    model = Ribasim.run(toml_path)
-    p = model.integrator.p
-    (; discrete_control) = p
-    (; record, greater_than) = discrete_control
-    level = Ribasim.get_storages_and_levels(model).level[1, :]
-    t = Ribasim.tsaves(model)
-
-    t_none_1 = discrete_control.record.time[2]
-    t_in = discrete_control.record.time[3]
-    t_none_2 = discrete_control.record.time[4]
-
-    level_min = greater_than[1][1]
-    setpoint = greater_than[1][2]
-
-    t_1_none_index = findfirst(>=(t_none_1), t)
-    t_in_index = findfirst(>=(t_in), t)
-    t_2_none_index = findfirst(>=(t_none_2), t)
-
-    @test record.control_state == ["out", "none", "in", "none"]
-    @test level[t_1_none_index] <= setpoint
-    @test level[t_in_index] >= level_min
-    @test level[t_2_none_index] <= setpoint
-end
-
 @testitem "Set PID target with DiscreteControl" begin
     using Ribasim: NodeID
 
diff --git a/docs/python/examples.ipynb b/docs/python/examples.ipynb
index d225579d3..7f7703ce3 100644
--- a/docs/python/examples.ipynb
+++ b/docs/python/examples.ipynb
@@ -433,327 +433,6 @@
     "ax.legend(bbox_to_anchor=(1.3, 1), title=\"Edge\")"
    ]
   },
-  {
-   "attachments": {},
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# Model with discrete control\n",
-    "\n",
-    "The model constructed below consists of a single basin which slowly drains trough a `TabulatedRatingCurve`, but is held within a range around a target level (setpoint) by two connected pumps. These two pumps behave like a reversible pump. When pumping can be done in only one direction, and the other direction is only possible under gravity, use an Outlet for that direction.\n"
-   ]
-  },
-  {
-   "attachments": {},
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Setup the basins:\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "model = Model(starttime=\"2020-01-01\", endtime=\"2021-01-01\", crs=\"EPSG:4326\")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "model.basin.add(\n",
-    "    Node(1, Point(0.0, 0.0)),\n",
-    "    [\n",
-    "        basin.Profile(area=[1000.0, 1000.0], level=[0.0, 1.0]),\n",
-    "        basin.State(level=[20.0]),\n",
-    "    ],\n",
-    ")"
-   ]
-  },
-  {
-   "attachments": {},
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Setup the discrete control:\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "model.discrete_control.add(\n",
-    "    Node(7, Point(1.0, 0.0)),\n",
-    "    [\n",
-    "        discrete_control.Variable(\n",
-    "            listen_node_id=[1],\n",
-    "            listen_node_type=\"Basin\",\n",
-    "            variable=\"level\",\n",
-    "            compound_variable_id=1,\n",
-    "        ),\n",
-    "        discrete_control.Condition(\n",
-    "            greater_than=[5.0, 10.0, 15.0], compound_variable_id=1\n",
-    "        ),\n",
-    "        discrete_control.Logic(\n",
-    "            truth_state=[\"FFF\", \"U**\", \"T*F\", \"**D\", \"TTT\"],\n",
-    "            control_state=[\"in\", \"in\", \"none\", \"out\", \"out\"],\n",
-    "        ),\n",
-    "    ],\n",
-    ")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "The above control logic can be summarized as follows:\n",
-    "\n",
-    "- If the level gets above the maximum, activate the control state \"out\" until the setpoint is reached;\n",
-    "- If the level gets below the minimum, active the control state \"in\" until the setpoint is reached;\n",
-    "- Otherwise activate the control state \"none\".\n"
-   ]
-  },
-  {
-   "attachments": {},
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Setup the pump:\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "model.pump.add(\n",
-    "    Node(2, Point(1.0, 1.0)),\n",
-    "    [pump.Static(control_state=[\"none\", \"in\", \"out\"], flow_rate=[0.0, 2e-3, 0.0])],\n",
-    ")\n",
-    "model.pump.add(\n",
-    "    Node(3, Point(1.0, -1.0)),\n",
-    "    [pump.Static(control_state=[\"none\", \"in\", \"out\"], flow_rate=[0.0, 0.0, 2e-3])],\n",
-    ")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "The pump data defines the following:\n",
-    "\n",
-    "| Control state | Pump #2 flow rate (m/s) | Pump #3 flow rate (m/s) |\n",
-    "| ------------- | ----------------------- | ----------------------- |\n",
-    "| \"none\"        | 0.0                     | 0.0                     |\n",
-    "| \"in\"          | 2e-3                    | 0.0                     |\n",
-    "| \"out\"         | 0.0                     | 2e-3                    |\n"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Setup the level boundary:\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "model.level_boundary.add(\n",
-    "    Node(4, Point(2.0, 0.0)), [level_boundary.Static(level=[10.0])]\n",
-    ")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Setup the rating curve:\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "model.tabulated_rating_curve.add(\n",
-    "    Node(5, Point(-1.0, 0.0)),\n",
-    "    [tabulated_rating_curve.Static(level=[2.0, 15.0], flow_rate=[0.0, 1e-3])],\n",
-    ")"
-   ]
-  },
-  {
-   "attachments": {},
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Setup the terminal:\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "model.terminal.add(Node(6, Point(-2.0, 0.0)))"
-   ]
-  },
-  {
-   "attachments": {},
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Setup edges:\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "model.edge.add(model.basin[1], model.pump[3])\n",
-    "model.edge.add(model.pump[3], model.level_boundary[4])\n",
-    "model.edge.add(model.level_boundary[4], model.pump[2])\n",
-    "model.edge.add(model.pump[2], model.basin[1])\n",
-    "model.edge.add(model.basin[1], model.tabulated_rating_curve[5])\n",
-    "model.edge.add(model.tabulated_rating_curve[5], model.terminal[6])\n",
-    "model.edge.add(model.discrete_control[7], model.pump[2])\n",
-    "model.edge.add(model.discrete_control[7], model.pump[3])"
-   ]
-  },
-  {
-   "attachments": {},
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Let’s take a look at the model:\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "model.plot()"
-   ]
-  },
-  {
-   "attachments": {},
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Listen edges are plotted with a dashed line since they are not present in the \"Edge / static\" schema but only in the \"Control / condition\" schema.\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "datadir = Path(\"data\")\n",
-    "model.write(datadir / \"level_setpoint_with_minmax/ribasim.toml\")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# | include: false\n",
-    "from subprocess import run\n",
-    "\n",
-    "run(\n",
-    "    [\n",
-    "        \"julia\",\n",
-    "        \"--project=../../core\",\n",
-    "        \"--eval\",\n",
-    "        f'using Ribasim; Ribasim.main(\"{datadir.as_posix()}/level_setpoint_with_minmax/ribasim.toml\")',\n",
-    "    ],\n",
-    "    check=True,\n",
-    ")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Now run the model with `ribasim level_setpoint_with_minmax/ribasim.toml`.\n",
-    "After running the model, read back the results:\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from matplotlib.dates import date2num\n",
-    "\n",
-    "df_basin = pd.read_feather(datadir / \"level_setpoint_with_minmax/results/basin.arrow\")\n",
-    "df_basin_wide = df_basin.pivot_table(\n",
-    "    index=\"time\", columns=\"node_id\", values=[\"storage\", \"level\"]\n",
-    ")\n",
-    "\n",
-    "ax = df_basin_wide[\"level\"].plot()\n",
-    "\n",
-    "greater_than = model.discrete_control.condition.df.greater_than\n",
-    "\n",
-    "ax.hlines(\n",
-    "    greater_than,\n",
-    "    df_basin.time[0],\n",
-    "    df_basin.time.max(),\n",
-    "    lw=1,\n",
-    "    ls=\"--\",\n",
-    "    color=\"k\",\n",
-    ")\n",
-    "\n",
-    "df_control = pd.read_feather(\n",
-    "    datadir / \"level_setpoint_with_minmax/results/control.arrow\"\n",
-    ")\n",
-    "\n",
-    "y_min, y_max = ax.get_ybound()\n",
-    "ax.fill_between(\n",
-    "    df_control.time[:2].to_numpy(), 2 * [y_min], 2 * [y_max], alpha=0.2, color=\"C0\"\n",
-    ")\n",
-    "ax.fill_between(\n",
-    "    df_control.time[2:4].to_numpy(), 2 * [y_min], 2 * [y_max], alpha=0.2, color=\"C0\"\n",
-    ")\n",
-    "\n",
-    "ax.set_xticks(\n",
-    "    date2num(df_control.time).tolist(),\n",
-    "    df_control.control_state.tolist(),\n",
-    "    rotation=50,\n",
-    ")\n",
-    "\n",
-    "ax.set_yticks(greater_than, [\"min\", \"setpoint\", \"max\"])\n",
-    "ax.set_ylabel(\"level\")\n",
-    "plt.show()"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "The highlighted regions show where a pump is active.\n"
-   ]
-  },
   {
    "cell_type": "markdown",
    "metadata": {},
diff --git a/python/ribasim/tests/conftest.py b/python/ribasim/tests/conftest.py
index 30939e521..4b4e1a7ba 100644
--- a/python/ribasim/tests/conftest.py
+++ b/python/ribasim/tests/conftest.py
@@ -39,11 +39,6 @@ def discrete_control_of_pid_control() -> ribasim.Model:
     return ribasim_testmodels.discrete_control_of_pid_control_model()
 
 
-@pytest.fixture()
-def level_setpoint_with_minmax() -> ribasim.Model:
-    return ribasim_testmodels.level_setpoint_with_minmax_model()
-
-
 @pytest.fixture()
 def trivial() -> ribasim.Model:
     return ribasim_testmodels.trivial_model()
diff --git a/python/ribasim/tests/test_io.py b/python/ribasim/tests/test_io.py
index e6507d77e..bd35fcc5c 100644
--- a/python/ribasim/tests/test_io.py
+++ b/python/ribasim/tests/test_io.py
@@ -92,41 +92,6 @@ def test_extra_columns():
         terminal.Static(meta_id=[-1, -2, -3], extra=[-1, -2, -3])
 
 
-def test_sort(level_setpoint_with_minmax, tmp_path):
-    model = level_setpoint_with_minmax
-    table = model.discrete_control.condition
-    edge = model.edge
-
-    # apply a wrong sort, then call the sort method to restore order
-    table.df.sort_values("greater_than", ascending=False, inplace=True)
-    assert table.df.iloc[0]["greater_than"] == 15.0
-    assert table._sort_keys == [
-        "node_id",
-        "compound_variable_id",
-        "greater_than",
-    ]
-    table.sort()
-    assert table.df.iloc[0]["greater_than"] == 5.0
-
-    # The edge table is not sorted
-    assert edge.df.iloc[1]["from_node_type"] == "Pump"
-    assert edge.df.iloc[1]["from_node_id"] == 3
-
-    # re-apply wrong sort, then check if it gets sorted on write
-    table.df.sort_values("greater_than", ascending=False, inplace=True)
-    model.write(tmp_path / "basic/ribasim.toml")
-    # write sorts the model in place
-    assert table.df.iloc[0]["greater_than"] == 5.0
-    model_loaded = ribasim.Model.read(filepath=tmp_path / "basic/ribasim.toml")
-    table_loaded = model_loaded.discrete_control.condition
-    edge_loaded = model_loaded.edge
-    assert table_loaded.df.iloc[0]["greater_than"] == 5.0
-    assert edge.df.iloc[1]["from_node_type"] == "Pump"
-    assert edge.df.iloc[1]["from_node_id"] == 3
-    __assert_equal(table.df, table_loaded.df)
-    __assert_equal(edge.df, edge_loaded.df)
-
-
 def test_roundtrip(trivial, tmp_path):
     model1 = trivial
     # set custom Edge index
diff --git a/python/ribasim_testmodels/ribasim_testmodels/__init__.py b/python/ribasim_testmodels/ribasim_testmodels/__init__.py
index 4f9ce86c3..736cdeab3 100644
--- a/python/ribasim_testmodels/ribasim_testmodels/__init__.py
+++ b/python/ribasim_testmodels/ribasim_testmodels/__init__.py
@@ -31,7 +31,6 @@
     compound_variable_condition_model,
     flow_condition_model,
     level_boundary_condition_model,
-    level_setpoint_with_minmax_model,
     pump_discrete_control_model,
     tabulated_rating_curve_control_model,
 )
@@ -79,7 +78,6 @@
     "leaky_bucket_model",
     "level_boundary_condition_model",
     "level_demand_model",
-    "level_setpoint_with_minmax_model",
     "linear_resistance_demand_model",
     "linear_resistance_model",
     "local_pidcontrolled_cascade_model",
diff --git a/python/ribasim_testmodels/ribasim_testmodels/discrete_control.py b/python/ribasim_testmodels/ribasim_testmodels/discrete_control.py
index 650e73876..7111c0d4b 100644
--- a/python/ribasim_testmodels/ribasim_testmodels/discrete_control.py
+++ b/python/ribasim_testmodels/ribasim_testmodels/discrete_control.py
@@ -323,99 +323,6 @@ def tabulated_rating_curve_control_model() -> Model:
     return model
 
 
-def level_setpoint_with_minmax_model() -> Model:
-    """
-    Set up a minimal model in which the level of a basin is kept within an acceptable range
-    around a setpoint while being affected by time-varying forcing.
-    This is done by bringing the level back to the setpoint once the level goes beyond this range.
-    """
-
-    model = Model(
-        starttime="2020-01-01",
-        endtime="2021-01-01",
-        crs="EPSG:28992",
-    )
-
-    model.basin.add(
-        Node(1, Point(0, 0)),
-        [
-            basin.Profile(area=1000.0, level=[0.0, 1.0]),
-            basin.State(level=[20.0]),
-        ],
-    )
-    model.pump.add(
-        Node(2, Point(1, 1)),
-        [pump.Static(control_state=["none", "in", "out"], flow_rate=[0.0, 2e-3, 0.0])],
-    )
-    model.pump.add(
-        Node(3, Point(1, -1)),
-        [pump.Static(control_state=["none", "in", "out"], flow_rate=[0.0, 0.0, 2e-3])],
-    )
-    model.level_boundary.add(
-        Node(4, Point(2, 0)), [level_boundary.Static(level=[10.0])]
-    )
-    model.tabulated_rating_curve.add(
-        Node(5, Point(-1, 0)),
-        [tabulated_rating_curve.Static(level=[2.0, 15.0], flow_rate=[0.0, 1e-3])],
-    )
-    model.terminal.add(Node(6, Point(-2, 0)))
-    model.discrete_control.add(
-        Node(7, Point(1, 0)),
-        [
-            discrete_control.Variable(
-                listen_node_type="Basin",
-                listen_node_id=[1],
-                variable="level",
-                compound_variable_id=1,
-            ),
-            discrete_control.Condition(
-                # min, setpoint, max
-                greater_than=[5.0, 10.0, 15.0],
-                compound_variable_id=1,
-            ),
-            discrete_control.Logic(
-                truth_state=["FFF", "U**", "T*F", "**D", "TTT"],
-                control_state=["in", "in", "none", "out", "out"],
-            ),
-        ],
-    )
-
-    model.edge.add(
-        model.basin[1],
-        model.pump[3],
-    )
-    model.edge.add(
-        model.pump[3],
-        model.level_boundary[4],
-    )
-    model.edge.add(
-        model.level_boundary[4],
-        model.pump[2],
-    )
-    model.edge.add(
-        model.pump[2],
-        model.basin[1],
-    )
-    model.edge.add(
-        model.basin[1],
-        model.tabulated_rating_curve[5],
-    )
-    model.edge.add(
-        model.tabulated_rating_curve[5],
-        model.terminal[6],
-    )
-    model.edge.add(
-        model.discrete_control[7],
-        model.pump[2],
-    )
-    model.edge.add(
-        model.discrete_control[7],
-        model.pump[3],
-    )
-
-    return model
-
-
 def compound_variable_condition_model() -> Model:
     """
     Set up a minimal model containing a condition on a compound variable

From 9c174eb302428b3d908bfb881634c460f80dbb95 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Tue, 16 Apr 2024 17:07:02 +0200
Subject: [PATCH 22/25] Bring back example model

---
 docs/python/examples.ipynb                    | 326 +++++++++++++++++-
 .../ribasim_testmodels/__init__.py            |   2 +
 .../ribasim_testmodels/discrete_control.py    |  94 +++++
 3 files changed, 419 insertions(+), 3 deletions(-)

diff --git a/docs/python/examples.ipynb b/docs/python/examples.ipynb
index 7f7703ce3..95fec279d 100644
--- a/docs/python/examples.ipynb
+++ b/docs/python/examples.ipynb
@@ -433,6 +433,328 @@
     "ax.legend(bbox_to_anchor=(1.3, 1), title=\"Edge\")"
    ]
   },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Model with discrete control\n",
+    "\n",
+    "The model constructed below consists of a single basin which slowly drains trough a `TabulatedRatingCurve`, but is held within a range by two connected pumps. These two pumps together behave like a reversible pump. When pumping can be done in only one direction, and the other direction is only possible under gravity, use an Outlet for that direction.\n"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Setup the basins:\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model = Model(starttime=\"2020-01-01\", endtime=\"2021-01-01\", crs=\"EPSG:4326\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model.basin.add(\n",
+    "    Node(1, Point(0.0, 0.0)),\n",
+    "    [\n",
+    "        basin.Profile(area=[1000.0, 1000.0], level=[0.0, 1.0]),\n",
+    "        basin.State(level=[20.0]),\n",
+    "        basin.Time(time=[\"2020-01-01\", \"2020-07-01\"], precipitation=[0.0, 3e-6]),\n",
+    "    ],\n",
+    ")"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Setup the discrete control:\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model.discrete_control.add(\n",
+    "    Node(7, Point(1.0, 0.0)),\n",
+    "    [\n",
+    "        discrete_control.Variable(\n",
+    "            compound_variable_id=1,\n",
+    "            listen_node_id=1,\n",
+    "            listen_node_type=[\"Basin\", \"Basin\"],\n",
+    "            variable=[\"level\", \"level\"],\n",
+    "        ),\n",
+    "        discrete_control.Condition(\n",
+    "            compound_variable_id=1,\n",
+    "            # min, max\n",
+    "            greater_than=[5.0, 15.0],\n",
+    "        ),\n",
+    "        discrete_control.Logic(\n",
+    "            truth_state=[\"FF\", \"TF\", \"TT\"],\n",
+    "            control_state=[\"in\", \"none\", \"out\"],\n",
+    "        ),\n",
+    "    ],\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "The above control logic can be summarized as follows:\n",
+    "\n",
+    "- If the level is above the maximum, activate the control state \"out\";\n",
+    "- If the level is below the minimum, active the control state \"in\";\n",
+    "- Otherwise activate the control state \"none\".\n"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Setup the pump:\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model.pump.add(\n",
+    "    Node(2, Point(1.0, 1.0)),\n",
+    "    [pump.Static(control_state=[\"none\", \"in\", \"out\"], flow_rate=[0.0, 2e-3, 0.0])],\n",
+    ")\n",
+    "model.pump.add(\n",
+    "    Node(3, Point(1.0, -1.0)),\n",
+    "    [pump.Static(control_state=[\"none\", \"in\", \"out\"], flow_rate=[0.0, 0.0, 2e-3])],\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "The pump data defines the following:\n",
+    "\n",
+    "| Control state | Pump #2 flow rate (m/s) | Pump #3 flow rate (m/s) |\n",
+    "| ------------- | ----------------------- | ----------------------- |\n",
+    "| \"none\"        | 0.0                     | 0.0                     |\n",
+    "| \"in\"          | 2e-3                    | 0.0                     |\n",
+    "| \"out\"         | 0.0                     | 2e-3                    |\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Setup the level boundary:\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model.level_boundary.add(\n",
+    "    Node(4, Point(2.0, 0.0)), [level_boundary.Static(level=[10.0])]\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Setup the rating curve:\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model.tabulated_rating_curve.add(\n",
+    "    Node(5, Point(-1.0, 0.0)),\n",
+    "    [tabulated_rating_curve.Static(level=[2.0, 15.0], flow_rate=[0.0, 2e-3])],\n",
+    ")"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Setup the terminal:\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model.terminal.add(Node(6, Point(-2.0, 0.0)))"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Setup edges:\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model.edge.add(model.basin[1], model.pump[3])\n",
+    "model.edge.add(model.pump[3], model.level_boundary[4])\n",
+    "model.edge.add(model.level_boundary[4], model.pump[2])\n",
+    "model.edge.add(model.pump[2], model.basin[1])\n",
+    "model.edge.add(model.basin[1], model.tabulated_rating_curve[5])\n",
+    "model.edge.add(model.tabulated_rating_curve[5], model.terminal[6])\n",
+    "model.edge.add(model.discrete_control[7], model.pump[2])\n",
+    "model.edge.add(model.discrete_control[7], model.pump[3])"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Let’s take a look at the model:\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model.plot()"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Listen edges are plotted with a dashed line since they are not present in the \"Edge / static\" schema but only in the \"Control / condition\" schema.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "datadir = Path(\"data\")\n",
+    "model.write(datadir / \"level_range/ribasim.toml\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# | include: false\n",
+    "from subprocess import run\n",
+    "\n",
+    "run(\n",
+    "    [\n",
+    "        \"julia\",\n",
+    "        \"--project=../../core\",\n",
+    "        \"--eval\",\n",
+    "        f'using Ribasim; Ribasim.main(\"{datadir.as_posix()}/level_range/ribasim.toml\")',\n",
+    "    ],\n",
+    "    check=True,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Now run the model with `ribasim level_range/ribasim.toml`.\n",
+    "After running the model, read back the results:\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from matplotlib.dates import date2num\n",
+    "\n",
+    "df_basin = pd.read_feather(datadir / \"level_range/results/basin.arrow\")\n",
+    "df_basin_wide = df_basin.pivot_table(\n",
+    "    index=\"time\", columns=\"node_id\", values=[\"storage\", \"level\"]\n",
+    ")\n",
+    "\n",
+    "ax = df_basin_wide[\"level\"].plot()\n",
+    "\n",
+    "greater_than = model.discrete_control.condition.df.greater_than\n",
+    "\n",
+    "ax.hlines(\n",
+    "    greater_than,\n",
+    "    df_basin.time[0],\n",
+    "    df_basin.time.max(),\n",
+    "    lw=1,\n",
+    "    ls=\"--\",\n",
+    "    color=\"k\",\n",
+    ")\n",
+    "\n",
+    "df_control = pd.read_feather(datadir / \"level_range/control.arrow\")\n",
+    "\n",
+    "y_min, y_max = ax.get_ybound()\n",
+    "ax.fill_between(\n",
+    "    df_control.time[:2].to_numpy(), 2 * [y_min], 2 * [y_max], alpha=0.2, color=\"C0\"\n",
+    ")\n",
+    "ax.fill_between(\n",
+    "    df_control.time[2:4].to_numpy(), 2 * [y_min], 2 * [y_max], alpha=0.2, color=\"C0\"\n",
+    ")\n",
+    "\n",
+    "ax.set_xticks(\n",
+    "    date2num(df_control.time).tolist(),\n",
+    "    df_control.control_state.tolist(),\n",
+    "    rotation=50,\n",
+    ")\n",
+    "\n",
+    "ax.set_yticks(greater_than, [\"min\", \"setpoint\", \"max\"])\n",
+    "ax.set_ylabel(\"level\")\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "The highlighted regions show where a pump is active.\n"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -865,12 +1187,10 @@
     "model.discrete_control.add(\n",
     "    Node(11, Point(4.5, 0.25), subnetwork_id=1),\n",
     "    [\n",
-    "        discrete_control.Variable(\n",
+    "        discrete_control.Condition(\n",
     "            listen_node_id=[5],\n",
     "            listen_node_type=[\"Basin\"],\n",
     "            variable=[\"level\"],\n",
-    "        ),\n",
-    "        discrete_control.Condition(\n",
     "            greater_than=[0.52],\n",
     "        ),\n",
     "        discrete_control.Logic(\n",
diff --git a/python/ribasim_testmodels/ribasim_testmodels/__init__.py b/python/ribasim_testmodels/ribasim_testmodels/__init__.py
index 736cdeab3..bae49d032 100644
--- a/python/ribasim_testmodels/ribasim_testmodels/__init__.py
+++ b/python/ribasim_testmodels/ribasim_testmodels/__init__.py
@@ -31,6 +31,7 @@
     compound_variable_condition_model,
     flow_condition_model,
     level_boundary_condition_model,
+    level_range_model,
     pump_discrete_control_model,
     tabulated_rating_curve_control_model,
 )
@@ -78,6 +79,7 @@
     "leaky_bucket_model",
     "level_boundary_condition_model",
     "level_demand_model",
+    "level_range_model",
     "linear_resistance_demand_model",
     "linear_resistance_model",
     "local_pidcontrolled_cascade_model",
diff --git a/python/ribasim_testmodels/ribasim_testmodels/discrete_control.py b/python/ribasim_testmodels/ribasim_testmodels/discrete_control.py
index 7111c0d4b..fcdd5b501 100644
--- a/python/ribasim_testmodels/ribasim_testmodels/discrete_control.py
+++ b/python/ribasim_testmodels/ribasim_testmodels/discrete_control.py
@@ -379,3 +379,97 @@ def compound_variable_condition_model() -> Model:
     model.edge.add(model.discrete_control[6], model.pump[4])
 
     return model
+
+
+def level_range_model() -> Model:
+    """
+    Set up a minimal model in which the level of a basin is kept within an acceptable range
+    around a setpoint while being affected by time-varying forcing.
+    This is done by bringing the level back to the setpoint once the level goes beyond this range.
+    """
+
+    model = Model(
+        starttime="2020-01-01",
+        endtime="2021-01-01",
+        crs="EPSG:28992",
+    )
+
+    model.basin.add(
+        Node(1, Point(0, 0)),
+        [
+            basin.Profile(area=1000.0, level=[0.0, 1.0]),
+            basin.State(level=[20.0]),
+            basin.Time(time=["2020-01-01", "2020-07-01"], precipitation=[0.0, 3e-6]),
+        ],
+    )
+    model.pump.add(
+        Node(2, Point(1, 1)),
+        [pump.Static(control_state=["none", "in", "out"], flow_rate=[0.0, 2e-3, 0.0])],
+    )
+    model.pump.add(
+        Node(3, Point(1, -1)),
+        [pump.Static(control_state=["none", "in", "out"], flow_rate=[0.0, 0.0, 2e-3])],
+    )
+    model.level_boundary.add(
+        Node(4, Point(2, 0)), [level_boundary.Static(level=[10.0])]
+    )
+    model.tabulated_rating_curve.add(
+        Node(5, Point(-1, 0)),
+        [tabulated_rating_curve.Static(level=[2.0, 15.0], flow_rate=[0.0, 2e-3])],
+    )
+    model.terminal.add(Node(6, Point(-2, 0)))
+    model.discrete_control.add(
+        Node(7, Point(1, 0)),
+        [
+            discrete_control.Variable(
+                listen_node_type="Basin",
+                listen_node_id=[1],
+                variable="level",
+                compound_variable_id=1,
+            ),
+            discrete_control.Condition(
+                # min, max
+                greater_than=[5.0, 15.0],
+                compound_variable_id=1,
+            ),
+            discrete_control.Logic(
+                truth_state=["FF", "TF", "TT"],
+                control_state=["in", "none", "out"],
+            ),
+        ],
+    )
+
+    model.edge.add(
+        model.basin[1],
+        model.pump[3],
+    )
+    model.edge.add(
+        model.pump[3],
+        model.level_boundary[4],
+    )
+    model.edge.add(
+        model.level_boundary[4],
+        model.pump[2],
+    )
+    model.edge.add(
+        model.pump[2],
+        model.basin[1],
+    )
+    model.edge.add(
+        model.basin[1],
+        model.tabulated_rating_curve[5],
+    )
+    model.edge.add(
+        model.tabulated_rating_curve[5],
+        model.terminal[6],
+    )
+    model.edge.add(
+        model.discrete_control[7],
+        model.pump[2],
+    )
+    model.edge.add(
+        model.discrete_control[7],
+        model.pump[3],
+    )
+
+    return model

From 2dcefd303950ebbe76a8c7805454569896afcd9d Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Tue, 16 Apr 2024 17:18:46 +0200
Subject: [PATCH 23/25] Bring back io tests

---
 docs/python/examples.ipynb       | 27 +++++-------------------
 python/ribasim/tests/conftest.py |  5 +++++
 python/ribasim/tests/test_io.py  | 35 ++++++++++++++++++++++++++++++++
 3 files changed, 45 insertions(+), 22 deletions(-)

diff --git a/docs/python/examples.ipynb b/docs/python/examples.ipynb
index 95fec279d..13cbb99fd 100644
--- a/docs/python/examples.ipynb
+++ b/docs/python/examples.ipynb
@@ -496,8 +496,8 @@
     "        discrete_control.Variable(\n",
     "            compound_variable_id=1,\n",
     "            listen_node_id=1,\n",
-    "            listen_node_type=[\"Basin\", \"Basin\"],\n",
-    "            variable=[\"level\", \"level\"],\n",
+    "            listen_node_type=[\"Basin\"],\n",
+    "            variable=[\"level\"],\n",
     "        ),\n",
     "        discrete_control.Condition(\n",
     "            compound_variable_id=1,\n",
@@ -707,8 +707,6 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from matplotlib.dates import date2num\n",
-    "\n",
     "df_basin = pd.read_feather(datadir / \"level_range/results/basin.arrow\")\n",
     "df_basin_wide = df_basin.pivot_table(\n",
     "    index=\"time\", columns=\"node_id\", values=[\"storage\", \"level\"]\n",
@@ -727,23 +725,7 @@
     "    color=\"k\",\n",
     ")\n",
     "\n",
-    "df_control = pd.read_feather(datadir / \"level_range/control.arrow\")\n",
-    "\n",
-    "y_min, y_max = ax.get_ybound()\n",
-    "ax.fill_between(\n",
-    "    df_control.time[:2].to_numpy(), 2 * [y_min], 2 * [y_max], alpha=0.2, color=\"C0\"\n",
-    ")\n",
-    "ax.fill_between(\n",
-    "    df_control.time[2:4].to_numpy(), 2 * [y_min], 2 * [y_max], alpha=0.2, color=\"C0\"\n",
-    ")\n",
-    "\n",
-    "ax.set_xticks(\n",
-    "    date2num(df_control.time).tolist(),\n",
-    "    df_control.control_state.tolist(),\n",
-    "    rotation=50,\n",
-    ")\n",
-    "\n",
-    "ax.set_yticks(greater_than, [\"min\", \"setpoint\", \"max\"])\n",
+    "ax.set_yticks(greater_than, [\"min\", \"max\"])\n",
     "ax.set_ylabel(\"level\")\n",
     "plt.show()"
    ]
@@ -752,7 +734,8 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "The highlighted regions show where a pump is active.\n"
+    "We see that in Januari the level of the basin is too high and thus water is pumped out until the maximum level of the desried range is reached. Then until May water flows out of the basin freely trough the tabulated rating curve until the minimum leavel is reached. From \n",
+    "this point until the start of July water is pumped into the basin in short bursts to stay within the desired range. at the Start of July rain starts falling on the basin, which causes the basin level to rise until the maximum level. From this point onward water is pumped out of the basin in short bursts to stay within the desired range."
    ]
   },
   {
diff --git a/python/ribasim/tests/conftest.py b/python/ribasim/tests/conftest.py
index 4b4e1a7ba..d77d287e1 100644
--- a/python/ribasim/tests/conftest.py
+++ b/python/ribasim/tests/conftest.py
@@ -39,6 +39,11 @@ def discrete_control_of_pid_control() -> ribasim.Model:
     return ribasim_testmodels.discrete_control_of_pid_control_model()
 
 
+@pytest.fixture()
+def level_range() -> ribasim.Model:
+    return ribasim_testmodels.level_range_model()
+
+
 @pytest.fixture()
 def trivial() -> ribasim.Model:
     return ribasim_testmodels.trivial_model()
diff --git a/python/ribasim/tests/test_io.py b/python/ribasim/tests/test_io.py
index bd35fcc5c..2fb6e50de 100644
--- a/python/ribasim/tests/test_io.py
+++ b/python/ribasim/tests/test_io.py
@@ -92,6 +92,41 @@ def test_extra_columns():
         terminal.Static(meta_id=[-1, -2, -3], extra=[-1, -2, -3])
 
 
+def test_sort(level_range, tmp_path):
+    model = level_range
+    table = model.discrete_control.condition
+    edge = model.edge
+
+    # apply a wrong sort, then call the sort method to restore order
+    table.df.sort_values("greater_than", ascending=False, inplace=True)
+    assert table.df.iloc[0]["greater_than"] == 15.0
+    assert table._sort_keys == [
+        "node_id",
+        "compound_variable_id",
+        "greater_than",
+    ]
+    table.sort()
+    assert table.df.iloc[0]["greater_than"] == 5.0
+
+    # The edge table is not sorted
+    assert edge.df.iloc[1]["from_node_type"] == "Pump"
+    assert edge.df.iloc[1]["from_node_id"] == 3
+
+    # re-apply wrong sort, then check if it gets sorted on write
+    table.df.sort_values("greater_than", ascending=False, inplace=True)
+    model.write(tmp_path / "basic/ribasim.toml")
+    # write sorts the model in place
+    assert table.df.iloc[0]["greater_than"] == 5.0
+    model_loaded = ribasim.Model.read(filepath=tmp_path / "basic/ribasim.toml")
+    table_loaded = model_loaded.discrete_control.condition
+    edge_loaded = model_loaded.edge
+    assert table_loaded.df.iloc[0]["greater_than"] == 5.0
+    assert edge.df.iloc[1]["from_node_type"] == "Pump"
+    assert edge.df.iloc[1]["from_node_id"] == 3
+    __assert_equal(table.df, table_loaded.df)
+    __assert_equal(edge.df, edge_loaded.df)
+
+
 def test_roundtrip(trivial, tmp_path):
     model1 = trivial
     # set custom Edge index

From f2517edf46e58bfb978c5b5ad4c26ef135b5287e Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Tue, 16 Apr 2024 18:27:25 +0200
Subject: [PATCH 24/25] Fix examples notebook

---
 docs/python/examples.ipynb | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/docs/python/examples.ipynb b/docs/python/examples.ipynb
index 13cbb99fd..2569156fa 100644
--- a/docs/python/examples.ipynb
+++ b/docs/python/examples.ipynb
@@ -1170,10 +1170,14 @@
     "model.discrete_control.add(\n",
     "    Node(11, Point(4.5, 0.25), subnetwork_id=1),\n",
     "    [\n",
-    "        discrete_control.Condition(\n",
+    "        discrete_control.Variable(\n",
+    "            compound_variable_id=1,\n",
     "            listen_node_id=[5],\n",
     "            listen_node_type=[\"Basin\"],\n",
     "            variable=[\"level\"],\n",
+    "        ),\n",
+    "        discrete_control.Condition(\n",
+    "            compound_variable_id=1,\n",
     "            greater_than=[0.52],\n",
     "        ),\n",
     "        discrete_control.Logic(\n",

From f00f459eb12783229ab99df9305b1a84bdef184c Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Wed, 17 Apr 2024 11:56:23 +0200
Subject: [PATCH 25/25] Comments adressed

---
 docs/python/examples.ipynb      | 4 ++--
 python/ribasim/ribasim/model.py | 1 -
 2 files changed, 2 insertions(+), 3 deletions(-)

diff --git a/docs/python/examples.ipynb b/docs/python/examples.ipynb
index 6d33f2fc8..5e1eb45ce 100644
--- a/docs/python/examples.ipynb
+++ b/docs/python/examples.ipynb
@@ -734,8 +734,8 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "We see that in Januari the level of the basin is too high and thus water is pumped out until the maximum level of the desried range is reached. Then until May water flows out of the basin freely trough the tabulated rating curve until the minimum leavel is reached. From \n",
-    "this point until the start of July water is pumped into the basin in short bursts to stay within the desired range. at the Start of July rain starts falling on the basin, which causes the basin level to rise until the maximum level. From this point onward water is pumped out of the basin in short bursts to stay within the desired range."
+    "We see that in Januari the level of the basin is too high and thus water is pumped out until the maximum level of the desired range is reached. Then until May water flows out of the basin freely through the tabulated rating curve until the minimum level is reached. From \n",
+    "this point until the start of July water is pumped into the basin in short bursts to stay within the desired range. At the start of July rain starts falling on the basin, which causes the basin level to rise until the maximum level. From this point onward water is pumped out of the basin in short bursts to stay within the desired range."
    ]
   },
   {
diff --git a/python/ribasim/ribasim/model.py b/python/ribasim/ribasim/model.py
index 966a65888..2b4f195ef 100644
--- a/python/ribasim/ribasim/model.py
+++ b/python/ribasim/ribasim/model.py
@@ -326,7 +326,6 @@ def plot_control_listen(self, ax):
             to_add = df_variable[
                 ["node_id", "listen_node_id", "listen_node_type"]
             ].drop_duplicates()
-            to_add = to_add[to_add["listen_node_type"] != "compound"]
             to_add.columns = ["control_node_id", "listen_node_id", "listen_node_type"]
             to_add["control_node_type"] = "DiscreteControl"
             df_listen_edge = pd.concat([df_listen_edge, to_add])