Allow partial basin state input in Ribasim core

core/src/bmi.jl

@@ -73,15 +73,9 @@ function BMI.initialize(T::Type{Model}, config::Config)::Model
     end
     @debug "Read database into memory."
 
-    storage = if isempty(state)
-        # default to nearly empty basins, perhaps make required input
-        fill(1.0, n)
-    else
-        storages, errors = get_storages_from_levels(parameters.basin, state.level)
-        if errors
-            error("Encountered errors while parsing the initial levels of basins.")
-        end
-        storages
+    storage, errors = get_storages_from_levels(parameters.basin, state.node_id, state.level)
+    if errors
+        error("Encountered errors while parsing the initial levels of basins.")
     end
     @assert length(storage) == n "Basin / state length differs from number of Basins"
     # Integrals for PID control

core/src/utils.jl

@@ -376,12 +376,29 @@ end
 """Compute the storages of the basins based on the water level of the basins."""
 function get_storages_from_levels(
     basin::Basin,
-    levels::Vector,
+    state_node_id::Vector{Int},
+    state_level::Vector{Float64},
 )::Tuple{Vector{Float64}, Bool}
-    storages = Float64[]
-
-    for (i, level) in enumerate(levels)
-        push!(storages, get_storage_from_level(basin, i, level))
+    (; node_id) = basin
+
+    storages = fill(1.0, length(node_id))
+    n_specified_states = length(state_node_id)
+
+    if n_specified_states > 0
+        basin_state_index = 1
+        basin_state_node_id = state_node_id[1]
+
+        for (i, id) in enumerate(node_id)
+            if basin_state_node_id == id.value
+                storages[i] =
+                    get_storage_from_level(basin, i, state_level[basin_state_index])
+                basin_state_index += 1
+                if basin_state_index > n_specified_states
+                    break
+                end
+                basin_state_node_id = state_node_id[basin_state_index]
+            end
+        end
     end
     return storages, any(isnan.(storages))
 end

core/test/utils_test.jl

@@ -110,24 +110,26 @@ end
     ]
     storage = Ribasim.profile_storage(level, area)
     basin = Ribasim.Basin(
-        Indices(Ribasim.NodeID[1]),
-        zeros(1),
-        zeros(1),
-        zeros(1),
-        zeros(1),
-        zeros(1),
-        zeros(1),
-        [area],
-        [level],
-        [storage],
+        Indices(Ribasim.NodeID[1, 2]),
+        zeros(2),
+        zeros(2),
+        zeros(2),
+        zeros(2),
+        zeros(2),
+        zeros(2),
+        [area, area],
+        [level, level],
+        [storage, storage],
         StructVector{Ribasim.BasinTimeV1}(undef, 0),
     )
 
     logger = TestLogger()
     with_logger(logger) do
-        storages, errors = Ribasim.get_storages_from_levels(basin, [-1.0])
+        storages, errors = Ribasim.get_storages_from_levels(basin, [1], [-1.0])
         @test isnan(storages[1])
         @test errors
+        # Storage for basin with unspecified level is set to 1.0
+        @test storages[2] == 1.0
     end
 
     @test length(logger.logs) == 1