State and time functions for coupling

For model coupling (OpenDA) additonal functions for loading and saving state information and start time in Unix time format are required. For this, the set states functionality from the initialization function is moved to a separate `set_states` function for each Model type.

src/bmi.jl

@@ -375,3 +375,22 @@ function BMI.get_grid_edge_nodes(model::Model, grid::Int, edge_nodes::Vector{Int
         error("unknown grid type $grid")
     end
 end
+
+# Extension of BMI functions (state handling and start time), required for OpenDA coupling.
+# May also be useful for other external software packages.
+function load_state(model::Model)
+    model = set_states(model)
+    return model
+end
+
+function save_state(model::Model)
+    @unpack config, writer, clock = model
+    if haskey(config, "state") && haskey(config.state, "path_output")
+        @info "Write output states to NetCDF file `$(model.writer.state_nc_path)`."
+    end
+    write_netcdf_timestep(model, writer.state_dataset, writer.state_parameters)
+end
+
+function get_start_unix_time(model::Model)
+    datetime2unix(DateTime(model.config.starttime))
+end

src/flextopo_model.jl

@@ -12,7 +12,6 @@ function initialize_flextopo_model(config::Config)
     clock = Clock(config, reader)
     Δt = clock.Δt
 
-    reinit = get(config.model, "reinit", true)::Bool
     do_reservoirs = get(config.model, "reservoirs", false)::Bool
     do_lakes = get(config.model, "lakes", false)::Bool
     do_pits = get(config.model, "pits", false)::Bool
@@ -661,23 +660,7 @@ function initialize_flextopo_model(config::Config)
         FlextopoModel(),
     )
 
-    # read and set states in model object if reinit=true
-    if reinit == false
-        instate_path = input_path(config, config.state.path_input)
-        state_ncnames = ncnames(config.state)
-        set_states(instate_path, model, state_ncnames; type = Float, dimname = :classes)
-        # update kinematic wave volume for river and land domain
-        @unpack lateral = model
-        lateral.land.volume .= lateral.land.h .* lateral.land.width .* lateral.land.dl
-        lateral.river.volume .= lateral.river.h .* lateral.river.width .* lateral.river.dl
-
-        if do_lakes
-            # storage must be re-initialized after loading the state with the current
-            # waterlevel otherwise the storage will be based on the initial water level
-            lakes.storage .=
-                initialize_storage(lakes.storfunc, lakes.area, lakes.waterlevel, lakes.sh)
-        end
-    end
+    model = set_states(model)
 
     return model
 end
@@ -732,3 +715,27 @@ function update(model::Model{N,L,V,R,W,T}) where {N,L,V,R,W,T<:FlextopoModel}
 
     return model
 end
+
+function set_states(model::Model{N,L,V,R,W,T}) where {N,L,V,R,W,T<:FlextopoModel}
+    @unpack lateral, config = model
+    reinit = get(config.model, "reinit", true)::Bool
+    # read and set states in model object if reinit=true
+    if reinit == false
+        instate_path = input_path(config, config.state.path_input)
+        state_ncnames = ncnames(config.state)
+        set_states(instate_path, model, state_ncnames; type = Float, dimname = :classes)
+
+        # update kinematic wave volume for river and land domain
+        lateral.land.volume .= lateral.land.h .* lateral.land.width .* lateral.land.dl
+        lateral.river.volume .= lateral.river.h .* lateral.river.width .* lateral.river.dl
+
+        if do_lakes
+            # storage must be re-initialized after loading the state with the current
+            # waterlevel otherwise the storage will be based on the initial water level
+            lakes = lateral.river.lake
+            lakes.storage .=
+                initialize_storage(lakes.storfunc, lakes.area, lakes.waterlevel, lakes.sh)
+        end
+    end
+    return model
+end

src/hbv_model.jl

@@ -12,7 +12,6 @@ function initialize_hbv_model(config::Config)
     clock = Clock(config, reader)
     Δt = clock.Δt
 
-    reinit = get(config.model, "reinit", true)::Bool
     do_reservoirs = get(config.model, "reservoirs", false)::Bool
     do_lakes = get(config.model, "lakes", false)::Bool
     do_pits = get(config.model, "pits", false)::Bool
@@ -381,33 +380,7 @@ function initialize_hbv_model(config::Config)
         HbvModel(),
     )
 
-    # read and set states in model object if reinit=true
-    if reinit == false
-        instate_path = input_path(config, config.state.path_input)
-        @info "Set initial conditions from state file `$instate_path`."
-        state_ncnames = ncnames(config.state)
-        set_states(instate_path, model, state_ncnames; type = Float)
-        # update kinematic wave volume for river and land domain
-        @unpack lateral = model
-        # makes sure land cells with zero flow width are set to zero q and h
-        for i in eachindex(lateral.land.width)
-            if lateral.land.width[i] <= 0.0
-                lateral.land.q[i] = 0.0
-                lateral.land.h[i] = 0.0
-            end
-        end
-        lateral.land.volume .= lateral.land.h .* lateral.land.width .* lateral.land.dl
-        lateral.river.volume .= lateral.river.h .* lateral.river.width .* lateral.river.dl
-
-        if do_lakes
-            # storage must be re-initialized after loading the state with the current
-            # waterlevel otherwise the storage will be based on the initial water level
-            lakes.storage .=
-                initialize_storage(lakes.storfunc, lakes.area, lakes.waterlevel, lakes.sh)
-        end
-    else
-        @info "Set initial conditions from default values."
-    end
+    model = set_states(model)
 
     return model
 end
