Add weir

core/src/bmi.jl

@@ -25,13 +25,14 @@ function BMI.initialize(T::Type{Model}, config::Config)::Model
             error("Invalid discrete control state definition(s).")
         end
 
-        (; pid_control, connectivity, basin, pump, fractional_flow) = parameters
+        (; pid_control, connectivity, basin, pump, weir, fractional_flow) = parameters
         if !valid_pid_connectivity(
             pid_control.node_id,
             pid_control.listen_node_id,
             connectivity.graph_flow,
             connectivity.graph_control,
             basin.node_id,
+            pump.node_id,
         )
             error("Invalid PidControl connectivity.")
         end
@@ -45,9 +46,14 @@ function BMI.initialize(T::Type{Model}, config::Config)::Model
         end
 
         for id in pid_control.node_id
-            id_pump = only(outneighbors(connectivity.graph_control, id))
-            pump_idx = findsorted(pump.node_id, id_pump)
-            pump.is_pid_controlled[pump_idx] = true
+            id_controlled = only(outneighbors(connectivity.graph_control, id))
+            pump_idx = findsorted(pump.node_id, id_controlled)
+            if isnothing(pump_idx)
+                weir_idx = findsorted(weir.node_id, id_controlled)
+                weir.is_pid_controlled[weir_idx] = true
+            else
+                pump.is_pid_controlled[pump_idx] = true
+            end
         end
 
         # tstops for transient flow_boundary

core/src/create.jl

@@ -4,11 +4,19 @@ function parse_static(
     nodetype::String,
     defaults::NamedTuple,
 )::NamedTuple
+    # E.g. `PumpStatic`
     static_type = eltype(static)
     columnnames_static = collect(fieldnames(static_type))
     mask = [symb ∉ [:node_id, :control_state] for symb in columnnames_static]
+
+    # The names of the variables that can define a control state
     columnnames_variables = columnnames_static[mask]
+
+    # The types of the variables that can define a control state
     columntypes_variables = collect(fieldtypes(static_type))[mask]
+
+    # A vector of vectors, for each variable the (initial) values for all nodes
+    # of the current type
     vals = []
 
     node_ids = get_ids(db, nodetype)
@@ -273,8 +281,6 @@ function Pump(db::DB, config::Config)::Pump
     static = load_structvector(db, config, PumpStaticV1)
     defaults = (; min_flow_rate = 0.0, max_flow_rate = NaN, active = true)
     static_parsed = parse_static(static, db, "Pump", defaults)
-
-    # TODO: use this in formulate_jac! for pump
     is_pid_controlled = falses(length(static_parsed.node_id))
 
     return Pump(
@@ -288,6 +294,23 @@ function Pump(db::DB, config::Config)::Pump
     )
 end
 
+function Weir(db::DB, config::Config)::Weir
+    static = load_structvector(db, config, WeirStaticV1)
+    defaults = (; min_flow_rate = 0.0, max_flow_rate = NaN, active = true)
+    static_parsed = parse_static(static, db, "Weir", defaults)
+    is_pid_controlled = falses(length(static_parsed.node_id))
+
+    return Weir(
+        static_parsed.node_id,
+        static_parsed.active,
+        static_parsed.flow_rate,
+        static_parsed.min_flow_rate,
+        static_parsed.max_flow_rate,
+        static_parsed.control_mapping,
+        is_pid_controlled,
+    )
+end
+
 function Terminal(db::DB, config::Config)::Terminal
     static = load_structvector(db, config, TerminalStaticV1)
     return Terminal(static.node_id)
@@ -407,6 +430,7 @@ function Parameters(db::DB, config::Config)::Parameters
     level_boundary = LevelBoundary(db, config)
     flow_boundary = FlowBoundary(db, config)
     pump = Pump(db, config)
+    weir = Weir(db, config)
     terminal = Terminal(db, config)
     discrete_control = DiscreteControl(db, config)
     pid_control = PidControl(db, config)
@@ -424,6 +448,7 @@ function Parameters(db::DB, config::Config)::Parameters
         level_boundary,
         flow_boundary,
         pump,
