Add basin_state.arrow results (#1626)

Fixes #1358.
Includes a test to ensure that we can continue a simulation with this.

core/src/util.jl

@@ -22,7 +22,7 @@ function get_storage_from_level(basin::Basin, state_idx::Int, level::Float64)::F
 end
 
 """Compute the storages of the basins based on the water level of the basins."""
-function get_storages_from_levels(basin::Basin, levels::Vector)::Vector{Float64}
+function get_storages_from_levels(basin::Basin, levels::AbstractVector)::Vector{Float64}
     errors = false
     state_length = length(levels)
     basin_length = length(basin.storage_to_level)

core/src/write.jl

@@ -9,6 +9,11 @@ function write_results(model::Model)::Model
     compress = get_compressor(results)
     remove_empty_table = model.integrator.t != 0
 
+    # state
+    table = basin_state_table(model)
+    path = results_path(config, RESULTS_FILENAME.basin_state)
+    write_arrow(path, table, compress; remove_empty_table)
+
     # basin
     table = basin_table(model)
     path = results_path(config, RESULTS_FILENAME.basin)
@@ -44,6 +49,7 @@ function write_results(model::Model)::Model
 end
 
 const RESULTS_FILENAME = (
+    basin_state = "basin_state.arrow",
     basin = "basin.arrow",
     flow = "flow.arrow",
     control = "control.arrow",
@@ -77,6 +83,19 @@ function get_storages_and_levels(
     return (; time = tsteps, node_id, storage, level)
 end
 
+"Create the basin state table from the saved data"
+function basin_state_table(
+    model::Model,
+)::@NamedTuple{node_id::Vector{Int32}, level::Vector{Float64}}
+    (; storage) = model.integrator.u
+    (; basin) = model.integrator.p
+
+    # ensure the levels are up-to-date
+    set_current_basin_properties!(basin, storage)
+
+    return (; node_id = Int32.(basin.node_id), level = get_tmp(basin.current_level, 0))
+end
+
 "Create the basin result table from the saved data"
 function basin_table(
     model::Model,
@@ -335,17 +354,19 @@ function write_arrow(
     # At the start of the simulation, write an empty table to ensure we have permissions
     # and fail early.
     # At the end of the simulation, write all non-empty tables, and remove existing empty ones.
-    if isempty(table.time) && remove_empty_table
+    if haskey(table, :time) && isempty(table.time) && remove_empty_table
         try
             rm(path; force = true)
         catch
             @warn "Failed to remove results, file may be locked." path
         end
         return nothing
     end
-    # ensure DateTime is encoded in a compatible manner
-    # https://github.com/apache/arrow-julia/issues/303
-    table = merge(table, (; time = convert.(Arrow.DATETIME, table.time)))
+    if haskey(table, :time)
+        # ensure DateTime is encoded in a compatible manner
+        # https://github.com/apache/arrow-julia/issues/303
+        table = merge(table, (; time = convert.(Arrow.DATETIME, table.time)))
+    end
     metadata = ["ribasim_version" => string(pkgversion(Ribasim))]
     mkpath(dirname(path))
     try

core/test/io_test.jl

@@ -138,3 +138,39 @@ end
     @test Arrow.getmetadata(tbl) ===
           Base.ImmutableDict("ribasim_version" => ribasim_version)
 end
+
+@testitem "warm state" begin
+    using IOCapture: capture
+    using Ribasim: solve!, write_results
+    import TOML
+
+    model_path_src = normpath(@__DIR__, "../../generated_testmodels/basic/")
+
+    # avoid changing the original model for other tests
+    model_path = normpath(@__DIR__, "../../generated_testmodels/basic_warm/")
+    cp(model_path_src, model_path; force = true)
+    toml_path = normpath(model_path, "ribasim.toml")
+
+    config = Ribasim.Config(toml_path)
+    model = Ribasim.Model(config)
+    storage1_begin = copy(model.integrator.u.storage)
+    solve!(model)
+    storage1_end = model.integrator.u.storage
+    @test storage1_begin != storage1_end
+
+    # copy state results to input
+    write_results(model)
+    state_path = Ribasim.results_path(config, Ribasim.RESULTS_FILENAME.basin_state)
+    cp(state_path, Ribasim.input_path(config, "warm_state.arrow"))
+
+    # point TOML to the warm state
+    toml_dict = TOML.parsefile(toml_path)
+    toml_dict["basin"] = Dict("state" => "warm_state.arrow")
+    open(toml_path, "w") do io
+        TOML.print(io, toml_dict)
+    end
+
+    model = Ribasim.Model(toml_path)
+    storage2_begin = model.integrator.u.storage
+    @test storage1_end ≈ storage2_begin
+end

docs/reference/node/basin.qmd

@@ -34,18 +34,26 @@ The only difference is that a time column is added.
 The table must by sorted by time, and per time it must be sorted by `node_id`.
 At the given timestamps the values are set in the simulation, such that the timeseries can be seen as forward filled.
 
-## State
+## State {#sec-state}
 
-The state table aims to capture the full state of the Basin, such that it can be used as an
-initial condition, potentially the outcome of an earlier simulation. Currently only the
-Basin node types have state.
+The state table gives the initial water levels of all Basins.
 
 column    | type    | unit         | restriction
 --------- | ------- | ------------ | -----------
 node_id   | Int32   | -            | sorted
 level     | Float64 | $m$          | $\ge$ basin bottom
 
 Each Basin ID needs to be in the table.
+To use the final state of an earlier simulation as an initial condition, copy [`results/basin_state.arrow`](/reference/usage.qmd#sec-state) over to the `input_dir`, and point the TOML to it:
+
+```toml
+[basin]
+state = "basin_state.arrow"
+```
+
+This will start of the simulation with the same water levels as the end of the earlier simulation.
+Since there is no time information in this state, the user is responsible to ensure that the earlier `endtime` matches the current `starttime`.
+This only applies when the user wishes to continue an existing simulation as if it was one continuous simulation.
 
 ## Profile
 

docs/reference/usage.qmd

@@ -237,6 +237,17 @@ The table is sorted by time, and per time the same `edge_id` order is used, thou
 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.
 
+## State - `basin_state.arrow` {#sec-state}
+
+The Basin state table contains the water levels in each Basin at the end of the simulation.
+
+column    | type    | unit
+--------- | ------- | ------------
+node_id   | Int32   | -
+level     | Float64 | $m$
+
+To use this result as the initial condition of another simulation, see the [Basin / state](/reference/node/basin.qmd#sec-state) table reference.
+
 ## DiscreteControl - `control.arrow`
 
 The control table contains a record of each change of control state: when it happened, which control node was involved, to which control state it changed and based on which truth state.