Fix many tests

core/src/graph.jl

@@ -166,28 +166,6 @@ function Base.iterate(iter::OutNeighbors, state = 1)
     return label_out, state
 end
 
-function set_flow!(graph::MetaGraph, edge_metadata::EdgeMetadata, q::Number, du)::Nothing
-    set_flow!(graph, edge_metadata.flow_idx, q, du)
-    return nothing
-end
-
-function set_flow!(graph, flow_idx::Int, q::Number, du)::Nothing
-    (; flow) = graph[]
-    flow[parent(du)][flow_idx] = q
-    return nothing
-end
-
-"""
-Get the flow over the given edge (du is needed for LazyBufferCache from ForwardDiff.jl).
-"""
-function get_flow(graph, edge_metadata::EdgeMetadata, du)::Number
-    return get_flow(graph, edge_metadata.flow_idx, du)
-end
-
-function get_flow(graph::MetaGraph, flow_idx::Int, du)
-    return graph[].flow[parent(du)][flow_idx]
-end
-
 """
 Get the inneighbor node IDs of the given node ID (label)
 over the given edge type in the graph.
@@ -271,3 +249,36 @@ function metadata_from_edge(graph::MetaGraph, edge::Edge{Int})::EdgeMetadata
     label_dst = label_for(graph, edge.dst)
     return graph[label_src, label_dst]
 end
+
+function get_flow(
+    flow::ComponentVector,
+    p::Parameters,
+    t::Number,
+    edge::Tuple{NodeID, NodeID};
+    boundary_flow = nothing,
+)
+    (; flow_boundary) = p
+    from_id, to_id = edge
+    from_node_type = snake_case(Symbol(from_id.type))
+    to_node_type = snake_case(Symbol(to_id.type))
+    if from_node_type in keys(flow)
+        getproperty(flow, from_node_type)[from_id.idx]
+    elseif to_node_type in keys(flow)
+        getproperty(flow, to_node_type)[to_id.idx]
+    elseif from_node_type == :flow_boundary
+        isnothing(boundary_flow) ? flow_boundary.flow_rate[from_id.idx](t) :
+        boundary_flow[from_id.idx]
+    elseif from_node_type == :user_demand
+        flow.user_demand_outflow[from_id.idx]
+    elseif to_node_type == :user_demand
+        flow.user_demand_inflow[to_id.idx]
+    else
+        error("$from_id, $to_id")
+    end
+end
+
+function get_influx(du::ComponentVector, id::NodeID)
+    @assert id.type == NodeType.Basin
+    return du.precipitation[id.idx] + drainage[id.idx] - du.evaporation[id.idx] -
+           du.infiltration[id.idx]
+end

core/src/parameter.jl

@@ -579,15 +579,20 @@ node_id: node ID of the Terminal node
 end
 
 """