+        weir,
         terminal,
         discrete_control,
         pid_control,

core/src/jac.jl

@@ -212,17 +212,17 @@
 end
 
 """
-The contributions of Pump nodes to the Jacobian.
+The contributions of Pump and Weir nodes to the Jacobian.
 """
 function formulate_jac!(
-    pump::Pump,
+    node::Union{Pump, Weir},
     J::SparseMatrixCSC{Float64, Int64},
     u::ComponentVector{Float64},
     p::Parameters,
     t::Float64,
 )::Nothing
-    (; basin, fractional_flow, connectivity, pid_control) = p
-    (; active, node_id, flow_rate, is_pid_controlled) = pump
+    (; basin, fractional_flow, connectivity) = p
+    (; active, node_id, flow_rate, is_pid_controlled) = node
 
     (; graph_flow) = connectivity
 
@@ -357,7 +357,7 @@
 end
 
 """
-The contributions of TabulatedRatingCurve nodes to the Jacobian.
+The contributions of PidControl nodes to the Jacobian.
 """
 function formulate_jac!(
     pid_control::PidControl,
@@ -366,8 +366,7 @@
     p::Parameters,
     t::Float64,
 )::Nothing
-    # TODO: update after pid_control equation test merge
-    (; basin, connectivity, pump) = p
+    (; basin, connectivity, pump, weir) = p
     (; node_id, active, listen_node_id, proportional, integral, derivative, error) =
         pid_control
     (; min_flow_rate, max_flow_rate) = pump
@@ -391,18 +390,40 @@
             continue
         end
 
-        id_pump = only(outneighbors(graph_control, id))
+        controlled_node_id = only(outneighbors(graph_control, id))
+        controls_pump = (controlled_node_id in pump.node_id)
+
         listened_node_id = listen_node_id[i]
         _, listened_node_idx = id_index(basin.node_id, listened_node_id)
         listen_area = basin.current_area[listened_node_idx]
-        storage_listened_basin = u.storage[listened_node_idx]
-        area = basin.current_area[listened_node_idx]
-        reduction_factor = min(storage_listened_basin, 10.0) / 10.0
+
+        if controls_pump
+            controlled_node_idx = findsorted(pump.node_id, controlled_node_id)
+
+            listened_basin_storage = u.storage[listened_node_idx]
+            reduction_factor = min(listened_basin_storage, 10.0) / 10.0
+        else
+            controlled_node_idx = findsorted(weir.node_id, controlled_node_id)
+
+            # Upstream node of weir does not have to be a basin
+            upstream_node_id = only(inneighbors(graph_flow, controlled_node_id))
+            has_upstream_index, upstream_basin_idx =
+                id_index(basin.node_id, upstream_node_id)
+            if has_upstream_index
+                upstream_basin_storage = u.storage[upstream_basin_idx]
+                reduction_factor = min(upstream_basin_storage, 10.0) / 10.0
+            else
+                reduction_factor = 1.0
+            end
+        end
 
         K_d = derivative[i]
         if !isnan(K_d)
-            # dlevel/dstorage = 1/area
-            D = 1.0 - K_d * reduction_factor / area
+            if controls_pump
+                D = 1.0 - K_d * reduction_factor / listen_area
+            else
+                D = 1.0 + K_d * reduction_factor / listen_area
+            end
         else
             D = 1.0
         end
@@ -422,22 +443,22 @@
         if !isnan(K_d)
             dtarget_level = 0.0
             du_listened_basin_old = du.storage[listened_node_idx]
-            E += K_d * (dtarget_level - du_listened_basin_old / area)
+            E += K_d * (dtarget_level - du_listened_basin_old / listen_area)
         end
 
         # Clip values outside pump flow rate bounds
         flow_rate = reduction_factor * E / D
         was_clipped = false
 
-        if flow_rate < min_flow_rate[i]
+        if flow_rate < min_flow_rate[controlled_node_idx]
             was_clipped = true
-            flow_rate = min_flow_rate[i]
+            flow_rate = min_flow_rate[controlled_node_idx]
         end
 
