Pass more tests

core/src/allocation_optim.jl

@@ -306,12 +306,12 @@ function get_basin_data(
     u::ComponentVector,
     node_id::NodeID,
 )
-    (; graph, allocation) = p
+    (; graph, allocation, basin) = p
     (; Δt_allocation) = allocation_model
     (; mean_input_flows) = allocation
     @assert node_id.type == NodeType.Basin
     influx = mean_input_flows[(node_id, node_id)][]
-    storage_basin = u.storage[node_id.idx]
+    storage_basin = basin.current_storage[parent(u)][node_id.idx]
     control_inneighbors = inneighbor_labels_type(graph, node_id, EdgeType.control)
     if isempty(control_inneighbors)
         level_demand_idx = 0

core/src/callback.jl

@@ -386,8 +386,8 @@ end
 
 "Solve the allocation problem for all demands and assign allocated abstractions."
 function update_allocation!(integrator)::Nothing
-    (; p, t, u) = integrator
-    (; allocation) = p
+    (; p, t, u, sol) = integrator
+    (; allocation, flow_boundary) = p
     (; allocation_models, mean_input_flows, mean_realized_flows) = allocation
 
     # Don't run the allocation algorithm if allocation is not active
@@ -397,11 +397,21 @@ function update_allocation!(integrator)::Nothing
     end
 
     (; Δt_allocation) = allocation_models[1]
-
-    # Divide by the allocation Δt to obtain the mean input flows
-    # from the integrated flows
-    for key in keys(mean_input_flows)
-        mean_input_flows[key] /= Δt_allocation
+    if t > 0
+        for edge in keys(mean_input_flows)
+            mean_flow = if edge[1] == edge[2]
+                (get_influx(sol(t), edge[1]) - get_influx(sol(t - Δt_allocation), edge[1])) / Δt_allocation
+            elseif edge[1].type == NodeType.FlowBoundary
+                # TODO: This is not correct if the flow boundary has been inactive
+                integral(flow_boundary.flow_rate[edge[1].idx], t - Δt_allocation, t) /
+                Δt_allocation
+            else
+                flow_idx = flow_index(u, edge)
+                (sol(t; idxs = flow_idx) - sol(t - Δt_allocation; idxs = flow_idx)) /
+                Δt_allocation
+            end
+            mean_input_flows[edge] = mean_flow
+        end
     end
 
     # Divide by the allocation Δt to obtain the mean realized flows

core/src/graph.jl

@@ -240,16 +240,6 @@ function inflow_id(graph::MetaGraph, id::NodeID)::NodeID
     return only(inflow_ids(graph, id))
 end
 
-"""
-Get the metadata of an edge in the graph from an edge of the underlying
-DiGraph.
-"""
-function metadata_from_edge(graph::MetaGraph, edge::Edge{Int})::EdgeMetadata
-    label_src = label_for(graph, edge.src)
-    label_dst = label_for(graph, edge.dst)
-    return graph[label_src, label_dst]
-end
-
 function get_flow(
     flow::ComponentVector,
     p::Parameters,
@@ -279,6 +269,6 @@ 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] -
+    return du.precipitation[id.idx] + du.drainage[id.idx] - du.evaporation[id.idx] -
            du.infiltration[id.idx]
 end

core/src/model.jl

@@ -212,7 +212,6 @@ function SciMLBase.step!(model::Model, dt::Float64)::Model
     if round(ntimes) ≈ ntimes
         update_allocation!(integrator)
     end
-    set_previous_flows!(integrator)
     step!(integrator, dt, true)
     return model
 end

core/src/util.jl

