Comments adressed

core/src/allocation_optim.jl

@@ -327,9 +327,7 @@ function get_basin_data(
     vertical_flux = get_tmp(vertical_flux, 0)
     _, basin_idx = id_index(basin.node_id, node_id)
     # NOTE: Instantaneous
-    influx =
-        vertical_flux.precipitation[basin_idx] - vertical_flux.evaporation[basin_idx] +
-        vertical_flux.drainage[basin_idx] - vertical_flux.infiltration[basin_idx]
+    influx = get_influx(basin, node_id)
     _, basin_idx = id_index(basin.node_id, node_id)
     storage_basin = u.storage[basin_idx]
     control_inneighbors = inneighbor_labels_type(graph, node_id, EdgeType.control)

core/src/bmi.jl

@@ -36,9 +36,9 @@ function BMI.get_value_ptr(model::Model, name::AbstractString)
         model.integrator.u.storage
     elseif name == "basin.level"
         get_tmp(model.integrator.p.basin.current_level, 0)
-    elseif name == "basin.infiltration_instantaneous"
+    elseif name == "basin.infiltration"
         get_tmp(model.integrator.p.basin.vertical_flux, 0).infiltration
-    elseif name == "basin.drainage_instantaneous"
+    elseif name == "basin.drainage"
         get_tmp(model.integrator.p.basin.vertical_flux, 0).drainage
     elseif name == "basin.infiltration_integrated"
         model.integrator.p.basin.vertical_flux_bmi.infiltration

core/src/callback.jl

@@ -138,11 +138,9 @@ end
 "Compute the average flows over the last saveat interval and write
 them to SavedValues"
 function save_flow(u, t, integrator)
-    (; p) = integrator
-    (; graph) = p
-    (; flow_integrated) = graph[]
+    (; flow_integrated) = integrator.p.graph[]
 
-    Δt = get_Δt(integrator, graph)
+    Δt = get_Δt(integrator)
     flow_mean = copy(flow_integrated)
     flow_mean ./= Δt
     fill!(flow_integrated, 0.0)
@@ -153,11 +151,10 @@ end
 "Compute the average vertical fluxes over the last saveat interval and write
 them to SavedValues"
 function save_vertical_flux(u, t, integrator)
-    (; p) = integrator
-    (; basin, graph) = p
+    (; basin) = integrator.p
     (; vertical_flux_integrated) = basin
 
-    Δt = get_Δt(integrator, graph)
+    Δt = get_Δt(integrator)
     vertical_flux_mean = copy(vertical_flux_integrated)
     vertical_flux_mean ./= Δt
     fill!(vertical_flux_integrated, 0.0)
@@ -497,7 +494,7 @@ function update_basin(integrator)::Nothing
         update_vertical_flux!(basin, storage, i)
     end
 
-    # Forget about vertical fluxes before basin update
+    # Forget about vertical fluxes to handle discontinuous forcing from basin_update
     copyto!(vertical_flux_prev, vertical_flux)
     return nothing
 end

core/src/parameter.jl

@@ -176,8 +176,8 @@ struct Basin{T, C, V1, V2, V3} <: AbstractParameterNode
         vertical_flux_from_input::V1,
         vertical_flux::V2,
         vertical_flux_prev::V3,
-        vertical_flux_integrated,
-        vertical_flux_bmi,
+        vertical_flux_integrated::V3,
+        vertical_flux_bmi::V3,
         current_level::T,
         current_area::T,
         area,

core/src/solve.jl

@@ -51,7 +51,7 @@ end
 Smoothly let the evaporation flux go to 0 when at small water depths
 Currently at less than 0.1 m.
 """
-function update_vertical_flux!(basin::Basin, storage::AbstractVector, i::Int)::Number
+function update_vertical_flux!(basin::Basin, storage::AbstractVector, i::Int)::Nothing
     (; current_level, current_area, vertical_flux_from_input, vertical_flux) = basin
     current_level = get_tmp(current_level, storage)
     current_area = get_tmp(current_area, storage)
@@ -75,8 +75,7 @@ function update_vertical_flux!(basin::Basin, storage::AbstractVector, i::Int)::N
     vertical_flux.drainage[i] = drainage
     vertical_flux.infiltration[i] = infiltration
 
-    influx = precipitation - evaporation + drainage - infiltration
-    return influx
+    return nothing
 end
 
 function formulate_basins!(
@@ -86,8 +85,8 @@ function formulate_basins!(
 )::Nothing
     for (i, id) in enumerate(basin.node_id)
         # add all precipitation that falls within the profile
-        influx = update_vertical_flux!(basin, storage, i)
-        du.storage[i] += influx
+        update_vertical_flux!(basin, storage, i)
+        du.storage[i] += get_influx(basin, i)
     end
     return nothing
 end

core/src/util.jl

@@ -711,9 +711,9 @@ end
 """
 Get the time interval between (flow) saves
 """
-function get_Δt(integrator, graph)::Float64
-    (; t, dt) = integrator
-    (; saveat) = graph[]
+function get_Δt(integrator)::Float64
+    (; p, t, dt) = integrator
+    (; saveat) = p.graph[]
     if iszero(saveat)
         dt
     elseif isinf(saveat)
@@ -729,3 +729,19 @@ function get_Δt(integrator, graph)::Float64
         end
     end
 end
+
+function get_influx(basin::Basin, node_id::NodeID)::Float64
+    has_index, basin_idx = id_index(basin.node_id, node_id)
+    if !has_index
+        error("Sum of vertical fluxes requested for non-basin $id.")
+    end
+    return get_influx(basin, basin_idx)
+end
+
+function get_influx(basin::Basin, basin_idx::Int)::Float64
+    (; vertical_flux) = basin
+    vertical_flux = get_tmp(vertical_flux, 0)
+    (; precipitation, evaporation, drainage, infiltration) = vertical_flux
+    return precipitation[basin_idx] - evaporation[basin_idx] + drainage[basin_idx] -
+           infiltration[basin_idx]
+end

core/src/write.jl

@@ -84,10 +84,10 @@ function basin_table(
     node_id::Vector{Int32},
     storage::Vector{Float64},
     level::Vector{Float64},
-    precipitation::Vector{Float64},
-    evaporation::Vector{Float64},
-    drainage::Vector{Float64},
-    infiltration::Vector{Float64},
+    precipitation::Vector{Union{Missing, Float64}},
+    evaporation::Vector{Union{Missing, Float64}},
+    drainage::Vector{Union{Missing, Float64}},
+    infiltration::Vector{Union{Missing, Float64}},
 }
     (; saved) = model
     (; vertical_flux) = saved
@@ -100,10 +100,10 @@ function basin_table(
     ntsteps = length(data.time)
     nrows = nbasin * ntsteps
 
-    precipitation = zeros(nrows)
-    evaporation = zeros(nrows)
-    drainage = zeros(nrows)
-    infiltration = zeros(nrows)
+    precipitation = Vector{Union{Missing, Float64}}(missing, nrows)
+    evaporation = Vector{Union{Missing, Float64}}(missing, nrows)
+    drainage = Vector{Union{Missing, Float64}}(missing, nrows)
+    infiltration = Vector{Union{Missing, Float64}}(missing, nrows)
 
     idx_row = 0
 

core/test/bmi_test.jl

@@ -61,8 +61,8 @@ end
     for name in [
         "basin.storage",
         "basin.level",
-        "basin.infiltration_instantaneous",
-        "basin.drainage_instantaneous",
+        "basin.infiltration",
+        "basin.drainage",
         "basin.infiltration_integrated",
         "basin.drainage_integrated",
         "basin.subgrid_level",

core/test/run_models_test.jl

@@ -54,7 +54,16 @@
                 :drainage,
                 :infiltration,
             ),
-            (DateTime, Int32, Float64, Float64, Float64, Float64, Float64, Float64),
+            (
+                DateTime,
+                Int32,
+                Float64,
+                Float64,
+                Union{Missing, Float64},
+                Union{Missing, Float64},
+                Union{Missing, Float64},
+                Union{Missing, Float64},
+            ),
         )
         @test Tables.schema(control) == Tables.Schema(
             (:time, :control_node_id, :truth_state, :control_state),

core/test/time_test.jl

@@ -24,43 +24,65 @@
 end
 
 @testitem "vertical_flux_means" begin
-    using DataFrames: DataFrame
+    using DataFrames: DataFrame, transform!, ByRow
 
     toml_path =
         normpath(@__DIR__, "../../generated_testmodels/basic_transient/ribasim.toml")
     @test ispath(toml_path)
     model = Ribasim.run(toml_path)
     basin = model.integrator.p.basin
+    n_basin = length(basin.node_id)
     basin_table = DataFrame(Ribasim.basin_table(model))
-    filter!(:node_id => id -> id == 1, basin_table)
-    basin_table[!, "time"] .=
-        Ribasim.seconds_since.(basin_table.time, model.config.starttime)
 
     # No vertical flux data for last saveat
     t_end = last(basin_table).time
     data_end = filter(:time => t -> t == t_end, basin_table)
-    @test all(data_end.precipitation .== 0)
-    @test all(data_end.evaporation .== 0)
-    @test all(data_end.drainage .== 0)
-    @test all(data_end.infiltration .== 0)
+    @test all(ismissing.(data_end.precipitation))
+    @test all(ismissing.(data_end.evaporation))
+    @test all(ismissing.(data_end.drainage))
+    @test all(ismissing.(data_end.infiltration))
 
     time_table = DataFrame(basin.time)
-    filter!(:node_id => id -> id == 1, time_table)
-    time_table[!, "time"] .= Ribasim.seconds_since.(time_table.time, model.config.starttime)
-    fixed_area = basin.area[1][end]
-    area = [
-        Ribasim.get_area_and_level(basin, 1, storage)[1] for storage in basin_table.storage
+    time_table[!, "basin_idx"] = [
+        Ribasim.id_index(basin.node_id, node_id)[2] for
+        node_id in Ribasim.NodeID.(:Basin, time_table.node_id)
     ]
-    area = (area[2:end] .+ area[1:(end - 1)]) ./ 2
-    @test all(
-        basin_table.precipitation[1:(end - 1)] .≈
-        fixed_area .* time_table.precipitation[1:(end - 1)],
-    )
+    time_table[!, "area"] = [
+        Ribasim.get_area_and_level(basin, idx, storage)[1] for
+        (idx, storage) in zip(time_table.basin_idx, basin_table.storage)
+    ]
+    # Mean areas are sufficient to compute the mean flows
+    # (assuming the saveats coincide with the solver timepoints),
+    # as the potential evaporation is constant over the saveat intervals
+    time_table[!, "mean_area"] .= 0.0
+    n_basins = length(basin.node_id)
+    n_times = length(unique(time_table.time)) - 1
+    for basin_idx in 1:n_basins
+        for time_idx in 1:n_times
+            idx_1 = n_basins * (time_idx - 1) + basin_idx
+            idx_2 = n_basins * time_idx + basin_idx
+            mean_area = (time_table.area[idx_1] + time_table.area[idx_2]) / 2
+            time_table.mean_area[idx_1] = mean_area
+        end
+    end
+
+    filter!(:time => t -> t !== t_end, basin_table)
+    filter!(:time => t -> t !== t_end, time_table)
+
     @test all(
         isapprox(
-            basin_table.evaporation[1:(end - 1)],
-            area .* time_table.potential_evaporation[1:(end - 1)];
+            basin_table.evaporation,
+            time_table.mean_area .* time_table.potential_evaporation;
             rtol = 1e-4,
         ),
     )
+
+    fixed_area = Dict(
+        id.value => basin.area[Ribasim.id_index(basin.node_id, id)[2]][end] for
+        id in basin.node_id
+    )
+    transform!(time_table, :node_id => ByRow(id -> fixed_area[id]) => :fixed_area)
+    @test all(
+        basin_table.precipitation .≈ time_table.fixed_area .* time_table.precipitation,
+    )
 end

docs/core/usage.qmd

@@ -708,20 +708,20 @@ derivative       | Float64  | -        | -
 The basin table contains:
 
 - results of the storage and level of each basin, which are instantaneous values;
-- results of the vertical fluxes on each basin, which are mean values over the `saveat` intervals. In the time column the start of the period is indicated. This means that for the final time in the table the mean vertical fluxes are not computed, and thus left $0$.
+- results of the vertical fluxes on each basin, which are mean values over the `saveat` intervals. In the time column the start of the period is indicated. This means that for the final time in the table the mean vertical fluxes are not computed, and thus are `missing`.
 
 The initial condition is also written to the file.
 
-column        | type     | unit
-------------- | -------- | ----
-time          | DateTime | -
-node_id       | Int32    | -
-storage       | Float64  | $m^3$
-level         | Float64  | $m$
-precipitation | Float64  | $m^3 s^{-1}$
-evaporation   | Float64  | $m^3 s^{-1}$
-drainage      | Float64  | $m^3 s^{-1}$
-infiltration  | Float64  | $m^3 s^{-1}$
+column        | type                     | unit
+------------- | ------------------------ | ----
+time          | DateTime                 | -
+node_id       | Int32                    | -
+storage       | Float64                  | $m^3$
+level         | Float64                  | $m$
+precipitation | Union{Float64, Missing}  | $m^3 s^{-1}$
+evaporation   | Union{Float64, Missing}  | $m^3 s^{-1}$
+drainage      | Union{Float64, Missing}  | $m^3 s^{-1}$
+infiltration  | Union{Float64, Missing}  | $m^3 s^{-1}$
 
 
 The table is sorted by time, and per time it is sorted by `node_id`.
@@ -734,7 +734,7 @@ In the time column the start of the period is indicated.
 column         | type                  | unit
 -------------- | --------------------- | ----
 time           | DateTime              | -
-edge_id        | Union{Int32, Missing} | -
+edge_id        | Int32                 | -
 from_node_type | String                | -
 from_node_id   | Int32                 | -
 to_node_type   | String                | -
@@ -743,6 +743,7 @@ flow_rate      | Float64               | $m^3 s^{-1}$
 
 The table is sorted by time, and per time the same `edge_id` order is used, though not sorted.
 The `edge_id` value is the same as the `fid` written to the Edge table, and can be used to directly look up the Edge geometry.
+Flows from the "from" to the "to" node have a positive sign, and if the flow is reversed it will be negative.
 
 ## DiscreteControl - `control.arrow`