Repository: Deltares/Wflow.jl (language: python)

Issue:
BMI get_value and grid type
### Wflow version checks

- [X] I have checked that this issue has not already been reported.

- [X] I have checked that this bug exists on the latest version of Wflow.


### Reproducible Example

Request of grid type seems wrong for some variables. As an example, for `lateral.subsurface.to_river`: request of grid type is incorrect. With BMI.get_var_grid, grid=2 is returned. However, the Vector that is returned with BMI.get_value has the size of grid=4.

### Current behaviour

BMI.get_value deviates from the BasicModelInterface.jl implementation (`dest` argument is missing).
Grid type for some variables seems wrong (see also Reproducible Example).
BMI.get_value should always return a flattened array, this is not the case for variables with an extra dimension (e.g. `layer`), then a Vector of SVectors is returned.


### Desired behaviour

BMI.get_value should have `dest` as argument.
Grid type should be correct for each variable. This has already been resolved in the zmq_server branch (coupling with OpenDA), but good to include in a separate branch that addresses this issue.
BMI.get_value should always return a flattened array (probably need to make use of get_grid_z).

### Additional Context

These BMI issues were reported during an update of eWaterCycle to support Wflow.jl.
`bmi.get_current_time` and `bmi.get_end_time` are broken.
### Wflow version checks

- [X] I have checked that this issue has not already been reported.

- [X] I have checked that this bug exists on the latest version of Wflow.


### Reproducible Example

After bmi model initialization:
```jl
BMI.get_current_time(model)
BMI.get_end_time(model)
```

### Current behaviour

Gives `MethodError: no method matching value -(::String)` (or similar)

### Desired behaviour

Gives the end/current time.

### Additional Context

The subtractions don't work because `model.config.starttime` and `model.config.endtime` are strings:

```jl
function BMI.get_current_time(model::Model)
    0.001 * Dates.value(model.clock.time - model.config.starttime)
end
```

```jl
function BMI.get_end_time(model::Model)
    0.001 * Dates.value(model.config.endtime - model.config.starttime)
end
```


Discussion / hints:
Thanks for tackling these issues.


I have found one more (possible) deviation from BMI in Wflow.jl: `get_var_type` usually returns the type of the elements of the array/vector. 

In the bmi-example-julia this is implemented as:
```julia
repr(eltype(m.temperature))
```
However, in Wflow.jl this is implemented as
```julia
string(typeof(param(model, name)))
```
So instead of `Float64`, the returned string will be `Vector{Float64}`.
Thanks for spotting this! 
Yes, this should be `eltype` instead of `typeof` in Wflow.jl.


Fix patch:
diff --git a/docs/src/changelog.md b/docs/src/changelog.md
index 91d02e4a6..5b4ccfaeb 100644
--- a/docs/src/changelog.md
+++ b/docs/src/changelog.md
@@ -11,11 +11,25 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Documentation: add leakage term to the wflow\_sbm figure, document external input
   parameter `ksathorfrac` and fix description of adding external `inflow` to the kinematic
   wave.
+- Fixed BMI functions (e.g. `BMI.get_value`) that deviated from BMI specifications
+  (BasicModelInterface.jl), including function arguments, return types and the BMI
+  specification that arrays are always flattened (this was not the case for variables stored
+  as 2-dimensional arrays or as vector of SVectors).
 - Bump compat for NCDatasets to 0.13.
 
 ### Changed
 - For cyclic parameters different cyclic time inputs are supported (only one common cyclic
   time (for example daily or monthly) was allowed).
+- BMI: 1) added grid information (type and location) and whether a variable can be exchanged
+  to metadata Structs, 2) extend model grid functions for Wflow components that store
+  variables on `edges` (local inertial model) with `get_grid_edge_count` and
+  `get_grid_edge_nodes`.
+
+### Added
+- Functions for loading and saving states and getting the `starttime` in Unix time format.
+  This is convenient for coupling with OpenDA (and other external) software. The set states
+  functionality from the initialization function has been moved to a separate `set_states`
+  function for each `Model` type, to support loading states independent of initialization.
 
 ## v0.7.2 - 2023-09-27
 