@@ -656,8 +656,7 @@ function get_variable_ref(
                 Ref(Float64[], 0)
             else
                 # Index in the state vector
-                flow_idx =
-                    only(getfield(u, :axes))[snake_case(Symbol(node_id.type))].idx[node_id.idx]
+                flow_idx = flow_index(u, node_id)
                 PreallocationRef(cache(1), flow_idx; from_du = true)
             end
         else
@@ -671,6 +670,23 @@ function get_variable_ref(
     return ref, errors
 end
 
+function flow_index(
+    u::ComponentVector{Float64, Vector{Float64}, <:Tuple{<:Axis{NT}}},
+    node_id::NodeID,
+)::Union{Int, Nothing} where {NT}
+    node_type = snake_case(Symbol(node_id.type))
+    if node_type in keys(NT)
+        only(getfield(u, :axes))[node_type].idx[node_id.idx]
+    else
+        nothing
+    end
+end
+
+function flow_index(u::ComponentVector, edge::Tuple{NodeID, NodeID})::Int
+    idx = flow_index(u, edge[1])
+    isnothing(idx) ? flow_index(u, edge[2]) : idx
+end
+
 """
 Set references to all variables that are listened to by discrete/continuous control
 """
@@ -833,23 +849,6 @@ function basin_areas(basin::Basin, state_idx::Int)
     return basin.level_to_area[state_idx].u
 end
 
-"""
-Just before setting a timestep, call water_balance! again
-to get a correct value of the flows for integrating
-"""
-function set_previous_flows!(integrator)
-    (; p, u, t) = integrator
-    (; flow, flow_prev) = p.graph[]
-    (; vertical_flux, vertical_flux_prev) = p.basin
-    du = get_du(integrator)
-    water_balance!(du, u, p, t)
-
-    flow = flow[parent(u)]
-    vertical_flux = vertical_flux[parent(u)]
-    copyto!(flow_prev, flow)
-    copyto!(vertical_flux_prev, vertical_flux)
-end
-
 """
 Split the single forcing vector into views of the components
 precipitation, evaporation, drainage, infiltration

core/test/allocation_test.jl

@@ -246,6 +246,7 @@ end
     import JuMP
     using Ribasim: NodeID
     using DataFrames: DataFrame
+    using OrdinaryDiffEqCore: get_du
     using DataInterpolations: LinearInterpolation, integral
 
     toml_path = normpath(@__DIR__, "../../generated_testmodels/level_demand/ribasim.toml")
@@ -262,14 +263,15 @@ end
 
     storage = Ribasim.get_storages_and_levels(model).storage[1, :]
     t = Ribasim.tsaves(model)
+    du = get_du(model.integrator)
 
     d = user_demand.demand_itp[1][2](0)
     ϕ = 1e-3 # precipitation
     q = Ribasim.get_flow(
-        graph,
-        Ribasim.NodeID(:FlowBoundary, 1, p),
-        Ribasim.NodeID(:Basin, 2, p),
-        Float64[],
+        du,
+        p,
+        0.0,
+        (Ribasim.NodeID(:FlowBoundary, 1, p), Ribasim.NodeID(:Basin, 2, p)),
     )
     A = Ribasim.basin_areas(basin, 1)[1]
     l_max = level_demand.max_level[1](0)

core/test/run_models_test.jl

@@ -140,19 +140,23 @@ end
 
 @testitem "leaky bucket model" begin
     using SciMLBase: successful_retcode
+    using OrdinaryDiffEqCore: get_du
     import BasicModelInterface as BMI
 
     toml_path = normpath(@__DIR__, "../../generated_testmodels/leaky_bucket/ribasim.toml")
     @test ispath(toml_path)
     model = Ribasim.Model(toml_path)
     @test model isa Ribasim.Model
 
-    stor = model.integrator.u.storage
-    vertical_flux = Ribasim.wrap_forcing(model.integrator.p.basin.vertical_flux[Float64[]])
-    prec = vertical_flux.precipitation
-    evap = vertical_flux.evaporation
-    drng = vertical_flux.drainage
-    infl = vertical_flux.infiltration
+    (; integrator) = model
+    du = get_du(integrator)
+    (; u, p, t) = integrator
+    Ribasim.water_balance!(du, u, p, t)
+    stor = integrator.p.basin.current_storage[parent(du)]
+    prec = du.precipitation
+    evap = du.evaporation
+    drng = du.drainage
+    infl = du.infiltration
     # The dynamic data has missings, but these are not set.
     @test prec == [0.0]
     @test evap == [0.0]
@@ -199,8 +203,8 @@ end
 
     @test successful_retcode(model)
     @test allunique(Ribasim.tsaves(model))
-    @test model.integrator.u ≈ Float32[803.7093, 803.68274, 495.241, 1318.3053] skip =
-        Sys.isapple() atol = 1.5
+    @test model.integrator.p.basin.current_storage[Float64[]] ≈
+          Float32[803.7093, 803.68274, 495.241, 1318.3053] skip = Sys.isapple() atol = 1.5
 
     @test length(logger.logs) > 10
     @test logger.logs[1].level == Debug
@@ -224,6 +228,7 @@ end
 
 @testitem "basic transient model" begin
     using SciMLBase: successful_retcode
+    using OrdinaryDiffEqCore: get_du
 
     toml_path =
         normpath(@__DIR__, "../../generated_testmodels/basic_transient/ribasim.toml")
@@ -232,13 +237,11 @@ end
     @test model isa Ribasim.Model
     @test successful_retcode(model)
     @test allunique(Ribasim.tsaves(model))
-    precipitation =
-        Ribasim.wrap_forcing(
-            model.integrator.p.basin.vertical_flux[Float64[]],
-        ).precipitation
+    du = get_du(model.integrator)
+    precipitation = du.precipitation
     @test length(precipitation) == 4
-    @test model.integrator.u ≈ Float32[698.22736, 698.2014, 421.20447, 1334.4354] atol = 2.0 skip =
-        Sys.isapple()
+    @test model.integrator.p.basin.current_storage[parent(du)] ≈
+          Float32[698.22736, 698.2014, 421.20447, 1334.4354] atol = 2.0 skip = Sys.isapple()
 end
 
 @testitem "Allocation example model" begin
@@ -291,7 +294,8 @@ end
     model = Ribasim.run(toml_path)
     @test model isa Ribasim.Model
     @test successful_retcode(model)
-    @test model.integrator.u ≈ Float32[368.31558, 365.68442] skip = Sys.isapple()
+    @test model.integrator.p.basin.current_storage[Float64[]] ≈
+          Float32[368.31558, 365.68442] skip = Sys.isapple()
     # the highest level in the dynamic table is updated to 1.2 from the callback
     @test model.integrator.p.tabulated_rating_curve.table[end].t[end] == 1.2
 end
@@ -391,18 +395,25 @@ end
     using SciMLBase: successful_retcode
     using Dates
     using DataFrames: DataFrame
+    using Ribasim: formulate_storages!
 
     toml_path = normpath(@__DIR__, "../../generated_testmodels/user_demand/ribasim.toml")
     @test ispath(toml_path)
     model = Ribasim.run(toml_path)
     @test successful_retcode(model)
 
-    seconds_in_day = 86400.0
-    @test only(model.integrator.sol(0seconds_in_day)) == 1000.0
+    (; u, p, t, sol) = model.integrator
+    current_storage = p.basin.current_storage[Float64[]]
+
+    day = 86400.0
+    formulate_storages!(current_storage, sol(0day), p, 0day)
+    @test only(current_storage) == 1000.0
     # constant UserDemand withdraws to 0.9m or 900m3 due to min level = 0.9
-    @test only(model.integrator.sol(150seconds_in_day)) ≈ 900 atol = 5
+    formulate_storages!(current_storage, sol(150day), p, 150day)
+    @test only(current_storage) ≈ 900 atol = 5
     # dynamic UserDemand withdraws to 0.5m or 500m3 due to min level = 0.5
-    @test only(model.integrator.sol(220seconds_in_day)) ≈ 500 atol = 1
+    formulate_storages!(current_storage, sol(220day), p, 220day)
+    @test only(current_storage) ≈ 500 atol = 1
 
     # Trasient return factor
     flow = DataFrame(Ribasim.flow_table(model))
@@ -422,6 +433,7 @@ end
 @testitem "ManningResistance" begin
     using PreallocationTools: get_tmp
     using SciMLBase: successful_retcode
+    using OrdinaryDiffEqCore: get_du
     using Ribasim: NodeID
 
     """
@@ -483,9 +495,9 @@ end
     model = Ribasim.run(toml_path)
     @test successful_retcode(model)
 
-    u = model.integrator.sol.u[end]
-    p = model.integrator.p
-    h_actual = p.basin.current_level[parent(u)][1:50]
+    du = get_du(model.integrator)
+    (; p, t) = model.integrator
+    h_actual = p.basin.current_level[parent(du)][1:50]
     x = collect(10.0:20.0:990.0)
     h_expected = standard_step_method(x, 5.0, 1.0, 0.04, h_actual[end], 1.0e-6)
 
@@ -495,18 +507,13 @@ end
     # https://www.hec.usace.army.mil/confluence/rasdocs/ras1dtechref/latest/theoretical-basis-for-one-dimensional-and-two-dimensional-hydrodynamic-calculations/1d-steady-flow-water-surface-profiles/friction-loss-evaluation
     @test all(isapprox.(h_expected, h_actual; atol = 0.02))
     # Test for conservation of mass, flow at the beginning == flow at the end
-    n_self_loops = length(u)
-    @test Ribasim.get_flow(
-        p.graph,
-        NodeID(:FlowBoundary, 1, p),
-        NodeID(:Basin, 2, p),
-        parent(u),
-    ) ≈ 5.0 atol = 0.001 skip = Sys.isapple()
+    @test Ribasim.get_flow(du, p, t, (NodeID(:FlowBoundary, 1, p), NodeID(:Basin, 2, p))) ≈
+          5.0 atol = 0.001 skip = Sys.isapple()
     @test Ribasim.get_flow(
-        p.graph,
-        NodeID(:ManningResistance, 101, p),
-        NodeID(:Basin, 102, p),
-        parent(u),
+        du,
+        p,
+        t,
+        (NodeID(:ManningResistance, 101, p), NodeID(:Basin, 102, p)),
     ) ≈ 5.0 atol = 0.001 skip = Sys.isapple()
 end
 

core/test/utils_test.jl

@@ -171,10 +171,7 @@ end
 
     p = Ribasim.Parameters(db, cfg)
     t0 = 0.0
-    u0 = ComponentVector{Float64}(;
-        storage = zeros(length(p.basin.node_id)),
-        integral = Float64[],
-    )
+    u0 = Ribasim.build_state_vector(p)
     du0 = copy(u0)
     jac_prototype = Ribasim.get_jac_prototype(du0, u0, p, t0)