From 249b79a2941b9dd72cb8a5c3657ea888aed70bc4 Mon Sep 17 00:00:00 2001
From: verseve <willem.vanverseveld@gmail.com>
Date: Fri, 29 Sep 2023 13:46:28 +0200
Subject: [PATCH 01/20] Fix deviation from BMI for `get_value` and
 `get_value_at_indices` Both were missing the required `dest` argument.

---
 src/bmi.jl  | 15 +++++++++++++++
 test/bmi.jl |  6 ++++++
 2 files changed, 21 insertions(+)

diff --git a/src/bmi.jl b/src/bmi.jl
index 5fb80fd8f..0672e7ef1 100644
--- a/src/bmi.jl
+++ b/src/bmi.jl
@@ -187,6 +187,11 @@ function BMI.get_time_step(model::Model)
     Float64(model.config.timestepsecs)
 end
 
+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(model::Model, name::String)
     copy(BMI.get_value_ptr(model, name))
 end
@@ -208,6 +213,16 @@ function BMI.get_value_at_indices(model::Model, name::String, inds::Vector{Int})
     BMI.get_value_ptr(model, name)[inds]
 end
 
+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
+
 """
     BMI.set_value(model::Model, name::String, src::Vector{T}) where T<:AbstractFloat
 
diff --git a/test/bmi.jl b/test/bmi.jl
index 1affd67da..b0d3be184 100644
--- a/test/bmi.jl
+++ b/test/bmi.jl
@@ -50,12 +50,18 @@ 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
+            dest = zeros(Float, size(model.vertical.zi))
+            dest = BMI.get_value(model, "vertical.zi", dest)
+            @test mean(dest) ≈ 276.3767651555451
             @test BMI.get_value_at_indices(model, "lateral.river.q", [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
             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]
+            dest = zeros(Float, 2)
+            dest = BMI.get_value_at_indices(model, "vertical.zi", dest, [1, 2])
+            @test dest == [250.0, 300.0]
         end
 
         @testset "model grid functions" begin

From ebd658c8f8056b5583dcf19884cfc424c54b0101 Mon Sep 17 00:00:00 2001
From: verseve <willem.vanverseveld@gmail.com>
Date: Fri, 29 Sep 2023 14:07:55 +0200
Subject: [PATCH 02/20] Fix `BMI.get_var_type` Should return the datatype of a
 variable (with `eltype`) and not the type of a complete Array (`typeof`).

---
 src/bmi.jl  | 2 +-
 test/bmi.jl | 4 +---
 2 files changed, 2 insertions(+), 4 deletions(-)

diff --git a/src/bmi.jl b/src/bmi.jl
index 0672e7ef1..7f390ca21 100644
--- a/src/bmi.jl
+++ b/src/bmi.jl
@@ -143,7 +143,7 @@ function BMI.get_var_grid(model::Model, name::String)
 end
 
 function BMI.get_var_type(model::Model, name::String)
-    string(typeof(param(model, name)))
+    repr(eltype(param(model, name)))
 end
 
 function BMI.get_var_units(model::Model, name::String)
diff --git a/test/bmi.jl b/test/bmi.jl
index b0d3be184..ed16aa2c8 100644
--- a/test/bmi.jl
+++ b/test/bmi.jl
@@ -33,9 +33,7 @@ tomlpath = joinpath(@__DIR__, "sbm_config.toml")
             @test BMI.get_var_grid(model, "lateral.river.h") == 2
             @test BMI.get_var_grid(model, "lateral.river.reservoir.inflow") == 0
             @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, "lateral.river.reservoir.inflow") == "$Float"
             @test BMI.get_var_type(model, "vertical.n") == "Int64"
             @test BMI.get_var_units(model, "vertical.θₛ") == "-"
             @test BMI.get_var_itemsize(model, "lateral.subsurface.ssf") == sizeof(Float)

From b4a5fa41ebd59d04bb550da47f8a2687c1040962 Mon Sep 17 00:00:00 2001
From: Willem van Verseveld <willem.vanverseveld@gmail.com>
Date: Wed, 18 Oct 2023 09:39:59 +0200
Subject: [PATCH 03/20] Fix and update BMI - fix functions that deviate from
 BasicModelInterface.jl (e.g. get_grid_x). - add metadata to Structs, for
 example grid location (node or edge), that can be used by BMI (and Wflow).

---
 src/Wflow.jl                           |   6 +-
 src/bmi.jl                             | 164 +++++++++------
 src/flextopo.jl                        |  28 +--
 src/flow.jl                            | 280 ++++++++++++-------------
 src/groundwater/aquifer.jl             |   6 +-
 src/groundwater/boundary_conditions.jl |  14 +-
 src/hbv.jl                             |  54 ++---
 src/reservoir_lake.jl                  |   8 +-
 src/sbm.jl                             |   8 +-
 src/sediment.jl                        |  14 +-
 test/bmi.jl                            |  62 +++---
 11 files changed, 339 insertions(+), 305 deletions(-)

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 7f390ca21..b7239e948 100644
--- a/src/bmi.jl
+++ b/src/bmi.jl
@@ -1,19 +1,16 @@
 # 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")
+grids = Dict{Int,Tuple}(
+    1 => (:reservoir,),
+    2 => (:lake,),
+    3 => (:drain,),
+    4 => (:river,),
+    5 => (:land,),
+    6 => (:land,),
+)
 
 """
     BMI.initialize(::Type{<:Wflow.Model}, config_file)
