More accurate interpolation

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

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

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))

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
 

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

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

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)

core/test/utils.jl

@@ -18,25 +18,24 @@ 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],
         [2.0, 3.0],
         [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

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,