fix level extrapolation, simplify tests

core/src/utils.jl

@@ -76,12 +76,14 @@ function get_area_and_level(
     # storage_idx: smallest index such that storage_discrete[storage_idx] >= storage
     storage_idx = searchsortedfirst(storage_discrete, storage)
 
-    if storage_idx == 1 # If the lowest discrete_storage level 0, this can only happen if the storage is 0 since storage is never negative
+    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
-        # These values yield linear extrapolation of area(level) based on the last 2 values
+        # 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]
@@ -95,9 +97,8 @@ function get_area_and_level(
         if area_diff ≈ 0
             # Constant area means linear interpolation of level
             area = area_lower
-
             level =
-                area_higher +
+                level_higher +
                 level_diff * (storage - storage_higher) / (storage_higher - storage_lower)
         else
             area = sqrt(

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")
@@ -44,41 +45,49 @@ end
     @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
-    area_discrete = [0.0, 100.0, 100.0]
-    level_discrete = [0.0, 10.0, 15.0]
-    storage_discrete = Ribasim.profile_storage(level_discrete, area_discrete)
+    A = [0.0, 100.0, 100.0]
+    h = [0.0, 10.0, 15.0]
+    profile = (; A, h, S = Ribasim.profile_storage(h, A))
 
-    area, level = zip(
-        [
-            Ribasim.get_area_and_level(storage_discrete, area_discrete, level_discrete, s) for s in storage
-        ]...,
-    )
+    # On profile points we reproduce the profile
+    for (; S, A, h) in Tables.rows(profile)
+        @test lookup(profile, S) == (A, h)
+    end
 
-    level_expected =
-        ifelse.(storage .< 500.0, sqrt.(storage ./ 5), 10.0 .+ (storage .- 500.0) ./ 100.0)
+    # Robust to negative storage
+    @test lookup(profile, -1.0) == (profile.A[1], profile.h[1])
 
-    @test all(level .≈ level_expected)
-    area_expected = min.(10.0 * level_expected, 100.0)
-    @test all(area .≈ area_expected)
+    # On the first segment
+    S = 100.0
+    A, h = lookup(profile, S)
+    @test h ≈ sqrt(S / 5)
+    @test A ≈ 10 * h
 
-    # Covers extrapolation for non-constant area
-    area_discrete = [0.0, 100.0]
-    level_discrete = [0.0, 10.0]
-    storage_discrete = Ribasim.profile_storage(level_discrete, area_discrete)
+    # 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
 
-    area, level = zip(
-        [
-            Ribasim.get_area_and_level(storage_discrete, area_discrete, level_discrete, s) for s in storage
-        ]...,
-    )
+    # 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))
 
-    level_expected = sqrt.(storage ./ 5)
-    @test all(level .≈ level_expected)
-    area_expected = 10.0 * level_expected
-    @test all(area .≈ area_expected)
+    S = 500.0 + 100.0
+    A, h = lookup(profile, S)
+    @test h ≈ sqrt(S / 5)
+    @test A ≈ 10 * h
 end

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