Pid equation test (#474)

Fixes #473.

---------

Co-authored-by: Bart de Koning <[email protected]>

core/src/bmi.jl

@@ -21,7 +21,7 @@ function BMI.initialize(T::Type{Model}, config::Config)::Model
             error("Invalid number of connections for certain node types.")
         end
 
-        (; pid_control, connectivity, basin) = parameters
+        (; pid_control, connectivity, basin, pump) = parameters
         if !valid_pid_connectivity(
             pid_control.node_id,
             pid_control.listen_node_id,
@@ -32,6 +32,12 @@ function BMI.initialize(T::Type{Model}, config::Config)::Model
             error("Invalid PidControl connectivity.")
         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
+        end
+
         # use state
         state = load_structvector(db, config, BasinStateV1)
         n = length(get_ids(db, "Basin"))

core/src/create.jl

@@ -252,13 +252,17 @@ function Pump(db::DB, config::Config)::Pump
     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(
         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
 
@@ -292,8 +296,6 @@ function Basin(db::DB, config::Config)::Basin
     # If not specified, target_level = NaN
     target_level = coalesce.(static.target_level, NaN)
 
-    dstorage = zero(target_level)
-
     return Basin(
         Indices(node_id),
         precipitation,
@@ -307,7 +309,6 @@ function Basin(db::DB, config::Config)::Basin
         storage,
         target_level,
         time,
-        dstorage,
     )
 end
 

core/src/solve.jl

@@ -86,8 +86,6 @@ struct Basin{C} <: AbstractParameterNode
     target_level::Vector{Float64}
     # data source for parameter updates
     time::StructVector{BasinForcingV1, C, Int}
-    # Storage derivative for use in PID controller
-    dstorage::Vector{Float64}
 
     function Basin(
         node_id,
@@ -102,7 +100,6 @@ struct Basin{C} <: AbstractParameterNode
         storage,
         target_level,
         time::StructVector{BasinForcingV1, C, Int},
-        dstorage,
     ) where {C}
         errors = valid_profiles(node_id, level, area)
         if isempty(errors)
@@ -119,7 +116,6 @@ struct Basin{C} <: AbstractParameterNode
                 storage,
                 target_level,
                 time,
-                dstorage,
             )
         else
             @error join(errors, "\n")
@@ -252,6 +248,7 @@ struct Pump <: AbstractParameterNode
     min_flow_rate::Vector{Float64}
     max_flow_rate::Vector{Float64}
     control_mapping::Dict{Tuple{Int, String}, NamedTuple}
+    is_pid_controlled::BitVector
 end
 
 """
@@ -434,14 +431,19 @@ end
 
 function get_error!(pid_control::PidControl, p::Parameters)
     (; basin) = p
-    (; listen_node_id, error) = pid_control
+    (; listen_node_id) = pid_control
+
+    pid_error = pid_control.error
 
     for i in eachindex(listen_node_id)
         listened_node_id = listen_node_id[i]
         has_index, listened_node_idx = id_index(basin.node_id, listened_node_id)
         @assert has_index "Listen node $listened_node_id is not a Basin."
         target_level = basin.target_level[listened_node_idx]
-        error[i] = target_level - basin.current_level[listened_node_idx]
+        if isnan(target_level)
+            error("No target level specified for listen basin #$listened_node_id.")
+        end
+        pid_error[i] = target_level - basin.current_level[listened_node_idx]
     end
 end
 
@@ -454,10 +456,9 @@ function continuous_control!(
 )::Nothing
     # TODO: Also support being able to control weir
     # TODO: also support time varying target levels
-    (; connectivity, pump, basin) = p
+    (; connectivity, pump, basin, fractional_flow) = p
     (; min_flow_rate, max_flow_rate) = pump
-    (; dstorage) = basin
-    (; graph_control) = connectivity
+    (; graph_control, graph_flow, flow) = connectivity
     (; node_id, active, proportional, integral, derivative, listen_node_id, error) =
         pid_control
 
@@ -478,34 +479,74 @@ function continuous_control!(
 
         listened_node_id = listen_node_id[i]
         _, listened_node_idx = id_index(basin.node_id, listened_node_id)
+        storage_listened_basin = u.storage[listened_node_idx]
+        reduction_factor = min(storage_listened_basin, 10.0) / 10.0
 
         flow_rate = 0.0
 
-        if !isnan(proportional[i])
-            flow_rate += proportional[i] * error[i]
+        K_p = proportional[i]
+        if !isnan(K_p)
+            flow_rate += reduction_factor * K_p * error[i]
         end
 
-        if !isnan(derivative[i])
-            # dlevel/dstorage = 1/area
-            area = basin.current_area[listened_node_idx]
-
-            error_deriv = -dstorage[listened_node_idx] / area
-            flow_rate += derivative[i] * error_deriv
+        K_i = integral[i]
+        if !isnan(K_i)
+            flow_rate += reduction_factor * K_i * integral_value[i]
         end
 
-        if !isnan(integral[i])
-            # coefficient * current value of integral
-            flow_rate += integral[i] * integral_value[i]
+        K_d = derivative[i]
+        if !isnan(K_d)
+            # dlevel/dstorage = 1/area
+            area = basin.current_area[listened_node_idx]
+            dtarget_level = 0.0
+            du_listened_basin_old = du.storage[listened_node_idx]
+            # The expression below is the solution to an implicit equation for
+            # du_listened_basin. This equation results from the fact that if the derivative
+            # term in the PID controller is used, the controlled pump flow rate depends on itself.
+            du_listened_basin =
+                (
+                    du_listened_basin_old - flow_rate -
+                    reduction_factor * K_d * dtarget_level
+                ) / (1 - reduction_factor * K_d / area)
+            flow_rate = du_listened_basin_old - du_listened_basin
         end
 
         # Clip values outside pump flow rate bounds
-        flow_rate = max(flow_rate, min_flow_rate[i])
+        flow_rate = clamp(flow_rate, min_flow_rate[i], max_flow_rate[i])
 
-        if !isnan(max_flow_rate[i])
-            flow_rate = min(flow_rate, max_flow_rate[i])
+        pump.flow_rate[controlled_node_idx] = flow_rate
+        du.storage[listened_node_idx] -= flow_rate
+
+        # Set flow for connected edges
+        src_id = only(inneighbors(graph_flow, controlled_node_id))
+        dst_id = only(outneighbors(graph_flow, controlled_node_id))
+
+        flow[src_id, controlled_node_id] = flow_rate
+        flow[controlled_node_id, dst_id] = flow_rate
+
+        # Below du.storage is updated. This is normally only done
+        # in formulate!(du, connectivity, basin), but in this function
+        # flows are set so du has to be updated too.
+
+        has_index, dst_idx = id_index(basin.node_id, dst_id)
+        if has_index
+            du.storage[dst_idx] += flow_rate
         end
 
-        pump.flow_rate[controlled_node_idx] = flow_rate
+        for id in outneighbors(graph_flow, controlled_node_id)
+            if id in fractional_flow.node_id
+                after_ff_id = only(outneighbours(graph_flow, id))
+                ff_idx = findsorted(fractional_flow, id)
+                flow_rate_fraction = fractional_flow.fraction[ff_idx] * flow_rate
+                flow[id, after_ff_id] = flow_rate_fraction
+
+                has_index, basin_idx = id_index(basin.node_id, after_ff_id)
+
+                if has_index
+                    du.storage[basin_idx] += flow_rate_fraction
+                end
+            end
+        end
     end
     return nothing
 end
@@ -700,8 +741,9 @@ end
 function formulate!(pump::Pump, p::Parameters, storage::AbstractVector{Float64})::Nothing
     (; connectivity, basin, level_boundary) = p
     (; graph_flow, flow) = connectivity
-    (; node_id, active, flow_rate) = pump
-    for (id, isactive, rate) in zip(node_id, active, flow_rate)
+    (; node_id, active, flow_rate, is_pid_controlled) = pump
+    for (id, isactive, rate, pid_controlled) in
+        zip(node_id, active, flow_rate, is_pid_controlled)
         @assert rate >= 0 "Pump flow rate must be positive, found $rate for Pump #$id"
 
         src_id = only(inneighbors(graph_flow, id))
@@ -713,6 +755,10 @@ function formulate!(pump::Pump, p::Parameters, storage::AbstractVector{Float64})
             continue
         end
 
+        if pid_controlled
+            continue
+        end
+
         hasindex, basin_idx = id_index(basin.node_id, src_id)
 
         if hasindex
@@ -783,23 +829,21 @@ function water_balance!(
     storage = u.storage
     integral = u.integral
 
-    basin.dstorage .= du.storage
-
     du .= 0.0
     nonzeros(connectivity.flow) .= 0.0
 
     # ensures current_level is current
     formulate!(du, basin, storage, t)
 
-    # PID control (does not set flows)
-    continuous_control!(u, du, pid_control, p, integral)
-
     # First formulate intermediate flows
     formulate_flows!(p, storage)
 
     # Now formulate du
     formulate!(du, connectivity, basin)
 
+    # PID control (changes the du of PID controlled basins)
+    continuous_control!(u, du, pid_control, p, integral)
+
     # Negative storage musn't decrease, based on Shampine's et. al. advice
     # https://docs.sciml.ai/DiffEqCallbacks/stable/step_control/#DiffEqCallbacks.PositiveDomain
     for i in eachindex(u.storage)

core/src/utils.jl

@@ -531,30 +531,33 @@ function update_jac_prototype!(
     node::PidControl,
 )::Nothing
     (; basin, connectivity) = p
-    (; graph_control) = connectivity
+    (; graph_control, graph_flow) = connectivity
 
     n_basins = length(basin.node_id)
 
     for (pid_idx, (listen_node_id, id)) in enumerate(zip(node.listen_node_id, node.node_id))
-        id_out = only(outneighbors(graph_control, id))
-        id_out_in = only(inneighbors(graph_control, id_out))
+        id_pump = only(outneighbors(graph_control, id))
+        id_pump_out = only(inneighbors(graph_flow, id_pump))
 
         _, listen_idx = id_index(basin.node_id, listen_node_id)
 
+        # Controlled basin affects itself
+        jac_prototype[listen_idx, listen_idx] = 1.0
+
         # PID control integral state
         pid_state_idx = n_basins + pid_idx
-        jac_prototype[pid_state_idx, listen_idx] = 1.0
         jac_prototype[listen_idx, pid_state_idx] = 1.0
+        jac_prototype[pid_state_idx, listen_idx] = 1.0
 
-        # The basin upstream of the pump
-        has_index, idx_out_in = id_index(basin.node_id, id_out_in)
+        # The basin downstream of the pump
+        has_index, idx_out_out = id_index(basin.node_id, id_pump_out)
 
         if has_index
-            jac_prototype[pid_state_idx, idx_out_in] = 1.0
-            jac_prototype[idx_out_in, pid_state_idx] = 1.0
+            # The basin downstream of the pump PID control integral state
+            jac_prototype[pid_state_idx, idx_out_out] = 1.0
 
-            # The basin upstream of the pump also depends on the controlled basin
-            jac_prototype[listen_idx, idx_out_in] = 1.0
+            # The basin downstream of the pump also depends on the controlled basin
+            jac_prototype[listen_idx, idx_out_out] = 1.0
         end
     end
     return nothing

core/test/equations.jl

@@ -149,6 +149,51 @@ end
     @test all(isapprox.(LHS, RHS; rtol = 0.005)) # Fails with '≈'
 end
 
+# The second order linear inhomogeneous ODE for this model is derived by
+# differentiating the equation for the storage of the controlled basin
+# once to time to get rid of the integral term.
+@testset "PID control" begin
+    toml_path =
+        normpath(@__DIR__, "../../data/pid_control_equation/pid_control_equation.toml")
+    @test ispath(toml_path)
+    model = Ribasim.run(toml_path)
+    @test model.integrator.sol.retcode == Ribasim.ReturnCode.Success
+    p = model.integrator.p
+    (; basin, pid_control) = p
+
+    storage = Ribasim.get_storages_and_levels(model).storage[:]
+    t = Ribasim.timesteps(model)
+
+    K_p = pid_control.proportional[1]
+    K_i = pid_control.integral[1]
+    K_d = pid_control.derivative[1]
+    storage_min = 50
+    level_min = basin.level[1][2]
+    SP = basin.target_level[1]
+    storage0 = storage[1]
+    area = basin.area[1][2]
+    level0 = level_min + (storage0 - storage_min) / area
+
+    α = 1 - K_d / area
+    β = -K_p / area
+    γ = -K_i / area
+    δ = -K_i * (SP - level_min + storage_min / area)
+
+    λ_1 = (-β + sqrt(β^2 - 4 * α * γ)) / (2 * α)
+    λ_2 = (-β - sqrt(β^2 - 4 * α * γ)) / (2 * α)
+
+    c_1 = storage0 - δ / γ
+    c_2 = -K_p * (SP - level0) / (1 - K_d / area)
+
+    Δλ = λ_2 - λ_1
+    k_1 = (λ_2 * c_1 - c_2) / Δλ
+    k_2 = (-λ_1 * c_1 + c_2) / Δλ
+
+    storage_predicted = @. k_1 * exp(λ_1 * t) + k_2 * exp(λ_2 * t) + δ / γ
+
+    @test all(isapprox.(storage, storage_predicted; rtol = 0.001))
+end
+
 # Simple solutions:
 # storage1 = storage1(t0) + (t-t0)*(frac*q_boundary - q_pump)
 # storage2 = storage2(t0) + (t-t0)*q_pump

core/test/utils.jl

@@ -24,7 +24,6 @@ end
     level = [[0.0, 1.0], [4.0, 5.0]]
     storage = Ribasim.profile_storage.(level, area)
     target_level = [0.0, 0.0]
-    dstorage = target_level
     basin = Ribasim.Basin(
         Indices([5, 7]),
         [2.0, 3.0],
@@ -38,7 +37,6 @@ end
         storage,
         target_level,
         StructVector{Ribasim.BasinForcingV1}(undef, 0),
-        dstorage,
     )
 
     @test basin.level[2][1] === 4.0
@@ -128,7 +126,7 @@ end
 
     @test jac_prototype.m == 2
     @test jac_prototype.n == 2
-    @test jac_prototype.colptr == [1, 2, 3]
-    @test jac_prototype.rowval == [2, 1]
-    @test jac_prototype.nzval == ones(2)
+    @test jac_prototype.colptr == [1, 3, 4]
+    @test jac_prototype.rowval == [1, 2, 1]
+    @test jac_prototype.nzval == ones(3)
 end

core/test/validation.jl

@@ -60,6 +60,7 @@ end
         [0.0, 0.0],
         [1.0, 1.0],
         Dict{Tuple{Int, String}, NamedTuple}(),
+        falses(2),
     )
 
     errors = Ribasim.valid_n_neighbors(graph_flow, pump)