@@ -438,3 +411,40 @@ function update(model::Model{N,L,V,R,W,T}) where {N,L,V,R,W,T<:HbvModel}
 
     return model
 end
+
+function set_states(model::Model{N,L,V,R,W,T}) where {N,L,V,R,W,T<:HbvModel}
+    @unpack lateral, config = model
+
+    reinit = get(config.model, "reinit", true)::Bool
+    do_lakes = get(config.model, "lakes", false)::Bool
+    # read and set states in model object if reinit=true
+    if reinit == false
+        instate_path = input_path(config, config.state.path_input)
+        @info "Set initial conditions from state file `$instate_path`."
+        state_ncnames = ncnames(config.state)
+        set_states(instate_path, model, state_ncnames; type = Float)
+        # update kinematic wave volume for river and land domain
+        @unpack lateral = model
+        # makes sure land cells with zero flow width are set to zero q and h
+        for i in eachindex(lateral.land.width)
+            if lateral.land.width[i] <= 0.0
+                lateral.land.q[i] = 0.0
+                lateral.land.h[i] = 0.0
+            end
+        end
+        lateral.land.volume .= lateral.land.h .* lateral.land.width .* lateral.land.dl
+        lateral.river.volume .= lateral.river.h .* lateral.river.width .* lateral.river.dl
+
+        if do_lakes
+            # storage must be re-initialized after loading the state with the current
+            # waterlevel otherwise the storage will be based on the initial water level
+            lakes = lateral.river.lake
+            lakes.storage .=
+                initialize_storage(lakes.storfunc, lakes.area, lakes.waterlevel, lakes.sh)
+        end
+    else
+        @info "Set initial conditions from default values."
+    end
+
+    return model
+end

src/sbm_gwf_model.jl

@@ -22,7 +22,6 @@ function initialize_sbm_gwf_model(config::Config)
     clock = Clock(config, reader)
     Δt = clock.Δt
 
-    reinit = get(config.model, "reinit", true)::Bool
     do_reservoirs = get(config.model, "reservoirs", false)::Bool
     do_lakes = get(config.model, "lakes", false)::Bool
     do_drains = get(config.model, "drains", false)::Bool
@@ -384,33 +383,7 @@ function initialize_sbm_gwf_model(config::Config)
         SbmGwfModel(),
     )
 
-    # read and set states in model object if reinit=false
-    if reinit == false
-        instate_path = input_path(config, config.state.path_input)
-        @info "Set initial conditions from state file `$instate_path`."
-        state_ncnames = ncnames(config.state)
-        set_states(instate_path, model, state_ncnames, type = Float, dimname = :layer)
-        # update kinematic wave volume for river and land domain
-        @unpack lateral = model
-        # makes sure land cells with zero flow width are set to zero q and h
-        for i in eachindex(lateral.land.width)
-            if lateral.land.width[i] <= 0.0
-                lateral.land.q[i] = 0.0
-                lateral.land.h[i] = 0.0
-            end
-        end
-        lateral.land.volume .= lateral.land.h .* lateral.land.width .* lateral.land.dl
-        lateral.river.volume .= lateral.river.h .* lateral.river.width .* lateral.river.dl
-
-        if do_lakes
-            # storage must be re-initialized after loading the state with the current
-            # waterlevel otherwise the storage will be based on the initial water level
-            lakes.storage .=
-                initialize_storage(lakes.storfunc, lakes.area, lakes.waterlevel, lakes.sh)
-        end
-    else
-        @info "Set initial conditions from default values."
-    end
+    model = set_states(model)
 
     return model
 end
@@ -488,3 +461,39 @@ function update(model::Model{N,L,V,R,W,T}) where {N,L,V,R,W,T<:SbmGwfModel}
 
     return model
 end
+
+function set_states(model::Model{N,L,V,R,W,T}) where {N,L,V,R,W,T<:SbmGwfModel}
+    @unpack lateral, config = model
+
+    reinit = get(config.model, "reinit", true)::Bool
+    do_lakes = get(config.model, "lakes", false)::Bool
+    # read and set states in model object if reinit=false
+    if reinit == false
+        instate_path = input_path(config, config.state.path_input)
+        @info "Set initial conditions from state file `$instate_path`."
+        state_ncnames = ncnames(config.state)
+        set_states(instate_path, model, state_ncnames, type = Float, dimname = :layer)
+        # update kinematic wave volume for river and land domain
+        @unpack lateral = model
+        # makes sure land cells with zero flow width are set to zero q and h
+        for i in eachindex(lateral.land.width)
+            if lateral.land.width[i] <= 0.0
+                lateral.land.q[i] = 0.0
+                lateral.land.h[i] = 0.0
+            end
+        end
+        lateral.land.volume .= lateral.land.h .* lateral.land.width .* lateral.land.dl
+        lateral.river.volume .= lateral.river.h .* lateral.river.width .* lateral.river.dl
+
+        if do_lakes
+            # storage must be re-initialized after loading the state with the current
+            # waterlevel otherwise the storage will be based on the initial water level
+            lakes = lateral.river.lake
+            lakes.storage .=
+                initialize_storage(lakes.storfunc, lakes.area, lakes.waterlevel, lakes.sh)
+        end
+    else
+        @info "Set initial conditions from default values."
+    end
+    return model
+end