@@ -33,6 +30,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 +105,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,45 +139,60 @@ 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))
+        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
+        else
+            6
         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)
-    repr(eltype(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
 
@@ -192,27 +221,30 @@ function BMI.get_value(model::Model, name::String, dest::Vector{T}) where {T<:Ab
     return dest
 end
 
-function BMI.get_value(model::Model, name::String)
-    copy(BMI.get_value_ptr(model, name))
-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
+    key = symbols(first(split(name, "-")))
+    if exchange(param(model, key[1:end-1]), key[end]) == 1
+        n = length(active_indices(network, key))
+        s = split(name, ".")
+        if tryparse(Int, s[end]) !== nothing
+            var = split(s[end], "-")
+            ind = tryparse(Int, var[end])
+            if eltype(param(model, key)) <: SVector
+                value = @view getindex.(param(model, key), ind)[1:n]
+                return value
+            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]
-end
-
 function BMI.get_value_at_indices(
     model::Model,
     name::String,
@@ -249,43 +281,47 @@ 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 range(1, 3)
+        "points"
+    elseif grid in range(4, 6)
+        "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 range(1, 6)
+        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
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/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/reservoir_lake.jl b/src/reservoir_lake.jl
index 303c98d14..bcacdbdd9 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                                     # 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                             # 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..bb5744a8e 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/sediment.jl b/src/sediment.jl
index 620d0ce63..ad4ebea83 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/test/bmi.jl b/test/bmi.jl
index ed16aa2c8..59e5e0b87 100644
--- a/test/bmi.jl
+++ b/test/bmi.jl
@@ -18,26 +18,24 @@ tomlpath = joinpath(@__DIR__, "sbm_config.toml")
 
         @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.θₛ") == 6
+            @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")
             @test BMI.get_var_type(model, "lateral.river.reservoir.inflow") == "$Float"
-            @test BMI.get_var_type(model, "vertical.n") == "Int64"
             @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"
@@ -47,33 +45,29 @@ 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
             dest = zeros(Float, size(model.vertical.zi))
-            dest = BMI.get_value(model, "vertical.zi", dest)
-            @test mean(dest) ≈ 276.3767651555451
-            @test BMI.get_value_at_indices(model, "lateral.river.q", [1, 100, 5617]) ≈
+            BMI.get_value(model, "vertical.zi", dest)
+            @test mean(dest) ≈ 276.3767651555451            
+            @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]
-            dest = zeros(Float, 2)
-            dest = BMI.get_value_at_indices(model, "vertical.zi", dest, [1, 2])
-            @test dest == [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_rank(model, 0) == 0
+            @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
         end
 
         @testset "update until and finalize" begin
@@ -105,12 +99,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")

From b03c33ea1c9329d5a46e1b4fca21d19659bbf8d8 Mon Sep 17 00:00:00 2001
From: Willem van Verseveld <willem.vanverseveld@gmail.com>
Date: Wed, 18 Oct 2023 11:44:02 +0200
Subject: [PATCH 04/20] Fix `BMI.get_value_ptr` (`SVector` and 2D Array )
 SVector type and 2D Array variables are exposed per third dimension
 coordinate, by adding "-i" (i refers to index) to the variable name (to
 minimize the number of required grid types). Extracting index i from the
 variable name has been fixed.

---
 src/bmi.jl  |  7 +++----
 test/bmi.jl | 35 ++++++++++++++++++++++++++---------
 2 files changed, 29 insertions(+), 13 deletions(-)

diff --git a/src/bmi.jl b/src/bmi.jl
index b7239e948..1f7673048 100644
--- a/src/bmi.jl
+++ b/src/bmi.jl
@@ -223,13 +223,12 @@ end
 
 function BMI.get_value_ptr(model::Model, name::String)
     @unpack network = model
-    key = symbols(first(split(name, "-")))
+    s = split(name, "-")
+    key = symbols(first(s))
     if exchange(param(model, key[1:end-1]), key[end]) == 1
         n = length(active_indices(network, key))
-        s = split(name, ".")
         if tryparse(Int, s[end]) !== nothing
-            var = split(s[end], "-")
-            ind = tryparse(Int, var[end])
+            ind = tryparse(Int, s[end])
             if eltype(param(model, key)) <: SVector
                 value = @view getindex.(param(model, key), ind)[1:n]
                 return value
diff --git a/test/bmi.jl b/test/bmi.jl
index 59e5e0b87..494d1c06f 100644
--- a/test/bmi.jl
+++ b/test/bmi.jl
@@ -47,13 +47,26 @@ tomlpath = joinpath(@__DIR__, "sbm_config.toml")
             @test BMI.get_current_time(model) == 86400.0
             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, "lateral.river.q", zeros(Float, 3), [1, 100, 5617]) ≈
-                  [0.6211503865184697, 5.219305686635002, 0.026163746306482282]
+            @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, 2, 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", zeros(Float, size(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", zeros(Float, 2), [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
@@ -63,10 +76,14 @@ tomlpath = joinpath(@__DIR__, "sbm_config.toml")
             @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 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
         end
 

From 4a8d17cd888b0bbdf1d69b14bc060602566c1cd3 Mon Sep 17 00:00:00 2001
From: Willem van Verseveld <willem.vanverseveld@gmail.com>
Date: Wed, 18 Oct 2023 17:27:28 +0200
Subject: [PATCH 05/20] Add BMI functions for grids with edges

---
 src/bmi.jl  | 53 ++++++++++++++++++++++++++++++++++++++++++++++++++---
 test/bmi.jl |  7 +++++--
 2 files changed, 55 insertions(+), 5 deletions(-)

diff --git a/src/bmi.jl b/src/bmi.jl
index 1f7673048..2a9ab6101 100644
--- a/src/bmi.jl
+++ b/src/bmi.jl
@@ -10,6 +10,7 @@ grids = Dict{Int,Tuple}(
     4 => (:river,),
     5 => (:land,),
     6 => (:land,),
+    7 => (:land,),
 )
 
 """
@@ -156,8 +157,10 @@ function BMI.get_var_grid(model::Model, name::String)
             4
         elseif type <: ShallowWaterLand && occursin("x", s[end])
             5
-        else
+        elseif type <: ShallowWaterLand && occursin("y", s[end])
             6
+        else
+            7
         end
     else
         error("$name not listed as variable for BMI exchange")
@@ -284,7 +287,7 @@ function BMI.get_grid_type(model::Model, grid::Int)
         "scalar"
     elseif grid in range(1, 3)
         "points"
-    elseif grid in range(4, 6)
+    elseif grid in range(4, 7)
         "unstructured"
     else
         error("unknown grid type $grid")
@@ -294,7 +297,7 @@ end
 function BMI.get_grid_rank(model::Model, grid::Int)
     if grid == 0
         0
-    elseif grid in range(1, 6)
+    elseif grid in range(1, 7)
         2
     else
         error("unknown grid type $grid")
@@ -324,3 +327,47 @@ end
 function BMI.get_grid_node_count(model::Model, grid::Int)
     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 range(0, 3) || grid == 7
+        warn("edges are not defined for grid type $grid")
+    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)
+    if grid == 4
+        nodes_at_edge = adjacent_nodes_at_link(network.river.graph)
+        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.staggered_indices.xu
+        m = length(xu)
+        edge_nodes[range(1, n, step = 2)] = range(1, m)
+        xu[xu.==m+1] .= -999
+        edge_nodes[range(2, n, step = 2)] = xu
+        return edge_nodes
+    elseif grid == 6
+        yu = network.staggered_indices.yu
+        m = length(yu)
+        edge_nodes[range(1, n, step = 2)] = range(1, m)
+        yu[yu.==m+1] .= -999
+        edge_nodes[range(2, n, step = 2)] = yu
+        return edge_nodes
+    elseif grid in range(0, 3) || grid == 7
+        warn("edges are not defined for grid type $grid")
+    else
+        error("unknown grid type $grid")
+    end
+end
diff --git a/test/bmi.jl b/test/bmi.jl
index 494d1c06f..5863dd647 100644
--- a/test/bmi.jl
+++ b/test/bmi.jl
@@ -29,7 +29,7 @@ tomlpath = joinpath(@__DIR__, "sbm_config.toml")
         end
 
         @testset "variable information functions" begin
-            @test BMI.get_var_grid(model, "vertical.θₛ") == 6
+            @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")
@@ -53,7 +53,7 @@ tomlpath = joinpath(@__DIR__, "sbm_config.toml")
                 "vertical.vwc-1",
                 zeros(Float, 3),
                 [1, 2, 3],
-            ) ≈ getindex.(model.vertical.vwc, 1)[[1, 2, 3]]
+            ) ≈ getindex.(model.vertical.vwc, 1)[1:3]
             @test BMI.get_value_at_indices(
                 model,
                 "lateral.river.q",
@@ -85,6 +85,9 @@ tomlpath = joinpath(@__DIR__, "sbm_config.toml")
             @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

From 0c8b6814166f9c9f01acff9bd55f8f69afb7b5ae Mon Sep 17 00:00:00 2001
From: Willem van Verseveld <willem.vanverseveld@gmail.com>
Date: Thu, 19 Oct 2023 10:52:03 +0200
Subject: [PATCH 06/20] Set inactive nodes at -999 in BMI This is done for
 boundary/ghost nodes.

---
 src/bmi.jl | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/src/bmi.jl b/src/bmi.jl
index 2a9ab6101..366ce0ea3 100644
--- a/src/bmi.jl
+++ b/src/bmi.jl
@@ -337,7 +337,7 @@ function BMI.get_grid_edge_count(model::Model, grid::Int)
     elseif grid == 6
         return length(network.land.staggered_indices.yu)
     elseif grid in range(0, 3) || grid == 7
-        warn("edges are not defined for grid type $grid")
+        warn("edges are not provided for grid type $grid (variables are located at nodes)")
     else
         error("unknown grid type $grid")
     end
@@ -346,27 +346,28 @@ 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.staggered_indices.xu
-        m = length(xu)
         edge_nodes[range(1, n, step = 2)] = range(1, m)
         xu[xu.==m+1] .= -999
         edge_nodes[range(2, n, step = 2)] = xu
         return edge_nodes
     elseif grid == 6
         yu = network.staggered_indices.yu
-        m = length(yu)
         edge_nodes[range(1, n, step = 2)] = range(1, m)
         yu[yu.==m+1] .= -999
         edge_nodes[range(2, n, step = 2)] = yu
         return edge_nodes
     elseif grid in range(0, 3) || grid == 7
-        warn("edges are not defined for grid type $grid")
+        warn("edges are not provided for grid type $grid (variables are located at nodes)")
     else
         error("unknown grid type $grid")
     end

From 7969e0899c3e1fc68cd658957de54f7be4c17b14 Mon Sep 17 00:00:00 2001
From: Willem van Verseveld <willem.vanverseveld@gmail.com>
Date: Thu, 19 Oct 2023 15:03:21 +0200
Subject: [PATCH 07/20] =?UTF-8?q?Change=20BMI=20variable=20name=20indexing?=
 =?UTF-8?q?=20(3rd=20dim)=20For=20SVector=20type=20and=202D=20Array=20vari?=
 =?UTF-8?q?ables=20use=20"[=C3=AC]"=20instead=20of=20"-i"?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/bmi.jl  | 16 ++++++++--------
 test/bmi.jl |  2 +-
 2 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/src/bmi.jl b/src/bmi.jl
index 366ce0ea3..90de8e81d 100644
--- a/src/bmi.jl
+++ b/src/bmi.jl
@@ -115,11 +115,11 @@ function BMI.get_input_var_names(model::Model)
                 model_var = param(model, var)
                 if eltype(model_var) <: SVector
                     for i = 1:length(first(model_var))
-                        push!(var_names, string(var, "-", i))
+                        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))
+                        push!(var_names, string(var, "[", i, "]"))
                     end
                 else
                     push!(var_names, var)
@@ -140,7 +140,7 @@ function BMI.get_output_var_names(model::Model)
 end
 
 function BMI.get_var_grid(model::Model, name::String)
-    s = split(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))
@@ -173,7 +173,7 @@ function BMI.get_var_type(model::Model, name::String)
 end
 
 function BMI.get_var_units(model::Model, name::String)
-    key = symbols(first(split(name, "-")))
+    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
@@ -191,7 +191,7 @@ function BMI.get_var_nbytes(model::Model, name::String)
 end
 
 function BMI.get_var_location(model::Model, name::String)
-    key = symbols(first(split(name, "-")))
+    key = symbols(first(split(name, "[")))
     if exchange(param(model, key[1:end-1]), key[end]) == 1
         return grid_location(model, key)
     else
@@ -226,12 +226,12 @@ end
 
 function BMI.get_value_ptr(model::Model, name::String)
     @unpack network = model
-    s = split(name, "-")
+    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 tryparse(Int, s[end]) !== nothing
-            ind = tryparse(Int, s[end])
+        if occursin("[", name)
+            ind = tryparse(Int, split(s[end], "]")[1])
             if eltype(param(model, key)) <: SVector
                 value = @view getindex.(param(model, key), ind)[1:n]
                 return value
diff --git a/test/bmi.jl b/test/bmi.jl
index 5863dd647..d99c1bf08 100644
--- a/test/bmi.jl
+++ b/test/bmi.jl
@@ -50,7 +50,7 @@ tomlpath = joinpath(@__DIR__, "sbm_config.toml")
             @test mean(dest) ≈ 276.3767651555451
             @test BMI.get_value_at_indices(
                 model,
-                "vertical.vwc-1",
+                "vertical.vwc[1]",
                 zeros(Float, 3),
                 [1, 2, 3],
             ) ≈ getindex.(model.vertical.vwc, 1)[1:3]

From 8fb571ce5e46adf92200c5a3faf820349a19e167 Mon Sep 17 00:00:00 2001
From: Willem van Verseveld <willem.vanverseveld@gmail.com>
Date: Thu, 19 Oct 2023 17:19:05 +0200
Subject: [PATCH 08/20] Update BMI docs

---
 docs/src/user_guide/additional_options.md | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

diff --git a/docs/src/user_guide/additional_options.md b/docs/src/user_guide/additional_options.md
index 7d66f29bb..84a476a1f 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:

From 3c8e8152efcf3cbad7ac999fe4f0cfa03ca26ffc Mon Sep 17 00:00:00 2001
From: Willem van Verseveld <willem.vanverseveld@gmail.com>
Date: Mon, 23 Oct 2023 08:59:28 +0200
Subject: [PATCH 09/20] `BMI.get_value_ptr` and `SVector` type Replace getindex
 (values are retrieved from the array (copy)), with reinterpret (a copy is not
 created).

---
 src/bmi.jl  | 7 +++++--
 test/bmi.jl | 6 ++++++
 2 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/src/bmi.jl b/src/bmi.jl
index 90de8e81d..cfc2c99f6 100644
--- a/src/bmi.jl
+++ b/src/bmi.jl
@@ -233,8 +233,11 @@ function BMI.get_value_ptr(model::Model, name::String)
         if occursin("[", name)
             ind = tryparse(Int, split(s[end], "]")[1])
             if eltype(param(model, key)) <: SVector
-                value = @view getindex.(param(model, key), ind)[1:n]
-                return value
+                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
diff --git a/test/bmi.jl b/test/bmi.jl
index d99c1bf08..85de7a17e 100644
--- a/test/bmi.jl
+++ b/test/bmi.jl
@@ -54,6 +54,12 @@ tomlpath = joinpath(@__DIR__, "sbm_config.toml")
                 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",

From e607bc503ec1f82d93c022d4fa1aea575bb0e196 Mon Sep 17 00:00:00 2001
From: Willem van Verseveld <willem.vanverseveld@gmail.com>
Date: Mon, 23 Oct 2023 11:07:55 +0200
Subject: [PATCH 10/20] Update changelog

---
 docs/src/changelog.md | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/docs/src/changelog.md b/docs/src/changelog.md
index e8e7de35a..96e15db0b 100644
--- a/docs/src/changelog.md
+++ b/docs/src/changelog.md
@@ -5,6 +5,20 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [unreleased]
+
+### Fixed
+- 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).
+
+### Changed
+- 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`.
+
 ## v0.7.2 - 2023-09-27
 
 ### Fixed

From cfafd2f2c3cb32dfdc12d722741a0cccdfc6edda Mon Sep 17 00:00:00 2001
From: Willem van Verseveld <willem.vanverseveld@gmail.com>
Date: Tue, 24 Oct 2023 13:31:47 +0200
Subject: [PATCH 11/20] 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            |  19 +++++
 src/flextopo_model.jl |  43 ++++++-----
 src/hbv_model.jl      |  66 +++++++++--------
 src/sbm_gwf_model.jl  |  65 +++++++++--------
 src/sbm_model.jl      | 162 +++++++++++++++++++++++-------------------
 src/sediment_model.jl |  27 ++++---
 6 files changed, 223 insertions(+), 159 deletions(-)

diff --git a/src/bmi.jl b/src/bmi.jl
index cfc2c99f6..34a971942 100644
--- a/src/bmi.jl
+++ b/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
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/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/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_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

From 26464f618f7f4f780c5e79e1e6923ee34d11bdd5 Mon Sep 17 00:00:00 2001
From: Willem van Verseveld <willem.vanverseveld@gmail.com>
Date: Tue, 24 Oct 2023 14:09:53 +0200
Subject: [PATCH 12/20] Update changelog

---
 docs/src/changelog.md | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/docs/src/changelog.md b/docs/src/changelog.md
index 96e15db0b..bb642c810 100644
--- a/docs/src/changelog.md
+++ b/docs/src/changelog.md
@@ -19,6 +19,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   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
 
 ### Fixed

From 5dc79b1786082672a1845d337483ceefa8349447 Mon Sep 17 00:00:00 2001
From: Willem van Verseveld <willem.vanverseveld@gmail.com>
Date: Wed, 25 Oct 2023 08:34:42 +0200
Subject: [PATCH 13/20] Add more tests for (extended) BMI

---
 src/bmi.jl  |  8 ++++----
 test/bmi.jl | 38 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 42 insertions(+), 4 deletions(-)

diff --git a/src/bmi.jl b/src/bmi.jl
index 34a971942..32ed5be4a 100644
--- a/src/bmi.jl
+++ b/src/bmi.jl
@@ -144,7 +144,7 @@ function BMI.get_var_grid(model::Model, name::String)
     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))
+        type = typeof(param(model, key[1:end-1]))
         if gridtype == "scalar"
             0
         elseif :reservoir in key
@@ -358,19 +358,19 @@ function BMI.get_grid_edge_nodes(model::Model, grid::Int, edge_nodes::Vector{Int
         edge_nodes[range(2, n, step = 2)] = nodes_at_edge.dst
         return edge_nodes
     elseif grid == 5
-        xu = network.staggered_indices.xu
+        xu = network.land.staggered_indices.xu
         edge_nodes[range(1, n, step = 2)] = range(1, m)
         xu[xu.==m+1] .= -999
         edge_nodes[range(2, n, step = 2)] = xu
         return edge_nodes
     elseif grid == 6
-        yu = network.staggered_indices.yu
+        yu = network.land.staggered_indices.yu
         edge_nodes[range(1, n, step = 2)] = range(1, m)
         yu[yu.==m+1] .= -999
         edge_nodes[range(2, n, step = 2)] = yu
         return edge_nodes
     elseif grid in range(0, 3) || grid == 7
-        warn("edges are not provided for grid type $grid (variables are located at nodes)")
+        @warn("edges are not provided for grid type $grid (variables are located at nodes)")
     else
         error("unknown grid type $grid")
     end
diff --git a/test/bmi.jl b/test/bmi.jl
index 85de7a17e..a8d9d5dc2 100644
--- a/test/bmi.jl
+++ b/test/bmi.jl
@@ -45,6 +45,7 @@ tomlpath = joinpath(@__DIR__, "sbm_config.toml")
 
         @testset "update and get and set functions" begin
             @test BMI.get_current_time(model) == 86400.0
+            @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
@@ -77,7 +78,12 @@ tomlpath = joinpath(@__DIR__, "sbm_config.toml")
 
         @testset "model grid functions" begin
             @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
@@ -105,6 +111,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)
@@ -150,4 +173,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

From 74ffd1a063d9b7c6b7269e843cecccc30db53cf5 Mon Sep 17 00:00:00 2001
From: Willem van Verseveld <willem.vanverseveld@gmail.com>
Date: Wed, 25 Oct 2023 10:31:34 +0200
Subject: [PATCH 14/20] Make `range` compatible with Julia 1.6

---
 src/bmi.jl | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/src/bmi.jl b/src/bmi.jl
index 32ed5be4a..1d17e9128 100644
--- a/src/bmi.jl
+++ b/src/bmi.jl
@@ -288,9 +288,9 @@ end
 function BMI.get_grid_type(model::Model, grid::Int)
     if grid == 0
         "scalar"
-    elseif grid in range(1, 3)
+    elseif grid in 1:3
         "points"
-    elseif grid in range(4, 7)
+    elseif grid in 4:7
         "unstructured"
     else
         error("unknown grid type $grid")
@@ -300,7 +300,7 @@ end
 function BMI.get_grid_rank(model::Model, grid::Int)
     if grid == 0
         0
-    elseif grid in range(1, 7)
+    elseif grid in 1:7
         2
     else
         error("unknown grid type $grid")
@@ -339,7 +339,7 @@ function BMI.get_grid_edge_count(model::Model, grid::Int)
         return length(network.land.staggered_indices.xu)
     elseif grid == 6
         return length(network.land.staggered_indices.yu)
-    elseif grid in range(0, 3) || grid == 7
+    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")
@@ -359,17 +359,17 @@ function BMI.get_grid_edge_nodes(model::Model, grid::Int, edge_nodes::Vector{Int
         return edge_nodes
     elseif grid == 5
         xu = network.land.staggered_indices.xu
-        edge_nodes[range(1, n, step = 2)] = range(1, m)
+        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)] = range(1, m)
+        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 range(0, 3) || grid == 7
+    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")

From b16a0dcd403682ec12949d23a01d155d53f6d822 Mon Sep 17 00:00:00 2001
From: Willem van Verseveld <willem.vanverseveld@gmail.com>
Date: Wed, 25 Oct 2023 17:10:29 +0200
Subject: [PATCH 15/20] Fix BMI `starttime` and `endtime` (conversion)
 config.startime and config.endtime should be converted to `DateTime`, because
 Wflow accepts also Strings here (TOML).

---
 src/bmi.jl  | 11 +++++++++--
 test/bmi.jl |  2 ++
 2 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/src/bmi.jl b/src/bmi.jl
index 1d17e9128..323b441e0 100644
--- a/src/bmi.jl
+++ b/src/bmi.jl
@@ -200,7 +200,10 @@ function BMI.get_var_location(model::Model, name::String)
 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)
@@ -208,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)
diff --git a/test/bmi.jl b/test/bmi.jl
index a8d9d5dc2..bb89cdf57 100644
--- a/test/bmi.jl
+++ b/test/bmi.jl
@@ -14,6 +14,8 @@ 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

From 99b39968501fa5cb6430b6c9aec3d8fbac44c0d2 Mon Sep 17 00:00:00 2001
From: Martijn Visser <mgvisser@gmail.com>
Date: Tue, 14 Nov 2023 14:49:20 +0100
Subject: [PATCH 16/20] remove trailing whitespace

---
 docs/src/user_guide/additional_options.md | 2 +-
 src/sbm.jl                                | 6 +++---
 src/sediment.jl                           | 6 +++---
 3 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/docs/src/user_guide/additional_options.md b/docs/src/user_guide/additional_options.md
index 84a476a1f..7862a5f7c 100644
--- a/docs/src/user_guide/additional_options.md
+++ b/docs/src/user_guide/additional_options.md
@@ -223,7 +223,7 @@ Variables with a third dimension, for example `layer` as part of the vertical `S
 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.    
+`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.
diff --git a/src/sbm.jl b/src/sbm.jl
index bb5744a8e..a85254bf7 100644
--- a/src/sbm.jl
+++ b/src/sbm.jl
@@ -1,10 +1,10 @@
 @get_units @exchange @grid_type @grid_location @with_kw struct SBM{T,N,M}
     # Model time step [s]
-    Δt::T | "s" | 0 | "none" | "none" 
+    Δt::T | "s" | 0 | "none" | "none"
     # Maximum number of soil layers
-    maxlayers::Int | "-" | 0 | "none" | "none" 
+    maxlayers::Int | "-" | 0 | "none" | "none"
     # number of cells
-    n::Int | "-" | 0 | "none" | "none" 
+    n::Int | "-" | 0 | "none" | "none"
     # Number of soil layers
     nlayers::Vector{Int} | "-"
     # Number of unsaturated soil layers
diff --git a/src/sediment.jl b/src/sediment.jl
index ad4ebea83..783c4e987 100644
--- a/src/sediment.jl
+++ b/src/sediment.jl
@@ -1,7 +1,7 @@
 ### Soil erosion ###
 @get_units @exchange @grid_type @grid_location @with_kw struct LandSediment{T}
     # number of cells
-    n::Int | "-" | 0 | "none" | "none" 
+    n::Int | "-" | 0 | "none" | "none"
     ### Soil erosion part ###
     # length of cells in y direction [m]
     yl::Vector{T} | "m"
@@ -595,7 +595,7 @@ end
 ### Sediment transport in overland flow ###
 @get_units @exchange @grid_type @grid_location @with_kw struct OverlandFlowSediment{T}
     # number of cells
-    n::Int | "-" | 0 | "none" | "none" 
+    n::Int | "-" | 0 | "none" | "none"
     # Filter with river cells
     rivcell::Vector{T} | "-"
     # Total eroded soil [ton Δt⁻¹]
@@ -685,7 +685,7 @@ end
     # number of cells
     n::Int | "-" | 0 | "none" | "none"
     # Timestep [s]
-    Δt::T | "s" | 0 | "none" | "none" 
+    Δt::T | "s" | 0 | "none" | "none"
     # River geometry (slope [-], length [m], width [m])
     sl::Vector{T} | "m"
     dl::Vector{T} | "m"

From 345fc287eee0b9c99b78f2e32bd171a99a7b6180 Mon Sep 17 00:00:00 2001
From: Willem van Verseveld <willem.vanverseveld@gmail.com>
Date: Tue, 14 Nov 2023 16:44:32 +0100
Subject: [PATCH 17/20] =?UTF-8?q?fix=20=CE=94t=20grid=20info=20(Reservoir?=
 =?UTF-8?q?=20and=20Lake)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/reservoir_lake.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/reservoir_lake.jl b/src/reservoir_lake.jl
index e0fac32b6..410745eef 100644
--- a/src/reservoir_lake.jl
+++ b/src/reservoir_lake.jl
@@ -1,5 +1,5 @@
 @get_units @exchange @grid_type @grid_location @with_kw struct SimpleReservoir{T}
-    Δt::T | "s" | 0                                     # Model time step [s]
+    Δ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⁻¹]
@@ -204,7 +204,7 @@ function update(res::SimpleReservoir, i, inflow, timestepsecs)
 end
 
 @get_units @exchange @grid_type @grid_location @with_kw struct Lake{T}
-    Δt::T | "s" | 0                             # Model time step [s]
+    Δ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³]

From 7552ba46ae2e432829508404846a116c77337d45 Mon Sep 17 00:00:00 2001
From: Willem van Verseveld <willem.vanverseveld@gmail.com>
Date: Tue, 14 Nov 2023 17:21:18 +0100
Subject: [PATCH 18/20] Explicit return

Co-authored-by: Martijn Visser <mgvisser@gmail.com>
---
 src/bmi.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/bmi.jl b/src/bmi.jl
index 323b441e0..52b80ea17 100644
--- a/src/bmi.jl
+++ b/src/bmi.jl
@@ -145,7 +145,7 @@ function BMI.get_var_grid(model::Model, name::String)
     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]))
-        if gridtype == "scalar"
+        return if gridtype == "scalar"
             0
         elseif :reservoir in key
             1

From 58708075f134ddf594220f4daa739e3e63f12691 Mon Sep 17 00:00:00 2001
From: Willem van Verseveld <willem.vanverseveld@gmail.com>
Date: Wed, 15 Nov 2023 09:54:17 +0100
Subject: [PATCH 19/20] Type grids (BMI)

Co-authored-by: Martijn Visser <mgvisser@gmail.com>
---
 src/bmi.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/bmi.jl b/src/bmi.jl
index 52b80ea17..d9ec5eba6 100644
--- a/src/bmi.jl
+++ b/src/bmi.jl
@@ -3,7 +3,7 @@
 
 # Mapping of grid identifier to a key, to get the active indices of the model domain.
 # See also function active_indices(network, key::Tuple).
-grids = Dict{Int,Tuple}(
+grids = Dict{Int,Tuple{Symbol}}(
     1 => (:reservoir,),
     2 => (:lake,),
     3 => (:drain,),

From 8cb00df93d7781b0f43ff46926f3b3e39921ebc4 Mon Sep 17 00:00:00 2001
From: Willem van Verseveld <willem.vanverseveld@gmail.com>
Date: Wed, 15 Nov 2023 13:16:34 +0100
Subject: [PATCH 20/20] Revert back to use of `const` for `grids` (BMI)

---
 src/bmi.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/bmi.jl b/src/bmi.jl
index d9ec5eba6..3f0c71979 100644
--- a/src/bmi.jl
+++ b/src/bmi.jl
@@ -3,7 +3,7 @@
 
 # Mapping of grid identifier to a key, to get the active indices of the model domain.
 # See also function active_indices(network, key::Tuple).
-grids = Dict{Int,Tuple{Symbol}}(
+const grids = Dict{Int,Tuple{Symbol}}(
     1 => (:reservoir,),
     2 => (:lake,),
     3 => (:drain,),