Add tests for discrete control times (#384)

Fixes #300. --------- Co-authored-by: Bart de Koning <[email protected]> Co-authored-by: Martijn Visser <[email protected]>
Deltares · Jul 4, 2023 · ed4fbef · ed4fbef
1 parent defeb53
commit ed4fbef
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 7 deletions.
diff --git a/core/src/io.jl b/core/src/io.jl
@@ -116,26 +116,43 @@ function parsefile(config_path::AbstractString)::Config
     return from_toml(Config, config_path; relative_dir = dirname(normpath(config_path)))
 end
 
-function write_basin_output(model::Model)
+"Get the storage and level of all basins as matrices of nbasin × ntime"
+function get_storages_and_levels(
+    model::Model,
+)::NamedTuple{
+    (:time, :node_id, :storage, :level),
+    Tuple{Vector{Dates.DateTime}, Vector{Int64}, Matrix{Float64}, Matrix{Float64}},
+}
     (; config, integrator) = model
     (; sol, p) = integrator
 
-    basin_id = p.basin.node_id.values::Vector{Int}
-    nbasin = length(basin_id)
+    node_id = p.basin.node_id.values::Vector{Int}
+    nbasin = length(node_id)
     tsteps = datetime_since.(timesteps(model), config.starttime)
     ntsteps = length(tsteps)
 
-    time = convert.(Arrow.DATETIME, repeat(tsteps; inner = nbasin))
-    node_id = repeat(basin_id; outer = ntsteps)
-
     storage = reshape(vec(sol), nbasin, ntsteps)
     level = zero(storage)
     for (i, basin_storage) in enumerate(eachrow(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))
+    return (; time = tsteps, node_id, storage, level)
+end
+
+function write_basin_output(model::Model)
+    (; config, integrator) = model
+    (; p) = integrator
+
+    data = get_storages_and_levels(model)
+    nbasin = length(data.node_id)
+    ntsteps = length(data.time)
+
+    time = convert.(Arrow.DATETIME, repeat(data.time; inner = nbasin))
+    node_id = repeat(data.node_id; outer = ntsteps)
+
+    basin = (; time, node_id, storage = vec(data.storage), level = vec(data.level))
     path = output_path(config, config.output.basin)
     mkpath(dirname(path))
     Arrow.write(path, basin; compress = :lz4)

diff --git a/core/test/control.jl b/core/test/control.jl
@@ -8,6 +8,7 @@ toml_path = normpath(@__DIR__, "../../data/pump_control/pump_control.toml")
     p = model.integrator.p
     control = model.integrator.p.control
 
+    # Control input
     pump_control_mapping::Dict{Tuple{Int64, String}, NamedTuple} =
         Dict((4, "off") => (flow_rate = 0.0,), (4, "on") => (flow_rate = 1.0e-5,))
 
@@ -16,6 +17,20 @@ toml_path = normpath(@__DIR__, "../../data/pump_control/pump_control.toml")
 
     @test p.pump.control_mapping == pump_control_mapping
     @test p.control.logic_mapping == logic_mapping
+
+    # Control result
     @test control.record.truth_state == ["TF", "FF", "FT"]
     @test control.record.control_state == ["off", "on", "off"]
+
+    level = Ribasim.get_storages_and_levels(model).level
+    timesteps = Ribasim.timesteps(model)
+
+    # Control times
+    t_1 = control.record.time[2]
+    t_1_index = findfirst(timesteps .≈ t_1)
+    @test level[1, t_1_index] ≈ control.greater_than[1]
+
+    t_2 = control.record.time[3]
+    t_2_index = findfirst(timesteps .≈ t_2)
+    @test level[2, t_2_index] ≈ control.greater_than[2]
 end