diff --git a/core/src/bmi.jl b/core/src/bmi.jl index d3a32698e..57d0ddf83 100644 --- a/core/src/bmi.jl +++ b/core/src/bmi.jl @@ -144,7 +144,8 @@ function get_value(p::Parameters, node_id::Int, variable::String, u) # NOTE: Getting the level with get_level does NOT work since water_balance! # is not called during rootfinding for callback hasindex, basin_idx = id_index(basin.node_id, node_id) - value = basin.level[basin_idx](u[basin_idx]) + _, level = get_area_and_level(basin, basin_idx, u[basin_idx]) + value = level else throw(ValueError("Unsupported condition variable $variable.")) end diff --git a/core/src/create.jl b/core/src/create.jl index 5bb16ffe9..8a186f1e5 100644 --- a/core/src/create.jl +++ b/core/src/create.jl @@ -110,7 +110,7 @@ function Basin(db::DB, config::Config)::Basin infiltration = fill(NaN, length(node_id)) table = (; precipitation, potential_evaporation, drainage, infiltration) - area, level = create_storage_tables(db, config) + area, level, storage = create_storage_tables(db, config) # both static and forcing are optional, but we need fallback defaults static = load_structvector(db, config, BasinStaticV1) @@ -129,6 +129,7 @@ function Basin(db::DB, config::Config)::Basin current_level, area, level, + storage, time, ) end diff --git a/core/src/io.jl b/core/src/io.jl index 7b141f76a..e8ab71b90 100644 --- a/core/src/io.jl +++ b/core/src/io.jl @@ -131,7 +131,8 @@ function write_basin_output(model::Model) storage = reshape(vec(sol), nbasin, ntsteps) level = zero(storage) for (i, basin_storage) in enumerate(eachrow(storage)) - level[i, :] = p.basin.level[i].(basin_storage) + level[i, :] = + [get_area_and_level(p.basin, i, storage)[2] for storage in basin_storage] end basin = (; time, node_id, storage = vec(storage), level = vec(level)) diff --git a/core/src/solve.jl b/core/src/solve.jl index 7105dccb7..2191d2352 100644 --- a/core/src/solve.jl +++ b/core/src/solve.jl @@ -69,9 +69,10 @@ struct Basin{C} <: AbstractParameterNode infiltration::Vector{Float64} # cache this to avoid recomputation current_level::Vector{Float64} - # f(storage) - area::Vector{Interpolation} - level::Vector{Interpolation} + # Discrete values for interpolation + area::Vector{Vector{Float64}} + level::Vector{Vector{Float64}} + storage::Vector{Vector{Float64}} # data source for parameter updates time::StructVector{BasinForcingV1, C, Int} end @@ -243,11 +244,10 @@ Currently at less than 0.1 m. function formulate!(du::AbstractVector, basin::Basin, u::AbstractVector, t::Real)::Nothing for i in eachindex(du) storage = u[i] - area = basin.area[i](storage) - level = basin.level[i](storage) + area, level = get_area_and_level(basin, i, storage) basin.current_level[i] = level - bottom = basin_bottom_index(basin, i) - fixed_area = median(basin.area[i].u) + bottom = basin.level[i][1] + fixed_area = median(basin.area[i]) depth = max(level - bottom, 0.0) reduction_factor = min(depth, 0.1) / 0.1 diff --git a/core/src/utils.jl b/core/src/utils.jl index 255311c33..8fa03fac4 100644 --- a/core/src/utils.jl +++ b/core/src/utils.jl @@ -34,23 +34,104 @@ function profile_storage(levels::Vector{Float64}, areas::Vector{Float64})::Vecto return storages end -"Read the Basin / profile table and return all A(S) and h(S) functions" +"Read the Basin / profile table and return all area and level and computed storage values" function create_storage_tables( db::DB, config::Config, -)::Tuple{Vector{Interpolation}, Vector{Interpolation}} +)::Tuple{Vector{Vector{Float64}}, Vector{Vector{Float64}}, Vector{Vector{Float64}}} profiles = load_structvector(db, config, BasinProfileV1) - area = Interpolation[] - level = Interpolation[] + area = Vector{Vector{Float64}}() + level = Vector{Vector{Float64}}() + storage = Vector{Vector{Float64}}() + for group in IterTools.groupby(row -> row.node_id, profiles) group_area = getproperty.(group, :area) group_level = getproperty.(group, :level) group_storage = profile_storage(group_level, group_area) - area_itp = LinearInterpolation(group_area, group_storage) - level_itp = LinearInterpolation(group_level, group_storage) - push!(area, area_itp) - push!(level, level_itp) + push!(area, group_area) + push!(level, group_level) + push!(storage, group_storage) end + return area, level, storage +end + +function get_area_and_level( + basin::Basin, + state_idx::Int, + storage::Float64, +)::Tuple{Float64, Float64} + storage_discrete = basin.storage[state_idx] + area_discrete = basin.area[state_idx] + level_discrete = basin.level[state_idx] + + return get_area_and_level(storage_discrete, area_discrete, level_discrete, storage) +end + +function get_area_and_level( + storage_discrete::Vector{Float64}, + area_discrete::Vector{Float64}, + level_discrete::Vector{Float64}, + storage::Float64, +)::Tuple{Float64, Float64} + # storage_idx: smallest index such that storage_discrete[storage_idx] >= storage + storage_idx = searchsortedfirst(storage_discrete, storage) + + if storage_idx == 1 + # This can only happen if the storage is 0 + level = level_discrete[1] + area = area_discrete[1] + + elseif storage_idx == length(storage_discrete) + 1 + # With a storage above the profile, use a linear extrapolation of area(level) + # based on the last 2 values. + area_lower = area_discrete[end - 1] + area_higher = area_discrete[end] + level_lower = level_discrete[end - 1] + level_higher = level_discrete[end] + storage_lower = storage_discrete[end - 1] + storage_higher = storage_discrete[end] + + area_diff = area_higher - area_lower + level_diff = level_higher - level_lower + + if area_diff ≈ 0 + # Constant area means linear interpolation of level + area = area_lower + level = + level_higher + + level_diff * (storage - storage_higher) / (storage_higher - storage_lower) + else + area = sqrt( + area_higher^2 + 2 * (storage - storage_higher) * area_diff / level_diff, + ) + level = level_lower + level_diff * (area - area_lower) / area_diff + end + + else + area_lower = area_discrete[storage_idx - 1] + area_higher = area_discrete[storage_idx] + level_lower = level_discrete[storage_idx - 1] + level_higher = level_discrete[storage_idx] + storage_lower = storage_discrete[storage_idx - 1] + storage_higher = storage_discrete[storage_idx] + + area_diff = area_higher - area_lower + level_diff = level_higher - level_lower + + if area_diff ≈ 0 + # Constant area means linear interpolation of level + area = area_lower + level = + level_lower + + level_diff * (storage - storage_lower) / (storage_higher - storage_lower) + + else + area = + sqrt(area_lower^2 + 2 * (storage - storage_lower) * area_diff / level_diff) + level = level_lower + level_diff * (area - area_lower) / area_diff + end + end + return area, level end @@ -206,23 +287,15 @@ function id_index(ids::Indices{Int}, id::Int) return hasindex, idx end -"Return the bottom elevation of the basin with index i" -function basin_bottom_index(basin::Basin, i::Int)::Float64 - # get level(storage) interpolation function - itp = basin.level[i] - # and return the first level in the underlying table, which represents the bottom - return first(itp.u) -end - "Return the bottom elevation of the basin with index i, or nothing if it doesn't exist" function basin_bottom(basin::Basin, node_id::Int)::Union{Float64, Nothing} basin = Dictionary(basin.node_id, basin.level) hasindex, token = gettoken(basin, node_id) return if hasindex # get level(storage) interpolation function - itp = gettokenvalue(basin, token) - # and return the first level in the underlying table, which represents the bottom - first(itp.u) + level_discrete = gettokenvalue(basin, token) + # and return the first level in this vector, representing the bottom + first(level_discrete) else nothing end diff --git a/core/test/basin.jl b/core/test/basin.jl index ca2494d11..1fc59efa4 100644 --- a/core/test/basin.jl +++ b/core/test/basin.jl @@ -2,6 +2,7 @@ using Test using Ribasim import BasicModelInterface as BMI using SciMLBase +import Tables @testset "trivial model" begin toml_path = normpath(@__DIR__, "../../data/trivial/trivial.toml") @@ -17,7 +18,7 @@ end model = Ribasim.run(toml_path) @test model isa Ribasim.Model @test model.integrator.sol.retcode == Ribasim.ReturnCode.Success - @test model.integrator.sol.u[end] ≈ Float32[476.8749, 476.85898, 1.8706497, 786.96686] skip = + @test model.integrator.sol.u[end] ≈ Float32[452.9688, 453.0431, 1.8501105, 1238.0144] skip = Sys.isapple() end @@ -28,7 +29,7 @@ end @test model isa Ribasim.Model @test model.integrator.sol.retcode == Ribasim.ReturnCode.Success @test length(model.integrator.p.basin.precipitation) == 4 - @test model.integrator.sol.u[end] ≈ Float32[480.5936, 480.58762, 1.4178488, 793.6097] skip = + @test model.integrator.sol.u[end] ≈ Float32[428.06897, 428.07315, 1.3662858, 1249.2343] skip = Sys.isapple() end @@ -39,7 +40,54 @@ end model = Ribasim.run(toml_path) @test model isa Ribasim.Model @test model.integrator.sol.retcode == Ribasim.ReturnCode.Success - @test model.integrator.sol.u[end] ≈ Float32[27.273159, 340.68964] skip = Sys.isapple() + @test model.integrator.sol.u[end] ≈ Float32[1.4875988, 366.4851] 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.tables[end].t[end] == 1.2 end + +"Shorthand for Ribasim.get_area_and_level" +function lookup(profile, S) + Ribasim.get_area_and_level(profile.S, profile.A, profile.h, S) +end + +@testset "Profile" begin + n_interpolations = 100 + storage = range(0.0, 1000.0, n_interpolations) + + # Covers interpolation for constant and non-constant area, extrapolation for constant area + A = [0.0, 100.0, 100.0] + h = [0.0, 10.0, 15.0] + profile = (; A, h, S = Ribasim.profile_storage(h, A)) + + # On profile points we reproduce the profile + for (; S, A, h) in Tables.rows(profile) + @test lookup(profile, S) == (A, h) + end + + # Robust to negative storage + @test lookup(profile, -1.0) == (profile.A[1], profile.h[1]) + + # On the first segment + S = 100.0 + A, h = lookup(profile, S) + @test h ≈ sqrt(S / 5) + @test A ≈ 10 * h + + # On the second segment and extrapolation + for S in [500.0 + 100.0, 1000.0 + 100.0] + S = 500.0 + 100.0 + A, h = lookup(profile, S) + @test h ≈ 10.0 + (S - 500.0) / 100.0 + @test A == 100.0 + end + + # Covers extrapolation for non-constant area + A = [0.0, 100.0] + h = [0.0, 10.0] + profile = (; A, h, S = Ribasim.profile_storage(h, A)) + + S = 500.0 + 100.0 + A, h = lookup(profile, S) + @test h ≈ sqrt(S / 5) + @test A ≈ 10 * h +end diff --git a/core/test/control.jl b/core/test/control.jl index 3401a9263..b37e8a80b 100644 --- a/core/test/control.jl +++ b/core/test/control.jl @@ -1,6 +1,6 @@ import Ribasim -toml_path = normpath(@__DIR__, "../../data/pump_control/control.toml") +toml_path = normpath(@__DIR__, "../../data/pump_control/pump_control.toml") @testset "pump control" begin @test ispath(toml_path) diff --git a/core/test/utils.jl b/core/test/utils.jl index a2824f5b8..5c6cb5c4f 100644 --- a/core/test/utils.jl +++ b/core/test/utils.jl @@ -18,6 +18,10 @@ end end @testset "bottom" begin + # create two basins with different bottoms/levels + area = [[0.0, 1.0], [0.0, 1.0]] + level = [[0.0, 1.0], [4.0, 5.0]] + storage = Ribasim.profile_storage.(level, area) basin = Ribasim.Basin( Indices([5, 7]), [2.0, 3.0], @@ -25,18 +29,13 @@ end [2.0, 3.0], [2.0, 3.0], [2.0, 3.0], - [ # area - LinearInterpolation([1.0, 1.0], [0.0, 1.0]), - LinearInterpolation([1.0, 1.0], [0.0, 1.0]), - ], - [ # level - LinearInterpolation([0.0, 1.0], [0.0, 1.0]), - LinearInterpolation([4.0, 3.0], [0.0, 1.0]), - ], + area, + level, + storage, StructVector{Ribasim.BasinForcingV1}(undef, 0), ) - @test Ribasim.basin_bottom_index(basin, 2) === 4.0 + @test basin.level[2][1] === 4.0 @test Ribasim.basin_bottom(basin, 5) === 0.0 @test Ribasim.basin_bottom(basin, 7) === 4.0 @test Ribasim.basin_bottom(basin, 6) === nothing diff --git a/python/ribasim_testmodels/ribasim_testmodels/control.py b/python/ribasim_testmodels/ribasim_testmodels/control.py index 94cd71f8f..374fea6fe 100644 --- a/python/ribasim_testmodels/ribasim_testmodels/control.py +++ b/python/ribasim_testmodels/ribasim_testmodels/control.py @@ -129,7 +129,7 @@ def pump_control_model() -> ribasim.Model: # Setup a model: model = ribasim.Model( - modelname="control", + modelname="pump_control", node=node, edge=edge, basin=basin,