PID control fix

core/src/solve.jl

@@ -222,29 +222,26 @@ function continuous_control!(
             max_flow_rate[controlled_node_idx],
         )
 
-        # 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.
-        if controls_pump
-            pump_flow_rate[controlled_node_idx] = flow_rate
-            du.storage[listened_node_idx] -= flow_rate
-        else
-            outlet_flow_rate[controlled_node_idx] = flow_rate
-            du.storage[listened_node_idx] += flow_rate
-        end
-
         # Set flow for connected edges
         src_id = inflow_id(graph, controlled_node_id)
         dst_id = outflow_id(graph, controlled_node_id)
 
         set_flow!(graph, src_id, controlled_node_id, flow_rate)
         set_flow!(graph, 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
 
+        has_index, src_idx = id_index(basin.node_id, src_id)
+        if has_index
+            du.storage[src_idx] -= flow_rate
+        end
+
         # When the controlled pump flows out into fractional flow nodes
         if controls_pump
             for id in outflow_ids(graph, controlled_node_id)

python/ribasim_testmodels/ribasim_testmodels/doc_example.py

@@ -17,7 +17,7 @@ def doc_example_model():
             "level": 6 * [0.0, 3.0],
         }
     )
-    basin_level = [1.0, 1.5, 1.0, 0.5, 0.3, 0.1]
+    basin_level = [5.0, 1.5, 1.0, 0.5, 0.3, 0.1]
 
     state = pd.DataFrame(data={"node_id": [1, 4, 6, 8, 10, 12], "level": basin_level})