-A variant on `Base.Ref` where the source array is a vector that is possibly wrapped in a ForwardDiff.DiffCache.
+A variant on `Base.Ref` where the source array is a vector that is possibly wrapped in a ForwardDiff.LazyBufferCache.
 Retrieve value with get_value(ref::PreallocationRef, val) where `val` determines the return type.
 """
 struct PreallocationRef
     vector::Cache
     idx::Int
+    from_du::Bool
+    function PreallocationRef(vector::Cache, idx::Int; from_du = false)
+        new(vector, idx, from_du)
+    end
 end
 
-get_value(ref::PreallocationRef, du) = ref.vector[parent(du)][ref.idx]
+get_value(ref::PreallocationRef, du) =
+    ref.from_du ? du[ref.idx] : ref.vector[parent(du)][ref.idx]
 
 function set_value!(ref::PreallocationRef, value, du)::Nothing
     ref.vector[parent(du)][ref.idx] = value

core/src/solve.jl

@@ -38,9 +38,6 @@ function water_balance!(
     # Formulate intermediate flow (controlled by PID control)
     formulate_flows!(du, p, t; continuous_control_type = ContinuousControlType.PID)
 
-    # Formulate du (controlled by PidControl)
-    formulate_du_pid_controlled!(du, graph, pid_control)
-
     return nothing
 end
 
@@ -265,7 +262,7 @@ function formulate_pid_control!(
 
         if !iszero(K_d)
             dlevel_demand = derivative(target[i], t)
-            du_listened_basin_old = du.storage[listened_node_id.idx]
+            du_listened_basin_old = formulate_dstorage(du, p, t, listened_node_id)#du.storage[listened_node_id.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.
@@ -278,6 +275,26 @@ function formulate_pid_control!(
     return nothing
 end
 
+function formulate_dstorage(du::ComponentVector, p::Parameters, t::Number, node_id::NodeID)
+    (; basin) = p
+    (; inflow_ids, outflow_ids) = basin
+    @assert node_id.type == NodeType.Basin
+    dstorage = 0.0
+    for inflow_id in inflow_ids[node_id.idx]
+        dstorage += get_flow(du, p, t, (inflow_id, node_id))
+    end
+    for outflow_id in outflow_ids[node_id.idx]
+        dstorage -= get_flow(du, p, t, (node_id, outflow_id))
+    end
+
+    dstorage += du.precipitation[node_id.idx]
+    dstorage += du.drainage[node_id.idx]
+    dstorage -= du.evaporation[node_id.idx]
+    dstorage -= du.infiltration[node_id.idx]
+
+    dstorage
+end
+
 function formulate_flow!(
     user_demand::UserDemand,
     du::AbstractVector,
@@ -294,6 +311,7 @@ function formulate_flow!(
         demand_itp,
         demand,
         allocated,
+        return_factor,
         min_level,
         demand_from_timeseries,
     ) in zip(
@@ -305,6 +323,7 @@ function formulate_flow!(
         # TODO permute these so the nodes are the last dimension, for performance
         eachrow(user_demand.demand),
         eachrow(user_demand.allocated),
+        user_demand.return_factor,
         user_demand.min_level,
         user_demand.demand_from_timeseries,
     )
@@ -341,7 +360,8 @@ function formulate_flow!(
         Δsource_level = source_level - min_level
         factor_level = reduction_factor(Δsource_level, 0.1)
         q *= factor_level
-        du.user_demand[id.idx] = q
+        du.user_demand_inflow[id.idx] = q
+        du.user_demand_outflow[id.idx] = q * return_factor(t)
     end
     return nothing
 end
@@ -641,23 +661,6 @@ function formulate_flow!(
     return nothing
 end
 
-function formulate_du_pid_controlled!(
-    du::ComponentVector,
-    graph::MetaGraph,
-    pid_control::PidControl,
-)::Nothing
-    for id in pid_control.controlled_basins
-        du[id.idx] = zero(eltype(du))
-        for id_in in inflow_ids(graph, id)
-            du[id.idx] += get_flow(graph, id_in, id, du)
-        end
-        for id_out in outflow_ids(graph, id)
-            du[id.idx] -= get_flow(graph, id, id_out, du)
-        end
-    end
-    return nothing
-end
-
 function formulate_flows!(
     du::AbstractVector,
     p::Parameters,

core/src/util.jl

@@ -594,9 +594,9 @@ function set_initial_allocation_mean_flows!(integrator)::Nothing
 
     for edge in keys(mean_input_flows)
         if edge[1] == edge[2]
-            q = get_influx(p.basin, edge[1])
+            q = get_influx(du, edge[1])
         else
-            q = get_flow(graph, edge..., du)
+            q = get_flow(du, p, t, edge)
         end
         # Multiply by Δt_allocation as averaging divides by this factor
         # in update_allocation!
@@ -609,7 +609,7 @@ function set_initial_allocation_mean_flows!(integrator)::Nothing
         if edge[1] == edge[2]
             mean_realized_flows[edge] = -u[edge[1].idx]
         else
-            q = get_flow(graph, edge..., du)
+            q = get_flow(du, p, t, edge)
             mean_realized_flows[edge] = q * Δt_allocation
         end
     end
@@ -643,6 +643,9 @@ function get_variable_ref(
     (; basin, graph) = p
     errors = false
 
+    # Only built here because it is needed to obtain indices
+    u = build_state_vector(p)
+
     ref = if node_id.type == NodeType.Basin && variable == "level"
         PreallocationRef(basin.current_level, node_id.idx)
     elseif variable == "flow_rate" && node_id.type != NodeType.FlowBoundary
@@ -652,17 +655,18 @@ function get_variable_ref(
                 @error "Cannot listen to flow_rate of $node_id, the node type must be one of $conservative_node_types"
                 Ref(Float64[], 0)
             else
-                id_in = inflow_id(graph, node_id)
-                (; flow_idx) = graph[id_in, node_id]
-                PreallocationRef(graph[].flow, flow_idx)
+                # Index in the state vector
+                flow_idx =
+                    only(getfield(u, :axes))[snake_case(Symbol(node_id.type))].idx[node_id.idx]
+                PreallocationRef(cache(1), flow_idx; from_du = true)
             end
         else
             node = getfield(p, snake_case(Symbol(node_id.type)))
             PreallocationRef(node.flow_rate, node_id.idx)
         end
     else
         # Placeholder to obtain correct type
-        PreallocationRef(graph[].flow, 0)
+        PreallocationRef(cache(1), 0)
     end
     return ref, errors
 end

core/src/write.jl

@@ -225,7 +225,7 @@ function flow_table(
 }
     (; config, saved, integrator) = model
     (; t, saveval) = saved.flow
-    (; p, u) = integrator
+    (; p) = integrator
     (; graph) = p
     (; flow_edges) = graph[]
 
@@ -247,26 +247,10 @@ function flow_table(
     flow_rate = zeros(nflow * ntsteps)
 
     for (i, flow_edge) in enumerate(flow_edges)
-        from_id, to_id = flow_edge.edge
-        from_node_type = snake_case(Symbol(from_id.type))
-        to_node_type = snake_case(Symbol(to_id.type))
-        component_name, index = if from_node_type in keys(u)
-            from_node_type, from_id.idx
-        elseif to_node_type in keys(u)
-            to_node_type, to_id.idx
-        elseif from_node_type == :flow_boundary
-            :flow_boundary, from_id.idx
-        else
-            error("$from_id, $to_id")
-        end
-
         for (j, cvec) in enumerate(saveval)
             (; flow, flow_boundary) = cvec
-
-            component =
-                component_name == (:flow_boundary) ? flow_boundary :
-                getproperty(flow, component_name)
-            flow_rate[i + (j - 1) * nflow] = component[index]
+            flow_rate[i + (j - 1) * nflow] =
+                get_flow(flow, p, 0.0, flow_edge.edge; boundary_flow = flow_boundary)
         end
     end
 

core/test/control_test.jl

@@ -1,5 +1,6 @@
 @testitem "Pump discrete control" begin
     using Ribasim: NodeID
+    using OrdinaryDiffEqCore: get_du
     using Dates: DateTime
     import Arrow
     import Tables
@@ -56,7 +57,7 @@
     t_2_index = findfirst(>=(t_2), t)
     @test level[2, t_2_index] >= discrete_control.compound_variables[1][2].greater_than[1]
 
-    flow = graph[].flow[Float64[]]
+    flow = get_du(model.integrator)[(:linear_resistance, :pump)]
     @test all(iszero, flow)
 end
 

core/test/run_models_test.jl

@@ -123,17 +123,18 @@ end
 
 @testitem "bucket model" begin
     using SciMLBase: successful_retcode
+    using OrdinaryDiffEqCore: get_du
 
     toml_path = normpath(@__DIR__, "../../generated_testmodels/bucket/ribasim.toml")
     @test ispath(toml_path)
     model = Ribasim.run(toml_path)
     @test model isa Ribasim.Model
-    @test model.integrator.u.storage ≈ [1000]
-    vertical_flux = Ribasim.wrap_forcing(model.integrator.p.basin.vertical_flux[Float64[]])
-    @test vertical_flux.precipitation == [0.0]
-    @test vertical_flux.evaporation == [0.0]
-    @test vertical_flux.drainage == [0.0]
-    @test vertical_flux.infiltration == [0.0]
+    @test model.integrator.p.basin.current_storage[Float64[]] ≈ [1000]
+    du = get_du(model.integrator)
+    @test du.precipitation == [0.0]
+    @test du.evaporation == [0.0]
+    @test du.drainage == [0.0]
+    @test du.infiltration == [0.0]
     @test successful_retcode(model)
 end
 

core/test/utils_test.jl

@@ -246,27 +246,12 @@ end
 
     node_id_1 = NodeID(:Basin, 5, 1)
     node_id_2 = NodeID(:Basin, 6, 2)
-    @test low_storage_factor_resistance_node(
-        (; storage = [3.0, 3.0]),
-        1.0,
-        node_id_1,
-        node_id_2,
-        2.0,
-    ) == 1.0
-    @test low_storage_factor_resistance_node(
-        (; storage = [1.0, 3.0]),
-        1.0,
-        node_id_1,
-        node_id_2,
-        2.0,
-    ) == 0.5
-    @test low_storage_factor_resistance_node(
-        (; storage = [1.0, 3.0]),
-        -1.0,
-        node_id_1,
-        node_id_2,
-        2.0,
-    ) == 1.0
+    @test low_storage_factor_resistance_node([3.0, 3.0], 1.0, node_id_1, node_id_2, 2.0) ==
+          1.0
+    @test low_storage_factor_resistance_node([1.0, 3.0], 1.0, node_id_1, node_id_2, 2.0) ==
+          0.5
+    @test low_storage_factor_resistance_node([1.0, 3.0], -1.0, node_id_1, node_id_2, 2.0) ==
+          1.0
 end
 
 @testitem "constraints_from_nodes" begin