From 932ecaf84fde2db9a4935c7ee462a79c630ac5a9 Mon Sep 17 00:00:00 2001 From: Bart de Koning Date: Thu, 12 Sep 2024 21:54:06 +0200 Subject: [PATCH] Pass more tests --- core/src/allocation_optim.jl | 4 +- core/src/callback.jl | 24 ++++++++---- core/src/graph.jl | 12 +----- core/src/model.jl | 1 - core/src/util.jl | 37 +++++++++--------- core/test/allocation_test.jl | 10 +++-- core/test/run_models_test.jl | 73 ++++++++++++++++++++---------------- core/test/utils_test.jl | 5 +-- 8 files changed, 85 insertions(+), 81 deletions(-) diff --git a/core/src/allocation_optim.jl b/core/src/allocation_optim.jl index 915a0139e..dc9613688 100644 --- a/core/src/allocation_optim.jl +++ b/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 diff --git a/core/src/callback.jl b/core/src/callback.jl index 0866e1980..bcd7aeb6f 100644 --- a/core/src/callback.jl +++ b/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 diff --git a/core/src/graph.jl b/core/src/graph.jl index e6ef710e4..5d64b6af7 100644 --- a/core/src/graph.jl +++ b/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 diff --git a/core/src/model.jl b/core/src/model.jl index cb0793dcb..f5ad346c8 100644 --- a/core/src/model.jl +++ b/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 diff --git a/core/src/util.jl b/core/src/util.jl index 20ff623ad..108e14a8b 100644 --- a/core/src/util.jl +++ b/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 diff --git a/core/test/allocation_test.jl b/core/test/allocation_test.jl index bbda92964..03c1ea004 100644 --- a/core/test/allocation_test.jl +++ b/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) diff --git a/core/test/run_models_test.jl b/core/test/run_models_test.jl index 1a92821ea..665fb3f46 100644 --- a/core/test/run_models_test.jl +++ b/core/test/run_models_test.jl @@ -140,6 +140,7 @@ 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") @@ -147,12 +148,15 @@ end 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 diff --git a/core/test/utils_test.jl b/core/test/utils_test.jl index f4ce58fd4..96e33eb15 100644 --- a/core/test/utils_test.jl +++ b/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)