diff --git a/docs/src/user_guide/additional_options.md b/docs/src/user_guide/additional_options.md
index 7d66f29bb..7862a5f7c 100644
--- a/docs/src/user_guide/additional_options.md
+++ b/docs/src/user_guide/additional_options.md
@@ -191,10 +191,10 @@ in the [Julia programming language](https://julialang.org/), makes use of the fo
 version.
 
 For the BMI implementation of Wflow all grids are defined as [unstructured
-grids](https://bmi-spec.readthedocs.io/en/latest/model_grids.html#unstructured-grids). While
-the input (forcing and model parameters) is structured (uniform rectilinear), internally
-wflow works with one dimensional arrays based on the active grid cells of the 2D model
-domain.
+grids](https://bmi-spec.readthedocs.io/en/latest/model_grids.html#unstructured-grids),
+including the special cases `scalar` and `points`. While the input (forcing and model
+parameters) is structured (uniform rectilinear), internally wflow works with one dimensional
+arrays based on the active grid cells of the 2D model domain.
 
 ### Configuration
 The variables that Wflow can exchange through BMI are based on the different model
@@ -219,6 +219,12 @@ Wflow.BMI.initialize
 Wflow.BMI.get_input_var_names
 ```
 
+Variables with a third dimension, for example `layer` as part of the vertical `SBM` concept,
+are exposed as two-dimensional grids through the Wflow BMI implementation. For these
+variables the index of this third dimension is required, by adding `[k]` to the variable
+name (`k` refers to the index of the third dimension). For example, the variable
+`vertical.vwc[1]` refers to the first soil layer of the vertical `SBM` concept.
+
 ### Couple to a groundwater model
 For the coupling of wflow\_sbm (SBM + kinematic wave) with a groundwater model (e.g.
 MODFLOW) it is possible to run:
diff --git a/src/Wflow.jl b/src/Wflow.jl
index 30ef7ad12..3bf4b03d8 100644
--- a/src/Wflow.jl
+++ b/src/Wflow.jl
@@ -22,7 +22,11 @@ using Polyester
 using LoopVectorization
 using IfElse
 
-@metadata get_units "mm Δt-1"
+@metadata get_units "mm Δt-1" String
+# metadata for BMI grid
+@metadata exchange 1 Integer
+@metadata grid_type "unstructured" String
+@metadata grid_location "node" String
 
 const BMI = BasicModelInterface
 const Float = Float64
diff --git a/src/bmi.jl b/src/bmi.jl
index 5fb80fd8f..3f0c71979 100644
--- a/src/bmi.jl
+++ b/src/bmi.jl
@@ -1,19 +1,17 @@
 # Basic Model Interface (BMI) implementation based on
 # https://github.com/Deltares/BasicModelInterface.jl
 
-# BMI grid type based on grid identifier
-const gridtype = Dict{Int,String}(
-    0 => "unstructured",
-    1 => "unstructured",
-    2 => "unstructured",
-    3 => "unstructured",
-    4 => "unstructured",
-)
-
 # Mapping of grid identifier to a key, to get the active indices of the model domain.
 # See also function active_indices(network, key::Tuple).
-const grids =
-    Dict{Int,String}(0 => "reservoir", 1 => "lake", 2 => "river", 3 => "drain", 4 => "land")
+const grids = Dict{Int,Tuple{Symbol}}(
+    1 => (:reservoir,),
+    2 => (:lake,),
+    3 => (:drain,),
+    4 => (:river,),
+    5 => (:land,),
+    6 => (:land,),
+    7 => (:land,),
+)
 
 """
     BMI.initialize(::Type{<:Wflow.Model}, config_file)
@@ -33,6 +31,8 @@ function BMI.initialize(::Type{<:Model}, config_file)
         initialize_hbv_model(config)
     elseif modeltype == "sediment"
         initialize_sediment_model(config)
+    elseif modeltype == "flextopo"
+        initialize_flextopo_model(config)
     else
         error("unknown model type")
     end
@@ -106,10 +106,25 @@ function BMI.get_input_var_names(model::Model)
     if haskey(config, "API")
         var_names = Vector{String}()
         for c in config.API.components
-            append!(
-                var_names,
-                collect(string.(c, ".", fieldnames(typeof(param(model, c))))),
-            )
+            type = typeof(param(model, c))
+            inds = findall(x -> x != 0, exchange(type))
+            field_names = fieldnames(type)[inds]
+
+            for name in field_names
+                var = string(c, ".", name)
+                model_var = param(model, var)
+                if eltype(model_var) <: SVector
+                    for i = 1:length(first(model_var))
+                        push!(var_names, string(var, "[", i, "]"))
+                    end
+                elseif ndims(model_var) > 1
+                    for i = 1:length(first(model_var))
+                        push!(var_names, string(var, "[", i, "]"))
+                    end
+                else
+                    push!(var_names, var)
+                end
+            end
         end
         return var_names
     else
@@ -125,50 +140,70 @@ function BMI.get_output_var_names(model::Model)
 end
 
 function BMI.get_var_grid(model::Model, name::String)
-    if name in BMI.get_input_var_names(model)
-        if occursin("reservoir", name)
+    s = split(name, "[")
+    key = symbols(first(s))
+    if exchange(param(model, key[1:end-1]), key[end]) == 1
+        gridtype = grid_type(param(model, key))
+        type = typeof(param(model, key[1:end-1]))
+        return if gridtype == "scalar"
             0
-        elseif occursin("lake", name)
+        elseif :reservoir in key
             1
-        elseif occursin("river", name)
+        elseif :lake in key
             2
-        elseif occursin("drain", name)
+        elseif :drain in key
             3
-        else
+        elseif :river in key
             4
+        elseif type <: ShallowWaterLand && occursin("x", s[end])
+            5
+        elseif type <: ShallowWaterLand && occursin("y", s[end])
+            6
+        else
+            7
         end
     else
-        error("Model variable $name not listed as input of output variable")
+        error("$name not listed as variable for BMI exchange")
     end
 end
 
 function BMI.get_var_type(model::Model, name::String)
-    string(typeof(param(model, name)))
+    value = BMI.get_value_ptr(model, name)
+    repr(eltype(first(value)))
 end
 
 function BMI.get_var_units(model::Model, name::String)
-    s = split(name, ".")
-    get_units(param(model, join(s[1:end-1], ".")), Symbol(s[end]))
+    key = symbols(first(split(name, "[")))
+    if exchange(param(model, key[1:end-1]), key[end]) == 1
+        get_units(param(model, key[1:end-1]), key[end])
+    else
+        error("$name not listed as variable for BMI exchange")
+    end
 end
 
 function BMI.get_var_itemsize(model::Model, name::String)
-    sizeof(param(model, name)[1])
+    value = BMI.get_value_ptr(model, name)
+    sizeof(eltype(first(value)))
 end
 
 function BMI.get_var_nbytes(model::Model, name::String)
-    sizeof(param(model, name))
+    sizeof(BMI.get_value_ptr(model, name))
 end
 
 function BMI.get_var_location(model::Model, name::String)
-    if name in BMI.get_input_var_names(model)
-        return "node"
+    key = symbols(first(split(name, "[")))
+    if exchange(param(model, key[1:end-1]), key[end]) == 1
+        return grid_location(model, key)
     else
-        error("$name not in get_input_var_names")
+        error("$name not listed as variable for BMI exchange")
     end
 end
 
 function BMI.get_current_time(model::Model)
-    0.001 * Dates.value(model.clock.time - model.config.starttime)
+    @unpack config = model
+    calendar = get(config, "calendar", "standard")::String
+    starttime = cftime(config.starttime, calendar)
+    return 0.001 * Dates.value(model.clock.time - starttime)
 end
 
 function BMI.get_start_time(model::Model)
@@ -176,7 +211,11 @@ function BMI.get_start_time(model::Model)
 end
 
 function BMI.get_end_time(model::Model)
-    0.001 * Dates.value(model.config.endtime - model.config.starttime)
+    @unpack config = model
+    calendar = get(config, "calendar", "standard")::String
+    starttime = cftime(config.starttime, calendar)
+    endtime = cftime(config.endtime, calendar)
+    return 0.001 * Dates.value(endtime - starttime)
 end
 
 function BMI.get_time_units(model::Model)
@@ -187,25 +226,45 @@ function BMI.get_time_step(model::Model)
     Float64(model.config.timestepsecs)
 end
 
-function BMI.get_value(model::Model, name::String)
-    copy(BMI.get_value_ptr(model, name))
+function BMI.get_value(model::Model, name::String, dest::Vector{T}) where {T<:AbstractFloat}
+    dest .= copy(BMI.get_value_ptr(model, name))
+    return dest
 end
 
 function BMI.get_value_ptr(model::Model, name::String)
-    if name in BMI.get_input_var_names(model)
-        return param(model, name)
+    @unpack network = model
+    s = split(name, "[")
+    key = symbols(first(s))
+    if exchange(param(model, key[1:end-1]), key[end]) == 1
+        n = length(active_indices(network, key))
+        if occursin("[", name)
+            ind = tryparse(Int, split(s[end], "]")[1])
+            if eltype(param(model, key)) <: SVector
+                model_vals = param(model, key)
+                el_type = eltype(first(model_vals))
+                dim = length(first(model_vals))
+                value = reshape(reinterpret(el_type, model_vals), dim, :)
+                return @view value[ind, 1:n]
+            else
+                value = @view param(model, key)[ind, 1:n]
+                return value
+            end
+        else
+            return @view(param(model, key)[1:n])
+        end
     else
-        error("$name not defined as an output BMI variable")
+        error("$name not listed as variable for BMI exchange")
     end
 end
 
-"""
-    BMI.get_value_at_indices(model::Model, name::String, inds::Vector{Int})
-
-Returns values of a model variable `name` at indices `inds`.
-"""
-function BMI.get_value_at_indices(model::Model, name::String, inds::Vector{Int})
-    BMI.get_value_ptr(model, name)[inds]
+function BMI.get_value_at_indices(
+    model::Model,
+    name::String,
+    dest::Vector{T},
+    inds::Vector{Int},
+) where {T<:AbstractFloat}
+    dest .= BMI.get_value_ptr(model, name)[inds]
+    return dest
 end
 
 """
@@ -234,43 +293,111 @@ function BMI.set_value_at_indices(
 end
 
 function BMI.get_grid_type(model::Model, grid::Int)
-    gridtype[grid]
+    if grid == 0
+        "scalar"
+    elseif grid in 1:3
+        "points"
+    elseif grid in 4:7
+        "unstructured"
+    else
+        error("unknown grid type $grid")
+    end
 end
 
 function BMI.get_grid_rank(model::Model, grid::Int)
-    if grid in keys(gridtype)
-        return 2
+    if grid == 0
+        0
+    elseif grid in 1:7
+        2
     else
         error("unknown grid type $grid")
     end
 end
 
-function BMI.get_grid_shape(model::Model, grid::Int)
-    size(active_indices(model.network, symbols(grids[grid])), 1)
-end
-
-function BMI.get_grid_size(model::Model, grid::Int)
-    BMI.get_grid_shape(model, grid)
-end
-
-function BMI.get_grid_x(model::Model, grid::Int)
+function BMI.get_grid_x(model::Model, grid::Int, x::Vector{T}) where {T<:AbstractFloat}
     @unpack reader, config = model
     @unpack dataset = reader
-    sel = active_indices(model.network, symbols(grids[grid]))
+    sel = active_indices(model.network, grids[grid])
     inds = [sel[i][1] for i in eachindex(sel)]
     x_nc = read_x_axis(dataset)
-    return x_nc[inds]
+    x .= x_nc[inds]
+    return x
 end
 
-function BMI.get_grid_y(model::Model, grid::Int)
+function BMI.get_grid_y(model::Model, grid::Int, y::Vector{T}) where {T<:AbstractFloat}
     @unpack reader, config = model
     @unpack dataset = reader
-    sel = active_indices(model.network, symbols(grids[grid]))
+    sel = active_indices(model.network, grids[grid])
     inds = [sel[i][2] for i in eachindex(sel)]
     y_nc = read_y_axis(dataset)
-    return y_nc[inds]
+    y .= y_nc[inds]
+    return y
 end
 
 function BMI.get_grid_node_count(model::Model, grid::Int)
-    BMI.get_grid_size(model, grid)
+    return length(active_indices(model.network, grids[grid]))
+end
+
+function BMI.get_grid_edge_count(model::Model, grid::Int)
+    @unpack network = model
+    if grid == 4
+        return ne(network.river.graph)
+    elseif grid == 5
+        return length(network.land.staggered_indices.xu)
+    elseif grid == 6
+        return length(network.land.staggered_indices.yu)
+    elseif grid in 0:3 || grid == 7
+        warn("edges are not provided for grid type $grid (variables are located at nodes)")
+    else
+        error("unknown grid type $grid")
+    end
+end
+
+function BMI.get_grid_edge_nodes(model::Model, grid::Int, edge_nodes::Vector{Int})
+    @unpack network = model
+    n = length(edge_nodes)
+    m = div(n, 2)
+    # inactive nodes (boundary/ghost points) are set at -999
+    if grid == 4
+        nodes_at_edge = adjacent_nodes_at_link(network.river.graph)
+        nodes_at_edge.dst[nodes_at_edge.dst.==m+1] .= -999
+        edge_nodes[range(1, n, step = 2)] = nodes_at_edge.src
+        edge_nodes[range(2, n, step = 2)] = nodes_at_edge.dst
+        return edge_nodes
+    elseif grid == 5
+        xu = network.land.staggered_indices.xu
+        edge_nodes[range(1, n, step = 2)] = 1:m
+        xu[xu.==m+1] .= -999
+        edge_nodes[range(2, n, step = 2)] = xu
+        return edge_nodes
+    elseif grid == 6
+        yu = network.land.staggered_indices.yu
+        edge_nodes[range(1, n, step = 2)] = 1:m
+        yu[yu.==m+1] .= -999
+        edge_nodes[range(2, n, step = 2)] = yu
+        return edge_nodes
+    elseif grid in 0:3 || grid == 7
+        @warn("edges are not provided for grid type $grid (variables are located at nodes)")
+    else
+        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
diff --git a/src/flextopo.jl b/src/flextopo.jl
index 466213d04..d02d79713 100644
--- a/src/flextopo.jl
+++ b/src/flextopo.jl
@@ -1,22 +1,22 @@
-@get_units @with_kw struct FLEXTOPO{T,N}
+@get_units @exchange @grid_type @grid_location @with_kw struct FLEXTOPO{T,N}
     # Model time step [s]
-    Δt::T | "s"
+    Δt::T | "s" | 0 | "none" | "none"
     # Number of classes
-    nclass::Int | "-"
+    nclass::Int | "-" | 0 | "none" | "none"
     # Number of cells
-    n::Int | "-"
+    n::Int | "-" | 0 | "none" | "none"
     #dictionary with all possible functions for each store
-    dic_function::Dict
+    dic_function::Dict |"-"| 0 | "none" | "none"
     #current class
-    kclass::Vector{Int64}
-    classes::Vector{String}
-    select_snow::Vector{String}
-    select_interception::Vector{String}
-    select_hortonponding::Vector{String}
-    select_hortonrunoff::Vector{String}
-    select_rootzone::Vector{String}
-    select_fast::Vector{String}
-    select_slow::Vector{String}
+    kclass::Vector{Int64} |"-"| 0 | "none" | "none"
+    classes::Vector{String} |"-"| 0 | "none" | "none"
+    select_snow::Vector{String} |"-"| 0 | "none" | "none"
+    select_interception::Vector{String} |"-"| 0 | "none" | "none"
+    select_hortonponding::Vector{String} |"-"| 0 | "none" | "none"
+    select_hortonrunoff::Vector{String} |"-"| 0 | "none" | "none"
+    select_rootzone::Vector{String} |"-"| 0 | "none" | "none"
+    select_fast::Vector{String} |"-"| 0 | "none" | "none"
+    select_slow::Vector{String} |"-"| 0 | "none" | "none"
 
     #fraction of each class
     hrufrac::Vector{SVector{N,T}} | "-"
diff --git a/src/flextopo_model.jl b/src/flextopo_model.jl
index 60db2233a..048359e7e 100644
--- a/src/flextopo_model.jl
+++ b/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
\ No newline at end of file
diff --git a/src/flow.jl b/src/flow.jl
index bb66feb5f..4aaa05e1c 100644
--- a/src/flow.jl
+++ b/src/flow.jl
@@ -1,33 +1,33 @@
 
 abstract type SurfaceFlow end
 
-@get_units @with_kw struct SurfaceFlowRiver{T,R,L} <: SurfaceFlow
-    β::T | "-"                              # constant in Manning's equation
-    sl::Vector{T} | "m m-1"                 # Slope [m m⁻¹]
-    n::Vector{T} | "s m-1/3"                # Manning's roughness [s m⁻⅓]
-    dl::Vector{T} | "m"                     # Drain length [m]
-    q::Vector{T} | "m3 s-1"                 # Discharge [m³ s⁻¹]
-    qin::Vector{T} | "m3 s-1"               # Inflow from upstream cells [m³ s⁻¹]
-    q_av::Vector{T} | "m3 s-1"              # Average discharge [m³ s⁻¹]
-    qlat::Vector{T} | "m2 s-1"              # Lateral inflow per unit length [m² s⁻¹]
-    inwater::Vector{T} | "m3 s-1"           # Lateral inflow [m³ s⁻¹]
-    inflow::Vector{T} | "m3 s-1"            # External inflow (abstraction/supply/demand) [m³ s⁻¹]
-    volume::Vector{T} | "m3"                # Kinematic wave volume [m³] (based on water level h)
-    h::Vector{T} | "m"                      # Water level [m]
-    h_av::Vector{T} | "m"                   # Average water level [m]
-    bankfull_depth::Vector{T} | "m"         # Bankfull water level [m]
-    Δt::T | "s"                             # Model time step [s]
-    its::Int | "-"                          # Number of fixed iterations
-    width::Vector{T} | "m"                  # Flow width [m]
-    alpha_pow::T | "-"                      # Used in the power part of α
-    alpha_term::Vector{T} | "-"             # Term used in computation of α
-    α::Vector{T} | "s3/5 m1/5"              # Constant in momentum equation A = αQᵝ, based on Manning's equation
-    cel::Vector{T} | "m s-1"                # Celerity of the kinematic wave
-    reservoir_index::Vector{Int} | "-"      # map cell to 0 (no reservoir) or i (pick reservoir i in reservoir field)
-    lake_index::Vector{Int} | "-"           # map cell to 0 (no lake) or i (pick lake i in lake field)
-    reservoir::R                            # Reservoir model struct of arrays
-    lake::L                                 # Lake model struct of arrays
-    kinwave_it::Bool                        # Boolean for iterations kinematic wave
+@get_units @exchange @grid_type @grid_location @with_kw struct SurfaceFlowRiver{T,R,L} <: SurfaceFlow
+    β::T | "-" | _ | "scalar"                    # constant in Manning's equation
+    sl::Vector{T} | "m m-1"                      # Slope [m m⁻¹]
+    n::Vector{T} | "s m-1/3"                     # Manning's roughness [s m⁻⅓]
+    dl::Vector{T} | "m"                          # Drain length [m]
+    q::Vector{T} | "m3 s-1"                      # Discharge [m³ s⁻¹]
+    qin::Vector{T} | "m3 s-1"                    # Inflow from upstream cells [m³ s⁻¹]
+    q_av::Vector{T} | "m3 s-1"                   # Average discharge [m³ s⁻¹]
+    qlat::Vector{T} | "m2 s-1"                   # Lateral inflow per unit length [m² s⁻¹]
+    inwater::Vector{T} | "m3 s-1"                # Lateral inflow [m³ s⁻¹]
+    inflow::Vector{T} | "m3 s-1"                 # External inflow (abstraction/supply/demand) [m³ s⁻¹]
+    volume::Vector{T} | "m3"                     # Kinematic wave volume [m³] (based on water level h)
+    h::Vector{T} | "m"                           # Water level [m]
+    h_av::Vector{T} | "m"                        # Average water level [m]
+    bankfull_depth::Vector{T} | "m"              # Bankfull water level [m]
+    Δt::T | "s" | 0 | "none" | "none"            # Model time step [s]
+    its::Int | "-" | 0 | "none" | "none"         # Number of fixed iterations
+    width::Vector{T} | "m"                       # Flow width [m]
+    alpha_pow::T | "-" | _ | "scalar"            # Used in the power part of α
+    alpha_term::Vector{T} | "-"                  # Term used in computation of α
+    α::Vector{T} | "s3/5 m1/5"                   # Constant in momentum equation A = αQᵝ, based on Manning's equation
+    cel::Vector{T} | "m s-1"                     # Celerity of the kinematic wave
+    reservoir_index::Vector{Int} | "-"           # map cell to 0 (no reservoir) or i (pick reservoir i in reservoir field)
+    lake_index::Vector{Int} | "-"                # map cell to 0 (no lake) or i (pick lake i in lake field)
+    reservoir::R | "-" | 0                       # Reservoir model struct of arrays
+    lake::L | "-" | 0                            # Lake model struct of arrays
+    kinwave_it::Bool | "-" | 0 | "none" | "none" # Boolean for iterations kinematic wave
 
     # TODO unclear why this causes a MethodError
     # function SurfaceFlow{T,R,L}(args...) where {T,R,L}
@@ -36,28 +36,28 @@ abstract type SurfaceFlow end
     # end
 end
 
-@get_units @with_kw struct SurfaceFlowLand{T} <: SurfaceFlow
-    β::T | "-"                              # constant in Manning's equation
-    sl::Vector{T} | "m m-1"                 # Slope [m m⁻¹]
-    n::Vector{T} | "s m-1/3"                # Manning's roughness [s m⁻⅓]
-    dl::Vector{T} | "m"                     # Drain length [m]
-    q::Vector{T} | "m3 s-1"                 # Discharge [m³ s⁻¹]
-    qin::Vector{T} | "m3 s-1"               # Inflow from upstream cells [m³ s⁻¹]
-    q_av::Vector{T} | "m3 s-1"              # Average discharge [m³ s⁻¹]
-    qlat::Vector{T} | "m2 s-1"              # Lateral inflow per unit length [m² s⁻¹]
-    inwater::Vector{T} | "m3 s-1"           # Lateral inflow [m³ s⁻¹]
-    volume::Vector{T} | "m3"                # Kinematic wave volume [m³] (based on water level h)
-    h::Vector{T} | "m"                      # Water level [m]
-    h_av::Vector{T} | "m"                   # Average water level [m]
-    Δt::T | "s"                             # Model time step [s]
-    its::Int | "-"                          # Number of fixed iterations
-    width::Vector{T} | "m"                  # Flow width [m]
-    alpha_pow::T | "-"                      # Used in the power part of α
-    alpha_term::Vector{T} | "-"             # Term used in computation of α
-    α::Vector{T} | "s3/5 m1/5"              # Constant in momentum equation A = αQᵝ, based on Manning's equation
-    cel::Vector{T} | "m s-1"                # Celerity of the kinematic wave
-    to_river::Vector{T} | "m3 s-1"          # Part of overland flow [m³ s⁻¹] that flows to the river
-    kinwave_it::Bool                        # Boolean for iterations kinematic wave
+@get_units @exchange @grid_type @grid_location @with_kw struct SurfaceFlowLand{T} <: SurfaceFlow
+    β::T | "-"  | _ | "scalar"                      # constant in Manning's equation
+    sl::Vector{T} | "m m-1"                         # Slope [m m⁻¹]
+    n::Vector{T} | "s m-1/3"                        # Manning's roughness [s m⁻⅓]
+    dl::Vector{T} | "m"                             # Drain length [m]
+    q::Vector{T} | "m3 s-1"                         # Discharge [m³ s⁻¹]
+    qin::Vector{T} | "m3 s-1"                       # Inflow from upstream cells [m³ s⁻¹]
+    q_av::Vector{T} | "m3 s-1"                      # Average discharge [m³ s⁻¹]
+    qlat::Vector{T} | "m2 s-1"                      # Lateral inflow per unit length [m² s⁻¹]
+    inwater::Vector{T} | "m3 s-1"                   # Lateral inflow [m³ s⁻¹]
+    volume::Vector{T} | "m3"                        # Kinematic wave volume [m³] (based on water level h)
+    h::Vector{T} | "m"                              # Water level [m]
+    h_av::Vector{T} | "m"                           # Average water level [m]
+    Δt::T | "s" | 0 | "none" | "none"               # Model time step [s]
+    its::Int | "-" | 0 | "none" | "none"            # Number of fixed iterations
+    width::Vector{T} | "m"                          # Flow width [m]
+    alpha_pow::T | "-" | _ | "scalar"               # Used in the power part of α
+    alpha_term::Vector{T} | "-"                     # Term used in computation of α
+    α::Vector{T} | "s3/5 m1/5"                      # Constant in momentum equation A = αQᵝ, based on Manning's equation
+    cel::Vector{T} | "m s-1"                        # Celerity of the kinematic wave
+    to_river::Vector{T} | "m3 s-1"                  # Part of overland flow [m³ s⁻¹] that flows to the river
+    kinwave_it::Bool | "-" | 0 |  "none" | "none"   # Boolean for iterations kinematic wave
 end
 
 function initialize_surfaceflow_land(nc, config, inds; sl, dl, width, iterate, tstep, Δt)
@@ -386,13 +386,13 @@ function stable_timestep(sf::S) where {S<:SurfaceFlow}
     return Δt, its
 end
 
-@get_units @with_kw struct LateralSSF{T}
+@get_units @exchange @grid_type  @grid_location @with_kw struct LateralSSF{T}
     kh₀::Vector{T} | "m d-1"               # Horizontal hydraulic conductivity at soil surface [m d⁻¹]
     f::Vector{T} | "m-1"                   # A scaling parameter [m⁻¹] (controls exponential decline of kh₀)
     soilthickness::Vector{T} | "m"         # Soil thickness [m]
     θₛ::Vector{T} | "-"                     # Saturated water content (porosity) [-]
     θᵣ::Vector{T} | "-"                    # Residual water content [-]
-    Δt::T | "d"                            # model time step [d]
+    Δt::T | "d" | 0 | "none" | "none"      # model time step [d]
     βₗ::Vector{T} | "m m-1"                 # Slope [m m⁻¹]
     dl::Vector{T} | "m"                    # Drain length [m]
     dw::Vector{T} | "m"                    # Flow width [m]
@@ -455,57 +455,57 @@ function update(ssf::LateralSSF, network, frac_toriver)
     end
 end
 
-@get_units @with_kw struct GroundwaterExchange{T}
-    Δt::T | "d"                         # model time step [d]
+@get_units @exchange @grid_type @grid_location @with_kw struct GroundwaterExchange{T}
+    Δt::T | "d" | 0 | "none" | "none"   # model time step [d]
     exfiltwater::Vector{T} | "m Δt-1"   # Exfiltration [m Δt⁻¹]  (groundwater above surface level, saturated excess conditions)
     zi::Vector{T} | "m"                 # Pseudo-water table depth [m] (top of the saturated zone)
     to_river::Vector{T} | "m3 d-1"      # Part of subsurface flow [m³ d⁻¹] that flows to the river
     ssf::Vector{T} | "m3 d-1"           # Subsurface flow [m³ d⁻¹]
 end
 
-@get_units @with_kw struct ShallowWaterRiver{T,R,L,F}
-    n::Int | "-"                            # number of cells
-    ne::Int | "-"                           # number of edges/links
-    active_n::Vector{Int} | "-"             # active nodes
-    active_e::Vector{Int} | "-"             # active edges/links
-    g::T | "m s-2"                          # acceleration due to gravity
-    α::T | "-"                              # stability coefficient (Bates et al., 2010)
-    h_thresh::T | "m"                       # depth threshold for calculating flow
-    Δt::T | "s"                             # model time step [s]
-    q::Vector{T} | "m3 s-1"                 # river discharge (subgrid channel)
-    q0::Vector{T} | "m3 s-1"                # river discharge (subgrid channel) at previous time step
-    q_av::Vector{T} | "m3 s-1"              # average river channel (+ floodplain) discharge [m³ s⁻¹]
-    q_channel_av::Vector{T} | "m3 s-1"      # average river channel discharge [m³ s⁻¹]
-    zb_max::Vector{T} | "m"                 # maximum channel bed elevation
-    mannings_n_sq::Vector{T} | "(s m-1/3)2" # Manning's roughness squared at edge/link
-    mannings_n::Vector{T} | "s m-1/3"       # Manning's roughness at node
-    h::Vector{T} | "m"                      # water depth
-    η_max::Vector{T} | "m"                  # maximum water elevation
-    η_src::Vector{T} | "m"                  # water elevation of source node of edge
-    η_dst::Vector{T} | "m"                  # water elevation of downstream node of edge
-    hf::Vector{T} | "m"                     # water depth at edge/link
-    h_av::Vector{T} | "m"                   # average water depth
-    dl::Vector{T} | "m"                     # river length
-    dl_at_link::Vector{T} | "m"             # river length at edge/link
-    width::Vector{T} | "m"                  # river width
-    width_at_link::Vector{T} | "m"          # river width at edge/link
-    a::Vector{T} | "m2"                     # flow area at edge/link
-    r::Vector{T} | "m"                      # wetted perimeter at edge/link
-    volume::Vector{T} | "m3"                # river volume
-    error::Vector{T} | "m3"                 # error volume
-    inwater::Vector{T} | "m3 s-1"           # lateral inflow [m³ s⁻¹]
-    inflow::Vector{T} | "m3 s-1"            # external inflow (abstraction/supply/demand) [m³ s⁻¹]
-    inflow_wb::Vector{T} | "m3 s-1"         # inflow waterbody (lake or reservoir model) from land part [m³ s⁻¹]
-    bankfull_volume::Vector{T} | "m3"       # bankfull volume
-    bankfull_depth::Vector{T} | "m"         # bankfull depth
-    zb::Vector{T} | "m"                     # river bed elevation
-    froude_limit::Bool | "-"                # if true a check is performed if froude number > 1.0 (algorithm is modified)
-    reservoir_index::Vector{Int} | "-"      # river cell index with a reservoir (each index of reservoir_index maps to reservoir i in reservoir field)
-    lake_index::Vector{Int} | "-"           # river cell index with a lake (each index of lake_index maps to lake i in lake field)
-    waterbody::Vector{Bool} | "-"           # water body cells (reservoir or lake)
-    reservoir::R                            # Reservoir model struct of arrays
-    lake::L                                 # Lake model struct of arrays
-    floodplain::F                           # Floodplain (1D) schematization
+@get_units @exchange @grid_type @grid_location @with_kw struct ShallowWaterRiver{T,R,L,F}
+    n::Int | "-" | 0 | "none" | "none"                      # number of cells
+    ne::Int | "-" | 0 | "none" | "none"                     # number of edges/links
+    active_n::Vector{Int} | "-"                             # active nodes
+    active_e::Vector{Int} | "-" | _ | "edge"                # active edges/links
+    g::T | "m s-2" | _ | "scalar"                           # acceleration due to gravity
+    α::T | "-" | _ | "scalar"                               # stability coefficient (Bates et al., 2010)
+    h_thresh::T | "m" | _ | "scalar"                        # depth threshold for calculating flow
+    Δt::T | "s" | 0 | "none" | "none"                       # model time step [s]
+    q::Vector{T} | "m3 s-1" | _ | "edge"                    # river discharge (subgrid channel)
+    q0::Vector{T} | "m3 s-1"| _ | "edge"                    # river discharge (subgrid channel) at previous time step
+    q_av::Vector{T} | "m3 s-1" | _ | "edge"                 # average river channel (+ floodplain) discharge [m³ s⁻¹]
+    q_channel_av::Vector{T} | "m3 s-1"                      # average river channel discharge [m³ s⁻¹]
+    zb_max::Vector{T} | "m"                                 # maximum channel bed elevation
+    mannings_n_sq::Vector{T} | "(s m-1/3)2" | _ | "edge"    # Manning's roughness squared at edge/link
+    mannings_n::Vector{T} | "s m-1/3"                       # Manning's roughness at node
+    h::Vector{T} | "m"                                      # water depth
+    η_max::Vector{T} | "m" | _ | "edge"                     # maximum water elevation at edge
+    η_src::Vector{T} | "m"                                  # water elevation of source node of edge
+    η_dst::Vector{T} | "m"                                  # water elevation of downstream node of edge
+    hf::Vector{T} | "m" | _ | "edge"                        # water depth at edge/link
+    h_av::Vector{T} | "m"                                   # average water depth
+    dl::Vector{T} | "m"                                     # river length
+    dl_at_link::Vector{T} | "m" | _ | "edge"                # river length at edge/link
+    width::Vector{T} | "m"                      	        # river width
+    width_at_link::Vector{T} | "m" | _ | "edge"             # river width at edge/link
+    a::Vector{T} | "m2" | _ | "edge"                        # flow area at edge/link
+    r::Vector{T} | "m" | _ | "edge"                         # wetted perimeter at edge/link
+    volume::Vector{T} | "m3"                                # river volume
+    error::Vector{T} | "m3"                                 # error volume
+    inwater::Vector{T} | "m3 s-1"                           # lateral inflow [m³ s⁻¹]
+    inflow::Vector{T} | "m3 s-1"                            # external inflow (abstraction/supply/demand) [m³ s⁻¹]
+    inflow_wb::Vector{T} | "m3 s-1"                         # inflow waterbody (lake or reservoir model) from land part [m³ s⁻¹]
+    bankfull_volume::Vector{T} | "m3"                       # bankfull volume
+    bankfull_depth::Vector{T} | "m"                         # bankfull depth
+    zb::Vector{T} | "m"                                     # river bed elevation
+    froude_limit::Bool | "-" | 0 | "none" | "none"          # if true a check is performed if froude number > 1.0 (algorithm is modified)
+    reservoir_index::Vector{Int} | "-"                      # river cell index with a reservoir (each index of reservoir_index maps to reservoir i in reservoir field)
+    lake_index::Vector{Int} | "-"                           # river cell index with a lake (each index of lake_index maps to lake i in lake field)
+    waterbody::Vector{Bool} | "-"                           # water body cells (reservoir or lake)
+    reservoir::R | "-" | 0                                  # Reservoir model struct of arrays
+    lake::L | "-" | 0                                       # Lake model struct of arrays
+    floodplain::F | "-" | 0                                 # Floodplain (1D) schematization
 end
 
 function initialize_shallowwater_river(
@@ -945,32 +945,32 @@ end
 # neigbors.
 const dirs = (:yd, :xd, :xu, :yu)
 
-@get_units @with_kw struct ShallowWaterLand{T}
-    n::Int | "-"                            # number of cells
-    xl::Vector{T} | "m"                     # cell length x direction
-    yl::Vector{T} | "m"                     # cell length y direction
-    xwidth::Vector{T} | "m"                 # effective flow width x direction (floodplain)
-    ywidth::Vector{T} | "m"                 # effective flow width y direction (floodplain)
-    g::T | "m2 s-1"                         # acceleration due to gravity
-    θ::T | "-"                              # weighting factor (de Almeida et al., 2012)
-    α::T | "-"                              # stability coefficient (de Almeida et al., 2012)
-    h_thresh::T | "m"                       # depth threshold for calculating flow
-    Δt::T | "s"                             # model time step [s]
-    qy0::Vector{T} | "m3 s-1"               # flow in y direction at previous time step
-    qx0::Vector{T} | "m3 s-1"               # flow in x direction at previous time step
-    qx::Vector{T} | "m3 s-1"                # flow in x direction
-    qy::Vector{T} | "m3 s-1"                # flow in y direction
-    zx_max::Vector{T} | "m"                 # maximum cell elevation (x direction)
-    zy_max::Vector{T} | "m"                 # maximum cell elevation (y direction)
-    mannings_n_sq::Vector{T} | "(s m-1/3)2" # Manning's roughness squared
-    volume::Vector{T} | "m3"                # total volume of cell (including river volume for river cells)
-    error::Vector{T} | "m3"                 # error volume
-    runoff::Vector{T} | "m3 s-1"            # runoff from hydrological model
-    h::Vector{T} | "m"                      # water depth of cell (for river cells the reference is the river bed elevation `zb`)
-    z::Vector{T} | "m"                      # elevation of cell
-    froude_limit::Bool | "-"                # if true a check is performed if froude number > 1.0 (algorithm is modified)
-    rivercells::Vector{Bool} | "-"          # river cells
-    h_av::Vector{T} | "m"                   # average water depth (for river cells the reference is the river bed elevation `zb`)
+@get_units @exchange @grid_type @grid_location @with_kw struct ShallowWaterLand{T}
+    n::Int | "-" | 0 | "none" | "none"                      # number of cells
+    xl::Vector{T} | "m"                                     # cell length x direction
+    yl::Vector{T} | "m"                                     # cell length y direction
+    xwidth::Vector{T} | "m" | _ | "edge"                    # effective flow width x direction (floodplain)
+    ywidth::Vector{T} | "m" | _ | "edge"                    # effective flow width y direction (floodplain)
+    g::T | "m2 s-1" | _ | "scalar"                          # acceleration due to gravity
+    θ::T | "-" | _ | "scalar"                               # weighting factor (de Almeida et al., 2012)
+    α::T | "-" | _ | "scalar"                               # stability coefficient (de Almeida et al., 2012)
+    h_thresh::T | "m" | _ | "scalar"                        # depth threshold for calculating flow
+    Δt::T | "s" | _ | "scalar"                              # model time step [s]
+    qy0::Vector{T} | "m3 s-1" | _ | "edge"                  # flow in y direction at previous time step
+    qx0::Vector{T} | "m3 s-1" | _ | "edge"                  # flow in x direction at previous time step
+    qx::Vector{T} | "m3 s-1" | _ | "edge"                   # flow in x direction
+    qy::Vector{T} | "m3 s-1" | _ | "edge"                   # flow in y direction
+    zx_max::Vector{T} | "m" | _ | "edge"                    # maximum cell elevation (x direction)
+    zy_max::Vector{T} | "m" | _ | "edge"                    # maximum cell elevation (y direction)
+    mannings_n_sq::Vector{T} | "(s m-1/3)2" | _ | "edge"    # Manning's roughness squared
+    volume::Vector{T} | "m3"                                # total volume of cell (including river volume for river cells)
+    error::Vector{T} | "m3"                                 # error volume
+    runoff::Vector{T} | "m3 s-1"                            # runoff from hydrological model
+    h::Vector{T} | "m"                                      # water depth of cell (for river cells the reference is the river bed elevation `zb`)
+    z::Vector{T} | "m"                                      # elevation of cell
+    froude_limit::Bool | "-" | 0  |"none" | "none"          # if true a check is performed if froude number > 1.0 (algorithm is modified)
+    rivercells::Vector{Bool} | "-"                          # river cells
+    h_av::Vector{T} | "m"                                   # average water depth (for river cells the reference is the river bed elevation `zb`)
 end
 
 function initialize_shallowwater_land(
@@ -1320,30 +1320,30 @@ Floodplain `volume` is a function of `depth` (flood depth intervals). Based on t
 cumulative floodplain `volume` a floodplain profile as a function of `flood_depth` is
 derived with floodplain area `a` (cumulative) and wetted perimeter radius `p` (cumulative).
 """
-@get_units @with_kw struct FloodPlainProfile{T,N}
-    depth::Vector{T} | "m"                    # Flood depth
+@get_units @exchange @grid_type @grid_location @with_kw struct FloodPlainProfile{T,N}
+    depth::Vector{T} | "m" | 0                # Flood depth
     volume::Array{T,2} | "m3"                 # Flood volume (cumulative)
     width::Array{T,2} | "m"                   # Flood width
     a::Array{T,2} | "m2"                      # Flow area (cumulative)
     p::Array{T,2} | "m"                       # Wetted perimeter (cumulative)
 end
 
-@get_units @with_kw struct FloodPlain{T,P}
-    profile::P                               # floodplain profile
-    mannings_n::Vector{T} | "s m-1/3"        # manning's roughness
-    mannings_n_sq::Vector{T} | "(s m-1/3)2"  # manning's roughness squared
-    volume::Vector{T} | "m3"                 # volume
-    h::Vector{T} | "m"                       # water depth
-    h_av::Vector{T} | "m"                    # average water depth
-    error::Vector{T} | "m3"                  # error volume
-    a::Vector{T} | "m2"                      # flow area
-    r::Vector{T} | "m"                       # hydraulic radius
-    hf::Vector{T} | "m"                      # water depth at edge/link
-    zb_max::Vector{T} | "m"                  # maximum bankfull elevation (edge/link)
-    q0::Vector{T} | "m3 s-1"                 # discharge at previous time step
-    q::Vector{T} | "m3 s-1"                  # discharge
-    q_av::Vector{T} | "m"                    # average river discharge
-    hf_index::Vector{Int} | "-"              # index with `hf` above depth threshold
+@get_units @exchange @grid_type @grid_location @with_kw struct FloodPlain{T,P}
+    profile::P | "-" | 0                                    # floodplain profile
+    mannings_n::Vector{T} | "s m-1/3"                       # manning's roughness
+    mannings_n_sq::Vector{T} | "(s m-1/3)2" | _ | "edge"    # manning's roughness squared
+    volume::Vector{T} | "m3"                                # volume
+    h::Vector{T} | "m"                                      # water depth
+    h_av::Vector{T} | "m"                                   # average water depth
+    error::Vector{T} | "m3"                                 # error volume
+    a::Vector{T} | "m2" | _ | "edge"                        # flow area
+    r::Vector{T} | "m" | _ | "edge"                         # hydraulic radius
+    hf::Vector{T} | "m" | _ | "edge"                        # water depth at edge/link
+    zb_max::Vector{T} | "m" | _ | "edge"                    # maximum bankfull elevation (edge/link)
+    q0::Vector{T} | "m3 s-1" | _ | "edge"                   # discharge at previous time step
+    q::Vector{T} | "m3 s-1" | _ | "edge"                    # discharge
+    q_av::Vector{T} | "m" | _ | "edge"                      # average river discharge
+    hf_index::Vector{Int} | "-" | _ | "edge"                # index with `hf` above depth threshold
 end
 
 "Determine the initial floodplain volume"
diff --git a/src/groundwater/aquifer.jl b/src/groundwater/aquifer.jl
index 64117ff56..43aa05438 100644
--- a/src/groundwater/aquifer.jl
+++ b/src/groundwater/aquifer.jl
@@ -86,7 +86,7 @@ NOTA BENE: **specific** storage is per m of aquifer (conf. specific weight).
 **Storativity** or (**storage coefficient**) is for the entire aquifer (conf.
 transmissivity).
 """
-@get_units struct ConfinedAquifer{T} <: Aquifer
+@get_units @exchange @grid_type @grid_location struct ConfinedAquifer{T} <: Aquifer
     head::Vector{T} | "m"  # hydraulic head [m]
     k::Vector{T} | "m d-1"  # horizontal conductivity [m d⁻¹]
     top::Vector{T} | "m" # top of groundwater layer [m]
@@ -109,7 +109,7 @@ aquifer will yield when all water drains and the pore volume is filled by air
 instead. Specific yield will vary roughly between 0.05 (clay) and 0.45 (peat)
 (Johnson, 1967).
 """
-@get_units struct UnconfinedAquifer{T} <: Aquifer
+@get_units @exchange @grid_type @grid_location struct UnconfinedAquifer{T} <: Aquifer
     head::Vector{T} | "m"  # hydraulic head [m]
     k::Vector{T} | "m d-1"  # reference horizontal conductivity [m d⁻¹]
     top::Vector{T} | "m" # top of groundwater layer [m]
@@ -307,7 +307,7 @@ function flux!(Q, aquifer, connectivity, conductivity_profile)
 end
 
 
-@get_units struct ConstantHead{T}
+@get_units @exchange @grid_type @grid_location struct ConstantHead{T}
     head::Vector{T} | "m"
     index::Vector{Int} | "-"
 end
diff --git a/src/groundwater/boundary_conditions.jl b/src/groundwater/boundary_conditions.jl
index 0aae19274..581bc9070 100644
--- a/src/groundwater/boundary_conditions.jl
+++ b/src/groundwater/boundary_conditions.jl
@@ -13,7 +13,7 @@ end
 check_flux(flux, aquifer::ConfinedAquifer, index::Int) = flux
 
 
-@get_units struct River{T} <: AquiferBoundaryCondition
+@get_units @exchange @grid_type @grid_location struct River{T} <: AquiferBoundaryCondition
     stage::Vector{T} | "m"
     infiltration_conductance::Vector{T} | "m2 d-1"
     exfiltration_conductance::Vector{T} | "m2 d-1"
@@ -41,11 +41,11 @@ function flux!(Q, river::River, aquifer)
 end
 
 
-@get_units struct Drainage{T} <: AquiferBoundaryCondition
-    elevation::Vector{T}
-    conductance::Vector{T}
-    flux::Vector{T}
-    index::Vector{Int}
+@get_units @exchange @grid_type @grid_location struct Drainage{T} <: AquiferBoundaryCondition
+    elevation::Vector{T} | "m"
+    conductance::Vector{T} | "m2 d-1"
+    flux::Vector{T} | "m3 d-1"
+    index::Vector{Int} | "-"
 end
 
 
@@ -79,7 +79,7 @@ function flux!(Q, headboundary::HeadBoundary, aquifer)
 end
 
 
-@get_units struct Recharge{T} <: AquiferBoundaryCondition
+@get_units @exchange @grid_type @grid_location struct Recharge{T} <: AquiferBoundaryCondition
     rate::Vector{T} | "m d-1"
     flux::Vector{T} | "m3 d-1"
     index::Vector{Int} | "-"
diff --git a/src/hbv.jl b/src/hbv.jl
index b5b376d62..df69ad78b 100644
--- a/src/hbv.jl
+++ b/src/hbv.jl
@@ -1,30 +1,30 @@
-@get_units @with_kw struct HBV{T}
-    Δt::T | "s"                     # Model time step [s]
-    n::Int | "-"                    # Number of cells
-    fc::Vector{T} | "mm"            # Field capacity [mm]
-    betaseepage::Vector{T} | "-"    # Exponent in soil runoff generation equation [-]
-    lp::Vector{T} | "-"             # Fraction of field capacity below which actual evaporation=potential evaporation [-]
-    threshold::Vector{T} | "mm"     # Threshold soilwater storage above which AE=PE [mm]
-    k4::Vector{T} | "Δt-1"          # Recession constant baseflow [Δt⁻¹]
-    kquickflow::Vector{T} | "Δt-1"  # Recession constant upper reservoir [Δt⁻¹]
-    suz::Vector{T} | "mm"           # Level over which k0 is used [mm]
-    k0::Vector{T} | "Δt-1"          # Recession constant upper reservoir [Δt⁻¹]
-    khq::Vector{T} | "Δt-1"         # Recession rate at flow hq [Δt⁻¹]
-    hq::Vector{T}                   # High flow rate hq for which recession rate of upper reservoir is known [mm Δt⁻¹]
-    alphanl::Vector{T}              # Measure of non-linearity of upper reservoir
-    perc::Vector{T}                 # Percolation from upper to lower zone [mm Δt⁻¹]
-    cfr::Vector{T} | "-"            # Refreezing efficiency constant in refreezing of freewater in snow [-]
-    pcorr::Vector{T} | "-"          # Correction factor for precipitation [-]
-    rfcf::Vector{T} | "-"           # Correction factor for rainfall [-]
-    sfcf::Vector{T} | "-"           # Correction factor for snowfall [-]
-    cflux::Vector{T}                # Maximum capillary rise from runoff response routine to soil moisture routine [mm Δt⁻¹]
-    icf::Vector{T} | "mm"           # Maximum interception storage (in forested and non-forested areas) [mm]
-    cevpf::Vector{T} | "-"          # Correction factor for potential evaporation [-]
-    epf::Vector{T} | "mm-1"         # Exponent of correction factor for evaporation on days with precipitation
-    ecorr::Vector{T} | "-"          # Evap correction [-]
-    tti::Vector{T} | "ᵒC"           # Critical temperature for snowmelt and refreezing [ᵒC]
-    tt::Vector{T} | "ᵒC"            # Defines interval in which precipitation falls as rainfall and snowfall [ᵒC]
-    ttm::Vector{T} | "ᵒC"           # Threshold temperature for snowmelt [ᵒC]
+@get_units @exchange @grid_type @grid_location @with_kw struct HBV{T}
+    Δt::T | "s" | 0 | "none" | "none"   # Model time step [s]
+    n::Int | "-" | 0 | "none" | "none"  # Number of cells
+    fc::Vector{T} | "mm"                # Field capacity [mm]
+    betaseepage::Vector{T} | "-"        # Exponent in soil runoff generation equation [-]
+    lp::Vector{T} | "-"                 # Fraction of field capacity below which actual evaporation=potential evaporation [-]
+    threshold::Vector{T} | "mm"         # Threshold soilwater storage above which AE=PE [mm]
+    k4::Vector{T} | "Δt-1"              # Recession constant baseflow [Δt⁻¹]
+    kquickflow::Vector{T} | "Δt-1"      # Recession constant upper reservoir [Δt⁻¹]
+    suz::Vector{T} | "mm"               # Level over which k0 is used [mm]
+    k0::Vector{T} | "Δt-1"              # Recession constant upper reservoir [Δt⁻¹]
+    khq::Vector{T} | "Δt-1"             # Recession rate at flow hq [Δt⁻¹]
+    hq::Vector{T}                       # High flow rate hq for which recession rate of upper reservoir is known [mm Δt⁻¹]
+    alphanl::Vector{T}                  # Measure of non-linearity of upper reservoir
+    perc::Vector{T}                     # Percolation from upper to lower zone [mm Δt⁻¹]
+    cfr::Vector{T} | "-"                # Refreezing efficiency constant in refreezing of freewater in snow [-]
+    pcorr::Vector{T} | "-"              # Correction factor for precipitation [-]
+    rfcf::Vector{T} | "-"               # Correction factor for rainfall [-]
+    sfcf::Vector{T} | "-"               # Correction factor for snowfall [-]
+    cflux::Vector{T}                    # Maximum capillary rise from runoff response routine to soil moisture routine [mm Δt⁻¹]
+    icf::Vector{T} | "mm"               # Maximum interception storage (in forested and non-forested areas) [mm]
+    cevpf::Vector{T} | "-"              # Correction factor for potential evaporation [-]
+    epf::Vector{T} | "mm-1"             # Exponent of correction factor for evaporation on days with precipitation
+    ecorr::Vector{T} | "-"              # Evap correction [-]
+    tti::Vector{T} | "ᵒC"               # Critical temperature for snowmelt and refreezing [ᵒC]
+    tt::Vector{T} | "ᵒC"                # Defines interval in which precipitation falls as rainfall and snowfall [ᵒC]
+    ttm::Vector{T} | "ᵒC"               # Threshold temperature for snowmelt [ᵒC]
     cfmax::Vector{T} | "mm ᵒC-1 Δt-1"   # Meltconstant in temperature-index [-]
     whc::Vector{T} | "-"                # Fraction of snow volume that can store water [-]
     g_tt::Vector{T} | "ᵒC"              # Threshold temperature for snowfall above glacier [ᵒC]
diff --git a/src/hbv_model.jl b/src/hbv_model.jl
index 3ab03d32c..a03298d71 100644
--- a/src/hbv_model.jl
+++ b/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
\ No newline at end of file
diff --git a/src/reservoir_lake.jl b/src/reservoir_lake.jl
index d599573ee..410745eef 100644
--- a/src/reservoir_lake.jl
+++ b/src/reservoir_lake.jl
@@ -1,5 +1,5 @@
-@get_units @with_kw struct SimpleReservoir{T}
-    Δt::T | "s"                                         # Model time step [s]
+@get_units @exchange @grid_type @grid_location @with_kw struct SimpleReservoir{T}
+    Δt::T | "s" | 0 | "none" | "none"                   # Model time step [s]
     maxvolume::Vector{T} | "m3"                         # maximum storage (above which water is spilled) [m³]
     area::Vector{T} | "m2"                              # reservoir area [m²]
     maxrelease::Vector{T} | "m3 s-1"                    # maximum amount that can be released if below spillway [m³ s⁻¹]
@@ -203,8 +203,8 @@ function update(res::SimpleReservoir, i, inflow, timestepsecs)
     return res
 end
 
-@get_units @with_kw struct Lake{T}
-    Δt::T | "s"                                 # Model time step [s]
+@get_units @exchange @grid_type @grid_location @with_kw struct Lake{T}
+    Δt::T | "s" | 0 | "none" | "none"           # Model time step [s]
     lowerlake_ind::Vector{Int} | "-"            # Index of lower lake (linked lakes)
     area::Vector{T} | "m2"                      # lake area [m²]
     maxstorage::Vector{Union{T,Missing}} | "m3" # lake maximum storage from rating curve 1 [m³]
diff --git a/src/sbm.jl b/src/sbm.jl
index 00631753a..a85254bf7 100644
--- a/src/sbm.jl
+++ b/src/sbm.jl
@@ -1,10 +1,10 @@
-@get_units @with_kw struct SBM{T,N,M}
+@get_units @exchange @grid_type @grid_location @with_kw struct SBM{T,N,M}
     # Model time step [s]
-    Δt::T | "s"
+    Δt::T | "s" | 0 | "none" | "none"
     # Maximum number of soil layers
-    maxlayers::Int | "-"
+    maxlayers::Int | "-" | 0 | "none" | "none"
     # number of cells
-    n::Int | "-"
+    n::Int | "-" | 0 | "none" | "none"
     # Number of soil layers
     nlayers::Vector{Int} | "-"
     # Number of unsaturated soil layers
diff --git a/src/sbm_gwf_model.jl b/src/sbm_gwf_model.jl
index dc54ca03d..f723bc5df 100644
--- a/src/sbm_gwf_model.jl
+++ b/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
\ No newline at end of file
diff --git a/src/sbm_model.jl b/src/sbm_model.jl
index f9274d0dd..907bbd9bc 100644
--- a/src/sbm_model.jl
+++ b/src/sbm_model.jl
@@ -16,7 +16,6 @@ function initialize_sbm_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
@@ -359,81 +358,9 @@ function initialize_sbm_model(config::Config)
         SbmModel(),
     )
 
-    # 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)
-        @warn string(
-            "The unit of `ssf` (lateral subsurface flow) is now m3 d-1. Please update your",
-            " input state file if it was produced with a Wflow version up to v0.5.2.",
-        )
-        set_states(instate_path, model, state_ncnames; type = Float, dimname = :layer)
-        @unpack lateral, vertical, network = model
-        # update zi for vertical sbm and kinematic wave volume for river and land domain
-        zi =
-            max.(
-                0.0,
-                vertical.soilthickness .-
-                vertical.satwaterdepth ./ (vertical.θₛ .- vertical.θᵣ),
-            )
-        vertical.zi .= zi
-        if land_routing == "kinematic-wave"
-            # make 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
-        elseif land_routing == "local-inertial"
-            @warn string(
-                "The reference level for the water depth `h` and `h_av` of overland flow ",
-                "(local inertial model) for cells containing a river has changed from river",
-                " bed elevation `zb` to cell elevation `z`. Please update the input state",
-                " file if it was produced with Wflow version v0.5.2.",
-            )
-            for i = 1:n
-                if lateral.land.rivercells[i]
-                    j = network.land.index_river[i]
-                    if lateral.land.h[i] > 0.0
-                        lateral.land.volume[i] =
-                            lateral.land.h[i] * lateral.land.xl[i] * lateral.land.yl[i] +
-                            lateral.river.bankfull_volume[j]
-                    else
-                        lateral.land.volume[i] =
-                            lateral.river.h[j] *
-                            lateral.river.width[j] *
-                            lateral.river.dl[j]
-                    end
-                else
-                    lateral.land.volume[i] =
-                        lateral.land.h[i] * lateral.land.xl[i] * lateral.land.yl[i]
-                end
-            end
-        end
-        # only set active cells for river (ignore boundary conditions/ghost points)
-        lateral.river.volume[1:nriv] .=
-            lateral.river.h[1:nriv] .* lateral.river.width[1:nriv] .*
-            lateral.river.dl[1:nriv]
-
-        if floodplain_1d
-            initialize_volume!(lateral.river, nriv)
-        end
-
-        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)
 
     @info "Initialized model"
-
     return model
 end
 
@@ -510,3 +437,90 @@ function update_after_subsurfaceflow(
 
     return model
 end
+
+function set_states(model::Model{N,L,V,R,W,T}) where {N,L,V,R,W,T<:SbmModel}
+    @unpack lateral, vertical, network, config = model
+
+    reinit = get(config.model, "reinit", true)::Bool
+    routing_options = ("kinematic-wave", "local-inertial")
+    land_routing =
+        get_options(config.model, "land_routing", routing_options, "kinematic-wave")::String
+    do_lakes = get(config.model, "lakes", false)::Bool
+    floodplain_1d = get(config.model, "floodplain_1d", false)::Bool
+
+    # read and set states in model object if reinit=false
+    if reinit == false
+        nriv = length(network.river.indices)
+        instate_path = input_path(config, config.state.path_input)
+        @info "Set initial conditions from state file `$instate_path`."
+        state_ncnames = ncnames(config.state)
+        @warn string(
+            "The unit of `ssf` (lateral subsurface flow) is now m3 d-1. Please update your",
+            " input state file if it was produced with a Wflow version up to v0.5.2.",
+        )
+        set_states(instate_path, model, state_ncnames; type = Float, dimname = :layer)
+        @unpack lateral, vertical, network = model
+        # update zi for vertical sbm and kinematic wave volume for river and land domain
+        zi =
+            max.(
+                0.0,
+                vertical.soilthickness .-
+                vertical.satwaterdepth ./ (vertical.θₛ .- vertical.θᵣ),
+            )
+        vertical.zi .= zi
+        if land_routing == "kinematic-wave"
+            # make 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
+        elseif land_routing == "local-inertial"
+            @warn string(
+                "The reference level for the water depth `h` and `h_av` of overland flow ",
+                "(local inertial model) for cells containing a river has changed from river",
+                " bed elevation `zb` to cell elevation `z`. Please update the input state",
+                " file if it was produced with Wflow version v0.5.2.",
+            )
+            for i = 1:n
+                if lateral.land.rivercells[i]
+                    j = network.land.index_river[i]
+                    if lateral.land.h[i] > 0.0
+                        lateral.land.volume[i] =
+                            lateral.land.h[i] * lateral.land.xl[i] * lateral.land.yl[i] +
+                            lateral.river.bankfull_volume[j]
+                    else
+                        lateral.land.volume[i] =
+                            lateral.river.h[j] *
+                            lateral.river.width[j] *
+                            lateral.river.dl[j]
+                    end
+                else
+                    lateral.land.volume[i] =
+                        lateral.land.h[i] * lateral.land.xl[i] * lateral.land.yl[i]
+                end
+            end
+        end
+        # only set active cells for river (ignore boundary conditions/ghost points)
+        lateral.river.volume[1:nriv] .=
+            lateral.river.h[1:nriv] .* lateral.river.width[1:nriv] .*
+            lateral.river.dl[1:nriv]
+
+        if floodplain_1d
+            initialize_volume!(lateral.river, nriv)
+        end
+    
+        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
diff --git a/src/sediment.jl b/src/sediment.jl
index 620d0ce63..783c4e987 100644
--- a/src/sediment.jl
+++ b/src/sediment.jl
@@ -1,7 +1,7 @@
 ### Soil erosion ###
-@get_units @with_kw struct LandSediment{T}
+@get_units @exchange @grid_type @grid_location @with_kw struct LandSediment{T}
     # number of cells
-    n::Int | "-"
+    n::Int | "-" | 0 | "none" | "none"
     ### Soil erosion part ###
     # length of cells in y direction [m]
     yl::Vector{T} | "m"
@@ -593,9 +593,9 @@ function tc_yalinpart(ols::LandSediment, i::Int, sinslope::Float, ts::Float)
 end
 
 ### Sediment transport in overland flow ###
-@get_units @with_kw struct OverlandFlowSediment{T}
+@get_units @exchange @grid_type @grid_location @with_kw struct OverlandFlowSediment{T}
     # number of cells
-    n::Int | "-"
+    n::Int | "-" | 0 | "none" | "none"
     # Filter with river cells
     rivcell::Vector{T} | "-"
     # Total eroded soil [ton Δt⁻¹]
@@ -681,11 +681,11 @@ function update(ols::OverlandFlowSediment, network, config)
 end
 
 ### River transport and processes ###
-@get_units @with_kw struct RiverSediment{T}
+@get_units @exchange @grid_type @grid_location @with_kw struct RiverSediment{T}
     # number of cells
-    n::Int | "-"
+    n::Int | "-" | 0 | "none" | "none"
     # Timestep [s]
-    Δt::T | "s"
+    Δt::T | "s" | 0 | "none" | "none"
     # River geometry (slope [-], length [m], width [m])
     sl::Vector{T} | "m"
     dl::Vector{T} | "m"
diff --git a/src/sediment_model.jl b/src/sediment_model.jl
index b823735ed..a1944263b 100644
--- a/src/sediment_model.jl
+++ b/src/sediment_model.jl
@@ -16,8 +16,6 @@ function initialize_sediment_model(config::Config)
     clock = Clock(config, reader)
     Δt = clock.Δt
 
-    reinit = get(config.model, "reinit", true)::Bool
-
     do_river = get(config.model, "runrivermodel", false)::Bool
 
     nc = NCDataset(static_path)
@@ -150,15 +148,7 @@ function initialize_sediment_model(config::Config)
         SedimentModel(),
     )
 
-    # 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)
-    else
-        @info "Set initial conditions from default values."
-    end
+    model = set_states(model)
 
     @info "Initialized model"
     return model
@@ -201,3 +191,18 @@ function update(model::Model{N,L,V,R,W,T}) where {N,L,V,R,W,T<:SedimentModel}
 
     return model
 end
+
+function set_states(model::Model{N,L,V,R,W,T}) where {N,L,V,R,W,T<:SedimentModel}
+    # read and set states in model object if reinit=false
+    @unpack config = model
+    reinit = get(config.model, "reinit", true)::Bool
+    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)
+    else
+        @info "Set initial conditions from default values."
+    end
+    return model
+end


Test patch:
diff --git a/test/bmi.jl b/test/bmi.jl
index 1affd67da..bb89cdf57 100644
--- a/test/bmi.jl
+++ b/test/bmi.jl
@@ -14,32 +14,30 @@ tomlpath = joinpath(@__DIR__, "sbm_config.toml")
             @test BMI.get_start_time(model) == 0.0
             @test BMI.get_current_time(model) == 0.0
             @test BMI.get_end_time(model) == 31 * 86400.0
+            model.config.endtime = "2000-02-01T00:00:00"
+            @test BMI.get_end_time(model) == 31 * 86400.0
         end
 
         @testset "model information functions" begin
             @test BMI.get_component_name(model) == "sbm"
-            @test BMI.get_input_item_count(model) == 174
-            @test BMI.get_output_item_count(model) == 174
-            @test BMI.get_input_var_names(model)[[1, 5, 120, 174]] == [
-                "vertical.Δt",
-                "vertical.n_unsatlayers",
-                "lateral.land.qlat",
-                "lateral.river.reservoir.evaporation",
+            @test BMI.get_input_item_count(model) == 183
+            @test BMI.get_output_item_count(model) == 183
+            @test BMI.get_input_var_names(model)[[1, 5, 150, 174]] == [
+                "vertical.nlayers",
+                "vertical.θᵣ",
+                "lateral.river.sl",
+                "lateral.river.reservoir.targetminfrac",
             ]
         end
 
         @testset "variable information functions" begin
-            @test BMI.get_var_grid(model, "vertical.θₛ") == 4
-            @test BMI.get_var_grid(model, "lateral.river.h") == 2
-            @test BMI.get_var_grid(model, "lateral.river.reservoir.inflow") == 0
+            @test BMI.get_var_grid(model, "vertical.θₛ") == 7
+            @test BMI.get_var_grid(model, "lateral.river.h") == 4
+            @test BMI.get_var_grid(model, "lateral.river.reservoir.inflow") == 1
             @test_throws ErrorException BMI.get_var_grid(model, "lateral.river.lake.volume")
-            # Vector{Float64} printing on Julia 1.6+
-            @test BMI.get_var_type(model, "lateral.river.reservoir.inflow") in
-                  ["Array{$Float,1}", "Vector{$Float}"]
-            @test BMI.get_var_type(model, "vertical.n") == "Int64"
+            @test BMI.get_var_type(model, "lateral.river.reservoir.inflow") == "$Float"
             @test BMI.get_var_units(model, "vertical.θₛ") == "-"
             @test BMI.get_var_itemsize(model, "lateral.subsurface.ssf") == sizeof(Float)
-            @test BMI.get_var_nbytes(model, "vertical.n") == 8
             @test BMI.get_var_nbytes(model, "lateral.river.q") ==
                   length(model.lateral.river.q) * sizeof(Float)
             @test BMI.get_var_location(model, "lateral.river.q") == "node"
@@ -49,27 +47,61 @@ tomlpath = joinpath(@__DIR__, "sbm_config.toml")
 
         @testset "update and get and set functions" begin
             @test BMI.get_current_time(model) == 86400.0
-            @test mean(BMI.get_value(model, "vertical.zi")) ≈ 276.3767651555451
-            @test BMI.get_value_at_indices(model, "lateral.river.q", [1, 100, 5617]) ≈
-                  [0.6211503865184697, 5.219305686635002, 0.026163746306482282]
+            @test_throws ErrorException BMI.get_value_ptr(model, "vertical.")
+            dest = zeros(Float, size(model.vertical.zi))
+            BMI.get_value(model, "vertical.zi", dest)
+            @test mean(dest) ≈ 276.3767651555451
+            @test BMI.get_value_at_indices(
+                model,
+                "vertical.vwc[1]",
+                zeros(Float, 3),
+                [1, 2, 3],
+            ) ≈ getindex.(model.vertical.vwc, 1)[1:3]
+            BMI.set_value_at_indices(
+                model,
+                "vertical.vwc[2]",
+                [1, 2, 3],
+                [0.10, 0.15, 0.20],
+            ) ≈ getindex.(model.vertical.vwc, 2)[1:3]
+            @test BMI.get_value_at_indices(
+                model,
+                "lateral.river.q",
+                zeros(Float, 3),
+                [1, 100, 5617],
+            ) ≈ [0.6211503865184697, 5.219305686635002, 0.026163746306482282]
             BMI.set_value(model, "vertical.zi", fill(300.0, length(model.vertical.zi)))
-            @test mean(BMI.get_value(model, "vertical.zi")) == 300.0
+            @test mean(
+                BMI.get_value(model, "vertical.zi", zeros(Float, size(model.vertical.zi))),
+            ) == 300.0
             BMI.set_value_at_indices(model, "vertical.zi", [1], [250.0])
-            @test BMI.get_value_at_indices(model, "vertical.zi", [1, 2]) == [250.0, 300.0]
+            @test BMI.get_value_at_indices(model, "vertical.zi", zeros(Float, 2), [1, 2]) ==
+                  [250.0, 300.0]
         end
 
         @testset "model grid functions" begin
-            @test BMI.get_grid_type(model, 0) == "unstructured"
-            @test BMI.get_grid_rank(model, 0) == 2
-            @test BMI.get_grid_size(model, 0) == 2
-            @test BMI.get_grid_size(model, 2) == 5809
-            @test BMI.get_grid_size(model, 4) == 50063
-            @test BMI.get_grid_shape(model, 4) == 50063
-            @test minimum(BMI.get_grid_x(model, 4)) ≈ 5.426666666666667f0
-            @test maximum(BMI.get_grid_x(model, 4)) ≈ 7.843333333333344f0
-            @test BMI.get_grid_x(model, 0) ≈ [5.760000000000002f0, 5.918333333333336f0]
-            @test BMI.get_grid_y(model, 0) ≈ [48.92583333333333f0, 49.909166666666664f0]
-            @test BMI.get_grid_node_count(model, 0) == 2
+            @test BMI.get_grid_type(model, 0) == "scalar"
+            @test BMI.get_grid_type(model, 2) == "points"
+            @test BMI.get_grid_type(model, 7) == "unstructured"
+            @test_throws ErrorException BMI.get_grid_rank(model, 8)
+            @test BMI.get_grid_rank(model, 0) == 0
+            @test BMI.get_grid_rank(model, 7) == 2
+            @test_throws ErrorException BMI.get_grid_rank(model, 8)
+            @test BMI.get_grid_node_count(model, 1) == 2
+            @test BMI.get_grid_node_count(model, 4) == 5809
+            @test BMI.get_grid_node_count(model, 5) == 50063
+            @test BMI.get_grid_node_count(model, 6) == 50063
+            @test minimum(BMI.get_grid_x(model, 5, zeros(Float, 50063))) ≈
+                  5.426666666666667f0
+            @test maximum(BMI.get_grid_x(model, 5, zeros(Float, 50063))) ≈
+                  7.843333333333344f0
+            @test BMI.get_grid_x(model, 1, zeros(Float, 2)) ≈
+                  [5.760000000000002f0, 5.918333333333336f0]
+            @test BMI.get_grid_y(model, 1, zeros(Float, 2)) ≈
+                  [48.92583333333333f0, 49.909166666666664f0]
+            @test BMI.get_grid_node_count(model, 1) == 2
+            @test BMI.get_grid_edge_count(model, 4) == 5808
+            @test BMI.get_grid_edge_nodes(model, 4, fill(0, 2 * 5808))[1:6] ==
+                  [1, 5, 2, 1, 3, 2]
         end
 
         @testset "update until and finalize" begin
@@ -81,6 +113,23 @@ tomlpath = joinpath(@__DIR__, "sbm_config.toml")
 
     end
 
+    @testset "BMI grid edges" begin
+        tomlpath = joinpath(@__DIR__, "sbm_swf_config.toml")
+        model = BMI.initialize(Wflow.Model, tomlpath)
+        @test BMI.get_var_grid(model, "lateral.land.qx") == 5
+        @test BMI.get_var_grid(model, "lateral.land.qy") == 6
+        @test BMI.get_grid_edge_count(model, 5) == 50063
+        @test BMI.get_grid_edge_count(model, 6) == 50063
+        @test BMI.get_grid_edge_nodes(model, 5, fill(0, 2 * 50063))[1:4] == [1, -999, 2, 3]
+        @test BMI.get_grid_edge_nodes(model, 6, fill(0, 2 * 50063))[1:4] == [1, 4, 2, 10]
+        @test_logs (
+            :warn,
+            "edges are not provided for grid type 3 (variables are located at nodes)",
+        ) BMI.get_grid_edge_nodes(model, 3, fill(0, 2 * 50063))
+        @test_throws ErrorException BMI.get_grid_edge_nodes(model, 8, fill(0, 2 * 50063))
+        BMI.finalize(model)
+    end
+
     @testset "BMI run SBM in parts" begin
         tomlpath = joinpath(@__DIR__, "sbm_gw.toml")
         model = BMI.initialize(Wflow.Model, tomlpath)
@@ -101,12 +150,12 @@ tomlpath = joinpath(@__DIR__, "sbm_config.toml")
         BMI.set_value(
             model,
             "lateral.subsurface.zi",
-            fill(0.25, BMI.get_grid_size(model, 4)),
+            fill(0.25, BMI.get_grid_node_count(model, 6)),
         )
         BMI.set_value(
             model,
             "lateral.subsurface.exfiltwater",
-            fill(1.0e-5, BMI.get_grid_size(model, 4)),
+            fill(1.0e-5, BMI.get_grid_node_count(model, 6)),
         )
         # update SBM after subsurface flow
         model = BMI.update(model, run = "sbm_after_subsurfaceflow")
@@ -126,4 +175,19 @@ tomlpath = joinpath(@__DIR__, "sbm_config.toml")
 
         BMI.finalize(model)
     end
+
+end
+
+@testset "BMI extension functions" begin
+
+    model = BMI.initialize(Wflow.Model, tomlpath)
+    @test Wflow.get_start_unix_time(model) == 9.466848e8
+    satwaterdepth = mean(model.vertical.satwaterdepth)
+    model.config.model.reinit = false
+    model = Wflow.load_state(model)
+    @test satwaterdepth ≠ mean(model.vertical.satwaterdepth)
+    @test_logs (
+        :info,
+        "Write output states to NetCDF file `$(model.writer.state_nc_path)`.",
+    ) Wflow.save_state(model)
 end