-        if !isnan(max_flow_rate[i])
-            if flow_rate > max_flow_rate[i]
+        if !isnan(max_flow_rate[controlled_node_idx])
+            if flow_rate > max_flow_rate[controlled_node_idx]
                 was_clipped = true
-                flow_rate = max_flow_rate[i]
+                flow_rate = max_flow_rate[controlled_node_idx]
             end
         end
 
@@ -453,51 +474,87 @@
         end
 
         # Only in this case the reduction factor has a non-zero derivative
-        reduction_factor_regime = (storage_listened_basin < 10.0)
+        reduction_factor_regime = if controls_pump
+            listened_basin_storage < 10.0
+        else
+            if has_upstream_index
+                upstream_basin_storage < 10.0
+            else
+                false
+            end
+        end
 
         # Computing D and E derivatives
         if !isnan(K_d)
             darea = basin.current_darea[listened_node_idx]
-            dD = reduction_factor * darea
+
+            dD_du_listen = reduction_factor * darea / (listen_area^2)
 
             if reduction_factor_regime
-                dD -= dreduction_factor / area
+                if controls_pump
+                    dD_du_listen -= 0.1 / darea
+                else
+                    dD_du_upstream = -0.1 * K_d / area
+                end
+            else
+                dD_du_upstream = 0.0
             end
 
-            dD *= K_d
+            dD_du_listen *= K_d
 
-            dE =
+            dE_du_listen =
                 -K_d * (
-                    area * J[listened_node_idx, listened_node_idx] -
+                    listen_area * J[listened_node_idx, listened_node_idx] -
                     du.storage[listened_node_idx] * darea
-                ) / (area^2)
+                ) / (listen_area^2)
         else
-            dD = 0.0
-            dE = 0.0
+            dD_du_listen = 0.0
+            dD_du_upstream = 0.0
+            dE_du_listen = 0.0
         end
 
         if !isnan(K_p)
-            dE -= K_p / area
+            dE_du_listen -= K_p / listen_area
+        end
+
+        dQ_du_listen = reduction_factor * (D * dE_du_listen - E * dD_du_listen) / (D^2)
+
+        if controls_pump && reduction_factor_regime
+            dQ_du_listen += 0.1 * E / D
         end
 
-        if reduction_factor_regime
-            dreduction_factor = 0.1
+        if controls_pump
+            J[listened_node_idx, listened_node_idx] -= dQ_du_listen
+
+            downstream_node_id = only(outneighbors(graph_flow, controlled_node_id))
+            has_downstream_index, downstream_node_idx =
+                id_index(basin.node_id, downstream_node_id)
 
-            dq = dreduction_factor * E / D
+            if has_downstream_index
+                J[listened_node_idx, downstream_node_idx] += dQ_du_listen
+            end
         else
-            dq = 0.0
+            J[listened_node_idx, listened_node_idx] += dQ_du_listen
+
+            if has_upstream_index
+                J[listened_node_idx, upstream_basin_idx] -= dQ_du_listen
+            end
         end
 
-        dq += reduction_factor * (D * dE - E * dD) / (D^2)
+        if !controls_pump
+            if !isnan(K_d) && has_upstream_index
+                dE_du_upstream = -K_d * J[upstream_basin_idx, listened_node_idx] / area
 
-        J[listened_node_idx, listened_node_idx] -= dq
+                dQ_du_upstream =
+                    reduction_factor * (D * dE_du_upstream - E * dD_du_upstream) / (D^2)
 
-        id_out = only(outneighbors(graph_flow, id_pump))
-        has_index, idx_out_out = id_index(basin.node_id, id_out)
+                if reduction_factor_regime
+                    dQ_du_upstream += 0.1 * E / D
+                end
 
-        if has_index
-            J[pid_state_idx, idx_out_out] += K_i * reduction_factor / D
-            J[listened_node_idx, idx_out_out] += dq
+                J[upstream_basin_idx, listened_node_idx] += dQ_du_upstream
+                J[upstream_basin_idx, upstream_basin_idx] -= dQ_du_upstream
+            end
         end
     end
     return nothing